Energy-Based Reward Models for Robust Language Model Alignment

Thank you for your thoughtful and constructive feedback. We're glad you found the paper interesting.

1. Seem to require access to RM's training data

EBRM does require access to the Base RM’s training data. This is a natural consequence of its design as a post-hoc refinement method. Training on the same preference data avoids distribution shift, ensures compatibility with the RM’s internal representations, and enables EBRM to refine the reward function rather than relearn it from scratch. We view removing this requirement as an interesting direction for future work.

2. Hyperparameter choice for noise can be tricky

We agree that tuning the noise hyperparameter can be tricky. Injecting noise into reward targets allows the model to better adapt across tasks with varying ambiguity. As shown in Section E, Figure 5, extreme values of $\beta$ (e.g., 0.01 or 0.5) reduce performance gain, but a moderate setting (0.05-0.2) yields consistent overall gains. This indicates that EBRM is robust to a wide range of reasonable settings. In practice, a simple validation sweep on tasks of differing ambiguity is sufficient to select a robust $\beta$ that generalizes well.

3. Clarification on the Base RM Choice

We would like to clarify a key misunderstanding: the Base RM used in our initial experiments was 70M-parameter
Pythia-based model - not 70B. Smaller RMs like the 70M variant are more susceptible to overconfidence, miscalibration, and reward hacking, making them particularly suitable for evaluating robustness and the effectiveness of post-hoc refinement methods such as EBRM.

Further, comparing against ensemble baselines (3× RMs) is computationally feasible only at small scales. Using a small Base RM also allows us to leverage stronger reward models (AlpacaFarm 7B Human Preference Reward Model) as "gold" RMs for RL experimentation.

Thank you for the detailed and insightful review. We are glad that you found our approach well-motivated, practical, and promising for broader adoption.

1. Inference Time

We believe EBRM offers a compelling cost-benefit trade-off, especially compared to alternatives like ensembling and retraining.

i. EBRM increases inference latency by only $\sim$ 2× compared to the Base RM. In contrast, ensembles typically incur $\sim$n× the inference time and parameter count, where n $\ge$ 3 in practice. Retraining full RMs is substantially more expensive in both compute and time. EBRM can be trained in just 465 seconds from a 70M Pythia model with only 3% the size of a standard RM.

ii. Across five alignment tasks, EBRM outperforms ensembles by 1–5% while being 34% faster and adding $<$3% in parameter count. These gains are non-trivial given that ensembles are a strong and widely adopted baseline.

iii. EBRM requires no retraining and is applied post hoc to any RM, making it a lightweight and practical solution for deployment scenarios where retraining or maintaining ensembles is infeasible.

2. Scaling Performance

We agree that benchmarking against stronger RMs is important and have addressed this by extending our experiments to include larger and competitive models.

In particular, we now evaluate on 2.6B Pythia Reward Model and Skywork-LLaMA3.1-8B RM [2]. The Skywork RM is ranked #11 on the RewardBench leaderboard and the strongest publicly available 8B reward model. Results are summarized below:

Win rates on RewardBench for the 2.6B Pythia Reward Model.

Win rates on Reward Model Benchmark (RMB) for the 2.6B Pythia Reward Model.

Win rates on RewardBench for the 8B Skywork Reward Model.

Win rates on Reward Model Benchmark (RMB) for the 8B Skywork Reward Model.

We observe a consistent positive gain in performance as the Base RM size increases. On the 2.6B Pythia-based RM, EBRM improves average scores on both RMB (+2.5%) and RewardBench (+1.2%) tasks. EBRM continues to provide measurable improvement even at the high end of the performance spectrum, delivering +0.29% on RewardBench and +0.17% on RMB over the already strong Skywork RM.

For comparison, the recent uncertainty-aware method (URM) [1] required full RM retraining, shifting from Bradley-Terry to multi-attribute regression loss and achieved a +0.4% gain on RewardBench. EBRM, by contrast, is plug-and-play: no retraining, no architecture changes, and still comparable gains. We also note that our Skywork evaluation uses the cleaned data version, unlike prior work that unknowingly used a version with RewardBench data leakage [2]. Due to the lack of open-source code, we are unable to reproduce a direct comparison with URM. We also omit ensemble baselines here due to compute constraints.

In summary, our new results show that EBRM adds value even to top-tier RMs, reinforcing its relevance for practical and scalable alignment refinement.

Thank you for the detailed and encouraging review. We are glad you found the method principled, practical, and well-motivated for real-world RLHF applications.

1. Theoretical Explanation and comparison with quantile regression and DPO

1.1 The paper lacks a strong theoretical explanation for why modeling with an energy function leads to improved robustness.

EBRM improves robustness as it replaces pointwise reward estimation with distributional modeling, enabling it to capture uncertainty in the reward signal.

• Point estimate vs. distribution modeling: Standard scalar reward models lack the ability to represent uncertainty, which is critical as human preference data is often noisy, inconsistent, or ambiguous. This limitation leads to overfitting, poor generalization, and susceptibility to reward hacking.

• Energy-based modeling enables uncertainty: EBRM defines a conditional reward distribution $p(r|e) \propto \exp(f_\theta(e,r))$. This allows the model to assign probability mass over a range of plausible rewards, rather than collapsing to a single value. As a result, it captures uncertainty.

• Uncertainty modeling improves robustness: EBRM is trained using NCE+ with soft positive labels and negative samples. This training objective encourages the model to learn well-calibrated reward regions, not just match noisy targets. At inference, EBRM refines the Base RM outputs via energy-guided optimization, pushing predictions toward high-likelihood regions of the learned distribution. This mitigates the impact of miscalibrated spurious Base RM scores and improves robustness.

