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 computation burden and PPO’s poor fit for LLM tasks motivate ReMax, which builds on REINFORCE with a greedy variance-reduction baseline. ReMax matches or surpasses PPO, and its Mistral-7B model achieved a 94.78% win rate against text-davinci-003 on AlpacaEval and a 7.739 score on MT-bench.

  • Problem

    RLHF’s heavy computation burden hinders its wide application, and PPO is not the best fit because LLM RLHF differs from general reinforcement learning tasks.

    Source: §1 · p. 1, §3.3 · p. 4

  • Method

    ReMax builds on REINFORCE and uses a greedy baseline for variance reduction to address stochastic-gradient variance in LLM RLHF.

    Source: §6 · p. 9

  • Results

    ReMax matches or surpasses PPO; its Mistral-7B model achieved a 94.78% win rate against text-davinci-003 on AlpacaEval and a 7.739 score on MT-bench.

    Source: §1.1 · p. 3, §6 · p. 9

  • Takeaways & Limitations

    ReMax facilitates easier application of RLHF in LLMs.

    Source: §6 · p. 9

    ReMax training took 1.3 times longer than DPO for one epoch due to online sampling.

    Source: §5.2 · p. 9

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

The paper addresses the computational burden of RLHF’s final fine-tuning stage by proposing ReMax, a simpler algorithm designed around properties of language-model RLHF. It reports improved efficiency and strong benchmark results for ReMax.

1. Introduction

  • RLHF uses supervised fine-tuning for initialization, reward-model learning from preference data, and RL fine-tuning on prompt-only data; its final stage is computationally burdensome with PPO.The introduction motivates focusing on that RL stage, where PPO’s value model and related training components impose heavy computation demands.Source: §1 · pp. 1–2

1.1. Our Contribution

  • ReMax exploits RLHF’s fast simulation, deterministic transitions, and trajectory-level rewards, using REINFORCE with an unbiased policy-gradient estimator and a new variance-reduction technique instead of a value model.The paper says PPO does not exploit these RLHF properties; removing the value model avoids its added memory and training-time costs.Source: §1.1 · p. 2
  • ReMax uses 6 lines of main code versus PPO’s 30+, and eliminates 4 PPO hyperparameters, including clipping, GAE coefficient, value-model learning rate, and off-policy training epochs.The authors argue that this simplifies implementation and reduces costly LLM hyperparameter tuning.Source: §1.1 · p. 2
  • ReMax runs about 1.6× as fast as PPO, and it trains Llama-2-7B on A800-80GB GPUs without optimizer offloading, which PPO requires.The paper attributes its faster iterations to removing value-model training and its higher throughput to memory savings that allow larger batch sizes.Source: Figure 2 · p. 2, §1.1 · p. 3
  • Mistral-7B (Jiang et al., 2023) achieves a 94.78% win rate against text-davinci-003 on AlpacaEval and 7.739 on MT-bench (Zheng et al., 2023b).The authors describe this as a new state of the art for open-source 7B models.Source: §1.1 · p. 3

2. Related Work

The section situates ReMax within prior RLHF research and highlights its distinction: leveraging properties specific to RLHF tasks while addressing a field in which PPO is often favored.

2. Related Work

3. Problem Formulation

The section formulates RLHF as maximizing a learned reward over generated responses and casts token generation as an MDP. It argues that RLHF’s properties make PPO’s value model unnecessary for estimating returns.

3. Problem Formulation

  • The LLM samples tokens from a finite vocabulary given prior context and stops at EOS or maximum length T, with responses treated as length T through padding.Source: §3 · p. 3
  • A reward model scores each complete response to a prompt, is learned from human preference data using binary classification, and supports reward maximization over the prompt distribution.The formulation omits KL regularization for simplicity, noting that its penalty can be incorporated into the reward.Source: §3 · p. 3

3.1. Reward Maximization by RL

  • The MDP represents each state as the prompt plus prior tokens, each action as the next token, and assigns zero reward to intermediate tokens before the response reward.Source: §3.1 · pp. 3–4

3.2. A Natural Solution to Problem (2): PPO

  • PPO uses importance sampling, clipping, a learned value model for advantage estimation, and GAE, adding trainable components and hyper-parameters.Source: §3.2 · p. 4

3.3. PPO is not the Best Fit for RLHF

  • For a 7B model, training PPO’s value model accounts for about 50% of GPU usage, and the model uses more than 4× the GPU memory during training than inference.Source: §3.3 · p. 4
  • RLHF offers fast simulation, deterministic transitions and rewards despite a randomized policy, and trajectory-level rewards that are zero before the response ends.Source: §3.3 · pp. 4–5
  • RLHF’s expected long-term return can be obtained computation- and sample-efficiently without PPO’s value model, which the authors argue does not fit the problem.Source: §3.3 · p. 5

