When Distance Distracts: Representation Distance Bias in BT-Loss for Reward Models

AI summary

Overview

Research area: Large Language Model alignment, specifically reward modeling with the Bradley-Terry (BT) loss within the RLHF pipeline, with a focus on per-sample gradient analysis and training-objective design.

Technical level: Advanced. The paper derives and interprets per-sample gradient decompositions for a linear reward head attached to an LLM backbone, then proposes a modified loss and validates it empirically.

One-sentence scope: The paper shows that the magnitude of a BT-loss update is driven partly by the representation distance between chosen and rejected responses rather than by prediction error alone, and proposes a lightweight pair-wise normalization (NormBT) that removes this bias.

What This Paper Is About

Reward models are trained with the Bradley-Terry loss to score a chosen response above a rejected one, and the intuition is that the strength of a parameter update should track how wrong the model currently is. The authors show that BT gradients actually depend on two coupled factors: prediction error and the distance between the chosen and rejected responses in the final-layer representation space — so close pairs get weak updates even when misranked, and far pairs get strong updates even when already ranked correctly. Their goal is to rescale the per-pair contribution so update magnitude is governed mainly by prediction error.

Key Contributions

1. A detailed per-sample gradient analysis of the BT loss, decomposing the gradient norm into a prediction-error term |σ(d) − 1| and a representation-distance term k‖h_w − h_l‖, with k = sqrt(1 + (L_g‖w_s‖)²).
2. Evidence that this coupling misaligns learning signals: small-distance pairs receive vanishingly weak updates even when misranked, while large-distance pairs receive disproportionately strong updates, undermining learning from fine-grained distinctions.
3. NormBT, an adaptive pair-wise normalization scheme (with an EMA-based stabilizer) that rescales each pair's gradient contribution by its representation distance, preserving the probabilistic BT foundation and incurring negligible overhead.
4. An empirical demonstration across two backbones and two datasets that NormBT consistently improves reward model performance, with notable gains on the Reasoning category of RewardBench and on Best-of-N response selection.

Main Findings

• Gradient decomposition: The BT gradient norm factorizes as |σ(d) − 1| × k‖h_w − h_l‖, where d = r_w − r_l is the predicted reward difference (prediction error) and ‖h_w − h_l‖ is the representation distance at the final layer (Eq. 7). The score-layer component of this distance term is exact, not merely an approximation.

• Distance dominates within fixed error bins: When preference pairs are grouped into bins with nearly identical prediction error (similar d), observed gradient norms still increase systematically with representation distance. The paper presents this as direct evidence that distance dominates the learning signal even when the model is "equally wrong."

• Category-level disparity on RewardBench: Under the BT objective, gradient norms vary widely by RewardBench category. The Reasoning category exhibits the smallest average gradient norm and also the smallest representation distance, while Safety and Chat pairs are much farther apart in representation space.

• Main results on RewardBench (Table 1): Across four settings (gemma-2b-it and Llama-3.2-3B-Instruct, each on Unified-Feedback 80K and Skywork-Reward-Preference-80K-v0.2):

  • gemma-2b-it / Unified-Feedback: BT 72.25 average → NormBT 73.57; Reasoning 75.41 → 80.71.
  • Llama-3.2-3B-Instruct / Unified-Feedback: BT 75.24 → NormBT 76.96; Reasoning 71.70 → 76.93.
  • gemma-2b-it / Skywork-80K-v0.2: BT 78.63 → NormBT 80.12; Reasoning 77.46 → 80.71.
  • Llama-3.2-3B-Instruct / Skywork-80K-v0.2: BT 80.31 → NormBT 81.48; Reasoning 67.05 → 74.60.
  • The paper reports gains of over 5% on the Reasoning category on average.

• Comparison to margin and label-smoothing baselines: BT + margin (Touvron et al., 2023) and BT + margin out (Wang et al., 2025) do not yield consistent improvements over vanilla BT and are described as possibly overfitting noisy margin annotations. BT + label smoothing (Liu et al., 2024a) also does not improve over vanilla BT; the largest drop appears in Reasoning (e.g., 75.41 → 72.28 for gemma-2b-it on Unified-Feedback). None of these vary the representation-distance factor.

• Best-of-N selection (Figure 3): Evaluated with reward models trained on gemma-2b-it and Unified-Feedback (80K) on a 1K held-out test set, sampling n responses per prompt with n ranging from 1 to 402 (roughly KL 0 to 5), scored by a 7B gold reward model (reward-model-Mistral-7B-instruct-Unified-Feedback). NormBT achieves higher gold scores than all BT baselines, and applying NormBT on top of each BT variant yields further improvements. The paper reports the gain is most pronounced for label smoothing.

• Where the gains come from (Figure 4): Analyzing RewardBench pairs where the BT baseline and NormBT disagree, binned by representation distance, the clearest gains appear in the small-distance regime; medium- and large-distance regimes show parity or modest gains. A small degradation is noted for large-distance pairs, e.g., 78.65 → 77.97 on Safety for gemma-2b-it on Unified-Feedback.

• Ablations (Table 2): Removing EMA drops the average to 67.78; replacing the last-token embedding L2 distance with average-pooled L2 gives 69.68; using cosine similarity gives 71.88; full NormBT gives 73.57. LayerNorm gives 72.25, gradient clipping 72.69, and batch-wise sum-to-zero regularization (BSR, Hong et al., 2025) gives 72.23 — with BSR reporting 74.40 versus NormBT's 80.71 on Reasoning.

