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

1. Problem Statement

Reward models (RMs) are central to Reinforcement Learning from Human Feedback (RLHF), serving as the proxy for human preferences that guide Large Language Model (LLM) alignment. The standard objective for training these models is the Bradley-Terry (BT) loss, which learns from pairwise data consisting of "chosen" ($y_w$) and "rejected" ($y_l$) responses.

The authors identify a critical flaw in the update dynamics of the standard BT loss: Representation Distance Bias. While the BT gradient is intended to be driven by prediction error (the magnitude of the mistake in ranking), the analysis reveals that the gradient norm scales with two distinct components:

1. Prediction Error: The difference in predicted rewards between the chosen and rejected responses.
2. Representation Distance: The distance between the two responses in the output space of the final layer ($\|h_w - h_l\|$).

This coupling creates a spurious learning signal where the magnitude of the update is dictated by how far apart the representations of the pair are, rather than solely by how wrong the model is.

• Small-Distance Pairs: Pairs that are semantically or structurally similar (e.g., a correct code snippet vs. one with a single logical error) often have small representation distances. Under standard BT loss, these pairs receive vanishingly weak updates, even when the model misranks them. This is problematic because fine-grained distinctions are often the most critical for tasks like reasoning and coding.
• Large-Distance Pairs: Pairs that are vastly different (e.g., a refusal vs. an answer to a harmful query) naturally have large representation distances. These pairs receive disproportionately strong updates, potentially overshadowing the learning signals from the more subtle, fine-grained pairs.

Consequently, the model under-trains on the difficult, fine-grained distinctions where it struggles most, while over-emphasizing superficial differences.

2. Methodology: NormBT

To address this bias, the authors propose NormBT, an adaptive pairwise normalization scheme that rescales the gradient contribution of each preference pair.

Theoretical Basis

The authors decompose the BT gradient norm $\|\nabla_\theta L_{BT}\|$ and show it is proportional to:
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
where $d$ is the reward difference and $\|h_w - h_l\|$ is the representation distance. To make the update size depend primarily on the prediction error term $|\sigma(d) - 1|$, the authors introduce a weight $w_i$ that inversely scales with the representation distance.

Implementation Details

Directly computing the exact gradient norm for weighting is computationally expensive. Instead, NormBT uses the final-layer representation distance ($\|h_w - h_l\|$) as a proxy, which is shown to be an exact proxy for the linear score head and a strong approximation for the full model.

The proposed loss function is:
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

Where the weight $\tilde{w}$ is defined as:
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

Key design choices include:

• Stop-Gradient: The distance term is computed with a stop-gradient operator to prevent the normalization weight itself from being optimized.
• EMA Stabilization: To prevent scale drift during training (as embedding magnitudes can shift), the authors use an Exponential Moving Average (EMA) of the batch mean representation distance ($\mu_t$) to normalize the weights. This ensures the weights remain stable and adaptive to distributional shifts.
• Lightweight Integration: NormBT requires no architectural changes, preserves the probabilistic foundation of BT, and incurs negligible computational overhead.

3. Key Contributions

1. Gradient Analysis: The paper provides a detailed theoretical and empirical analysis showing that BT gradient norms are jointly determined by prediction error and representation distance, leading to biased learning signals.
2. Identification of Bias: The authors demonstrate that this coupling causes small-distance pairs (often the most critical for fine-grained tasks) to receive weak updates, while large-distance pairs dominate the learning signal.
3. NormBT Proposal: They introduce NormBT, a drop-in modification that normalizes updates by representation distance, ensuring learning signals focus on prediction errors.
4. Empirical Validation: Extensive experiments across different LLM backbones (Gemma-2b-it, Llama-3.2-3B) and datasets (Unified-Feedback, Skywork-Reward) show consistent performance improvements.

4. Experimental Results

The authors evaluated NormBT against the standard BT baseline and several variants (BT + Margin, BT + Label Smoothing) on the RewardBench benchmark and downstream Best-of-N (BoN) selection tasks.

• RewardBench Performance: NormBT consistently outperformed the BT baseline across all settings. The most significant gains were observed in the Reasoning category, where accuracy improved by over 5% on average. This aligns with the hypothesis that Reasoning tasks rely on fine-grained distinctions (small representation distances) that are suppressed by standard BT.
• Downstream RLHF (BoN): In Best-of-N selection experiments, models trained with NormBT achieved higher "gold scores" (quality of selected responses) compared to BT baselines.
• Ablation Studies:
  • NormBT outperformed other training techniques like LayerNorm and Gradient Clipping, which failed to address the relative weighting issue between pairs.
  • The method proved complementary, improving performance even when applied on top of other BT variants (e.g., NormBT + Margin).
  • Removing the EMA stabilization or using alternative similarity metrics (cosine similarity, average-pooled embeddings) resulted in performance degradation, validating the specific design choices.

5. Significance and Claims

The paper claims that NormBT offers a principled and lightweight solution to a structural limitation in the standard BT loss. By decoupling the update magnitude from representation distance, NormBT ensures that the reward model learns more faithfully from fine-grained distinctions, which are often the most challenging yet critical for high-quality alignment.

The authors emphasize that their work does not address data quality or annotator bias directly but improves the efficiency and fidelity of the learning signals extracted from existing pairwise data. They position NormBT as a foundational improvement for the discriminative reward modeling paradigm, potentially leading to more reliable reward signals in downstream RLHF applications, particularly in domains requiring precise reasoning and subtle error detection.