PaperlayerView exact revision on arXiv

Source-linked AI summary

ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models

Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, Zhi-Quan Luo

arXiv:2310.10505v4cs.LG

TL;DR

RLHF’s PPO-based RL stage is computationally burdensome and requires extensive optimization machinery. ReMax builds on REINFORCE, exploiting RLHF’s fast simulation, deterministic transitions, and trajectory-level rewards while using a greedy baseline for variance reduction. It simplifies training, reduces memory and time, and reports strong open-source 7B results, although online sampling remains slower than DPO and variance can worsen in a worst case.

  • Problem

    RLHF’s RL stage is hindered by PPO’s computational burden, including value-model training and associated optimization components.

  • Method

    ReMax builds on REINFORCE, removes value-model training, exploits RLHF task properties, and introduces a greedy-baseline variance-reduction technique.

  • Results

    ReMax simplifies implementation, reduces memory usage, speeds training, and improves task performance, including a 94.78% AlpacaEval win rate and 7.739 MT-bench score for Mistral-7B.

  • Takeaways & Limitations

    ReMax facilitates easier application of RLHF while offering computation efficiency comparable to DPO and reported advantages over PPO.

  • Takeaways & Limitations

    ReMax is 1.3 times slower than DPO per epoch because of online sampling, and its variance may increase in the worst case.

Abstract

from arXiv · show

Reinforcement Learning from Human Feedback (RLHF) is key to aligning Large Language Models (LLMs), typically paired with the Proximal Policy Optimization (PPO) algorithm. While PPO is a powerful method designed for general reinforcement learning tasks, it is overly sophisticated for LLMs, leading to laborious hyper-parameter tuning and significant computation burdens. To make RLHF efficient, we present ReMax, which leverages 3 properties of RLHF: fast simulation, deterministic transitions, and trajectory-level rewards. These properties are not exploited in PPO, making it less suitable for RLHF. Building on the renowned REINFORCE algorithm, ReMax does not require training an additional value model as in PPO and is further enhanced with a new variance reduction technique. ReMax offers several benefits over PPO: it is simpler to implement, eliminates more than 4 hyper-parameters in PPO, reduces GPU memory usage, and shortens training time. ReMax can save about 46% GPU memory than PPO when training a 7B model and enables training on A800-80GB GPUs without the memory-saving offloading technique needed by PPO. Applying ReMax to a Mistral-7B model resulted in a 94.78% win rate on the AlpacaEval leaderboard and a 7.739 score on MT-bench, setting a new SOTA for open-source 7B models. These results show the effectiveness of ReMax while addressing the limitations of PPO in LLMs.

1. Introduction

RLHF’s RL stage is computationally demanding because PPO adds value-model training and related optimization components. ReMax targets this mismatch by exploiting RLHF properties, simplifying REINFORCE with variance reduction, and improving efficiency and reported performance.

  • Motivation: RLHF consists of supervised fine-tuning, reward-model learning, and RL-based language-model fine-tuning, with the RL stage presenting the main computational challenge.PPO is commonly used for the third stage but adds substantial training cost.
  • Motivation: PPO training can take 4× longer than the first two RLHF stages and increase GPU memory consumption by at least 2×.The extra value model and associated optimization components drive these demands.
  • Motivation: RLHF has fast simulation, deterministic transitions, and trajectory-level rewards, properties that PPO does not exploit.These differences motivate an algorithm tailored specifically to LLM alignment.
  • ReMax: ReMax builds on REINFORCE, avoids training a value model, and introduces variance reduction using a greedy baseline.Its estimator uses reward-weighted likelihood maximization while retaining unbiasedness without PPO’s value model.
  • ReMax: ReMax requires 6 main-code lines instead of PPO’s 30+ and eliminates 4 PPO hyper-parameters.The removed parameters include clipping, GAE, value-model learning-rate, and off-policy-epoch settings.
  • Results: ReMax operates about 1.6× as fast as PPO and achieves a 94.78% AlpacaEval win rate and 7.739 MT-bench score with Mistral-7B.The reported results establish a new open-source 7B model state-of-the-art.

2. Related Work

Prior work studies RLHF broadly and improves its optimization, implementation, reward modeling, and parameter-efficient training. ReMax differs by designing around RLHF’s task-specific properties.

  • RLHF research: Prior RLHF research covers alignment, reward-model overoptimization, sample complexity, optimization errors, and broader surveys of the field.These studies establish the surrounding research landscape for RLHF.
  • Optimization methods: Recent work improves the RL stage through NLP benchmarks, LoRA with PPO, PPO implementation analysis, reward ranking, and non-RL optimization methods.The cited approaches address efficiency, implementation, or objective design from different directions.
  • ReMax’s distinction: ReMax differs from prior studies by leveraging unique properties of RLHF tasks rather than treating them as generic reinforcement-learning problems.The paper positions this task-specific design as its central distinction.
  • Comparison baseline: PPO often outperforms alternatives and remains favored in open-source software, providing the principal comparison point for ReMax.The paper compares its method against an established and widely used baseline.

