PaperlayerView exact revision on arXiv

Source-linked AI summary

REBEL: Reinforcement Learning via Regressing Relative Rewards

Zhaolin Gao, Jonathan D. Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, J. Andrew Bagnell, Jason D. Lee, Wen Sun

arXiv:2404.16767v4cs.LGcs.CLcs.CV

TL;DR

PPO requires multiple stability heuristics and is sensitive to implementation details, creating complexity and memory demands for generative-model fine-tuning. REBEL reduces policy optimization to sequential squared-loss regression of reward differences between paired trajectories, using on-policy, offline, or hybrid datasets. REBEL empirically outperforms PPO and DPO in language generation, with faster convergence and similar asymptotic performance in image generation, plus competitive Llama-3-8B benchmark results.

  • Problem

    PPO requires multiple stability heuristics and is sensitive to implementation details, creating complexity and memory demands for generative-model fine-tuning.

  • Method

    REBEL reduces policy optimization to sequential squared-loss regression of reward differences between paired trajectories, using on-policy, offline, or hybrid datasets.

  • Results

    REBEL empirically outperforms PPO and DPO in language generation, with faster convergence and similar asymptotic performance in image generation, plus competitive Llama-3-8B benchmark results.

  • Takeaways & Limitations

    REBEL offers a simpler and more scalable alternative to policy-gradient methods while retaining strong theoretical guarantees and supporting offline data and intransitive preferences.

  • Takeaways & Limitations

    The analysis assumes a coverage condition in the general nonutility-based preference setting, whose necessity remains an open question.

Abstract

from arXiv · show

While originally developed for continuous control problems, Proximal Policy Optimization (PPO) has emerged as the work-horse of a variety of reinforcement learning (RL) applications, including the fine-tuning of generative models. Unfortunately, PPO requires multiple heuristics to enable stable convergence (e.g. value networks, clipping), and is notorious for its sensitivity to the precise implementation of these components. In response, we take a step back and ask what a minimalist RL algorithm for the era of generative models would look like. We propose REBEL, an algorithm that cleanly reduces the problem of policy optimization to regressing the relative reward between two completions to a prompt in terms of the policy, enabling strikingly lightweight implementation. In theory, we prove that fundamental RL algorithms like Natural Policy Gradient can be seen as variants of REBEL, which allows us to match the strongest known theoretical guarantees in terms of convergence and sample complexity in the RL literature. REBEL can also cleanly incorporate offline data and be extended to handle the intransitive preferences we frequently see in practice. Empirically, we find that REBEL provides a unified approach to language modeling and image generation with stronger or similar performance as PPO and DPO, all while being simpler to implement and more computationally efficient than PPO. When fine-tuning Llama-3-8B-Instruct, REBEL achieves strong performance in AlpacaEval 2.0, MT-Bench, and Open LLM Leaderboard.

1 Introduction

REBEL addresses the complexity and implementation sensitivity of PPO for generative-model fine-tuning by reducing policy optimization to relative-reward regression. It provides theoretical connections and guarantees while matching or sometimes exceeding PPO across language and image-generation applications.

  • Motivation and approach: REBEL replaces value functions and clipping with a sequence of squared-loss regressions predicting reward differences between trajectories.The approach uses iteratively collected datasets and treats the regressor as an implicit improved policy.
  • Empirical evaluation: When fine-tuning Llama-3-8B, REBEL outperforms PPO and DPO in language generation, converges faster than PPO to similar asymptotic performance in image generation, and achieves 30.1%, 8.16, and 68.2 on three benchmarks.
  • Data and preference settings: REBEL supports on-policy, offline, and hybrid data and can be extended to intransitive preferences.These capabilities broaden the training settings beyond purely online policy optimization.
  • Theoretical connections: REBEL generalizes Natural Policy Gradient, while first-order optimization avoids Fisher Information Matrix computation and reduces variance with finite data.Applying Gauss-Newton to REBEL's regression sequence recovers NPG.
  • Theoretical guarantees: REBEL's reduction-based analysis shows that accurate in-distribution regressions suffice to compete with policies covered by collected datasets.The result matches strong guarantees in the agnostic RL setting and relies on expressive policy-parameterized regressors.

2 REBEL: REgression to RElative REward Based RL

