PaperlayerView exact revision on arXiv

Source-linked AI summary

Statistical Rejection Sampling Improves Preference Optimization

Tianqi Liu, Yao Zhao, Rishabh Joshi, Misha Khalman, Mohammad Saleh, Peter J. Liu, Jialu Liu

arXiv:2309.06657v2cs.CL

TL;DR

The paper addresses the difficulty of obtaining preference pairs sampled from the optimal policy, which limits existing methods such as DPO and SLiC. It introduces RSO, which uses an explicit reward-ranking model and statistical rejection sampling, and reports consistent improvements over DPO and SLiC across diverse tasks and evaluators.

  • Problem

    DPO and SLiC do not obtain labeled preference pairs directly from the optimal policy: DPO uses data from other policies, while SLiC samples from the SFT policy.

  • Method

    RSO trains a pairwise reward-ranking model, uses the SFT policy and statistical rejection sampling to generate pairs from an approximated optimal policy, then fits a classifier on labeled pairs.

  • Results

    RSO consistently outperforms SLiC and DPO across three diverse tasks in evaluations by both LLMs and human raters.

  • Takeaways & Limitations

    The paper presents RSO as a simple, effective alternative to RLHF that combines improved preference-pair sampling with a unified view of DPO and SLiC losses.

  • Takeaways & Limitations

    The paper notes that subtle differences between the SLiC loss and hinge loss remain to be discussed in the Method section.

Abstract

from arXiv · show

Improving the alignment of language models with human preferences remains an active research challenge. Previous approaches have primarily utilized Reinforcement Learning from Human Feedback (RLHF) via online RL methods such as Proximal Policy Optimization (PPO). Recently, offline methods such as Sequence Likelihood Calibration (SLiC) and Direct Preference Optimization (DPO) have emerged as attractive alternatives, offering improvements in stability and scalability while maintaining competitive performance. SLiC refines its loss function using sequence pairs sampled from a supervised fine-tuned (SFT) policy, while DPO directly optimizes language models based on preference data, foregoing the need for a separate reward model. However, the maximum likelihood estimator (MLE) of the target optimal policy requires labeled preference pairs sampled from that policy. DPO's lack of a reward model constrains its ability to sample preference pairs from the optimal policy, and SLiC is restricted to sampling preference pairs only from the SFT policy. To address these limitations, we introduce a novel approach called Statistical Rejection Sampling Optimization (RSO) that aims to source preference data from the target optimal policy using rejection sampling, enabling a more accurate estimation of the optimal policy. We also propose a unified framework that enhances the loss functions used in both SLiC and DPO from a preference modeling standpoint. Through extensive experiments across three diverse tasks, we demonstrate that RSO consistently outperforms both SLiC and DPO on evaluations from both Large Language Model (LLM) and human raters.

1 INTRODUCTION

The paper identifies a sampling mismatch in preference optimization and introduces RSO to estimate the optimal policy using preference pairs generated through statistical rejection sampling. It also unifies DPO and SLiC as preference-fitting methods with different losses and emphasizes explicit reward modeling.

  • Motivation: DPO and SLiC estimate preference policies from data sampled from other policies rather than directly from the optimal policy π∗.This sampling mismatch means the resulting estimator is not strictly maximum likelihood for the preference model.
  • RSO framework: RSO trains a pairwise reward-ranking model, uses statistical rejection sampling with the SFT policy to generate pairs near π∗, labels them, and fits a classifier.The procedure constructs preference data from an approximation to the optimal policy before preference optimization.
  • RSO framework: RSO uses an explicit reward model because comparing responses through rewards is easier to learn than generating high-quality responses directly.The paper contrasts this design with DPO’s claim that the language model is implicitly a reward model.
  • Contributions: The proposed framework is presented as scalable and easy to implement, with experiments showing the importance of reward modeling over direct optimization on preference data.The contribution combines a preference-pair generation recipe with choices of loss functions and dataset construction.
  • Unified preference optimization: DPO is characterized as logistic regression on preference data, whereas SLiC is nearly equivalent to an SVM with hinge loss.The paper uses this statistical view to improve SLiC as the SVM counterpart of DPO.

2 PRELIMINARIES

This section formulates preference optimization through KL-constrained reward maximization and connects the optimal policy to Bradley-Terry preference modeling. It also describes the pairwise reward-ranking model used to estimate pointwise rewards.

  • KL-Constrained Reward Maximization Objective: DPO and RLHF optimize a KL-constrained reward-maximization objective, whose solution is an optimal policy governed by the reward function.The partition function normalizes the policy, while β controls the exploration–exploitation balance.
  • KL-Constrained Reward Maximization Objective: As β approaches infinity, the optimal policy becomes the SFT policy, whereas β approaching zero concentrates probability on the highest-reward response.
  • Preference Model: The Bradley-Terry model expresses preference probability as the sigmoid of the difference between two responses’ ground-truth rewards.This makes estimating the optimal policy from human preference data a density-estimation problem.
  • Reward Model: A pairwise T5-XXL reward-ranking model is trained on human preference data to approximate pairwise preference probabilities.It receives context and two candidate responses formatted differently for summarization and AI-assistant tasks.
  • Reward Model: The model’s probability of decoding “A” estimates preference probability, which can be converted into a pointwise reward using a baseline sequence with reward zero.The conversion uses rψ(x, y) = logit(ρψ(x, y, yb)).

