ReMax Explained: REINFORCE with a Greedy Baseline

Q: What is ReMax?

A REINFORCE-style RLHF algorithm from Li et al. (2023). For each prompt it samples one response and greedily decodes another, and uses the reward difference between them as the advantage for the sampled response. It trains no value network.

Q: Is the ReMax gradient unbiased?

Yes. The greedy response depends on the prompt and current parameters but not on the sampled response, so its reward is a valid baseline whose expected contribution to the gradient is zero.

Q: How much memory does ReMax save compared with PPO?

The paper reports about 46% less GPU memory than PPO for a 7B model, because there is no value model with its gradients and optimiser state. Our back-of-envelope estimate for a 7B policy and 7B reward model is 126 to 140 GB of model state versus about 252 GB for PPO.

Q: ReMax or RLOO?

Both remove the critic. At an equal budget of two generations per prompt, our toy run found RLOO with k = 2 had about 40% lower gradient variance than ReMax and slightly lower regret, because RLOO uses both generations as training samples. It is worth testing both on your task.

Q: Does ReMax work with binary rewards?

It works but wastes many updates: with a 0/1 reward the greedy baseline is 0 or 1, so when the greedy answer is right only failures get gradient and vice versa. Group methods such as GRPO, DAPO or RLOO with larger k are usually a better fit.

ReMax is a REINFORCE-style RL algorithm for language models that uses the reward of the model's own greedy response as the baseline. For each prompt it generates one response by sampling and one by greedy decoding, scores both with the reward model, and pushes up the sampled response's log-probability if it scored better than the greedy one and down if it scored worse. There is no value network and no learned baseline. It sits in the same slot as PPO and RLOO: after SFT and reward modeling, optimising the policy online against a reward. Its pitch is minimalism: PPO-level results with a fraction of PPO's memory and hyperparameters.

Where it came from#

ReMax was introduced by Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun and Zhi-Quan Luo in October 2023 (ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models). The paper starts from an observation about what makes RLHF different from the RL problems PPO was designed for. It identifies three properties: fast simulation (generating a full trajectory is just sampling from the model, cheap compared with a robot or a game engine), deterministic transitions (the next state is the old state plus the chosen token, with no environment noise), and trajectory-level rewards (the reward model scores the whole response once). A value network in PPO exists to estimate expected future reward from intermediate states under stochastic dynamics; with deterministic transitions and a single terminal reward, most of that machinery is unnecessary.

The greedy baseline itself is older. Rennie et al. (2017) used it for image captioning as self-critical sequence training: the model is "critical" of itself by comparing each sampled caption with its own test-time (greedy) caption. ReMax brings that idea to RLHF and analyses its variance. The paper reports that ReMax saves about 46% of GPU memory relative to PPO when training a 7B model, allows training on 80 GB GPUs without the offloading PPO needed, and, applied to Mistral-7B, reached a 94.78% win rate on the AlpacaEval leaderboard and 7.739 on MT-bench.

The objective and the estimator#

Objective

Let $π_{θ}$ be the policy, $x$ a prompt, $y = (y_{1}, \dots, y_{T})$ a response and $r (x, y)$ the reward of the full response (a reward model score, optionally minus a KL penalty to a reference model). ReMax maximises the expected reward

J (θ) = E_{x \sim D} E_{y \sim π_{θ} (\cdot ∣ x)} [r (x, y)],

whose gradient by the score-function identity is $E [r (x, y) \sum_{t} \nabla_{θ} lo g π_{θ} (y_{t} ∣ x, y_{< t})]$ .

The ReMax gradient

Define the greedy response $\overset{y}{ˉ} (x; θ)$ by $\overset{y}{ˉ}_{t} = ar g max_{v} π_{θ} (v ∣ x, \overset{y}{ˉ}_{< t})$ at each step. ReMax's estimator for one prompt is

\overset{g}{^}_{ReMax} = (r (x, y) - r (x, \overset{y}{ˉ} (x; θ))) t = 1 \sum T \nabla_{θ} lo g π_{θ} (y_{t} ∣ x, y_{< t}), y \sim π_{θ} (\cdot ∣ x) .

