PaperlayerView exact revision on arXiv

Source-linked AI summary

Iterative Preference Learning from Human Feedback: Bridging Theory and Practice for RLHF under KL-Constraint

Wei Xiong, Hanze Dong, Chenlu Ye, Ziqi Wang, Han Zhong, Heng Ji, Nan Jiang, Tong Zhang

arXiv:2312.11456v4cs.LGcs.AIstat.ML

TL;DR

RLHF implementations face ad-hoc tuning and distributional challenges, motivating a more rigorous treatment of alignment. The paper formulates RLHF as a reverse-KL regularized contextual bandit, develops algorithms across offline, online, and hybrid settings, and reports that practical implementations outperform DPO and RSO, including a 34.79% AlpacaEval2 win rate for Zephyr-SFT-7B.

  • Problem

    RLHF implementation involves ad-hoc tuning, reward misspecification, training instability, and distribution shift, while finite offline data may inadequately cover the prompt space.

  • Method

    The paper formulates RLHF as a reverse-KL regularized contextual bandit and designs statistically efficient offline, online, and hybrid algorithms using pessimism, optimism, and uncertainty-aware reward estimation.

  • Results

    34.79% win rate on AlpacaEval2 is achieved by the aligned Zephyr-SFT-7B policy, and the proposed implementations outperform DPO and RSO in real-world LLM alignment experiments.

  • Takeaways & Limitations

    The framework connects theoretical policy-improvement results with practical RLHF algorithms for offline, online, and hybrid alignment settings.

  • Takeaways & Limitations

    RSO’s rejection sampling can have very low efficiency; with reward gap r_x = 1 and η = 0.1, the acceptance rate is approximately 0.00004, potentially producing inferior performance under a fixed sampling budget.

Abstract

from arXiv · show

This paper studies the alignment process of generative models with Reinforcement Learning from Human Feedback (RLHF). We first identify the primary challenges of existing popular methods like offline PPO and offline DPO as lacking in strategical exploration of the environment. Then, to understand the mathematical principle of RLHF, we consider a standard mathematical formulation, the reverse-KL regularized contextual bandit for RLHF. Despite its widespread practical application, a rigorous theoretical analysis of this formulation remains open. We investigate its behavior in three distinct settings -- offline, online, and hybrid -- and propose efficient algorithms with finite-sample theoretical guarantees. Moving towards practical applications, our framework, with a robust approximation of the information-theoretical policy improvement oracle, naturally gives rise to several novel RLHF algorithms. This includes an iterative version of the Direct Preference Optimization (DPO) algorithm for online settings, and a multi-step rejection sampling strategy for offline scenarios. Our empirical evaluations on real-world alignment experiment of large language model demonstrate that these proposed methods significantly surpass existing strong baselines, such as DPO and Rejection Sampling Optimization (RSO), showcasing the connections between solid theoretical foundations and their potent practical implementations.

1 Introduction

The paper develops a theoretical foundation for RLHF that addresses limited coverage, reward over-optimization, and insufficient exploration in existing approaches. It analyzes offline, online, and hybrid settings and translates the results into practical algorithms that outperform strong baselines in LLM alignment experiments.

  • RLHF effectively aligns generative models with human values, producing systems described as helpful, harmless, honest, and controllable.
  • Existing RLHF implementations rely on ad-hoc practices and extensive tuning across preference collection, reward modeling, and policy optimization.
  • Finite preference datasets require uniform coverage of the exponentially large prompt-response space for existing methods to converge to the optimal policy.
  • The paper formulates RLHF as a reverse-KL regularized contextual bandit and analyzes offline, online, and hybrid learning settings.
  • Offline learning uses pessimistic reward estimation, with implementations derived for both PPO and DPO.
  • Online interactions with humans provide theoretical support for strategic exploration and continued policy improvement.
  • The framework builds on offline PPO, offline DPO, and InfoNCA, and its algorithms empirically outperform DPO and RSO in real-world LLM experiments.With Zephyr-SFT-7B as the initial model, the aligned policy achieves a 34.79% win-rate on AlpacaEval2.