4. Proposed Method

ReMax adds a greedy-response reward baseline to REINFORCE to address high gradient variance while retaining an unbiased estimator. Its theory establishes bounded variance and convergence to a stationary point, but its variance-reduction analysis is limited to 2-armed bandits and does not guarantee lower variance in every regime.

4.1. A Brief Review on REINFORCE

  • REINFORCE updates by reward-weighted likelihood maximization on rollout samples, unlike supervised learning with known responses and a fixed sample distribution.Its update is a single gradient-ascent step using the stochastic policy-gradient estimator.Source: §4.1 · p. 5
  • REINFORCE is computationally efficient because it estimates gradients with one language- and reward-model query and needs no value model, but its stochastic gradients have high variance.The paper attributes the remaining variance in LLM RLHF to policy token-generation randomness and widely varying reward scales across prompts and training.Source: §4.1 · pp. 5–6

4.2. From REINFORCE to ReMax

  • ReMax subtracts the reward of a greedily sampled response as a baseline, adapting reward normalization to prompts and training while reducing gradient variance.The baseline compares a random response’s reward with the greedy response’s reward; the resulting estimator is shown to remain unbiased.Source: §4.2 · p. 6
  • Unlike DPO (Rafailov et al., 2023), ReMax supports arbitrary reward models and prompt datasets without preference annotations, and updates online from current-model responses.The paper describes ReMax as comparable to DPO in computation efficiency; DPO instead uses a static preference dataset and a fixed reference-model baseline.Source: §4.2 · pp. 6–7

4.3. Theory of ReMax

  • ReMax’s estimator is unbiased with variance bounded by c × rmax^2 × T^2 × S^2/N, and learning rates ηk = O(1/k) yield convergence to a stationary point in expectation.The unbiasedness argument requires the baseline to be statistically independent of the sampled response.Source: §4.3 · p. 7
  • In 2-armed bandits, ReMax reduces variance when the optimal action is not dominant, but may increase it in the worst case, while variance remains bounded and convergence is unaffected.The authors restrict this analysis to 2-armed bandits and note that slower convergence in an over-optimized regime may be acceptable to avoid LLM overfitting.Source: §4.3 · p. 7

5. Experiments

Experiments evaluate ReMax’s effectiveness and efficiency on Llama-2-7B with full-hh-rlhf, then test its application to Mistral-7B on AlpacaEval and MT-Bench. Results cover response quality, training stability, compute costs, and sensitivity to prompt data.

5.1. Part I (a): On the Effectiveness of ReMax

  • ReMax achieved reward comparable with PPO and stable training gradients; its gradient norms were lower than DPO’s, which the authors attribute partly to ReMax’s token-normalized log-likelihood.They also explain that DPO’s gradient estimator calculates two score functions, while ReMax requires one; ReMax can normalize log-likelihood across tokens to stabilize training.Source: §5.1 · p. 8, Figure 5 · p. 8, Figure 6 · p. 8
  • ReMax achieved the largest improvement over SFT at 31.4 points, while DPO+ReMax reached the highest win rate at 84.7%.The authors interpret the DPO+ReMax result as evidence that DPO can initialize RL optimization and ReMax can address DPO’s out-of-distribution shift through online learning.Source: §5.1 · p. 8, Figure 7 · p. 8

5.2. Part I (b): On the Efficiency of ReMax

  • ReMax saves about 50% memory versus PPO and achieves about 1.6 times speed-up (≈2.9/1.8) with 4 GPUs, while supporting training without optimizer-state offloading.With optimizer-state offloading, ReMax supports a 1.4 times larger batch size; its efficiency is comparable to DPO, though one epoch takes 1.3 times longer than DPO’s 1.4 hours due to online sampling. Without offloading, DPO and ReMax have the same maximum batch size of 96.Source: §5.2 · pp. 8–9, Table 2 · p. 8

5.3. Part II: ReMax on LeaderBoard

  • ReMax improved performance across the tested cases, but performance deteriorated when the prompt-data size reached 40k, indicating over-optimization.The authors say regularization is needed and that effective prompt selection remains unresolved.Source: §5.3 · p. 9, Table 3 · p. 9
  • 20k sharegpt-en prompts yielded promising results against open-source and private models, setting a new standard for 7B models.The authors note PPO might achieve similar results but requires more hyper-parameter tuning and computational resources than were available to them.Source: §5.3 · p. 9, Table 4 · p. 9

