Reward Optimization Needs a Baseline

A reward of 1 for a correct answer and 0 for everything else leaves most gradients at zero. If the policy almost never lands a correct answer, the run barely moves. The repair is a baseline: sample the same prompt several times and center the rewards on the group mean. That is what turns an RL update into something you can read.

The goal of RL for a language model is small. Make the model pick answers that score well. The rewards are sparse and the variance is high. Inference is the expensive part.

Sparse rewards leave most gradients at zero

The state S is the prompt plus the response so far. Each generated token joins the state. An action picks the next token. With outcome rewards, treat the whole response as one action A.

The reward scores that whole response instead of each step inside it. Make it verifiable and deterministic. A program parses the final answer and compares it against a ground-truth key. Correct scores 1, and everything else scores 0.

Transitions here are known and dull. Each one appends the token the model chose. Robotics takes its transitions from the world, and they can be unknown.

Many robot states are physically impossible. A language model can reach every token sequence, so the hard part is finding the ones that earn a reward. The same freedom lets the model build scratchpads and chains of thought on the way there.

The policy π is a distribution over next tokens given the current state. Most runs start from a pretrained model and fine-tune it. A rollout starts from a prompt, samples a full response from π, and scores it once.

The objective is the expected reward.

S comes from the prompts, and you sample A from πθ. The gradient of that objective has one term.

One stochastic step follows directly.

• Sample a prompt S.
• Sample a response A from πθ.
• Compute the reward R(S,A).
• Update with R(S,A) times the gradient of log πθ(A|S).

The shape matches supervised learning with two changes. The label A comes from the model itself. The reward scales the update, so a better response pushes harder.

With a 0/1 reward, every wrong sample multiplies its gradient by zero. Only correct samples move the parameters.

A learned reward model changes the shape of the problem. RLHF trains one on preference data, and it returns a real-valued score for any response. The signal is smoother, so the same algorithm behaves differently and needs different tuning.

A baseline cuts the variance and leaves the expectation alone

Naive policy gradient is unbiased and noisy. You want the same expected gradient with less noise. Subtract a baseline B that depends only on the state.

For a fixed S, the baseline term sums to zero across actions. The expectation is untouched and only the variance moves.

Two states and two actions show the effect.

S1,A1 → 11, S1,A2 → 9
S2,A1 → 0,  S2,A2 → 2

Draw the single sample (S1,A2) with reward 9 and the update pushes toward A2. A1 is the better action in that state. The push is noise, and its size comes straight from the raw reward.

Now set B(S1) = 10 and B(S2) = 1.

S1,A1: 11-10 = 1,  S1,A2: 9-10 = -1
S2,A1: 0-1 = -1,   S2,A2: 2-1 = 1

The better action in each state carries +1 and the worse one carries -1. The effective rewards are smaller and centered, and the updates get less noisy.

Each baseline here sits at the average reward for its own state. That is the common choice, and it is the one that strips the most noise.

Name that quantity and the standard identities follow.

With outcome rewards on a full response, Q matches R for the response you sampled. R - V(S) is then an advantage estimate, which is what every method in this family builds.

Most of the family fits one form.

Δ is what changes between methods.

• Naive: Δ = R.
• Baseline: Δ = R - B(S).
• Advantage-based: Δ approximates A(S,A).
• GRPO-style: Δ is centered and normalized inside a group, usually with clipped ratios and an optional KL term.

GRPO gets its baseline by sampling the same prompt again

Estimating V(S) usually costs a second model. A language model offers a cheaper route, because you can reset to the same prompt and sample many answers. GRPO is a PPO-style method that builds its baseline out of that.

For one prompt, collect the rewards R1 through RK. Scoring a whole batch gives a reward matrix shaped [batch, num_samples]. The baseline is the mean inside each prompt.

Center each reward on it.

A response above its group average gets a positive Δ and a push up. A response below it gets a negative Δ and a push down. The group grades itself.

Divide by the group standard deviation to fix the scale.

The note on verifiable rewards objected to this divisor. The case for keeping it is narrow. The update size stops tracking the reward scale, which helps stability. In small runs the centered and normalized versions land in about the same place.

The group baseline brings its own failure mode. If every sample for a prompt earns the same reward, every Δ is zero and that prompt gives no gradient. A sparse reward and a weak policy produce exactly that.

A denser reward buys speed and sells loopholes

RL optimizes what you can measure, so the reward function decides what the model becomes. The toy that exposes this is sorting. The prompt is a fixed-length list of n numbers, and the model must return the same numbers in sorted order. Prompt length and response length are fixed.

The first reward is exact match. Score 1 when the output equals the sorted sequence and 0 otherwise. It matches the goal exactly, and a random policy almost always scores 0, so the run never starts.

The second reward counts positions that match the sorted truth.

truth: 0 1 2 3
0 1 2 3 → 4
2 1 3 0 → 1
1 0 2 3 → 1

