Technical Summary: Beyond RLHF: A Unified Theoretical Framework of Alignment

Problem Statement

Reinforcement Learning from Human Feedback (RLHF) has become the dominant paradigm for aligning Large Language Models (LLMs) with human preferences. However, the paper identifies significant theoretical gaps in the current understanding of RLHF and related alignment methods:

1. Lack of Justification: Existing theories often treat the RLHF objective as a tautology (converging to the population version of the objective itself) rather than providing a learning-theoretic justification for why it converges to the target language model.
2. Incompatible Frameworks: Different methods (e.g., DPO, RLHF, REBEL) are analyzed under disparate theoretical frameworks, preventing direct comparison of their convergence guarantees.
3. The RLHF Dilemma: The standard RLHF objective faces an asymptotic dilemma regarding the regularization strength $\beta$. If $\beta$ remains fixed, the model underfits (staying too close to the reference policy); if $\beta$ decays to zero as data increases, the solution may degenerate.
4. Empirical vs. Theoretical Gap: Empirically, on-policy objectives like RLHF often outperform likelihood-based methods like Direct Preference Optimization (DPO), yet existing theories cannot explain this superiority because they offer ad-hoc guarantees for each method individually.

The authors propose reframing alignment not as reward maximization, but as distribution learning from pairwise preferences. The core question is: Under what assumptions can we derive training objectives that provably converge to a target language model ($\pi^*$)?

Methodology

The paper posits the existence of a target (oracle) language model $\pi^*$ and explicitly models how preference feedback reveals information about $\pi^*$.

1. The Preference Model

The authors assume a Bradley-Terry (BT) model where the probability of preferring response $a$ over $b$ given prompt $x$ is determined directly by the target model's probabilities, rather than an external reward function:
$$P(a \succ b | x) = \frac{\pi^*(a | x)^\gamma}{\pi^*(a | x)^\gamma + \pi^*(b | x)^\gamma}$$
Here, $\gamma > 0$ is a hyperparameter. This assumption implies that the preference model is explicitly a language model, contrasting with DPO which relies on the "all-policy" assumption to derive a relationship between rewards and the LM.

2. Proposed Objectives

Based on this distribution learning perspective, the authors derive three principled alignment objectives:

• Preference Maximum Likelihood Estimation (PMLE):
  This objective maximizes the likelihood of the observed pairwise preferences under the BT assumption, regularized by a KL divergence to a reference policy $\pi_0$.
  $$L_{PMLE, \beta}(\pi) = \frac{1}{n} \sum -\ln \sigma\left(\gamma \ln \frac{\pi(a^+|x)}{\pi(a^-|x)}\right) + \beta KL(\pi \| \pi_0)$$
  Unlike DPO, which minimizes a likelihood ratio without an explicit regularizer (potentially leading to degenerate distributions), PMLE includes an explicit KL term to ensure stability.

• Preference Distillation:
  This approach first learns a reward model $\hat{R}$ (equivalent to learning a preference simulator) and then distills this preference distribution into the language model. Instead of using squared loss (as in REBEL), it minimizes the KL divergence between the binary preference distributions of the learned simulator and the target policy:
  $$L_{Distill}(\pi) = \frac{1}{n} \sum KL(\text{Bern}(P_{\hat{\pi}}(a^+ \succ a^-|x)) \| \text{Bern}(P_{\pi}(a^+ \succ a^-|x))) + \beta KL(\pi \| \pi_0)$$
  This formulation is explicitly derived from the BT model and avoids the statistical justification issues of squared loss.