2 Formulation and Existing Approaches

The paper models RLHF as a reverse-KL regularized contextual bandit with stochastic policies, then analyzes how reward estimation, policy improvement, and data coverage shape existing and proposed approaches. It highlights reward over-optimization and coverage limitations in classic PPO- and DPO-style pipelines.

  • 2.1 Formulation of RLHF: The framework treats prompts as contextual-bandit states, token sequences as actions, and LLMs as policies mapping prompts to action distributions.
  • 2.1 Formulation of RLHF: Human preferences are modeled with the Bradley-Terry model, where preference probabilities depend on exponentiated ground-truth rewards.
  • 2.1 Formulation of RLHF: The RLHF objective includes a KL penalty coefficient η that regularizes the policy relative to the pretrained policy.
  • 2.1 Formulation of RLHF: The reward is assumed linear in bounded features, with rθ(x, a) = ⟨θ, ϕ(x, a)⟩ and bounded parameter and feature norms.
  • 2.2 The Insufficiency of Classic Frameworks: Classic RLHF separates reward modeling from policy optimization, commonly estimating rewards by maximum likelihood and optimizing them with regularized PPO.
  • 2.3 Preliminary: The policy-improvement oracle produces a Gibbs policy, while KL constraints require separate control of policy-improvement error in addition to reward-estimation error.
  • 2.2 The Insufficiency of Classic Frameworks: DPO requires an effectively infinite coverage parameter D to cover the entire prompt-response space and converge to the optimal policy.
  • 2.2 The Insufficiency of Classic Frameworks: Reward over-optimization arises because proxy rewards align with gold rewards on in-distribution responses but can misalign after distribution shift toward out-of-distribution responses.

3 Offline Learning with Pessimism

The offline framework formulates RLHF as reverse-KL-regularized contextual bandits and uses pessimistic reward estimation to obtain sample-efficient policy learning from fixed preference data. Two pessimism implementations trade sharper theoretical bounds and robust improvement against practical approximability.

  • Problem setup: The offline setting learns a policy from a pre-collected preference dataset without further human interactions.The policy is induced through an information-theoretical improvement oracle.
  • Pessimistic learning: Pessimistic reward estimation addresses limited prompt-response coverage by constructing conservative lower-bound estimates of the true reward.The estimator subtracts an uncertainty bonus from the maximum-likelihood reward estimate.
  • Guarantees: The offline GSHF algorithm outputs a Gibbs policy based on a pessimistic reward and provides finite-sample theoretical guarantees.Its reference vector can be chosen from an available reference policy, often πref = π0.
  • Algorithms: The framework offers two pessimism implementations: an expected uncertainty estimator and a pointwise uncertainty estimator used with the policy-improvement oracle.The pointwise version defines ˆr(x, a) = rMLE(x, a) −β · Γ(x, a, ν, Doff).
  • Trade-offs: Option I achieves a sharper uncertainty bonus bound, while Option II is easier to approximate empirically through PPO or direct preference learning.Option I also guarantees robust improvement over πref regardless of offline-data coverage under the stated theorem.

4 Online Iterative Learning with Batch Exploration

The online iterative framework alternates exploitation by a main agent with strategically selected exploration by an enhancer, while optionally retaining offline data. Its guarantees show that explicit exploration can remove the need for broadly covered offline data, although uncertainty estimation remains difficult for neural networks.

  • Hybrid setting: Hybrid learning starts from an offline dataset when available and repeatedly collects fresh human preference data during training.Each iteration samples prompts, generates response pairs, queries preferences, and adds the resulting batch to the history.
  • Exploration design: The main technical challenge is balancing exploitation of learned rewards with exploration that improves information about uncertain policies.The non-symmetric design assigns exploitation to the main agent and exploratory behavior to an enhancer.
  • Option I: When offline data has good coverage, the enhancer can use a reference policy instead of actively exploring.This option is simpler algorithmically but depends on a diverse offline dataset.
  • Option II: When offline coverage is insufficient, the enhancer maximizes uncertainty relative to the main agent while maintaining a moderate KL divergence.The design requires the information gained to justify the cost of moving away from the main agent.
  • Practical trade-off: The two options are complementary: Option I favors simplicity with strong data coverage, whereas online learning reduces dependence on offline-data quality but makes uncertainty estimation challenging for neural networks.In practice, heuristic uncertainty estimators are used for neural-network policies.