6. Conclusion

ReMax adapts REINFORCE for LLM RLHF by using a greedy baseline to reduce stochastic-gradient variance. Compared with PPO, it simplifies implementation, reduces memory use and hyper-parameters, speeds training, and improves task performance.

6. Conclusion

  • ReMax uses a greedy baseline to reduce the stochastic-gradient variance that must be addressed when scaling REINFORCE to large models.The paper identifies REINFORCE as computationally suitable for RLHF in LLMs.Source: §6 · p. 9
  • Compared with PPO, ReMax simplifies implementation, reduces memory use and hyper-parameters, speeds training, and improves task performance.The paper frames these benefits as facilitating easier application of RLHF in LLMs.Source: §6 · p. 9
  • The work identifies fundamental differences between LLM RLHF and generic RL tasks and their implications for algorithm design.Source: §6 · p. 9

Broader Impacts … C.1. Proof of Proposition 1

The appendices situate ReMax within RLHF research, describe KL-regularized variants, and outline both potential benefits and limitations. They also provide a proof of Proposition 1.

Broader Impacts

  • ReMax aims to make LLM alignment through RLHF cheaper and easier, but misuse could lead to harmful or disrespectful responses.The paper hopes the technique will improve human interactions.Source: Broader Impacts · p. 10

A. Additional Related Work

  • REINFORCE is often inferior to actor-critic methods for stochastic transitions and dense rewards, yet the paper finds it computationally suitable for LLM RLHF.Prior theory studied baselines, regret, and convergence of expected-gradient ascent, including results by Dayan (1991), Weaver & Tao (2001), Greensmith et al. (2004), Zhang et al. (2021), and Mei et al. (2021; 2022).Source: A. Additional Related Work · p. 15
  • Unlike PPO, ReMax has a theoretical convergence guarantee, but it cannot replace PPO in general RL tasks where a value model may still be required.Best-of-n response selection can be effective but increases inference computation; Liu et al. (2023) describe statistical rejection as a way to mitigate this burden.Source: A. Additional Related Work · p. 15
  • ReMax’s variance-reduction method is related to control variates, but uses log πθ(a1:T |x) ∗ b(x), which is readily available rather than requiring a known expectation or many samples.The passage describes control variates as reducing variance when the estimated variable and auxiliary variable are correlated.Source: A. Additional Related Work · p. 15
  • ReMax addresses computational efficiency, not the broader challenges of inferring human-preference rewards or mitigating reward bias.The paper points to Skalse et al. (2022), Kirk et al. (2023), Pang et al. (2023), Munos et al. (2023), and Xiong et al. (2023) on these topics.Source: A. Additional Related Work · p. 15

B. Handling KL Regularization

  • KL regularization addresses reward-model bias through one-step or full-step penalties; full-step penalties more strongly discourage divergence but add stochastic noise that makes optimization harder.The penalty strength is controlled by β > 0, and both formulations are reported to work in practice.Source: B. Handling KL Regularization · pp. 15–16

C.1. Proof of Proposition 1

  • The proof of Proposition 1 uses the Bartlett identity and greedy sampling’s conditional independence to show the baseline-related term is zero.It then uses independence across samples and the triangle inequality to complete the second argument.Source: §C.1 · pp. 16–17

C.2. Proof of Proposition 2 … D. Technical Lemmas

These sections establish variance-reduction results for ReMax baselines, bound the smoothness of its policy objective, and outline the assumptions and limitations of the convergence analysis. The results distinguish globally variance-reducing optimal baselines from alternatives that reduce variance only under stated conditions.

C.2. Proof of Proposition 2

  • For bounded-reward autoregressive softmax policies, Proposition 2’s proof applies standard SGD analysis to ReMax’s update, using smoothness, variance bounds, and an unbiased gradient estimator.The stated analysis covers randomness in collecting prompts and responses.Source: §C.2 · pp. 17–18
  • Global-optimal convergence analysis for stochastic policy-gradient methods, including ReMax, remains open and would require regularity conditions such as the Polyak-Łojasiewicz inequality.The paper leaves this direction for future work.Source: §C.2 · pp. 17–18

C.3. Proof of Proposition 3

  • For the two-action case, Proposition 3’s baseline reduces variance when p ≤ 0.5 before the optimal action dominates and when p < 1 after it dominates.ReMax uses r2 before the optimal action dominates and r1 after it does.Source: §C.3 · p. 20

C.4. The Optimal Baseline Value

  • The optimal baseline minimizes stochastic-gradient variance and satisfies Var[eg(θ)] < Var[bg(θ)] for all θ, though computing it is impractical.Its computation requires expectations over the score function and its reward-weighted version.Source: §C.4 · p. 20