• EMA coefficient (Table 3): Varying β beyond the default 0.9 gives averages of 72.53 (β = 0.70), 73.20 (β = 0.80), 73.62 (β = 0.85), and 73.57 (β = 0.90). The paper states β in [0.8, 0.9] provides a robust option.

• Additional benchmarks (Table 4): On RM-Bench, BT scores 57.87 overall versus NormBT 58.98; on RewardBench-2, 30.08 versus 34.17; on Preference Proxy Evaluations (PPE), 50.82 versus 51.33. These benchmarks are described as generally more challenging than RewardBench.

Methodology in Plain English

The authors start from the standard setup where a reward model is a pretrained LLM backbone with a linear score head that outputs a scalar reward for a prompt-response pair. The Bradley-Terry loss pushes the chosen response's reward above the rejected response's reward. They write out the gradient of this loss for a single pair and split it into two multiplicative pieces. The first piece depends only on how wrong the current reward prediction is. The second piece depends on how different the two responses are in the backbone's final-layer representation — essentially, how far apart the two embeddings are.

They then verify this empirically with three analyses: binning pairs by prediction error to see if distance still predicts gradient size within bins; comparing gradient norms across RewardBench categories; and comparing representation distances across those same categories.

Their fix is a per-pair weight equal to the inverse of the representation distance, wrapped in a stop-gradient so it does not itself get differentiated. Since a running scalar is more stable than raw inverse-distance, the weight is computed relative to an exponential moving average (EMA) of the batch mean representation difference, with β = 0.9 by default, so the weight stays near unity on average and is robust to embedding scale drift during training. The result is a drop-in modification to the BT loss requiring no architecture changes and using only quantities already available from the forward pass.

Why This Matters

This work reframes a long-standing, largely treated-as-black-box component of the RLHF pipeline — the reward model's training objective — as an object with identifiable per-pair pathologies. It offers a mechanistic explanation for why reward models underperform on reasoning-heavy, fine-grained preference data, and it does so with a fix that is simple enough to be adopted as a one-line change.

Real-world applications:

• RLHF pipelines for general-purpose chat assistants, where the reward model's quality directly governs policy alignment.
• Code generation and verification, where chosen and rejected responses often differ by a single logical error and are nearly identical in surface form.
• Safety alignment, where preferred refusals and dispreferred harmful answers are stylistically far apart and, per this analysis, may already receive strong BT updates.
• Best-of-N response selection in deployment, where the reward model ranks candidate outputs and the paper reports improved gold scores from NormBT-trained reward models.

Industry relevance: reward modeling is a standard stage in commercial LLM alignment. NormBT requires no new annotations, no external signals (unlike margin-based methods), no architectural changes, and negligible computational overhead, which makes it a low-cost candidate for teams already training BT-based reward models.

Future Directions

• Extending the analysis beyond the linear-score-head, discriminative BT paradigm to generative reward models and to rationales, critiques, or verifier-style reward signals.
• Understanding whether representation distance is itself partly a learned artifact of the backbone and how it evolves over training, given the paper's finding that raw normalization weights drift and require EMA stabilization.
• Combining the distance normalization with uncertainty-aware or distribution-aware reward modeling, and with direct preference optimization methods that implicitly rely on BT-style pairwise signals.
• Determining when down-weighting large-distance pairs is undesirable, since the paper reports small degradations in Safety (78.65 → 77.97 for gemma-2b-it on Unified-Feedback) while gains concentrate in small-distance pairs.
• Further tuning and theoretical characterization of the EMA coefficient and the choice of distance proxy across other backbones and training scales.

Target Audience

Researchers and engineers working on RLHF and reward modeling who already understand the Bradley-Terry objective and want a gradient-level account of its behavior. The paper is best suited to readers comfortable with per-sample gradient derivations, Lipschitz-smoothness arguments, and standard LLM alignment terminology. Practitioners seeking an inexpensive drop-in improvement to reward model training will also find the empirical sections and ablation tables directly actionable, while readers new to alignment will benefit most from the qualitative framing in the Introduction and Figure 1.

Authors’ abstract

Reward models are central to Large Language Model (LLM) alignment within the framework of RLHF. The standard objective used in reward modeling is the Bradley-Terry (BT) loss, which learns from pairwise data consisting of chosen and rejected responses. In this work, we analyze the per-sample gradient of BT-loss and show spurious learning signals due to representation distance. In particular, BT gradient norm scales with two distinct components: (1) prediction error, reflected by the difference in predicted rewards between chosen and rejected responses, and critically, (2) representation distance between the pair measured in the output space of the final layer. While the first term captures the intended training signal, the second term can significantly impact the update magnitude and misalign learning. Specifically, pairs with small representation distance often receive vanishingly weak updates, even when misranked, while pairs with large distance receive disproportionately strong updates. This leads to gradients from large-distance pairs to overshadow those from small-distance pairs, where fine-grained distinctions are especially important. To overcome this limitation, we propose NormBT, an adaptive pair-wise normalization scheme that rescales updates to balance representation-driven effects and focuses learning signals on prediction error. NormBT is a lightweight, drop-in modification to BT loss with negligible overhead. Across various LLM backbones and datasets, NormBT improves reward model performance consistently, with notable gains of over 5% on the Reasoning category of RewardBench, which contains numerous fine-grained pairs.