5 Implementations of GSHF

The GSHF framework turns the paper’s theoretical policy-improvement procedures into practical RLHF implementations using existing optimization, preference-learning, and sampling methods. It also introduces pessimistic DPO updates and multi-step rejection sampling to address uncertainty and low acceptance rates in offline learning.

  • Framework implementation: GSHF combines existing algorithms such as PPO, DPO, and InfoNCA to approximate its computational policy-improvement oracle.The framework is presented as building on fixed-reward optimization methods to boost their performance.
  • PPO and DPO: PPO optimizes the regularized reward but is less stable and more implementation-sensitive than SFT, motivating direct preference-learning alternatives.DPO and related methods avoid reward modeling and attracted attention for stability and easy implementation.
  • Pessimistic DPO: DPO’s optimal policy matches the RLHF objective when the reward is the maximum-likelihood estimate, while pessimistic DPO adds uncertainty-aware margins to preference pairs.The uncertainty margin is based on the difference between the chosen and rejected responses’ uncertainty.
  • Uncertainty estimation: Ensemble-based uncertainty estimation trains five reward models and subtracts their variance to reduce reward overoptimization for PPO and best-of-n sampling.The paper describes this as a practical heuristic because general neural-network uncertainty estimation remains open.
  • Online exploration: Online exploration seeks enhancer policies that maximize information while maintaining moderate KL divergence from the main policy.The paper discusses model variants and other heuristic implementations for approximating this objective.
  • Multi-step RSO: Multi-step RSO progressively performs rejection sampling, raising each step’s acceptance probability above exp(-1) > 0.367 instead of relying on one inefficient rejection step.The method is motivated by RSO’s potentially very low acceptance rates and combines rejection sampling with iterative offline DPO.

6 Experiments

The experiments evaluate offline, hybrid, and online RLHF variants using gold rewards, KL divergence, GPT-4 comparisons, sampling temperatures, and output lengths. Iterative methods generally improve reward and robustness over strong baselines, while length bias remains a concern.

  • Experimental setup: The experiments compare SFT, offline DPO, RSO, Hybrid-GSHF-DPO, and Multi-step RSO using gold reward, KL divergence, and GPT-4 evaluations.The evaluation includes in-domain HH-RLHF and out-of-distribution UltraFeedback prompts.
  • Main results: Hybrid-GSHF-DPO outperforms DPO and RSO in gold reward and GPT-4 evaluations, while showing a smaller in-domain-to-OOD performance gap.The reported comparisons use reward-KL trade-off and evaluation results.
  • Online exploration: Each successive Hybrid-GSHF-DPO iteration improves the reward-KL frontier, making online data more efficient than additional offline data in the reported comparison.
  • Iterative pseudo-labelling: Multi-step RSO strictly dominates original RSO in reward-KL trade-off, and its third iteration achieves the highest ground-truth reward.
  • Robustness and bias: GSHF improvements remain stable across sampling temperatures, but later Hybrid-GSHF-DPO iterations produce longer outputs and reward gains may partly reflect length preference.Temperature 0.7 yields the highest gold reward for all models, while reward hacking motivates early stopping, replay, and validation.
  • Scaling-up experiments: Online-GSHF-DPO’s length-control win rate increases across iterations and eventually surpasses offline DPO and RAFT by a large margin, but UltraRM-13B favors longer responses.The three-iteration model has 30.49% unadjusted win rate versus 24.17% length-controlled win rate.
  • Scaling-up experiments: The scaling-up experiments report potential for iterative RLHF to produce appropriate, well-structured responses without using external closed-source model responses.The external information used is the relative preference signal.

