Technical Summary: Normalized Rewards for Preference Optimization

Problem Statement

Direct Alignment Algorithms (DAAs), such as Direct Preference Optimization (DPO) and SimPO, have become standard for aligning Large Language Models (LLMs) with human preferences due to their computational efficiency compared to Reinforcement Learning from Human Feedback (RLHF). However, recent work has identified a critical limitation: DAAs often suffer from over-optimization of their implicit reward models.

A primary manifestation of this is likelihood displacement, a phenomenon where the likelihoods of both preferred and non-preferred responses drop simultaneously during training. This occurs because standard DAA objectives (DPO, SimPO) optimize the difference in rewards between response pairs, causing the partition function $Z(x)$ to cancel out mathematically. Consequently, the model is not incentivized to maintain a calibrated estimate of the total probability mass over the response space. This lack of normalization leads to spurious probability shifts, particularly for "unseen" responses (those not in the training pair), potentially resulting in undesirable or harmful behaviors.

Existing hypotheses for over-optimization—such as high embedding similarity between pairs or insufficient regularization—fail to generalize across both reference-based (e.g., DPO) and reference-free (e.g., SimPO) methods. Furthermore, prior analyses often focus on single instances of DAAs rather than the shared mechanism of normalization.

Methodology

The authors propose that reward over-optimization stems from a lack of normalization in the implicit reward. To counteract this, they introduce a regularization term designed to conserve the total length-normalized probability mass of the chosen ($y_w$) and rejected ($y_l$) responses within a preference pair.

The Regularization Term

The core insight is that for a parameterized language model applying softmax, the output distribution should remain normalized (partition function = 1) throughout training. Since estimating the full partition function over all possible responses is infeasible, the authors enforce normalization on the seen responses.

They define a regularization penalty $R(\pi_\theta, \pi_{ref}, x, y_w, y_l)$:
$$R = \lambda \left( \log \frac{\bar{\pi}_\theta(y_w|x) + \bar{\pi}_\theta(y_l|x)}{\bar{\pi}_{ref}(y_w|x) + \bar{\pi}_{ref}(y_l|x)} \right)^2$$
Where:

• $\bar{\pi}$ denotes length-normalized probabilities (to mitigate length bias).
• $\lambda$ is a hyperparameter.
• The term minimizes the squared log-ratio of the total probability mass assigned to the pair by the current model versus the reference model.

This penalty is added to the standard objectives:

1. N-DPO: $L_{N-DPO} = L_{DPO} + R$
2. N-SimPO: $L_{N-SimPO} = L_{SimPO} + R$

Mechanistic Analysis

The paper provides a theoretical gradient analysis of this regularization term. The gradient of $R$ with respect to model parameters $\theta$ is inversely proportional to the likelihood of individual tokens.

• Outlier Tokens: The analysis reveals that likelihood displacement is driven by a small set of outlier tokens with very low likelihoods. These tokens dominate the gradient updates.
• Mitigation: When the total response probability decreases, the regularization term generates gradients that primarily increase the likelihood of these low-likelihood outlier tokens. This reshapes the reward distribution, preventing the massive spurious drops in probability associated with displacement, while making minimal changes to high-likelihood tokens relevant to generation.

Key Contributions

1. Identification of Normalization Deficit: The authors illustrate how the inability to account for the distribution of unseen responses (due to the cancellation of the partition function in standard objectives) leads to likelihood displacement.
2. Mechanistic Explanation: They connect likelihood displacement to the concentration of probability shifts on outlier tokens and demonstrate via gradient analysis that their regularizer mitigates these shifts by utilizing low-likelihood tokens.
3. Novel Objectives: They propose N-DPO and N-SimPO, which incorporate the proposed regularization term, applicable to both reference-based and reference-free alignment methods.
4. Empirical Validation: They demonstrate that these modified objectives improve the trade-off between generation quality and general benchmark capabilities.

Experimental Results

The authors evaluated N-DPO and N-SimPO on Mistral-7B-Instruct, Llama-3.1-8B-Instruct, and OLMo-7B-SFT using the UltraFeedback dataset.

• Generation Quality (AlpacaEval):

  • On Llama-3.1-8B-Instruct, N-DPO achieved a >20% relative increase in AlpacaEval2 (Length-Controlled) scores compared to standard DPO.
  • On OLMo-7B-SFT, N-SimPO showed a >75% relative increase in AlpacaEval2 (Raw Win-Rate) compared to SimPO.
  • The proposed methods generally outperformed other regularization baselines like DPO+SFT and DPOP.

• General Benchmark Performance:

  • Contrary to the typical trade-off where improved generation degrades general capabilities, N-DPO and N-SimPO maintained or improved scores on common sense and reasoning benchmarks (MMLU, ARC, HellaSwag, etc.).
  • For Llama-3.1-8B-Instruct, N-DPO improved average benchmark scores by over 3% while simultaneously improving generation quality.

• Reward Modeling:

  • Models trained with the regularization term showed improved or comparable reward modeling accuracy on held-out datasets (HH-RLHF, HelpSteer2) compared to standard DPO/SimPO, whereas DPO+SFT and DPOP showed significant drops in accuracy.

• Likelihood Dynamics:

  • Training curves confirmed that regularization mitigates likelihood displacement. In N-DPO, chosen response likelihoods remained stable, whereas in standard DPO, they dropped significantly.
  • Token-level analysis confirmed that the regularization specifically targets and corrects the distribution of minimum token rewards (outliers).

Significance and Claims

The paper claims that reward normalization is a critical, previously overlooked factor in the success of Direct Alignment Algorithms. By explicitly enforcing the conservation of probability mass for seen responses, the proposed method:

1. Mitigates Likelihood Displacement: It prevents the spurious probability shifts that can lead to harmful behaviors, particularly by correcting the behavior of outlier tokens.
2. Improves Trade-offs: It allows models to learn preferences more effectively without sacrificing general capabilities, addressing the "over-optimization" problem that has limited the reliability of DAAs.
3. Generalizes: The approach is effective across both reference-based (DPO) and reference-free (SimPO) frameworks, suggesting a unified mechanism for over-optimization.

The authors conclude that their analysis sheds light on why DAAs improve generation quality despite decreasing likelihoods (due to outlier token shifts) and demonstrates that utilizing low-likelihood tokens for regularization is an effective strategy for reshaping reward distributions. They note that while their current analysis focuses on offline settings, the findings motivate further research into the granular behavior of reward models in preference optimization.