On the Robustness of Reward Models for Language Model Alignment

We appreciate your suggestions on the typo and the section title. We will make sure to address them in the final version of our paper.

 

Q1 - Quantifying disjointness between datasets: We appreciate the reviewer’s question for disjointness quantification as it would make the splitting criteria clearer. We originally planned on splitting the datasets based on the cosine distances of the embeddings of the prompt and responses. However, due to the diverse split settings outlined in Section 2.2, we were unable to provide a strict quantification rule of thumb for disjointness quantifications.

Q3 - Over-confidence vs over-optimization: For clarity, we address the third question before the second one. We view the over-confidence issue, which is originally studied for the conventional multi-class classifiers, as a core cause of the reward model over-optimization problem in the context of reward modeling formulation. [1] shows that an overly growing magnitude of logits triggers classifiers to be over-confident due to the nature of the softmax function. In this paper, we carefully analyze how the backbone model’s hidden state norm is a core cause of the growth in reward magnitude (i.e., which is equivalent to the logit magnitude) as they are a unique type of classifier with two classes but with a single projection head, unlike conventional multi-class classifiers. We continue how over-confidence analysis in RMs connects to over-optimization in Q2.

Q2 - Clarifications on lines 184-187: This explanation intends to describe RMs’ formulation as a special case of conventional classifier models, where [1] studies the over-confidence problem, and builds how this over-confidence issue in RM formulation is connected to the over-optimization problem. Reward modeling loss $\mathcal{L}\_\mathrm{BT}$ can be written as a two-class softmax classification. While the over-confidence problem originates from growing logit magnitude in conventional classifiers [1], we are analyzing this problem in the RM context. As the magnitude of the reward (i.e., logit) comes largely from $||h(x, y_||$ by $||W_p|| \simeq 1$ even after training, as we showed, the over-confidence problem in RMs stems from the dispersion of hidden states. As the hidden dimension grows (which is common in larger language models), the norm tends to increase, thereby potentially exacerbating over-optimization issues. This is why these lines of explanation bridge the important connection between the over-confidence and over-optimization issues in RMs.hmm (briefly talked about utilizing embedding model + cosine distance for quantification)

Q4 - Clarification in Figures 5(c) and 5(d): Appendix A.3 explains that we used a linear decay scheduler during RLOO training. Consequently, the later phase of training operates with a lower learning rate, which leads to slower parameter updates and a tendency for the model to overfit the prompt over multiple epochs. Moreover, note that the y-axis scales differ between Figures 5(c) and 5(d): Figure 5(c) covers a range of [0.0, 1.0], while Figure 5(d) focuses on [0.85, 1.0] to allow a more fine-grained comparison between the policies trained with $RM_{\text{BT}}$ and $RM_{\text{BT-BSR}}$. In this context, the result can be interpreted as $RM_{\text{BT}}$ stagnating in the later phase of training, reaching about 95% of the maximum reward, while $RM_{\text{BT-BSR}}$ reaches that level by epoch 5. The transient dip and subsequent boost in the curve likely reflect the combined effects of the decaying learning rate and overfitting dynamics, resulting in non-monotonic updates to the proxy reward signal.

 

Reference

[1] Wei, Hongxin, et al. "Mitigating neural network overconfidence with logit normalization." International conference on machine learning. PMLR, 2022.

We thank the reviewer for the insightful comments and will add the suggested paper to the related works.

 

Q1 - Normalization techniques to implicit rewards of DAAs: DAAs use the language model as an implicit reward model, differing from our classifier-based reward models in that they use a sum of log ratios rather than a linear projection. As our normalization, BT-BSR, starts with the solid background on the over-confidence issue that classifiers are prone to, it may not be intuitive to directly apply batch-wise sum-to-zero constraint on implicit RMs. Applying a batch-wise sum-to-zero constraint directly to implicit rewards, which often exhibit negative scales [1], would be counterintuitive.

Nevertheless, over-optimization can also occur with implicit rewards [2], as evidenced by their divergence to negative infinity [3]. As such, divergence is similar to hidden state norm dispersion in our scenario, the straightforward solution could be Z-loss formulation [4,5]:

$$\mathcal{L}\_{IRM-BSR} = - \lambda \log^2 \sum\_i^{N} \left( \beta \log \frac{\pi\_\theta(y\_{i, w}|x)}{\pi\_\mathrm{ref}(y\_{i, w}|x)} + \beta \log \frac{\pi\_\theta(y\_{i, l}|x)}{\pi\_\mathrm{ref}(y\_{i, l}|x)} \right).$$

Z-loss is a well-established method for stabilizing large-scale language model pre-training by penalizing divergence in logits. Here, $\mathcal{L}\_\mathrm{IRM-BSR}$ prevents the implicit reward (log-ratio) from diverging to negative values, mirroring how Z-loss regulates the softmax normalizing constant.

We conducted a minimal experiment on Qwen2.5-1.5B trained with TULU3-SFT data [6] comparing plain DPO with DPO plus the $\mathcal{L}\_\mathrm{IRM-BSR}$ regularizer. As shown in Tables A, B, and C, the Z-loss-based regularizer substantially mitigates negative divergence in log-likelihood while preserving training preference accuracy.

 

Table A: Log-likelihood of chosen responses during training.

Table B: Log-likelihood of rejected responses during training.

Table C: Preference accuracy for each batch during training.

 

 

Reference

[1] Rafailov, Rafael, et al. "From $r$ to $Q^*$: Your Language Model is Secretly a Q-Function." First Conference on Language Modeling.

[2] Rafailov, Rafael, et al. "Scaling laws for reward model overoptimization in direct alignment algorithms." Advances in Neural Information Processing Systems 37 (2024): 126207-126242.

[3] Shi, Zhengyan, et al. "Understanding likelihood over-optimisation in direct alignment algorithms." arXiv preprint arXiv:2410.11677 (2024).

[4] Chowdhery, Aakanksha, et al. "Palm: Scaling language modeling with pathways." Journal of Machine Learning Research 24.240 (2023): 1-113.

[5] Wortsman, Mitchell, et al. "Small-scale proxies for large-scale Transformer training instabilities." The Twelfth International Conference on Learning Representations.

[6] Lambert, Nathan, et al. Tulu 3: Pushing frontiers in open language model post-training." arXiv preprint arXiv:2411.15124 (2024).

We thank the reviewer for the comment and would like to further discuss the addressed points.

 

W1 - Limited Experimental Results: Our experimental design can be streamlined into threefold: (1) assessing alignment between proxy RMs and gold RMs with different learning objectives, (2) propagation of over-optimization in RMs to their downstream usage for RLHF training, and (3) understanding real-world impact of such propagation with state-of-the-art model and data. This line of analysis is fully supported with Qwen2.5 and Llama3.2 series with 1B and 3B scale on carefully curated UltraFeedback, followed by Llama-3.1-8B model and 80k Skywork-Preferences dataset. While the clarity in writing could be a valid concern if the logical flow is not fully conveyed, we believe the paper encompasses a large range of models and datasets, especially aligning with practical usage in the research community (e.g., the majority of RMs in RewardBench [1] are based on Llama-3.1 or Llama-3.2 series and Skywork-Preferences dataset).

To ensure that the concerns on the width of experiments, we extend the experiment of Figure 5 (Section 5.2) by scaling the model to Qwen2.5-3B in Table A. This result further validates that the trend shown in Qwen2.5-1.5B on the propagation of RM robustness in the RLHF stage on a larger scale. If the concerns about the clarity of writing can be specified, we would happily provide further explanation and an update in the final version of the paper.

 

W2 - Novelty of Batch-wise Sum-to-Zero Constraint in Reward Modeling: We appreciate the concern on methodological novelty, but there seems to be a major confusion on the zero-centering aspect of BSR. While [1] is cited as the core reference to point out that the idea of centering rewards to zero is not novel, we do not find it to be connected to our approach:

(Excerpt from Discussion section of RewardBench according to reviewer’s pointer)...Few RMs are Gaussian in their scores across the REWARDBENCH datasets, fewer RMs are centered around 0 reward, and none we tested centered Gaussians. Future work should identify a preferred RM output distribution for downstream RL training.

This is the only part that we find in [1] that mentions centering. However, they intend that RMs are rarely centered through post-hoc analysis. To the best of our knowledge, [2] is the only work that adopts zero-centering per-prompt. However, their core motivation lies in the identifiability of RMs, which is distant from ours. We build a solid reason for adopting a batch-wise sum-to-zero constraint regarding the mitigation of reward model over-optimization. We would happily incorporate the discussion on [2] for our final version of the paper. However, we emphasize that our method remains novel from its theoretical background and empirical support with varying model sizes, types, and datasets.

 

Table A. Qwen2.5-3B-SFT trained with RLOO with each reward model for 5 epochs, evaluated by the gold reward model (ArmoRM). The reward scores are NOT normalized, unlike Figure 5. Unnormalized score of ArmoRM typically ranges around 0.1-0.2 [3]

 

Reference

[1] Lambert, Nathan, et al. "Rewardbench: Evaluating reward models for language modeling." arXiv preprint arXiv:2403.13787 (2024).

[2] Eisenstein, Jacob, et al. "Helping or Herding? Reward Model Ensembles Mitigate but do not Eliminate Reward Hacking." First Conference on Language Modeling.

[3] Wang, Haoxiang, et al. "Interpretable Preferences via Multi-Objective Reward Modeling and Mixture-of-Experts." Findings of the Association for Computational Linguistics: EMNLP 2024. 2024.

We appreciate the reviewer’s detailed comments. Below, we address each concern:

1. Claims/Evidence #1 – Notations in Section 4.2: We regret the confusion from the ambiguous explanations in lines 261–271 and acknowledge that lines 265–269 apply only to a specific case. In the final version, we will include a concise analysis of $\frac{\partial \mathcal{L}\_\mathrm{BSR}}{\partial h(x,y)}$, showing that it introduces a linear penalty on the growing reward scale, which indirectly limits the expansion of $\|h(x,y)\|$. Thus, $\mathcal{L}\_\mathrm{BSR}$ properly complements the BT loss with regularization.

2. Claims/Evidence #2 – Is BSR properly regularizing the hidden state norms?: Logit normalization (i.e., $\mathcal{L}\_\text{BT-Norm}$ in Section 4.3) decouples the reward margin from the hidden state norm by dividing the reward by its L2 norm, making the loss scale-invariant [1]. Our analysis in Section 4.1 shows that the reward scale is largely driven by the hidden state norm since $||W_p|| \simeq 1$. However, Figure 4 demonstrates that the strictly normalized loss (RM$_\text{BT-Norm}$) performs substantially worse, suggesting that completely discarding magnitude information removes a valuable discriminative signal for out-of-distribution generalization. In contrast, our BSR method softly regularizes the norm, mitigating over-optimization while preserving useful information. Moreover, Table A shows that even with BSR, the projection head’s norm remains around 1, confirming that the BSR penalty is effectively propagated without fully suppressing the norms.

3. Theoretical Claims #1 – Ambiguity in “unconstrained norm dispersion”: We acknowledge that “unconstrained” may not be ideal, given the sigmoid function’s gradient mechanism. Nonetheless, without a proper regularizer like BSR, norm dispersion can be amplified, especially under certain hyperparameter choices (e.g., learning rate), and we have fully demonstrated this throughout the paper. In the final version, we will revise the term to “excessive norm dispersion” or similar.

4. Strengths and Weaknesses #1 – Generalizability of the Method: We emphasize that the ultimate evaluation of reward models (RMs) is based on their performance in the RLHF stage. As shown in [2], downstream performance varies significantly with the RM used. Despite marginal improvements in RM benchmarks in some cases, BT-BSR outperforms in the RLHF stage while mitigating verbosity bias [3]—as evidenced by shorter responses and a higher win rate (Table 3). Combined with the improvements seen in RM benchmarks (Figure 4 and Table 2), our regularization method is both generalizable and practically valuable.

5. Comments or Suggestions #1 – Further Validation for Figures 1 and 2 with BT-BSR: Table A confirms that $\|W_p\| \simeq 1$ for BT-BSR checkpoints (as shown in Figure 1). Figure 2 empirically supports that the BT loss triggers norm dispersion, while the comparison between Figures 3(a) and 3(b) confirms that BSR effectively mitigates this phenomenon.

6. Comments or Suggestions #2 – Questions on Notations: Since $\mathbb{D}\_\mathrm{KL}$ is conditioned on the prompt $x$, both the reward and the KL penalty should be under $x \sim \mathcal{D}$. Additionally, lines 184–188 bridge the over-confidence issue in classic multi-class classifiers to RMs with shared projection heads. (For further details, please refer to our comment on “Q2” for reviewer YT72.) We will clarify these notations in the final version.

7. Question #1 – Drop in Easy Tasks: Appendix O of RM-Bench [4] shows that “hard accuracy” has the highest correlation with policy performance ($r=0.45$) while “easy accuracy” is near 0 ($r=0.07$). Although the drop in “easy accuracy” may stem from a compressed representation space due to BSR, it is crucial that “hard accuracy” aligns with actual policy performance after RLHF. Our experiments in Figure 5 and Table 3 confirm that BT-BSR improves outcomes in both controlled and practical settings.

 

Table A: The norm of the projection head for four different models after reward modeling with BT-BSR.

 

Reference

[1] Wei, Hongxin, et al. "Mitigating neural network overconfidence with logit normalization." International Conference on Machine Learning. PMLR, 2022.

[2] Meng, Yu, Mengzhou Xia, and Danqi Chen. "Simpo: Simple preference optimization with a reference-free reward." Advances in Neural Information Processing Systems 37 (2024): 124198–124235.

[3] Yann Dubois, et al. “Length-controlled alpacaeval: A simple way to debias automatic evaluators.” First Conference on Language Modeling.

[4] Liu, Yantao, et al. "RM-Bench: Benchmarking Reward Models of Language Models with Subtlety and Style." The Thirteenth International Conference on Learning Representations.