3 RSO APPROACH

RSO estimates the optimal policy by choosing preference-data distributions and loss functions, then using statistical rejection sampling to draw responses closer to that policy. The algorithm samples from the SFT policy, accepts candidates probabilistically, and repeats until enough responses are obtained.

  • 3.1 STATISTICAL ESTIMATION OF THE OPTIMAL POLICY π∗: RSO considers loss-function and preference-distribution choices because estimating the optimal policy requires deciding both how to fit preferences and how to obtain response triplets.
  • 3.1 STATISTICAL ESTIMATION OF THE OPTIMAL POLICY π∗: The rso-sample-rank setting samples from the reward-induced policy, which is closer to the optimal policy than direct fitting or SFT sampling.
  • 3.2 STATISTICAL REJECTION SAMPLING ALGORITHM: Statistical rejection sampling uses the SFT policy as a proposal distribution to generate samples from the reward-induced target policy.Each candidate is generated from πsft and evaluated using an acceptance probability based on πrψ and a scaling constant M.
  • 3.2 STATISTICAL REJECTION SAMPLING ALGORITHM: The algorithm accepts or rejects each SFT-generated response using a uniform random variable and repeats the process until the required sample set is large enough.In practice, M is estimated from 64 sequences sampled by the SFT policy rather than computed directly.
  • 3.2 STATISTICAL REJECTION SAMPLING ALGORITHM: As β approaches infinity, every SFT sample is accepted; as β approaches zero, only the highest-reward response is accepted.Thus β controls how much the procedure trusts the reward model.

4 RELATED WORK

Prior work studies preference optimization through RLHF and offline alternatives, while rejection sampling has been used to select high-reward candidates for later tuning. This paper frames those existing approaches as special cases of its broader rejection-sampling algorithm.

  • Preference Optimization: RLHF and offline preference-optimization methods fit models to human- or model-ranked response pairs, offering alternatives to reinforcement-learning-based training.
  • Rejection Sampling: AnthropicHH, ReST, and Llama2 use rejection sampling to select top-k candidates for further tuning, with Llama2 combining the approach with PPO.
  • Rejection Sampling: The paper identifies these existing rejection-sampling approaches as special cases of its proposed algorithm.

5 EXPERIMENTS

Experiments evaluate RSO across summarization and dialogue tasks, comparing sampling strategies, losses, model scales, and human judgments. RSO variants generally outperform DPO, SLiC, and related rejection-sampling baselines, while ablations identify settings that affect performance.

  • Tasks: RSO is evaluated on Reddit TL;DR summarization, AnthropicHH dialogue, and CNN/DailyMail cross-task generalization.The Reddit TL;DR and AnthropicHH datasets provide the main tasks; CNN/DailyMail is used for cross-task evaluation.
  • Method: The experiments use T5-large and T5-XXL policy models, a T5-XXL pairwise reward-ranking model, and nine combinations of losses and preference-data distributions.For pair construction, 64 SFT-policy responses are sampled and eight are sub-sampled; β = 0.5 and γ = 0.05 unless noted.
  • Performance comparison: RSO variants show significant gains over RAFT, ReST, DPO, and SLiC variants on both evaluated tasks.The rso-sample-rank strategy improves over direct and SFT-sample-rank construction, while hinge loss exhibits reward hacking on Reddit TL;DR.
  • RSO ablation: γ = 0.05 and β = 0.5 provide the optimal win rates in the reported Reddit TL;DR hyper-parameter studies.The study varies γ over 0.005, 0.05, and 0.5, and β over 0, 0.05, 0.5, and 5.
  • RSO ablation: Rso-8-sample-first-round-rank is optimal by AutoSxS, indicating that more sampled responses or tournament ranking does not always improve results.Tournament ranking consistently improves reward-model results but not AutoSxS win rate in the rso-8-sample case.
  • Scale-up: 1.1% and 33.1%: RSO improves AutoSxS over DPO on the two tasks with a T5-XXL policy.The comparison uses the sigmoid-norm loss while scaling the policy model.
  • Human evaluation: Human raters prefer rso-sample-rank over direct and sft-sample-rank across all evaluated losses and tasks, with RSO_sigmoid-norm preferred more than 2x as often as DPO.Each task was rated by three raters; 47 raters participated overall.

6 CONCLUSION