REBEL formulates generative-model reinforcement learning as relative-reward regression around a baseline distribution. Paired completions eliminate the partition function, allowing online, offline, or hybrid training through a simple squared-loss objective.

  • 2 REBEL: REgression to RElative REward Based RL: REBEL treats generative-model training as a contextual bandit, representing a response sequence as one action under deterministic transitions.The formulation applies to language and diffusion-model generation and can extend to deterministic MDPs.
  • 2 REBEL: REgression to RElative REward Based RL: At each iteration, REBEL samples a prompt, a response from the current policy, and a comparison response from a base distribution before solving a square-loss regression.The base distribution may be the current policy or an offline dataset.
  • 2 REBEL: REgression to RElative REward Based RL: The per-iteration objective optimizes reward while constraining the updated policy to remain close to a baseline policy in action distribution.This is expressed as a KL-constrained RL problem.
  • 2.1 Deriving REBEL: REgression to RElative REward Based RL: REBEL inverts the KL-constrained solution to express reward through the policy, then approximates that relationship with squared loss when exact equality is unavailable.The policy-parameterized regressor predicts relative reward differences between paired responses.
  • 2.1 Deriving REBEL: REgression to RElative REward Based RL: Regressing paired reward differences removes the prompt-dependent partition function, and perfect regression recovers the KL-constrained optimum.The resulting dataset may be fully online or hybrid depending on the chosen base distribution.

3 Understanding REBEL as an Adaptive Policy Gradient

This section places REBEL alongside Mirror Descent, NPG, and PPO, showing how its regression updates recover principled policy-improvement behavior while avoiding NPG’s scalability costs and PPO’s clipping heuristics.

  • 3.1 Adaptive Gradient Algorithms for Policy Optimization: Mirror Descent achieves 1/T convergence, bounded KL changes, and monotonic policy improvement in the tabular setting, providing the idealized update that NPG and PPO approximate.Its convergence depends only logarithmically on the response-space size under the stated bounded-reward condition.
  • 3.1 Adaptive Gradient Algorithms for Policy Optimization: NPG approximates the KL geometry with the Fisher Information Matrix, enabling adaptive parameter changes while preserving policy-space proximity, but its Fisher-matrix computation does not scale to modern generative models.The cited theory gives NPG a dimension-free 1/T convergence rate under softmax parameterization.
  • 3.1 Adaptive Gradient Algorithms for Policy Optimization: PPO replaces explicit KL geometry with clipped updates, but clipping neither guarantees policy closeness nor adapts its treatment across samples.Clipping also reduces the effective batch size by discarding useful gradient information.
  • 3.2 Connections between REBEL and MD / NPG: Under point-wise accurate regression, REBEL recovers the exact Mirror Descent update and inherits 1/T convergence, bounded policy KL, and monotonic policy improvement.The update has the form πθt+1(y|x) ∝ πθt(y|x) exp(ηr(x, y)).
  • 3.2 Connections between REBEL and MD / NPG: REBEL’s minimum-norm least-squares solution recovers an advantage-based NPG update, while Gauss-Newton optimization of its regression problems yields NPG as a special case.With finite data, relative-reward regression supplies variance reduction, and first-order optimization avoids explicit Fisher-matrix computation.
  • 3.2 Connections between REBEL and MD / NPG: REBEL addresses NPG’s scalability and PPO’s lack of conservativity or adaptivity without relying on the Fisher Information Matrix or unjustified clipping heuristics.It retains strong convergence and regret guarantees and can extend to preference-based settings with non-transitive preferences.

4 Theoretical Analysis

The theoretical analysis relaxes REBEL’s exact point-wise regression assumption to in-distribution generalization and derives regret guarantees against covered comparator policies.

  • 4 Theoretical Analysis: The regression assumption requires a policy-isomorphic regressor class to fit reward differences accurately, with error ϵ depending on policy complexity and dataset size.Rich log-policy-difference classes can drive ϵ toward zero as training data increase, without implying point-wise prediction accuracy.
  • 4 Theoretical Analysis: In linear feature settings with linear rewards, REBEL reduces to well-specified linear regression whose error typically scales as O(d/|Dt|).Kernelized extensions yield analogous poly(d/|Dt|) scaling in the effective dimension.
  • 4 Theoretical Analysis: Theorem 1 guarantees that among T learned policies, one achieves a regret bound combining a 1/T optimization term with regression error scaled by distribution shift.The shift cost is represented by the concentrability coefficient Cµ→π∗.
  • 4 Theoretical Analysis: REBEL can compete with any comparator policy covered by its training distributions when the regression problems achieve small in-distribution error.The analysis replaces the stronger point-wise accuracy requirement with a reduction from RL to supervised learning.
  • 4 Theoretical Analysis: This agnostic guarantee matches the strongest known results cited for policy-optimization methods including PSDP, CPI, NPG, and PC-PG.The comparator must be covered by the training distributions.
  • 4 Theoretical Analysis: REBEL’s framework extends to non-transitive preferences, with the extension and its analysis deferred to the appendices.The base distribution may be a human labeler, an SFT policy, or the current learned policy.
  • 4 Theoretical Analysis: Sampling at least one response on-policy helps avoid polynomial dependence on |Y|, while alternative covering distributions introduce an additional concentrability factor.Best-of-N and worst-of-N sampling are reported as design choices that improve coverage and reward gaps experimentally.