No gradient flows through $\overset{y}{ˉ}$ or its reward; it is a number computed with the current parameters.

Why it is unbiased

The greedy response depends on $x$ and $θ$ but not on the sampled $y$ . So $b (x) = r (x, \overset{y}{ˉ} (x; θ))$ is a legitimate baseline, and

E_{y \sim π_{θ}} [b (x) \nabla_{θ} lo g π_{θ} (y ∣ x)] = b (x) \nabla_{θ} y \sum π_{θ} (y ∣ x) = 0.

The estimator therefore has the same expectation as plain REINFORCE. Its variance is driven by $E [(r (x, y) - b (x))^{2} ∥ \nabla_{θ} lo g π_{θ} (y ∣ x) ∥^{2}]$ , which is small when the greedy reward is close to the rewards of typical samples. For the variance-minimising constant baseline you would want $b (x) \approx E_{y} [r (x, y)]$ ; ReMax bets that the greedy response is a good one-sample proxy for that expectation, which is often true for SFT models whose sampling distribution is concentrated around their greedy output.

When is the greedy reward a good baseline?

Ignore the gradient-norm weighting for a moment and look at the advantage's second moment, which drives the variance. For a fixed prompt write $μ = E_{y} [r (x, y)]$ and $σ^{2} = Var_{y} [r (x, y)]$ . For any baseline $b$ that is independent of $y$ ,

E_{y} [(r (x, y) - b)^{2}] = σ^{2} + E [(μ - b)^{2}] .

The first term is irreducible. The second is the baseline's error. ReMax's baseline is deterministic, so its error is $(μ - r (x, \overset{y}{ˉ}))^{2}$ : small when the greedy response is typical of what the policy samples, large when it is an outlier. A leave-one-out baseline from one other sample (RLOO with $k = 2$ ) has error $σ^{2}$ on average, so its single-term second moment is $2 σ^{2}$ , which looks worse; but RLOO gets two gradient terms from its two generations, and their advantages are exact negatives of each other, which cancels much of the noise when they are averaged. In our toy below, the mean squared gap between greedy reward and expected reward was 0.51 against a per-sample reward variance of 0.80. The greedy baseline was decent, better than a random single sample as an estimate of $μ$ , yet ReMax still ended up with higher gradient variance than RLOO at equal cost, because RLOO's second gradient term is worth more than ReMax's better baseline.

The practical reading: ReMax's baseline is good when the policy is concentrated and its greedy output is representative, which describes many SFT models on short-form tasks. It is weaker on open-ended prompts with many valid answers, where the greedy path is just one of many modes and can be unrepresentative (often blander, sometimes repetitive).

Side by side

Method	Baseline for sample $y_{i}$	Generations per gradient term	Trained networks
REINFORCE	none (or a moving average)	1	policy
ReMax	$r (x, \overset{y}{ˉ})$ , the greedy response's reward	2 (1 sample + 1 greedy)	policy
RLOO	$\frac{1}{k - 1} \sum_{j \neq = i} r (x, y_{j})$	1, in groups of $k$	policy
GRPO	group mean, scaled by group std	1, in groups of $G$	policy
PPO	learned $V_{ψ} (s_{t})$ per token	1	policy and value model

Method

Baseline for sample

y_{i}

Generations per gradient term

Trained networks

REINFORCE

none (or a moving average)

policy

ReMax

r (x, \overset{y}{ˉ})

, the greedy response's reward

2 (1 sample + 1 greedy)

policy

RLOO

\frac{1}{k - 1} \sum_{j \neq = i} r (x, y_{j})

1, in groups of

k

policy

GRPO

group mean, scaled by group std

1, in groups of

G

policy

PPO

learned

V_{ψ} (s_{t})

per token

policy and value model

ReMax is the $k = 1$ member of the family with a special extra sample. RLOO with $k = 2$ uses the same two generations per prompt but treats both as training samples, each serving as the other's baseline. Our toy run below compares the two at equal generation budget.

Worked example by hand#