7 Conclusion

The paper formulates RLHF as a reverse-KL regularized contextual bandit and analyzes offline, online, and hybrid learning. Its practical implementations outperform DPO and RSO, while the appendices provide proofs, notation, and experimental details.

  • Conclusion: RLHF is modeled as a reverse-KL regularized contextual bandit with a stochastic optimal policy that reflects generative-model alignment practices.
  • Conclusion: The paper develops statistically efficient algorithms with finite-sample guarantees for offline, online, and hybrid settings.
  • Conclusion: Practical implementations of the theoretical algorithms outperform strong DPO and RSO baselines in real-world LLM alignment experiments.
  • Appendices: The appendices cover rejection sampling, proofs for online, offline, and hybrid learning, DPO coverage, technical lemmas, and experimental details.

B Proof of Online Learning

The online-learning appendix analyzes batch and sequential policy updates under KL constraints, using uncertainty estimation and confidence sets to control suboptimality and regret. Theorems provide high-probability guarantees for the resulting online algorithms.

  • Online framework: The online framework targets policies with suboptimality below ϵ using samples polynomial in 1/ϵ, feature dimension d, and problem-dependent parameters.
  • Online framework: Different model variants and temperatures are recommended for collecting diverse response pairs during online preference learning.
  • Proof strategy: The proof controls online learning through covariance concentration, elliptical-potential bounds, Cauchy-Schwarz, Jensen, and confidence-set arguments.
  • Sequential online learning: The sequential setting seeks sublinear cumulative and average regret by having an enhancer maximize an uncertainty estimator.
  • Theoretical guarantee: Theorem 4 gives a high-probability regret guarantee for Algorithm 2 with online Option II when m = 1.The theorem uses λ = Ω(d log(T/δ)/(γ^2B^2)).
  • Proof strategy: The appendix establishes that the optimal policy remains in the confidence set at every iteration, supporting the subsequent regret analysis.

C Proof of the Offline Learning

The offline-learning appendix combines pessimistic reward estimation or pessimistic targets with KL-regularized policy optimization. Its analysis connects DPO, coverage, uncertainty bonuses, and finite-sample guarantees under offline and hybrid data.

  • Offline proof: The offline proof analyzes two options: pessimistic reward estimation followed by policy optimization, and point-wise pessimistic targets passed to the policy oracle.
  • DPO connection: The uncertainty bonus acts as an adaptive margin when connecting the KL-regularized target to the DPO objective.
  • Coverage: Single-policy coverage is weaker than uniform coverage and makes the hybrid coefficient depend mainly on online-to-offline data size.

D.3 Proof of Theorem 2

The proof controls exploration and suboptimality errors under the stated conditions, then combines parameter choices to establish Theorem 2. Its resulting guarantee depends on offline coverage and reward-function complexity, with online data potentially improving coverage.

  • The proof controls the exploration error P′t using the batch-online analysis.
  • Because the reference policy is directly available and collected the data, its uncertainty requires no optimism adjustment.
  • The remaining analysis handles the suboptimality source P1 using Cauchy–Schwarz and Lemma 8.
  • Under Assumption 2 and the stated parameter choices, the proof concludes Theorem 2.
  • When noff ≈ mT, α(mT, Doff) ≈ 1, as in the LLaMA2 project.
  • The convergence rate is jointly determined by offline data coverage and reward-function complexity, while online data may improve Ccov during training.

E Discussion on the Coverage Condition for Vanilla RLHF