5 Experiments

REBEL is evaluated across text summarization, general chat, large-batch preference optimization, and image generation, with comparisons emphasizing quality, efficiency, and sampling choices.

  • 5 Experiments: The experiments use reward-model evaluation across language and image generation, including GPT4 winrate, benchmark scores, and LAION aesthetic reward.The image-generation comparison reports inter-quartile means with 95% confidence intervals across three seeds.
  • 5.1 Summarization: REBEL outperforms all baselines on winrate in TL;DR summarization and achieves the highest GPT4 winrate against human references.It also leads reward-model score at 1.4B and 2.8B parameters, with slightly larger KL than PPO.
  • 5.1 Summarization: Compared with PPO and RLOO (k = 4), REBEL has shorter runtimes and lower peak memory usage while achieving comparable winrates.PPO updates policy and value networks, whereas RLOO (k = 4) requires four generations per prompt.
  • 5.2 General Chat: In general chat, REBEL outperforms APA on MT-Bench and AlpacaEval 2.0 while improving GSM8K and HellaSwag and maintaining the base model’s overall leaderboard average.Similar MMLU values indicate preservation of the pretrained model’s basic capability during RL fine-tuning.
  • 5.2.1 Ablation: batch size and data sampling distributions: With the entire dataset and best-of-5/worst-of-5 sampling, one REBEL iteration beats DPO on almost all datasets, especially AlpacaEval 2.0 and Arena Hard.Additional iterations generally do not help in this large-batch setting, possibly because of overfitting to the training dataset.
  • 5.3 Image Generation: REBEL optimizes the consistency model faster early in image-generation training and eventually reaches performance similar to PPO under equal reward-query and gradient-update budgets.Across seeds, REBEL generations also produced lush backgrounds compared with PPO’s generations.

6 Related Work

Related work spans policy-gradient methods, reward regression, and preference fine-tuning. REBEL combines iterative on-policy regression with preference handling while avoiding several approximations used by prior methods.

  • Policy Gradients: Policy-gradient methods include REINFORCE-based approaches and adaptive methods such as NPG, TRPO, and PPO.Adaptive methods precondition policy gradients to promote distribution-level rather than parameter-level updates.
  • Policy Gradients: REBEL avoids the Fisher-matrix approximations commonly used to scale adaptive policy-gradient methods to generative models.This gives REBEL a closer connection between theory and practice at scale.
  • Reward Regression: Unlike prior reward- or value-regression methods, REBEL parameterizes the reward regressor through the policy and collapses a two-stage procedure into one step.The reduction uses relative rewards rather than separately fitting a reward or value target and then extracting a policy.
  • Preference Fine-Tuning of Generative Models: Preference fine-tuning includes RLHF, ranking-loss methods such as DPO, and iterative online methods that collect additional preference feedback.The paper positions REBEL within the iterative online family while using regression targets and paired responses.
  • Preference Fine-Tuning of Generative Models: REBEL uses nuanced reward-model outputs as regression targets, supports mixed on-policy and off-policy data, and avoids some costly online-feedback requirements.Its paired responses also cancel the partition function, avoiding a constant-partition-function approximation.

7 Conclusion and Future Work

The paper concludes that REBEL offers a simple relative-reward regression approach with strong theoretical guarantees and competitive empirical performance. It also identifies open questions about losses, preference coverage, and broader offline RL settings.

  • Conclusion: REBEL reduces RL to sequential relative-reward least-squares problems, avoiding extra networks and clipping while matching or sometimes exceeding more complex methods across language and image generation.The conclusion emphasizes implementation simplicity, scalability, and strong agnostic-RL guarantees.
  • Future Work: Future work could test alternative losses, including log loss or cross-entropy, for improved practical performance or tighter theoretical bounds.The paper specifically mentions first-order and gap-dependent bounds as possible theoretical targets.
  • Future Work: The paper leaves open whether its preference-setting coverage condition is necessary and whether preference-model supervision reproduces improvements reported by related methods.These questions concern the general nonutility-based preference setting.
  • Future Work: A further direction is extending the primarily bandit-focused analysis to broader RL settings and using offline data to improve policy-optimization efficiency.The paper points to reset-based techniques as one possible route.