That reward produces a gradient from almost any output. It also hands the same score to two wrong answers with nothing in common. The model learns that one wrong answer is worth as much as the other.

The third reward adds inclusion and adjacency. Inclusion counts how many prompt tokens appear anywhere in the response. Adjacency counts how many neighboring pairs sit in sorted order. The signal gets denser, and some patterns score well without sorting anything.

Denser rewards keep the run moving, and they pull the model toward easy wrong strategies that pay. Reward design sits between those two failures and stays fragile in both directions.

One patch keeps Δ only for the top-scoring responses in each group and zeroes the rest. It slows the drift into mediocre partial-credit modes. It also changes the signal you optimize, so measure it before and after you turn it on.

The ratio collapses to one unless you detach the old policy

In code the objective becomes a loss.

Here log πθ(A|S) is the sum of the per-token log probabilities for that response. Run the model for logits, take log_softmax, then gather the log probability at each sampled index. The result is shaped [batch, num_samples, positions].

GRPO takes several gradient steps on the same sampled responses, and reuse is how you amortize the sampling cost. The policy that produced those responses goes stale in the meantime. PPO-style methods correct for that with a ratio against a snapshot.

Clipping keeps a reused batch from moving the policy too far.

• Compute r for each response.
• Clip it to the band from 1-ε to 1+ε.
• Take the minimum of rΔ and clipped_rΔ, then flip the sign for a minimization loss.

A KL term pulls the policy back toward a reference model.

KL is the expectation of log(P/Q) under P. The direct estimate of it is noisy, so many implementations use an unbiased form with lower variance.

Compute it per token, then average over the batch, the samples, and the positions.

The loop that holds all of this is short.

• Sample prompts, then sample several responses per prompt from the current policy or a frozen snapshot.
• Score every response, then compute Δ with the raw, centered, or normalized scheme.
• Compute logp_old for those responses, store it, and detach it.
• Compute log probs under π_ref when you use the KL term.
• Take several gradient steps on the same responses without resampling.
• Inside a step, recompute logp_current, form the ratios, clip them, and apply Δ.
• Add the KL term when you use it, then backprop and update θ.
• Refresh π_old to the current policy for the next batch, and update π_ref less often when it must move slowly.

Every update changes πθ, so the data distribution moves under you for the whole run. The loss curve is a weaker signal here than in supervised training. Read the reward on held-out prompts instead.

Outcome rewards assign credit bluntly

One reward covers the whole response, so the same Δ multiplies the log probability at every position. A response that nails one pivotal token and then drifts gets the same push on every token. Nothing in the method knows which token did the work.

Process rewards score the intermediate steps and can fix this. They are hard to design for language reasoning. Until you have one, every token in a rewarded response gets credit by association.

The cost is inference and model copies

Inference dominates the bill. Every prompt needs K full generations before one gradient step, and the policy must produce all of them.

The system also holds more than one model.

• The policy you update.
• A reward system, which is a verifiable checker or a trained reward model.
• An old snapshot for the ratios, or the stored old log probabilities in its place.
• A reference policy for the KL term.
• A critic or value model, in the methods that use one.

Full copies cost memory, which is why storing the old log probabilities is worth the bookkeeping. You also need distributed sampling, distributed reward computation, model synchronization across workers, and careful variance control.

That combination makes an RL run harder to operate than supervised pretraining. This unit prices a run before it starts. Price this one in sampled tokens, not in parameters.

The Builder Test

Pick one prompt with a checkable answer. Sample K responses from your current policy and score them all.

Compute Δ twice. Once with the raw rewards, once with the rewards centered on the group mean. Count how many responses carry a nonzero Δ in each version.

With a 0/1 reward, the raw version leaves every wrong response at zero. The centered version gives the wrong ones a negative push, and the count of usable deltas jumps.

Then run the same check on a prompt the policy already solves. Every sample comes back correct, the group mean sits at 1, and every centered Δ is zero. That decides which prompts are worth sampling. A group that agrees teaches nothing.

Name the proxy before you trust any improvement. Write down three things.

• What the reward pays for.
• What the reward ignores.
• The pattern that scores well without solving the task.

A model learns the shape of a helpful answer or the length a reward prefers without getting more reliable. If the proxy is incomplete, stronger optimization makes the model better at the proxy and worse at the goal. The evaluation note asked the same of a benchmark score. A reward is a benchmark you optimize against on purpose.

What Carries

One habit carries out of this unit. Name the binding constraint before you reach for a technique, then price it.

Here the constraint was variance under a sparse reward. The price was K samples of every prompt, paid in inference before one gradient step moved.

Each stage priced the one after it, from the tokenizer forward. The reward is the last stage in that chain, and it decides the behavior you get.

Every number here answers three questions: what produced it, what it cost, and how you catch it lying. Before your next run, write down what one gradient step costs you in sampled tokens. If you cannot write that number, do not start the run.