The conclusion presents RSO as a simple alternative to RLHF that improves preference-pair sampling and unifies DPO and SLiC loss perspectives. Experiments support its effectiveness across multiple tasks, while broader scaling and extensions remain future work.

  • Conclusion: RSO is proposed as a simple, effective alternative to RLHF with a better sampling strategy than DPO and SLiC.The authors describe the recipe as applicable to training language models from human feedback.
  • Conclusion: The paper unifies DPO and SLiC as logistic-regression and support-vector-machine approaches, respectively, from a preference-optimization perspective.The conclusion identifies the two loss-function perspectives as part of the proposed framework.
  • Conclusion: The authors report strong performance across multiple tasks based on comprehensive numerical experiments and analysis.The conclusion states that RSO is powerful on multiple tasks.
  • Future work: Future work includes larger decoding samples, other losses and generation tasks, online variants, and non-human feedback.These directions are listed as extensions of the reported recipe.

A Python Implementation

The implementation realizes statistical rejection sampling by using SFT responses and reward scores to accept a requested number of sequences. The appendix derives the scaling term used by the algorithm and presents its Python implementation.

  • Implementation: Algorithm 1 implements statistical rejection sampling over response candidates, rewards, a requested sample count, and β.The function returns rejection-sampled sequences from the estimated optimal policy.
  • Implementation: The procedure stores candidate–reward pairs, repeatedly finds the maximum reward, and removes or accepts candidates until num_samples are accepted.The implementation maintains accepted and to_remove collections during sampling.
  • Derivation: The derivation identifies M_Dx as the minimum scaling factor ensuring the scaled SFT policy covers the reward-based policy outside the accepted set.It is expressed through a maximum over responses not in D_x.
  • Derivation: Using a sample-based version of the maximum reward ratio, the derivation obtains Algorithm 1.The implementation follows the quantity derived from Equation (4).

A.2 PROOF OF THEOREM 1

The proof section analyzes the accepted sequence set during the rejection-sampling process and derives the corresponding acceptance rate. It frames the proof around the process specified in Algorithm 1.

  • Proof: The proof defines D_x as the accepted sequence set at the current step of Algorithm 1’s generation process.This set is used as the state of the rejection-sampling procedure.
  • Proof: Equation (13) is used to derive the rejection-sampling acceptance rate.The passage identifies the acceptance rate as the next result obtained in the proof.

A.3 QUALITATIVE EXAMPLES OF RSO COMPARING WITH OTHER APPROACHES

Figures 4 and 5 provide qualitative comparisons of summaries and dialogue responses generated by SFT, SLiC, DPO, and RSO. The examples highlight RSO’s concise, precise summaries and helpful, direct responses.

  • Figure 4 compares summaries from SFT, SLiC, DPO, and RSO on a Reddit post.
  • RSO produces the best Reddit summary by concisely and precisely covering key information.
  • Figure 5 compares responses from SFT, SLiC, DPO, and RSO on an AnthropicHH dialogue.
  • RSO gives the most helpful response by clearly explaining how to send a letter quickly through traditional mail.
  • SFT discusses email instead of answering the traditional-mail question, while SLiC and DPO are vague and repetitive.

A.4 PALM 2-L DETAILS AND FEW-SHOT SXS TEMPLATE

This appendix describes evaluation prompts, qualitative examples, cross-task transfer, baseline coverage, and the computational trade-offs of RSO. It also notes that RSO adds sampling and reward-ranking computation while remaining scalable and empirically efficient.

  • A.4.1 DETAILS: AutoSxS uses PaLM 2-L few-shot prompting, eight decoded samples, and four response-order flips to reduce reward-model hacking bias.
  • A.4.2 REDDIT TL;DR FEW-SHOT PROMPTS: Reddit TL;DR examples judge whether one of two summaries is more concise, informative, fluent, or equivalent.
  • A.4.3 ANTHROPICHH FEW-SHOT PROMPTS: AnthropicHH examples ask whether response A, response B, or neither is more helpful for a chatbot query.
  • A.7 CROSS-TASK ADAPTATION AND GENERALIZATION: RSO consistently improves over SLiC and DPO for cross-task transfer according to Table 6.
  • A.8 OTHER BASELINES: The study omits RRHF and RLHF baselines because comparable RRHF systems were unavailable and the authors lacked RLHF expertise.
  • A.10 COMPUTATIONAL EFFICIENCY: RSO’s additional sample-and-rank stages can be parallelized across model servers, while prompt caching and serving techniques can mitigate reward-inference costs.
  • A.10 COMPUTATIONAL EFFICIENCY: RSO adds sampling from the SFT policy and ranking by a pairwise reward model, but this extra cost is empirically minor compared with policy training.
  • A.10 COMPUTATIONAL EFFICIENCY: rso-sample-rank accounts for less than 10% of total training time.
Loading 2309.06657v2…