Research map / Large language models / Alignment and preferences
ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ReMax, a PPO-free RLHF method that adapts REINFORCE with a greedy-decoding reward baseline, cutting memory and training time while reaching SOTA 7B AlpacaEval/MT-Bench scores.
Problem
RLHF for LLMs with PPO is computationally expensive and hyper-parameter heavy because of the trainable value model, and REINFORCE is cheap but has high gradient variance at LLM scale.
Approach
Exploits RLHF properties (fast simulation, deterministic transitions, trajectory-level rewards) to drop the value model and use the reward of a greedily decoded response as a baseline in REINFORCE, giving an unbiased, lower-variance, online policy gradient.
Key results
- Mistral-7B with ReMax (sharegpt-en, 20k prompts) reaches 94.78% AlpacaEval win rate and 7.739 MT-Bench, above Zephyr-7B-beta (90.60%, 7.356).
- ReMax saves about 46% GPU memory versus PPO for a 7B model and runs without optimizer offloading on A800-80GB where PPO cannot.
- With 4 A800 GPUs, ReMax trains about 1.6x faster than PPO per epoch (1.8h vs 2.9h with offloading).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ReMax instead of the SFT baseline | Mistral-7B-Instruct-v0.2 · AlpacaEval | win rate: 92.78% -> 94.78% | Helped |
| Use ReMax instead of the SFT baseline | Mistral-7B-Instruct-v0.2 · MT-Bench | score: 7.516 -> 7.739 | Helped |
| Use ReMax instead of PPO | Llama-2-7B on full-hh-rlhf · AlpacaEval | win rate over SFT: +31.4 points | Helped |
| Initialize ReMax from a DPO-trained model | Llama-2-7B on full-hh-rlhf · AlpacaEval | win rate: 84.7% | Helped |
| Use ReMax instead of PPO with 4 GPUs and optimizer offload | Llama-2-7B · A800-80GB, 4 GPUs, optimizer offload | epoch training time: 2.9h -> 2.0h | Helped |
| Use ReMax instead of PPO with 4 GPUs and optimizer offload | Llama-2-7B · A800-80GB, 4 GPUs, optimizer offload | maximum batch size: 112 -> 152 | Helped |
| Increase ultrafeedback prompt count from 10k to 40k (ablation) | Mistral-7B ReMax · AlpacaEval | win rate: 94.29% -> 93.11% | Hurt |
| Increase ultrafeedback prompt count from 10k to 40k (ablation) | Mistral-7B ReMax · MT-Bench | score: 7.578 -> 7.538 | Hurt |
| Increase lmsys-chat-1m prompt count from 10k to 40k (ablation) | Mistral-7B ReMax · AlpacaEval | win rate: 94.40% -> 92.86% | Hurt |
| Increase sharegpt-en prompt count from 20k to 40k (ablation) | Mistral-7B ReMax · MT-Bench | score: 7.739 -> 7.534 | Hurt |
| Subtract the greedy-response reward as a baseline (ablation) | ReMax versus REINFORCE · OPT-1.3B fine-tuning | gradient norm | Helped |
| Subtract the greedy-response reward as a baseline (ablation) | ReMax versus REINFORCE · OPT-1.3B fine-tuning | evaluation reward | Helped |
Surprising findings
- Over-optimization: performance drops when Mistral training data grows to 40k prompts, indicating need for regularization.
- Theory shows the greedy baseline may increase variance when the optimal action already dominates, a drawback argued to be acceptable in RLHF.
Limitations
- Online greedy-baseline sampling adds generation cost, making it about 1.3x slower than DPO per epoch.
- Theory shows variance reduction only under conditions (e.g., optimal action not dominant); may increase variance in the worst case.
- Global convergence to optimum not established; only convergence to stationary points.
Experimental setup
Datasets: full-hh-rlhf ultrafeedback lmsys-chat-1m sharegpt-en
Benchmarks: AlpacaEval MT-Bench full-hh-rlhf evaluation reward
Metrics: win rate MT-Bench score evaluation reward GPU memory training time per epoch gradient norm
Baselines: PPO DPO SFT Llama-2-Chat Zephyr-7B-beta Mistral-7B-Instruct-v0.2 GPT-3.5-turbo GPT-4-turbo
Builds on: REINFORCE PPO REINFORCE with baseline DPO
Closest papers by meaning
- DPO Meets PPO: Reinforced Token Optimization for RLHF (ICML 2025 · similarity 0.90)
- Avoiding exp(R) scaling in RLHF through Preference-based Exploration (NeurIPS 2025 · similarity 0.90)
- Zeroth-Order Policy Gradient for Reinforcement Learning from Human Feedback without Reward Inference (ICLR 2025 · similarity 0.90)
- MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actions (ICLR 2025 · similarity 0.90)
- Leftover Lunch: Advantage-based Offline Reinforcement Learning for Language Models (ICLR 2024 · similarity 0.90)
- Towards Efficient Exact Optimization of Language Model Alignment (ICML 2024 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.