This discussion analyzes when vanilla DPO can recover the optimal policy and shows that offline preference data must cover the relevant prompt-response support. Uncovered outputs leave the loss indifferent, so finite offline datasets face a demanding coverage requirement.

  • DPO is derived from the reverse-KL-regularized contextual bandit formulation but skips part of its framework.
  • The loss in Equation (27) directly optimizes the language model on offline preference data, with aw chosen and al rejected.
  • Given x, aw, and al, the loss yields a maximum-likelihood estimate of preference probabilities, while accounting for generation distributions absent from original DPO.
  • With population loss, the preference model pθ uniquely minimizes the loss when pθ = p∗.
  • The minimizing policy preserves the optimal-policy probability ratio on actions covered by πoff.
  • For outputs generated by π∗ but absent from πoff, the loss does not constrain πθ, allowing arbitrary choices there.
  • Convergence to π∗ for every supported state-action pair requires no uncovered outputs or sufficient additional constraints, a difficult condition for finite offline datasets.

F Technical Lemma Proofs

The technical lemmas establish policy-oracle identities, loss optimality, concentration, and optimization tools used by the paper’s theoretical guarantees. Together they connect reward reparameterization and statistical estimation to KL-regularized policy solutions.

  • Calling Policy Improvement Oracle 1 with estimated reward induces the policy used in the proof.
  • The proof rewrites the reward function through a normalization constant and substitutes this form into the policy-optimization error.
  • The preference-learning loss can be reformulated as a KL divergence plus a constant, making p∗ = pθ the unique optimum.
  • Jensen’s inequality supplies the convexity relation used to compare a function at an expectation with the expected function value.
  • The MLE lemma assumes bounded feature and parameter norms and gives a high-probability in-sample error bound involving γ = 1/(2 + exp(−B) + exp(B)).
  • The elliptical potential lemma bounds accumulated contributions from bounded sequential vectors and positive-definite regularization.
  • The concentration lemma controls inverse covariance under bounded features when λ = Ω(d log(n/δ)).
  • KL-regularized optimization has a Gibbs-distribution minimizer proportional to π0(a|x) exp(·).

H More Experiment Details

The experiments use a controlled Gaussian-mixture reference-policy setup and report rejection rates across settings, alongside specified hardware and training infrastructure.

  • Experiments run on 8×A40 GPUs with 48G memory each, 600G RAM, and bf16 half-precision training.
  • Implementations use the open-source TRL and LMFlow packages, with hyperparameters summarized in Tables 8 and 9.
  • Figure 7 illustrates rejection rates across settings parameterized by {ηi}N.
  • The synthetic model uses a Gaussian-mixture π0, with ηi = Nη/i and reward r = [1, 0]⊤a for πr.

H.1 Implementation Details

The implementation details describe rejection-sampling and multi-step settings, hybrid-learning data construction, GPT-4 evaluation, and supporting example and hyper-parameter tables.

  • Rejection Sampling: Rejection sampling generates K responses per prompt and samples preference pairs from this pool, using 2n rejection samples to obtain n pairs.When multiple pairs are identified, the initial ranking round selects the appropriate pairs before sample order is randomized.
  • Multi-step approximation: The multi-step approximation divides the path into three steps with η ∈ {0.1, 0.3, 0.5} and uses 25K prompts per step.RSO uses a rejection-sampling coefficient selected from {0.5, 1, 2, 3}.
  • Hybrid learning: Hybrid GSHF derives preference signals from a gold reward function trained on UltraFeedback, Anthropic HH-RLHF, and other open-source datasets, using LLaMA2-13B.Anthropic HH-RLHF’s 75K training prompts are split across three online-algorithm iterations, with the initial iteration trained using DPO on an offline dataset.
  • GPT4 Evaluation: GPT-4 evaluation compares models against DPO on 100 prompts using GPT-4-turbo, with final wins counted as win + tie × 0.5.The test hyper-parameter is specified in Table 8.
  • Reward baseline: Absolute values in Table 2 use the initial model’s test reward as the reward baseline.The paper attributes this convention to prior work.
  • Examples and hyper-parameters: Tables 5–7 provide randomly sampled model outputs for in-domain HH-RLHF prompts and an out-of-distribution Ultra-Chat prompt, while Tables 8 and 9 report fine-tuning and auxiliary-training hyper-parameters.Table 8 defines SFT-RLHF as fine-tuning on preferred samples and indicates that multiple listed values represent hyper-parameter search spaces.
Loading 2312.11456v4…