A Proof of Claim 1

The proof of Claim 1 derives the Fisher information structure underlying REBEL's regression formulation. It uses least-squares solutions and the centered advantage representation to connect reward differences with policy-gradient quantities.

  • A Proof of Claim 1: The proof begins by deriving the Fisher information matrix from differences between policy score functions for paired trajectories.The displayed expression uses score-function differences under the current policy and a reference sampling distribution.
  • A Proof of Claim 1: The regression step is treated as an ordinary least-squares problem whose minimum-norm solution supplies the relevant update.The derivation relies on the vanishing cross terms obtained by completing the square.
  • A Proof of Claim 1: The proof rewrites rewards through the advantage A_π(x, y), defined as reward minus the policy-conditional expected reward.The score-function expectation identity removes terms independent of the sampled completion.

B Proof of Claim 2

The appendix connects REBEL's regression updates to policy-gradient quantities and extends the framework from scalar rewards to general, potentially intransitive preferences. The preference extension uses self-play and least-squares regression to target minimax-winner policies.

  • B Proof of Claim 2: The proof of Claim 2 approximates policy-log-probability changes with a first-order expansion around the current parameters.The approximation replaces the log-policy difference with a score-function inner product and ignores higher-order terms.
  • B Proof of Claim 2: The finite-data derivation estimates the Fisher information matrix and solves the resulting ordinary least-squares problem with a minimum-norm solution.These steps connect the regression formulation to the policy-update analysis.
  • C Extending REBEL to General Preferences: A single utility-based reward model cannot represent intransitive preferences such as a ≻ b, b ≻ c, and c ≻ a.The paper therefore replaces reward-model access with a preference model in the general preference setting.
  • C.1 A Game-Theoretic Perspective on Learning from Preferences: The preference model represents the probability that one completion is preferred to another and induces a skew-symmetric comparison function.Binary feedback is modeled as a Bernoulli draw with the corresponding preference probability.
  • C.1 A Game-Theoretic Perspective on Learning from Preferences: Preference aggregation is framed as finding a minimax winner, the Nash-equilibrium strategy of a symmetric two-player zero-sum game.Approximation quality is measured through the duality gap.
  • C.2 Self-Play Preference Optimization (SPO) with REBEL as Base Learner: REBEL extends to preferences through self-play: sample multiple completions from the current policy, query pairwise preferences, and use completion winrates as rewards.This instantiates the SPO reduction's stable self-play procedure.
  • C.2 Self-Play Preference Optimization (SPO) with REBEL as Base Learner: The preference regression targets differences in expected pairwise preference against the current policy, yielding an iteration-dependent reward.The resulting analysis uses fast convergence of mirror descent toward an approximate minimax winner while accounting for nonzero regression error.

D Proof of Theorem 1

The proof analyzes REBEL through reward-difference regression and mirror-descent-style policy updates, showing how regression error and policy optimization combine to yield a regret guarantee.

  • The learned regressor predicts rewards accurately under both the current policy and the data distribution up to response-independent terms.The proof defines the regression error and its expectations under π_t and μ, then uses these quantities to control advantage error.
  • REBEL updates the policy by exponentiating the estimated reward advantage, with reward-prediction errors adjusted by policy- and dataset-dependent baselines.The update can equivalently subtract the dataset baseline because it is independent of the response.
  • REBEL's policy sequence can be analyzed as mirror descent on estimated advantages, yielding a standard bounded-advantage regret bound with η = ln(|Y|)/(A^2T).The analysis assumes |A_t(x,y)| ≤ A and initializes π_0 uniformly over Y.
  • A telescoping argument converts the per-iteration mirror-descent inequalities into a cumulative guarantee over T iterations.The proof sums the inequalities and concludes using the selected learning rate.
  • The final comparison bounds comparator-policy performance against the uniform mixture of iterates using the estimated-versus-real advantage discrepancy and the concentrability coefficient.The proof explicitly uses the on-policy sample in controlling the discrepancy between estimated and real advantages.
  • Combining the regression and mirror-descent bounds yields the theorem's final regret guarantee.

E Extension of analysis to General Preferences