3. Problem Formulation

The paper formulates RLHF as reward maximization over token-generation trajectories and examines why PPO’s value model is poorly matched to this setting. It argues that fast, deterministic, trajectory-reward dynamics enable a more efficient alternative.

  • 3. Problem Formulation: An LLM generates tokens sequentially from a prompt according to πθ until an EOS token or maximum length T is reached.The state is the prompt plus previously generated tokens, and each action selects a vocabulary token.
  • 3. Problem Formulation: RLHF addresses LLMs’ difficulty aligning generated text with human preferences by fine-tuning against a reward model learned from preference data.The reward model assigns a scalar value to a complete response, with higher values generally preferred.
  • 3. Problem Formulation: The objective is to maximize expected reward over the prompt distribution, with KL regularization representable as a penalty added to the reward.The formulation can include training prompts and new prompts without preference annotations.
  • 3.1. Reward Maximization by RL: RLHF reward maximization can be cast as an MDP in which states are token histories, actions are vocabulary tokens, and transitions append the selected token.Intermediate rewards are zero while the completed response receives the reward, making the MDP equivalent to the original objective.
  • 3.2. A Natural Solution to Problem (2): PPO: PPO uses importance sampling, clipping, a value model, advantage estimation, and GAE-related hyper-parameters to optimize a surrogate objective.The value model and associated settings make PPO complicated to implement and tune.
  • 3.3. PPO is not the Best Fit for RLHF: For a 7B model, training the PPO value model consumes about 50% of GPU usage and more than 4× its inference memory.Its scale and training gradients, activations, and optimizer states substantially increase computational load.
  • 3.3. PPO is not the Best Fit for RLHF: RLHF offers fast simulation, deterministic transitions, and trajectory-level rewards, reducing the apparent need for PPO’s value-function machinery.A response reward requires one language-model and reward-model query, transitions append tokens deterministically, and intermediate rewards are zero.
  • 3.3. PPO is not the Best Fit for RLHF: The paper therefore claims that RLHF’s expected return can be estimated efficiently and sample-efficiently without a value model, motivating ReMax.This claim directly motivates the proposed algorithm in the following section.

4. Proposed Method

ReMax adapts REINFORCE to RLHF by using a greedy reward baseline that reduces gradient variance while avoiding PPO’s value model. The method targets RLHF’s deterministic, fast-simulation setting and retains computational efficiency with stronger practical performance.

  • From REINFORCE to ReMax: ReMax builds on REINFORCE but introduces a subtractive baseline for gradient estimation.The baseline is designed to improve REINFORCE’s effectiveness without adding a value model.
  • From REINFORCE to ReMax: The greedy baseline compares sampled and greedy-response rewards to reduce gradient variance and balance reward magnitudes across prompts.The estimator remains unbiased because the baseline is independent of the sampled response.
  • From REINFORCE to ReMax: ReMax’s baseline adapts to both prompts and training, unlike fixed or online reward normalization approaches.PPO can provide similar adaptation, but requires a heavy value model.
  • Theory of ReMax: The gradient estimator is unbiased, has bounded variance, and ReMax converges to a stationary point in expectation.The variance bound scales with reward magnitude, sequence length, score-function bound, and inverse sample count.
  • Theory of ReMax: ReMax reduces variance when the optimal action is not dominant, such as when πθ(a1|x) ≤0.5, but may increase worst-case variance.The paper states that this drawback does not affect convergence and may be acceptable in over-optimized regimes.

5. Experiments

Experiments evaluate ReMax on RLHF effectiveness, computational efficiency, and practical leaderboard performance. ReMax matches PPO’s reward, improves win rates, reduces memory demands, and achieves strong results on Mistral-7B, while prompt selection remains a limitation.

  • Part I (a): On the Effectiveness of ReMax: ReMax achieves a reward comparable with PPO and shows stable training without significant gradient variance.The gradient norms of PPO and ReMax are lower than DPO’s in the reported experiments.
  • Part I (a): On the Effectiveness of ReMax: 31.4 points was ReMax’s largest improvement over the SFT model under the same initialization.DPO+ReMax achieved the highest win rate of 84.7%.
  • Part I (b): On the Efficiency of ReMax: About 50% memory savings versus PPO allow ReMax to train without optimizer-state offloading on a 7B model.With optimizer-state offloading, ReMax supports a 1.4× larger batch size than PPO.
  • Part I (b): On the Efficiency of ReMax: 1.6× faster than PPO was ReMax’s reported evaluation speed-up with four GPUs.ReMax generates two responses but does not train a value model; generation is reported as faster than backward computation.
  • Part I (b): On the Efficiency of ReMax: ReMax’s computational efficiency is comparable to DPO, although one epoch takes 1.3 times longer because of online sampling.The reported maximum batch size is 96 for both methods without offloading on four GPUs.
  • Part II: ReMax on LeaderBoard: ReMax consistently improves performance across prompt datasets, but performance deteriorates at 40k prompts, indicating over-optimization and a need for regularization.The authors report that effective prompt selection remains unresolved.

