On the Robustness of Reward Models for Language Model Alignment

TL;DR

This work identifies excessive dispersion of hidden-state norms in reward models trained with the Bradley-Terry objective as a key driver of over-optimization and poor generalization. It introduces batch-wise sum-to-zero regularization (BSR), forming the BT-BSR objective $\,L_{BT-BSR} = L_{BT} + \lambda L_{BSR}$ with $L_{BSR} = ( \\frac{1}{2|\\mathcal{B}|} \\sum_{i=1}^{|\\mathcal{B}|} \\sum_{j \in {w,l}} r(x_i,y_{i,j}) )^2$, to cap extreme reward magnitudes and curb hidden-state dispersion. Across four generalization scenarios and multiple model families, BT-BSR yields superior robustness and better alignment of the policy $\,\\pi$ with the gold preference model $r^*$ during RLHF, with further gains on 8B-scale data, including a 40% reduction in generation length and a 7% rise in win rate in AlpacaEval 2.0. The paper demonstrates the practical impact of RM robustness on RLHF training and real-world tasks, and provides code, data, and models to facilitate adoption and further study.

Abstract

The Bradley-Terry (BT) model is widely practiced in reward modeling for reinforcement learning with human feedback (RLHF). Despite its effectiveness, reward models (RMs) trained with BT model loss are prone to over-optimization, losing generalizability to unseen input distributions. In this paper, we study the cause of over-optimization in RM training and its downstream effects on the RLHF procedure, accentuating the importance of distributional robustness of RMs in unseen data. First, we show that the excessive dispersion of hidden state norms is the main source of over-optimization. Then, we propose batch-wise sum-to-zero regularization (BSR) to enforce zero-centered reward sum per batch, constraining the rewards with extreme magnitudes. We assess the impact of BSR in improving robustness in RMs through four scenarios of over-optimization, where BSR consistently manifests better robustness. Subsequently, we compare the plain BT model and BSR on RLHF training and empirically show that robust RMs better align the policy to the gold preference model. Finally, we apply BSR to high-quality data and models, which surpasses state-of-the-art RMs in the 8B scale by adding more than 5% in complex preference prediction tasks. By conducting RLOO training with 8B RMs, AlpacaEval 2.0 reduces generation length by 40% while adding a 7% increase in win rate, further highlighting that robustness in RMs induces robustness in RLHF training. We release the code, data, and models: https://github.com/LinkedIn-XFACT/RM-Robustness.

Figures (6)

• Figure 1: $||W_p||$ distribution after reward modeling for four seeds each. $||W_p||$ generally stays around one after the training.
• Figure 2: Growth of $||h(x, y_w) - h(x, y_l)||$ throughout reward modeling with $\mathcal{L}_\text{BT}$. The variance of the hidden state difference grows incrementally with the right-skewed distribution. The width of the colored area indicates the standard deviation.
• Figure 3: Hidden state norm dispersion comparison between $\text{RM}_\text{BT}$ and $\text{RM}_\text{BT-BSR}$. Batch-wise sum-to-zero regularization (BSR) alleviates hidden state norm dispersion and demonstrates a consistent range of norms across different generalization scenarios in Section \ref{['subsec:setup']}. Reducing the variability of hidden state norms, $\mathcal{L}_\text{BSR}$ improves the robustness of RMs for unseen data.
• Figure 4: Assessing the robustness of each reward modeling objective on four generalization scenarios in Section \ref{['subsec:setup']}. By applying $\mathcal{L}_\text{BT-BSR}$, downstream RMs are more robust to unseen prompt sets or response generation model sets by being more aligned to the gold preference model $r^*$, measured through accuracy and Kendall's $\tau$ against the preference annotations of $r^*$. Full results are reported in Table \ref{['tab:rm-oop']}.
• Figure 5: Reward to KL divergence plots (Figures \ref{['subfig:kl-bt']} and \ref{['subfig:kl-BSR']}) and alignment to gold reward model $r^*$ (Figures \ref{['subfig:gold-rm-1st']} and \ref{['subfig:gold-rm']}) while fine-tuning Qwen2.5-1.5B SFT model as $\pi_\theta$ with RLOO using $\text{RM}_\text{BT}$ and $\text{RM}_\text{BT-BSR}$, respectively. We measure the alignment with the gold preference model $r^*$ to verify if reward maximization with each RM as a proxy is leading to maximizing $r^*$ scores in Figures \ref{['subfig:gold-rm-1st']} and \ref{['subfig:gold-rm']}. Using $\text{RM}_\text{BT-BSR}$ consistently improved the reward assessed by $r^*$, while $\text{RM}_\text{BT}$ stagnated in the last half of the training.