Thus, the robustness of EBRM arises from its ability to model and leverage uncertainty, enabled by the energy-based formulation that showcases the theoretical basis of our approach.

1.2 A deeper comparison with alternatives like quantile regression or DPO would clarify the unique benefits and mechanisms of EBRM.

Quantile regression [6] models uncertainty by predicting rewards at fixed quantile levels. While this captures distributional information, it learns a discrete set of targets and lacks a continuous density function over rewards. At test time it relies on a weighted average of quantile outputs to approximate the final score.

In contrast, EBRM learns a continuous conditional reward distribution via an energy function, enabling flexible uncertainty modeling and calibrated reward refinement at test time via gradient optimization. This allows EBRM to adjust predictions based on the local reward structure. We do not compare directly to quantile regression methods as the code is not publicly available.

Direct Preference Optimization (DPO), by design, lacks an explicit reward model and assumes scalar scores without modeling uncertainty. EBRM, in contrast, is a post-hoc refinement layer that improves the robustness of existing RMs by modeling uncertainty in scalar reward estimation.

2. Addressing Pseudo-Label Bias and External Validation

EBRM does not reinforce RM errors, as demonstrated by consistent gains on human-curated test sets. While it uses the Base RM outputs during training, it mitigates pseudo-label bias through two mechanisms:

i. Noise-aware training: Rather than regressing to scalar RM scores, EBRM learns a distribution over plausible rewards by injecting Gaussian noise and contrasting against negative samples. This forces the model to learn a smoother reward landscape based on relative plausibility, not the Base RM's raw score. As seen in Figure 5, tuning the noise scale affects performance across tasks: lower noise hurts ambiguous tasks (Chat), while higher noise degrades deterministic ones (Safety). This confirms that EBRM learns to generalize by modeling uncertainty, not memorizing RM bias.

ii. Filtering corrupted supervision: While EBRM does not rely on external supervision during training, we do leverage existing human preference labels for the training data to filter corrupted supervision. We discard preference pairs where the Base RM assigns a higher reward to the human-rejected response. These cases indicate unreliable or misleading supervision either due to annotation noise or RM bias and are excluded to prevent reinforcement of the Base RM errors.

Further, we validate EBRM externally at test time using two test sets - RewardBench and RMB. The performance gain (up to $\sim$ 6%) demonstrate that EBRM does not reinforce the Base RM limitations, but actively corrects them through distributional modeling and inference-time reward refinement.

Thank you for the comprehensive and thoughtful review. We appreciate your recognition of the method’s novelty, clarity, and practical relevance.

1. Cost-Benefit Trade-off

We believe EBRM offers a compelling cost-benefit trade-off, especially compared to alternatives like ensembling and retraining.

i. EBRM increases inference latency by only $\sim$ 2× compared to the Base RM. In contrast, ensembles typically incur $\sim$ n× the inference time and parameter count, where n $\ge$ 3 in practice. Retraining full RMs is substantially more expensive in both compute and time. EBRM can be trained in just 465 seconds from a 70M Pythia model with only 3% the size of a standard RM.

ii. Across five alignment tasks, EBRM outperforms ensembles by 1–5% while being 34% faster and adding $<$3% in parameter count. These gains are non-trivial given that ensembles are a strong and widely adopted baseline.

iii. EBRM requires no retraining and is applied post hoc to any RM, making it a lightweight and practical solution for deployment scenarios where retraining or maintaining ensembles is infeasible.

2. Drop in Chat Accuracy

The dip is not unique to EBRM. Similar declines appear in both uncertainty-aware [1,3] and standard methods [4]. This suggests that the issue arises from task ambiguity, not EBRM design. Open-ended tasks are inherently ambiguous, similar inputs can yield diverse valid preferences. EBRM’s noise-aware contrastive loss reflects this ambiguity by flattening the energy surface, leading to lower confidence but better calibrated uncertainty. To validate this interpretation, we analyzed the shape of EBRM’s reward distributions for a subset of each RewardBench task:

i. Kurtosis indicates a distribution's peakedness: $3$ indicates normal distribution, $>3$ (leptokurtic) implies sharper, more confident predictions, and $<3$ (platykurtic) indicates flatter, more uncertain ones.

ii. Variance captures the distribution’s spread.

Mean and standard deviation of the variance and kurtosis values.

Chat exhibits the highest variance and a platykurtic distribution, indicating that EBRM correctly models the ambiguity in these tasks. While this may reduce pairwise win rates, it improves the model’s calibration and mitigates the risk of overfitting.

3. Methodology: Circular Dependency and Filtering

Our data filtering strategy is not circular nor aggressive—it improves training signal quality by removing contradictory or corrupted supervision, not difficult cases.

• The filtered examples were excluded because they provide misleading signals that degrade EBRM's performance. As shown in Figure 6, without filtering, the performance of EBRM will slightly degrade. While it is possible to correct such examples, we are unaware of a principled way to do so, as it is often unclear whether the error lies with the Base RM, the human annotation, or both. Additionally, using high-quality external RMs for correction could introduce reward scale mismatches, leading to further inconsistency in EBRM’s training data.

• The 25% filtering only applies to the small RM (70M), where the capacity is not high. For the 1.3B RM, only 8 pairs were filtered. This clearly shows that filtering is not “aggressive”, but a step to avoid injecting harmful signals into training.

• EBRM does not blindly trust the Base RM. In fact, it treats the Base RM’s outputs as noisy priors and uses noise-aware training to learn the distribution rather than regressing to the Base RM. This allows EBRM to correct biases instead of reinforcing them.

In summary, filtering improves EBRM by training on high-quality, consistent signals. Combined with noise-aware training, it enables EBRM to identify and correct the Base RM errors instead of reproducing them.