6. Conclusion

ReMax is presented as a simpler and more efficient alternative to PPO for RLHF, using REINFORCE with variance reduction to improve practicality and performance.

  • 6. Conclusion: ReMax identifies RLHF-specific properties and uses them to guide algorithm design instead of directly applying generic reinforcement-learning methods.The method targets fast simulation, deterministic transitions, and trajectory-level rewards.
  • 6. Conclusion: ReMax builds on REINFORCE and introduces a greedy baseline for variance reduction, avoiding the value model required by PPO.This design is intended to reduce implementation complexity and computational requirements.
  • 6. Conclusion: Compared with PPO, ReMax simplifies implementation, reduces memory usage, minimizes hyper-parameters, speeds training, and improves task performance.The paper concludes that these benefits facilitate easier application of RLHF to LLMs.

Broader Impacts

The supplied passages frame RLHF as an approach for aligning LLM behavior, while emphasizing unresolved reward-modeling, bias, and broader-use concerns.

  • Broader Impacts: RLHF is used to improve LLM behavior toward human preferences, complementing instruction tuning as an alignment approach.Instruction tuning depends on substantial pre-prepared human-annotated data, whereas RLHF is presented as an alternative alignment procedure.
  • Broader Impacts: ReMax is intended to make RLHF cheaper and easier, but misuse may cause LLMs to generate harmful and disrespectful responses.This warning concerns potential consequences of applying the techniques outside their intended safe use.
  • Broader Impacts: The paper does not address how to infer reward functions effectively from human preferences or how to mitigate reward bias.These are identified as important open questions beyond the computational-efficiency focus of ReMax.
  • Broader Impacts: Because reward models may be imperfect and biased, the induced optimal response may not align with expectations; KL regularization is often incorporated to address this issue.The paper describes one-step and full-step KL penalties, with full-step regularization more strongly penalizing divergence but introducing additional stochastic noise.

C.2. Proof of Proposition 2

This section establishes theoretical properties of ReMax under autoregressive softmax policies, including unbiased variance-reduced gradients, smoothness, and convergence-related guarantees with explicit limitations.

  • C.2. Proof of Proposition 2: ReMax analyzes the objective R(θ) for autoregressive softmax policies under a bounded reward assumption and derives an update rule based on stochastic policy gradients.The analysis takes expectation over sampled prompts and responses.
  • C.2. Proof of Proposition 2: The convergence analysis does not establish global optimal convergence for stochastic policy gradients on these non-convex problems.The paper leaves this question for future work and notes that additional regularity conditions would be required.
  • C.2. Proof of Proposition 2: The greedy-baseline estimator is unbiased, relying on the Bartlett identity and conditional independence induced by greedy sampling.The proof shows that the baseline-dependent term has zero expectation.
  • C.2. Proof of Proposition 2: The optimal baseline value minimizes the stochastic-gradient variance, and its variance is globally lower than that of the unadjusted estimator for all θ.The baseline is derived by treating the objective as a quadratic function of the baseline.
  • C.2. Proof of Proposition 2: The expected-reward baseline does not provide global variance reduction, although Proposition 6 gives a condition under which variance is reduced.The condition depends on the policy probability and the two action rewards.
  • C.2. Proof of Proposition 2: The reward objective is shown to be 6-smooth with respect to θ under the stated softmax parameterization.The proof bounds Hessian eigenvalues and extends smoothness from partial parameters to the complete parameter vector.

E. Experiment Details