C.5. Variance Reduction of Using Expected Value

  • The expected-value baseline reduces variance when πθ(a1|x) < 2/3 + r(x, a2)/(3(r(x, a1) − r(x, a2))), but does not do so globally.Proposition 6 states that the expected-value baseline lacks the global variance-reduction property.Source: §C.5 · pp. 20–21

D. Technical Lemmas

  • R(θ) is 6-smooth with respect to the full parameter vector θ under autoregressive softmax policies.The proof first establishes smoothness for each timestep’s parameters, then extends it to the full parameter vector.Source: D. Technical Lemmas · pp. 21–22

E. Experiment Details … F. Additional Results

The appendix documents implementation and experiment settings, including a single-seed limitation for Mistral-7B. Its memory and training-time estimates show ReMax using fewer resources than PPO for Llama-2-7B.

E. Experiment Details

  • Llama-2-7B experiments use three seeds, but Mistral-7B experiments use one seed (1234), which the authors identify as a study limitation.Experiments run on four A800-80GB GPUs; the authors report no significant result differences across the Llama-2-7B seeds.Source: E. Experiment Details · p. 24

E.1. Llama-2-7B

  • Llama-2-7B SFT uses batch size 30 per device for two epochs, while reward-model training uses batch size 36 for two epochs and reaches 63% evaluation accuracy.Both use a 10^-5 learning rate, cosine scheduling without warm-up, and a maximum total length of 512; experiments start from pretrained Llama-2-7B rather than Vicuna-7B-1.5.Source: §E.1 · p. 25
  • Llama-2-7B RL uses a 10^-6 learning rate, one epoch, and KL penalty 0.1 for PPO and ReMax; DPO uses β = 0.05.The paper says β = 0.05 was best among 0.01, 0.05, and 0.1; trainable models avoid ZeRO-3 and parameter offloading because these make generation and training unacceptably slow.Source: §E.1 · p. 25

E.2. Mistral 7B

  • For Section 5.3, the experiments use Mistral-7B-instruct-v0.2 and UltraRM-13B; the authors’ Llama-2-13B reward model scores about 65% versus UltraRM-13B’s 71% on full-hh-rlhf (Cui et al., 2023).The paper says Mistral-7B was chosen for its pretrained and instruction-tuned benchmark performance, and describes UltraRM-13B as trained on several preference datasets.Source: §E.2 · p. 25
  • For Mistral-7B, prompts longer than 384 tokens are excluded, responses are capped at 384, and the resulting maximum total length is 784 tokens.Training uses a 5 × 10^-7 learning rate, reward clipping of 1.0, temperature 0.7, top-p 0.9, and full-step KL regularization; these hyperparameters were not tuned due to training cost.Source: §E.2 · p. 25

E.3. Details of Figures and Tables

  • ReMax requires 172.12 GB versus PPO’s 319.14 GB (54%); one Llama-2-7B training epoch takes 1.8 hours versus PPO’s 2.9 hours (1.6× slower).The memory estimate counts trainable and non-trainable models; the epoch-time comparison is based on Table 2 results.Source: §E.3 · p. 26
  • Figure 5 reports averages over five random seeds, using 20 evaluation intervals per epoch and 500 test prompts per evaluation iteration.The reward model scores generated responses at each evaluation iteration.Source: §E.3 · p. 26

F.1. Training Instability of REINFORCE

REINFORCE’s stochastic-gradient variance causes instability and poor evaluation rewards, motivating the authors to abandon it in the main experiments for variance-reduced ReMax. Larger-model training avoids divergence but remains unstable and suboptimal.

F.1. Training Instability of REINFORCE

  • Fine-tuning OPT-1.3B with REINFORCE tended to diverge, with rapidly increasing gradient norm and poor evaluation reward.The experiment used the full-hh-rlhf dataset.Source: §F.1 · p. 26
  • With Llama-2-7B, REINFORCE avoided divergence but remained suboptimal in evaluation reward, with variability in gradient norm, evaluation perplexity, and KL divergence.This experiment also used the full-hh-rlhf dataset.Source: §F.1 · p. 26
  • REINFORCE suffers high-variance stochastic gradients and inferior performance to ReMax, so the authors abandon it in their main experiments.They instead explore ReMax, their variance-reduced version of REINFORCE.Source: §F.1 · pp. 26–27, Figure 8 · p. 27

F.2. Ablation Study on KL Regularization · F.3. ReMax with Truncated Baseline Response