A prompt asks for a two-sentence product description. Greedy decoding produces a serviceable answer that the reward model scores at 0.60. A sampled response that adds a concrete detail scores 0.82, so its advantage is $0.82 - 0.60 = + 0.22$ : every token in it gets its log-probability nudged up with weight 0.22. On the next iteration another sample comes out vague and scores 0.35; its advantage is $0.35 - 0.60 = - 0.25$ , and its tokens are pushed down. Without the baseline, both samples would have been pushed up, by 0.82 and 0.35 respectively, and the policy would learn only from the difference between those two pushes across many updates. Note what happens if the sample happens to equal the greedy response, which is common for short answers from a confident model: the advantage is exactly zero and the prompt contributes nothing that step.

Figure 1. ReMax per prompt: one sampled response and one greedy response, both scored by the reward model. The reward difference is the advantage of the sampled response.

The algorithm#

inputs: SFT policy π_sft, reward model r_φ, prompts D, (optional) β and π_ref
π_θ ← π_sft
repeat:
    sample a batch of B prompts
    for each x (batched in practice):
        y    ~ π_θ(·|x)                  # temperature 1.0, keep token log-probs
        ȳ    = greedy_decode(π_θ, x)     # no grad
        R    = r_φ(x, y)    - β Σ_t log(π_θ(y_t|·)/π_ref(y_t|·))   # KL term optional
        R̄    = r_φ(x, ȳ)    - β Σ_t log(π_θ(ȳ_t|·)/π_ref(ȳ_t|·))
        A    = (R - R̄)                    # detached scalar
    loss = - mean_x( A · Σ_t log π_θ(y_t | x, y_


ReMax as written optimises the trajectory-level reward; in practice you should fold a KL penalty to the SFT reference into that reward, as RLOO does, or the policy will drift freely toward whatever the reward model over-scores. Apply the same penalty to the greedy response so the baseline and the sample are measured on the same scale.
Data. Prompts only, 10k to a few hundred thousand, from the distribution you care about. Each step produces  $2 B$  generations, of which  $B$  carry gradient. Batches of 64 to 512 prompts are typical.
What the loop costs. Greedy decoding is no cheaper per token than sampling (it is the same forward pass with a different selection rule), so ReMax pays two generations per gradient term. The training step is a single forward and backward pass through the policy on the sampled responses, plus a reference forward pass if you use KL. There is no critic, no GAE and no clipping, which is where the memory and hyperparameter savings come from.

Variations worth knowing
Several samples, one greedy baseline. Draw  $k$  samples per prompt and subtract the same greedy reward from each. That costs  $k + 1$  generations for  $k$  gradient terms, so the baseline's overhead shrinks as  $k$  grows. At that point you can also compare against the leave-one-out mean of the other samples, or average the two baselines; both remain unbiased because neither depends on the sample being scored.
Batching the two decodes. Sampled and greedy generation can share one batched call if your inference engine supports per-sequence sampling parameters, which avoids a second pass over the prompts' prefill. The reward model can then score all  $2 B$  responses in a single forward pass.
Reusing evaluation decodes. Many pipelines already produce greedy outputs on the training prompts for monitoring. Those decodes are the baseline; ReMax turns a logging cost into a variance-reduction step.


Hyperparameters that matter#
Knob Starting value Notes

Learning rate 5e-7 to 2e-6 (full), around 1e-5 (LoRA) The main stability lever, since nothing clips the update.
KL coefficient  $β$ 0.01 to 0.1 Same role as in PPO. Tune against judge win rate versus KL.
Prompts per batch 64 to 512 Each prompt yields one gradient term, so ReMax needs larger batches than RLOO for the same smoothness.
Sampling temperature 1.0 Lower temperatures make samples coincide with the greedy response more often (zero advantage) and bias the gradient.
Max new tokens Fit the task; apply to both decodes A truncated greedy baseline scores poorly and inflates every sample's advantage.
Gradient clipping (norm) 1.0 A cheap safety net against occasional large-advantage batches.

The list is short by design. The paper reports that ReMax eliminates more than four of PPO's hyperparameters; the ones that disappear are those tied to the critic and the trust region, such as the value-model learning rate, GAE  $λ$ , value clipping and the policy clip range. In exchange, the learning rate carries more of the burden of stability, so sweep it first on a short run (three values spaced by 2×) while watching KL.


Compute and memory#
ReMax trains one network. With roughly 16 bytes per trained parameter for mixed-precision Adam and 2 bytes per frozen parameter,
 $M_{ReMax} \approx 16 P + 2 P_{r} (+ 2 P for a KL reference), M_{PPO} \approx 16 P + 16 P_{v} + 2 P_{r} + 2 P .$ 
For a 7B policy and a 7B reward model that is about 126 GB (140 GB with the reference) against about 252 GB for PPO with a 7B critic, a saving of 44% to 50% that lines up with the paper's reported 46%. Per gradient term ReMax generates twice as many tokens as PPO. Per prompt it generates the same as RLOO with  $k = 2$ . Since generation usually dominates RLHF wall-clock, ReMax's speed advantage over PPO comes from dropping the critic's training passes and from fitting on fewer GPUs without offloading, not from cheaper rollouts.


A toy experiment: ReMax against RLOO at equal budget#
We built 200 "prompts", each a softmax policy over 16 responses with logits drawn from  $N (0, 1. 5^{2})$ . Rewards are  $μ_{p} + ε$ , with a per-prompt offset  $μ_{p} \sim N (0, 4)$  standing in for prompt difficulty and reward-model scale differences across prompts, and  $ε \sim N (0, 1)$  per response. Every method got exactly two generations per prompt per step. We measured per-prompt gradient variance at initialisation over 400 repeated draws, then trained each method for 150 steps (learning rate 1.0 on the logits, 5 seeds) and tracked regret: the best response's reward minus the policy's expected reward.
Estimator (2 generations per prompt) Gradient variance at init Regret at step 0 Step 50 Step 150

REINFORCE, 2 samples, no baseline 2.037 1.753 0.675 0.589
REINFORCE, 2 samples, global mean baseline 2.025 1.753 0.685 0.594
ReMax: 1 sample + greedy baseline 1.183 1.753 0.557 0.471
RLOO,  $k = 2$ 0.682 1.753 0.524 0.439

The global baseline does almost nothing (2.025 vs 2.037) because the per-prompt offsets dominate: a baseline has to be per-prompt to help, and both ReMax and RLOO are. ReMax cut variance by 42% relative to REINFORCE and reached lower regret throughout training. On this toy, though, RLOO with  $k = 2$  was better still at the same generation budget: variance 0.68 versus 1.18, and regret 0.439 versus 0.471 at step 150. The reason is visible in the estimator. RLOO turns both generations into gradient terms; ReMax spends one generation purely on the baseline.


Figure 2. Per-prompt gradient variance at initialisation (200 prompts, 400 draws), every method using two generations per prompt. A global baseline barely helps when prompts differ in reward level.

We also checked whether ReMax catches up as the policy becomes more peaked, since a greedy response is a better proxy for the mean of a concentrated distribution. Scaling the logits so the average top-response probability went from 0.14 to 0.69, the variance ratio of ReMax to RLOO fell only from 1.93 to 1.69. Both estimators improve as the policy concentrates, and RLOO stayed ahead at every level we tried. The toy has no notion of response length or of the reward model's errors, so it does not capture everything; but it does say that if you can afford two generations per prompt, it is worth testing RLOO with  $k = 2$  before settling on ReMax. Where ReMax keeps a practical edge is in implementations that already produce a greedy decode for evaluation or logging, and in its deterministic, easy-to-inspect baseline.


Figure 3. Mean regret (best reward minus expected reward) over 150 steps, 5 seeds. ReMax beats plain REINFORCE; RLOO with k = 2 at the same generation budget does slightly better.



Failure modes and how to detect them#
Sample equals greedy. For short or highly constrained outputs a confident policy often samples exactly its greedy response, giving zero advantage. Log the fraction of prompts where  $y = \overset{y}{ˉ}$  (or where  $∣ R - \overset{ˉ}{R} ∣$  is tiny). Above about 30%, those prompts are wasted generations; use temperature 1.0, filter such prompts or switch to RLOO with a larger group.
Unlucky greedy baselines. If the greedy response is unusually bad, every sample beats it and the update pushes all sampled responses up; if it is unusually good, everything is pushed down. In expectation this is fine, but it makes per-prompt gradients noisy. A histogram of advantages that is strongly skewed per prompt across steps is the sign.
Greedy degeneration. Greedy decoding is prone to repetition loops in some models. A looping baseline scores low, inflating sample advantages and teaching nothing useful. Check greedy outputs for repetition and length hitting the cap; apply the same truncation penalty you use for samples.
KL drift and reward hacking. With nothing clipping updates, too high a learning rate or too low a  $β$  lets KL climb quickly. As with any RLHF method, track KL to the reference and an independent judge's win rate, and stop when win rate flattens while reward still climbs.
Reward scale changes. ReMax's update size is proportional to raw reward differences. Swap in a reward model with twice the output scale and the effective learning rate doubles. Either normalise advantages by a running standard deviation of  $R - \overset{ˉ}{R}$  across batches, or re-tune the learning rate every time the reward model changes; log the mean absolute advantage so a scale change is visible immediately rather than discovered as a KL spike.
Entropy collapse. As the policy sharpens, more samples equal the greedy decode and learning slows on its own. Track token entropy; a steady fall alongside a rising share of zero-advantage prompts means the run is approaching a plateau.


When to use ReMax#
Situation Recommendation

Replacing PPO under a memory limit, dense reward model scores ReMax or RLOO; test RLOO  $k = 2$  at the same budget
Want the simplest possible RLHF loop with few hyperparameters ReMax
Binary verifiable rewards (pass/fail) GRPO, DAPO or RLOO with large  $k$ ; a single greedy baseline is 0 or 1 and wastes many updates
Long responses where token-level credit matters PPO with a good critic, or a process reward model
No reward model, fixed preference pairs DPO
Want online improvement without policy gradients Rejection sampling fine-tuning



Minimal implementation in PyTorch#
The sketch uses Hugging Face Transformers' generate for both decodes and a sequence-classification reward model. Batching, padding and KL are kept minimal for readability.
import torch

@torch.no_grad()
def decode(policy, tok, prompts, greedy, max_new=256):
    enc = tok(prompts, return_tensors="pt", padding=True).to(policy.device)
    out = policy.generate(**enc, do_sample=not greedy, temperature=1.0, top_k=0, top_p=1.0,
                          max_new_tokens=max_new, pad_token_id=tok.pad_token_id)
    return enc, out                       # out = prompt + response ids (left-padded prompts)

@torch.no_grad()
def reward(rm, rm_tok, prompts, texts):
    enc = rm_tok([p + t for p, t in zip(prompts, texts)], return_tensors="pt",
                 padding=True, truncation=True).to(rm.device)
    return rm(**enc).logits.squeeze(-1).float()

def response_logprob(policy, out, prompt_len, pad_id):
    logits = policy(input_ids=out, attention_mask=(out != pad_id)).logits[:, :-1].float()
    lp = torch.log_softmax(logits, -1).gather(-1, out[:, 1:].unsqueeze(-1)).squeeze(-1)
    mask = (out[:, 1:] != pad_id).float(); mask[:, : prompt_len - 1] = 0   # response tokens only
    return (lp * mask).sum(-1)

def remax_step(policy, tok, rm, rm_tok, prompts, opt):
    tok.padding_side = "left"
    enc, y_s = decode(policy, tok, prompts, greedy=False)
    _,   y_g = decode(policy, tok, prompts, greedy=True)
    L = enc["input_ids"].shape[1]
    txt_s = tok.batch_decode(y_s[:, L:], skip_special_tokens=True)
    txt_g = tok.batch_decode(y_g[:, L:], skip_special_tokens=True)
    adv = reward(rm, rm_tok, prompts, txt_s) - reward(rm, rm_tok, prompts, txt_g)   # [B]
    logp = response_logprob(policy, y_s, L, tok.pad_token_id)
    loss = -(adv.to(logp.device) * logp).mean()
    loss.backward(); torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
    opt.step(); opt.zero_grad()
    return {"adv_mean": adv.mean().item(), "adv_abs": adv.abs().mean().item(),
            "same_as_greedy": (adv.abs() < 1e-6).float().mean().item()}
Add the KL term by computing reference log-probabilities for both decodes and subtracting  $β$  times the log-ratio sum from each reward before taking the difference. In POFE, ReMax runs as an RL recipe beside RLOO and PPO with identical evaluation panels, which makes the equal-budget comparison above a one-line configuration change on your own prompts.


Further reading#
Li et al. (2023) for ReMax, its variance analysis and the memory measurements. Rennie et al. (2017) for self-critical sequence training, the origin of the greedy baseline. Williams (1992) for REINFORCE. Ahmadian et al. (2024) for RLOO and the broader case for REINFORCE-style RLHF. Schulman et al. (2017) for the PPO baseline ReMax was designed to replace.

      Frequently asked questions#
What is ReMax?
A REINFORCE-style RLHF algorithm from Li et al. (2023). For each prompt it samples one response and greedily decodes another, and uses the reward difference between them as the advantage for the sampled response. It trains no value network.
Is the ReMax gradient unbiased?
Yes. The greedy response depends on the prompt and current parameters but not on the sampled response, so its reward is a valid baseline whose expected contribution to the gradient is zero.
How much memory does ReMax save compared with PPO?
The paper reports about 46% less GPU memory than PPO for a 7B model, because there is no value model with its gradients and optimiser state. Our back-of-envelope estimate for a 7B policy and 7B reward model is 126 to 140 GB of model state versus about 252 GB for PPO.
ReMax or RLOO?
Both remove the critic. At an equal budget of two generations per prompt, our toy run found RLOO with k = 2 had about 40% lower gradient variance than ReMax and slightly lower regret, because RLOO uses both generations as training samples. It is worth testing both on your task.
Does ReMax work with binary rewards?
It works but wastes many updates: with a 0/1 reward the greedy baseline is 0 or 1, so when the greedy answer is right only failures get gradient and vice versa. Group methods such as GRPO, DAPO or RLOO with larger k are usually a better fit.
      
      References#
Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, Zhi-Quan Luo (2023). ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models. arxiv.org/abs/2310.10505
Steven J. Rennie, Etienne Marcheret, Youssef Mroueh, Jarret Ross, Vaibhava Goel (2017). Self-critical Sequence Training for Image Captioning. arxiv.org/abs/1612.00563
Ronald J. Williams (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. doi.org/10.1007/BF00992696
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, Sara Hooker (2024). Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. arxiv.org/abs/2402.14740
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, Oleg Klimov (2017). Proximal Policy Optimization Algorithms. arxiv.org/abs/1707.06347
      Related reading
REINFORCE and RLOOAlgorithms 101RLHF with PPOAlgorithms 101GRPO (Group Relative Policy Optimization) explainedAlgorithms 101Reward modeling with Bradley-TerryAlgorithms 101Post-training algorithms 101Hub

Estimator (2 generations per prompt)	Gradient variance at init	Regret at step 0	Step 50	Step 150
REINFORCE, 2 samples, no baseline	2.037	1.753	0.675	0.589
REINFORCE, 2 samples, global mean baseline	2.025	1.753	0.685	0.594
ReMax: 1 sample + greedy baseline	1.183	1.753	0.557	0.471
RLOO, $k = 2$	0.682	1.753	0.524	0.439

Where it came from#

The objective and the estimator#

Objective

The ReMax gradient

Why it is unbiased

When is the greedy reward a good baseline?

Side by side

Worked example by hand#

The algorithm#

Variations worth knowing

Hyperparameters that matter#

Compute and memory#

A toy experiment: ReMax against RLOO at equal budget#

Failure modes and how to detect them#

When to use ReMax#

Minimal implementation in PyTorch#

Further reading#

Frequently asked questions#

References#

Related reading

Bring a model and a metric.