The general-preference extension replaces single-policy coverage with a stronger unilateral coverage condition and uses regression generalization assumptions to establish approximate minimax-winner learning.

  • E Extension of analysis to General Preferences: General preferences require comparing a learned policy against all other policies, so single-policy coverage is insufficient for efficient learning.The analysis therefore adopts unilateral concentrability, represented here by C_μ.
  • E Extension of analysis to General Preferences: The extension assumes a small regression generalization error over all T iterations under the stronger coverage condition.
  • E Extension of analysis to General Preferences: Under these assumptions and a suitable learning rate, REBEL learns an approximate Minimax Winner after T iterations.The theorem hides problem-dependent constants independent of ϵ, C_μ, and T in its O-notation.
  • E Extension of analysis to General Preferences: The authors identify the stronger coverage requirement as a limitation of the general-preference extension and leave the necessary condition for future work.
  • E.1 Proof of Theorem 2: The proof reuses reward-difference regression and mirror-descent arguments, adapting the advantage analysis to skew-symmetric preference rewards.Skew symmetry makes the self-comparison reward zero and supports the corresponding advantage simplification.

F.1.1 Dataset Details

The experimental appendices specify datasets, models, baselines, reward shaping, evaluation prompts, and example generations across summarization, chat, and image generation.

  • F.1.2 Model Details: The TL;DR summarization experiments use SFT and reward models at 1.4B, 2.8B, and 6.9B scales, with LoRA for smaller models and full-parameter training for 6.9B.Training uses A6000 GPUs for the smaller models and H100 GPUs for the 6.9B model.
  • F.1.3 Baseline Implementation Details: REBEL is compared with PPO, DPO, RLOO, REINFORCE, and iterative DPO using specified batch sizes, learning rates, schedules, and training epochs.The appendix also describes iterative DPO as regressing pairwise preference signals rather than reward differences.
  • F.1.4 Reward Details: The reward includes a KL penalty to keep the policy near its initialization and a fixed penalty for generations exceeding the maximum length.The reward-model score is combined with these constraints during training.
  • F.1.6 Winrate Details: Winrate evaluation asks GPT-4 to compare summaries for accuracy, coverage, and coherence, then select the preferred summary.The example evaluation prefers Summary B because it better preserves the post's emotional conflict coherently.
  • F.2.2 Model Details: The general-chat experiments use OpenChat-3.5 with Starling-RM on Nectar and Meta-Llama-3-8B-Instruct with ArmoRM on UltraFeedback.Prompts longer than 1,024 tokens are filtered, and the two model settings use different training durations and parameter-update scopes.
  • F.3.1 Dataset Details: Image-generation experiments use 45 animal prompts with an LCM distillation of Dreamshaper v7 and compare PPO and REBEL under reported consistency-model settings.Each run uses four A6000 GPUs and requires 10 hours.
  • G.1 Summarization Generations from a 2.8B model: The appendix includes example summarization and chat generations, including a concise relationship-post summary and a detailed ice-hockey explanation.

H Ablation Analysis

The ablation identifies η as REBEL's key tuning parameter, while the reward–KL analysis shows that REBEL can obtain higher reward-model scores at the cost of larger divergence for its highest-scoring generations.

  • H Ablation Analysis: REBEL performs best when η = 1 on the 2.8B summarization ablation, while increasing or decreasing η reduces performance.The authors recommend η = 1.0 as a starting point because it performed best across their language-modeling and image-generation experiments.
  • I Trade-off between Reward Model Score and KL-divergence: During training, PPO exhibits lower KL divergence, whereas REBEL reaches larger reward-model scores; by the end, their KL and reward-model scores are similar.The comparison holds sample complexity per update constant across the algorithms.
  • I Trade-off between Reward Model Score and KL-divergence: At the final checkpoint, REBEL achieves higher reward-model scores for generations with small divergence but requires larger divergence for generations with the highest scores.The analysis bins generations by KL divergence into ten equally sized groups and averages reward-model scores within each bin.

J Regression Loss During Training

During summarization training, REBEL maintains low reward-difference prediction error despite an unbounded learned reward. On MT-Bench, it outperforms comparison models across six of eight evaluated dimensions.

  • J Regression Loss During Training: Figure 6 tracks both smoothed reward-difference prediction error and per-iteration loss during training.
  • J Regression Loss During Training: REBEL minimizes the regression loss throughout 6.9B-parameter summarization training, maintaining relatively low mean squared error.Observed rewards were mostly within [−10, 10], although the learned reward model is unbounded.
  • K Breakdown of MT-Bench: Figure 7 presents the MT-Bench results broken down across eight dimensions.
  • K Breakdown of MT-Bench: Six of eight MT-Bench dimensions favor REBEL over APA and base models, including writing, roleplay, math, extraction, STEM, and humanities.
Loading 2404.16767v4…