The experiments use specified model, optimizer, hardware, data-processing, and evaluation settings to compare ReMax with PPO and assess performance and efficiency.

  • E. Experiment Details: Experiments run on four A800-80GB GPUs, using ZeRO-2, flash attention, and gradient checkpointing unless otherwise stated.Trainable-model optimizer states are offloaded to CPU when offloading is used, while parameters remain on the GPU.
  • E. Experiment Details: The Llama-2-7B experiments use three random seeds, whereas the Mistral-7B experiments use a single seed because of computational burden.The authors explicitly identify the single-seed Mistral-7B evaluation as a study limitation.
  • E. Experiment Details: For Llama-2-7B, supervised fine-tuning and reward-model training each run for 2 epochs, while RL runs for 1 epoch with a learning rate of 10^-6.PPO and ReMax use a KL penalty coefficient of 0.1, temperature 1, and top-p 0.9 during generation.
  • E.2. Mistral 7B: The Mistral-7B evaluation uses Mistral-7B-instruct-v0.2 with UltraRM-13B as the reward model.The authors report that UltraRM-13B has higher full-hh-rlhf accuracy than their trained Llama-2-13B reward model: 71% versus about 65%.
  • E.2. Mistral 7B: Mistral-7B experiments exclude prompts longer than 384 tokens and use a maximum combined prompt-response length of 784 to reduce computational costs.The authors did not tune these Mistral-7B hyper-parameters because training was expensive.
  • E. Experiment Details: 54% of PPO’s GPU memory consumption is required by ReMax in the Llama-2-7B calculation: 172.12 GB versus 319.14 GB.The calculation attributes the difference to omitting PPO’s additional trainable value model.
  • E. Experiment Details: ReMax takes 1.8 hours per epoch compared with PPO’s 2.9 hours for Llama-2-7B, making PPO 1.6× slower.These times are reported for the one-epoch training comparison used in Figure 2.

F.1. Training Instability of REINFORCE

REINFORCE is computationally simple but suffers from high-variance stochastic gradients and training instability, motivating ReMax’s variance-reduced formulation.

  • REINFORCE tends to diverge when fine-tuning OPT-1.3B on full-hh-rlhf, as indicated by a rapidly increasing gradient norm.Its evaluation reward performance was poor in this experiment.
  • On Llama-2-7B with full-hh-rlhf, REINFORCE avoids divergence but still shows instability-related concerns.The experiment tracks evaluation reward and gradient norm to assess its behavior on the larger model.
  • REINFORCE exhibits high variance in stochastic gradients, which can cause training instability in practice.The paper therefore abandons REINFORCE in its main experiments and studies ReMax instead.
  • ReMax addresses REINFORCE’s variance problem with a proposed variance-reduction technique.This preserves the simplicity and computational efficiency motivating the REINFORCE-based approach.

F.2. Ablation Study on KL Regularization

The experiments examine KL regularization choices and a shortened greedy-sampling strategy for ReMax. Both KL variants work, while shortening greedy responses can improve training speed under suitable settings.

  • KL Regularization: Both one-step KL and full-step KL regularization work in ReMax.The paper reports greater effectiveness for full-step KL regularization and reduces β from 0.1 to 0.01 for comparable performance.
  • KL Regularization: The coefficient β must be tuned for the specific KL-regularization method and scenario.The paper presents β sweeps separately for one-step and full-step KL regularization.
  • Fast Greedy Sampling: Shorter responses can provide reward scores that closely approximate those of longer responses.The comparison evaluates a complete response against a version truncated to its initial half.
  • Fast Greedy Sampling: A greedy response length of 128 yields performance comparable to original ReMax, whereas length 64 significantly sacrifices performance.The normal generation length remains unchanged in this experiment.
  • Fast Greedy Sampling: 2.1 times: ReMax’s training speed-up over PPO when halving greedy response length.The reduction is attributed in principle to lower generation time from transformer self-attention’s quadratic complexity.

F.4. Fine-tuning GPT-2 on a Classical NLP Task

ReMax is evaluated beyond RLHF by fine-tuning GPT-2 on sentiment-controlled movie-review completion. It maximizes reward while using less memory and training faster than PPO in this setting.

  • ReMax fine-tunes GPT-2 to generate movie-review continuations that maintain positive sentiment.A sentiment classifier supplies the reward, using the IMDB dataset rather than human-preference data.
  • The experiment tests whether reward maximization can be applied beyond RLHF on a classical NLP task.The task uses partial movie reviews as inputs and asks the model to complete them while preserving sentiment.
  • 2.2 times: ReMax’s training speed-up while effectively maximizing rewards on the GPT-2 IMDB experiment.PPO uses an off-policy training epoch of 4 in this setup, compared with 1 in the RLHF experiments.

G.2. In-distribution Samples

The in-distribution examples use prompts from the training data to inspect how RLHF changes model behavior. The authors report substantially improved response quality in the ReMax fine-tuned model.

  • The in-distribution evaluation uses sharegpt-en prompts from the training data and compares Mistral-Instruct-v0.2 with its ReMax fine-tuned version.The fine-tuned model uses 20k training prompts, with examples shown in Figures 18–21.
  • The authors observe that response quality has greatly improved in the in-distribution samples.This conclusion is based on qualitative response examples rather than a reported numerical metric.
  • The sample set includes comparisons involving SFT and DPO outputs alongside ReMax-related examples.The provided passages explicitly show SFT and DPO responses for some prompts.
  • The examples cover questions about classification, therapists, dog breeds, Atlantis, and instruction-following behavior.The passages identify these prompts and indicate that ReMax responses are highlighted in the corresponding figures.
Loading 2310.10505v4…