• Reverse KL (RKL) Minimization:
  Recognizing that the goal is distribution learning, the authors propose minimizing the reverse KL divergence $E_x[KL(\hat{\pi}(x) \| \pi^*(x))]$. Since $\pi^*$ is unknown, they plug in an estimator $\hat{\pi}$ derived from the reward model.
  $$L_{RKL, \beta}(\pi) = \frac{1}{n} \sum \left( -E_{a \sim \pi(x)}[\hat{R}(x, a)] - \gamma H(\pi(x)) + \beta KL(\pi(x) \| \pi_0(x)) \right)$$
  This objective resembles RLHF but includes an additional entropy term ($-\gamma H(\pi(x))$). The authors argue this acts as a correction to RLHF, resolving the asymptotic dilemma where standard RLHF fails to converge to the true $\pi^*$ as $n \to \infty$.

3. Theoretical Guarantees

The paper provides non-asymptotic $O(1/n)$ convergence guarantees for all three methods. Specifically, they prove that the forward KL error $E_x[KL(\pi^*(x) \| \hat{\pi}(x))]$ converges to zero at a rate of $O(1/n)$, where $n$ is the dataset size.

• PMLE and Preference Distillation: Achieve bounds dependent on the coverage coefficient of the policy class $\Pi$.
• Reverse KL: Achieves a tighter bound that depends on the coverage coefficient of the reward-induced model class $\mathcal{P}_\gamma(\mathcal{R})$, which is often smaller than that of $\Pi$. This theoretically explains why on-policy methods (like RKL/RLHF) often empirically outperform likelihood-based methods (like PMLE/DPO).

Key Results

Theoretical Findings

• Justification for RLHF: The RKL objective provides a strong theoretical justification for the RLHF objective, showing it is a special case (with $\gamma \to 0$) of a consistent distribution learning objective, provided an entropy correction is applied.
• Resolution of the Dilemma: Theoretical analysis and toy experiments (Figure 1) demonstrate that while standard RLHF suffers from inconsistency (saturating at non-zero error) as $n \to \infty$ due to the regularization dilemma, the proposed RKL method converges to the true target model.
• Superiority of On-Policy Methods: The theory confirms that Reverse KL (similar to RLHF) enjoys a superior convergence guarantee compared to PMLE (similar to DPO) because it leverages the smaller complexity of the reward model class rather than the full policy class.

Empirical Findings

Experiments were conducted on two tasks:

1. TL;DR Summarization: Using Pythia-2.8B and Pythia-6.9B models.
   • Win-Rates: The proposed methods (PMLE, Preference Distillation, Reverse KL) consistently achieved higher GPT-4 win-rates compared to their baselines (DPO, REBEL, RLHF).
   • KL Divergence: PMLE achieved significantly lower KL divergence to the reference model compared to DPO, validating the benefit of the explicit regularizer.
2. General Chat (UltraFeedback): Using LLaMA-3-8B-Instruct.
   • Evaluated on MT-Bench, AlpacaEval 2.0, and Arena Hard.
   • PMLE and Preference Distillation outperformed DPO and REBEL across most benchmarks, particularly showing strong gains in AlpacaEval 2.0 and Arena Hard.

Significance and Claims

The paper claims to provide the first unified theoretical framework for alignment that treats the problem as distribution learning from pairwise preferences. Its primary contributions are:

1. Unified Perspective: It unifies existing methods (DPO, RLHF, REBEL) under a single probabilistic assumption, making their hidden assumptions explicit.
2. Theoretical Grounding: It offers rigorous, non-asymptotic $O(1/n)$ convergence guarantees for learning a target distribution from pairwise feedback, a result previously unavailable for this setting.
3. Explanation of Empirical Trends: It provides the first theoretical confirmation for the empirical observation that on-policy objectives (RLHF) often outperform likelihood-based objectives (DPO), attributing this to the coverage properties of the reward model class versus the policy class.
4. Correction to RLHF: It identifies a specific deficiency in the standard RLHF objective (the lack of an entropy term leading to inconsistency) and proposes a corrected objective (RKL) that retains the practical benefits of RLHF while ensuring theoretical consistency.

The authors conclude that while their methods are competitive with strong baselines, the primary value lies in the theoretical framework that explains why these methods work and how they can be improved. Future work suggested includes exploring mode-seeking vs. mode-covering objectives and extending the framework to alternative divergences.