ReMax’s KL regularization works with either one-step or full-step KL, while ReMax-fast shortens greedy baseline generation and improves training speed versus PPO. The supplied results also show that KL coefficient tuning matters and that shorter greedy responses can approximate full-response rewards, though overly short responses reduce performance.

F.2. Ablation Study on KL Regularization

  • Both one-step and full-step KL regularization work in ReMax, and full-step KL regularization is more effective.The paper reduces β from 0.1 to 0.01 to achieve comparable performance with full-step KL.Source: §F.2 · p. 27, Figure 9 · p. 27
  • Both one-step and full-step KL regularization are viable, provided β is tuned to the scenario.The paper presents coefficient adjustments for one-step KL in Figure 10 and full-step KL in Figure 11.Source: §F.2 · p. 28, Figure 10 · p. 28, Figure 11 · p. 28

F.3. ReMax with Truncated Baseline Response

  • ReMax-fast is presented as a practical solution for reducing time spent on generation.Its strategy constrains generation length during greedy sampling.Source: §F.3 · p. 28
  • A truncated response’s reward can reliably approximate the complete response’s reward, motivating ReMax-fast’s constrained greedy sampling without compromising unbiasedness or convergence guarantees.The method may cause a small rise in gradient variance.Source: §F.3 · p. 28, Figure 12 · p. 29
  • A greedy response length of 128 yields performance comparable to original ReMax, whereas length 64 significantly sacrifices performance.The normal generation length remains unchanged; Figure 13 uses full-step KL for training stability and does not offload optimizer states.Source: §F.3 · p. 29, Figure 13 · p. 29
  • 2.1 times: ReMax-fast improves training speed versus PPO from 1.6 times to 2.1 times.Halving greedy response length can reduce generation time by a factor of 0.75 in principle, due to transformers’ quadratic self-attention complexity.Source: §F.3 · p. 29
  • Halving greedy response length is a practical way to reduce generation time, with a theoretical reduction factor of 0.75 from quadratic self-attention complexity.The paper reports the speed comparison with PPO separately.Source: §F.3 · p. 29

F.4. Fine-tuning GPT-2 on a Classical NLP Task · G. Samples of Response · G.1. Out-of-distribution Samples

The appendices test ReMax on sentiment-guided GPT-2 generation beyond RLHF and describe its AlpacaEval responses to unseen prompts as content-rich; they also provide sample responses from several methods.

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

  • The 137M-parameter GPT-2 experiment tests reward maximization beyond RLHF by using a sentiment classifier that scores how positive generated text is.The classifier replaces direct learning from human preference data in this experiment.Source: §F.4 · pp. 29–30
  • With fast greedy sampling, ReMax achieves a 2.1 times speed-up over PPO when training Llama-2-7B on A800-80GB GPUs with 33k samples of length 512.The setting matches Table 2.Source: Table 5 · p. 30
  • ReMax reaches comparable GPT-2 IMDB reward with a 2.2 times speed-up and uses about 15GB GPU memory versus PPO’s 37GB, saving about 60%.PPO uses an off-policy training epoch of 4 in this setup, compared with 1 in the RLHF experiments; the authors attribute the acceleration to that larger epoch on the small model.Source: Figure 14 · p. 30, §F.4 · p. 30
  • On IMDB, GPT-2 receives up to 64 tokens of a movie review and generates 48 more to maintain positive sentiment and fluency.The dataset has 25,000 training, 5,000 validation, and 5,000 test examples of reviews labeled positive or negative.Source: §F.4 · p. 30

G. Samples of Response

  • Complete SFT, DPO, PPO, and ReMax responses are available in the authors’ anonymous repository, with samples presented in this section.The passage provides the repository location as an anonymous results repository.Source: G. Samples of Response · p. 30

G.1. Out-of-distribution Samples

  • On AlpacaEval prompts absent from training, ReMax outputs tend to be rich in content, as illustrated by Figures 15–17.The section frames these prompts as an out-of-distribution evaluation and contrasts them with in-distribution samples discussed later.Source: §G.1 · p. 31

G.2. In-distribution Samples

The section compares responses to training-dataset prompts from Mistral-Instruct-v0.2 and its version fine-tuned on 20k prompts, reporting that response quality greatly improved.

G.2. In-distribution Samples

  • The ShareGPT-en in-distribution samples show greatly improved response quality after fine-tuning Mistral-Instruct-v0.2 on 20k training prompts.The samples include responses to ShareGPT-en questions #32 and #83.Source: §G.2 · pp. 31, 35, Figure 18 · p. 35, Figure 19 · p. 35

Related papers

Loading 2310.10505v4…