Section 14

The baseline trick

A free fix for REINFORCE's noisy updates

REINFORCE works in principle, but in practice its updates are so noisy that learning crawls. The previous chapter ended on the reason: the gradient is multiplied by the raw reward, so when every response scores, say, $+9$ , the algorithm enthusiastically pushes up the probability of everything it just did, learning almost nothing about which responses were actually better. A $+9$ tells you nothing if $+9$ is just what this policy normally earns.

The fix is one subtraction. Instead of asking “was this response good?”, ask “was this response better than my typical response?” Grade on a curve. That reframing, from absolute score to relative score, is the single idea of this chapter.

See the problem, then the fix

Before any algebra, look at what the subtraction does to a real batch of updates. The top strip below is REINFORCE as chapter 13 left it: every reward is positive, so every sampled response gets a hefty upward push, and the update is dominated by which responses happened to get sampled rather than which were good. The bottom strip is the same batch with the batch’s mean reward subtracted.

The baseline, before and after

One batch of rollouts from a decent policy. Each dot is one sampled response's update weight: raw reward r on top, r − mean below.

batch mean 7.16

raw reward r — baseline OFFavg push size 7.16 · pushed up 24/24

r − mean — baseline ONavg push size 0.92 · pushed up 12/24

Without a baseline, every response gets a strong upward shove — the amber dots all sit far right of zero, so the policy reinforces everything it just did, and the noise in which responses got sampled dominates the update. Subtract the batch mean and the same batch turns into small, signed nudges: above-average responses up, below-average down, typical ones barely touched. No information was lost — the dots kept their left-to-right order — but the average push size drops several-fold. Draw a few batches: the raw strip jumps around; the centered strip stays put.

Nothing about the batch changed: same responses, same rewards, same ordering. But the pushes shrank several-fold and became signed: above-average responses up, below-average down, typical ones left alone. The noise collapsed; the signal survived.

Why the subtraction doesn’t bias the gradient

Formally, we subtract a value $b$ , called the baseline, from the reward inside the policy gradient. The one requirement is that $b$ must not depend on the sampled response $y$ itself:

\nabla_\theta J = \mathbb{E}_{y \sim \pi_\theta}\big[\, (r(y) - b) \, \nabla_\theta \log \pi_\theta(y) \,\big]

Read aloud: the gradient of the objective is still “sample responses, and for each one push up its log-probability”, but the strength of the push is now reward minus typical reward instead of the raw reward. A $+9$ in a sea of $+9$ s contributes nothing; a $+9$ among $+2$ s gets a strong upward push, and a $+2$ in that same batch gets pushed down.

The remarkable part: the subtraction leaves the expected gradient exactly unchanged. Averaged over many rollouts, training points where it pointed before; what changes is how much individual updates scatter around that average. A well-chosen $b$ , one close to the typical reward, shrinks that scatter substantially. (A badly chosen one, say a large arbitrary constant, would inflate it instead. The baseline is a lever on variance, and you still have to set it sensibly.)

The simplest useful baseline is just the mean reward of the current batch of rollouts. That alone helps enormously. (One honest footnote: the batch mean is computed from the very samples being scored, so it slips just outside the proof above and introduces a small bias that shrinks as the batch grows. Using the mean of the other samples for each rollout, a leave-one-out mean, restores exact unbiasedness.) As a preview, the batch-mean idea is essentially what GRPO (chapter 27) does to avoid training a separate network at all.

Try it: the fix, in the same sandbox

This is the bandit from the last chapter, with one addition: the baseline toggle now exists. First leave it OFF and pull arms for a while to re-feel the problem: every positive reward shoves its own arm up, and a few lucky pulls on a mediocre arm can send the policy chasing it. Then hit Reset, turn the baseline ON (it subtracts the running mean reward, the batch-mean idea from above), and run the same experiment. The update becomes reward minus typical reward: merely-OK pulls now barely move the policy, above-average pulls push up, below-average pulls push down, and the policy locks onto the best arm faster and far more steadily.

Policy gradient sandbox (REINFORCE bandit)

Click an arm to sample it. Each pull collects a noisy reward and applies one update: logit ← logit + η·(reward − baseline). Watch the policy shift.

steps 0avg reward 0.000

Policy π(arm) — click a bar to sample

Last sampled arm

—

Reward received

—

Advantage (reward − baseline)

—

REINFORCE pushes probability toward actions that beat expectation and away from those that underperform. With the baseline ON, the update uses reward − running mean, so a merely-OK arm produces a small signed nudge instead of a big positive shove — that is lower variance, and the policy locks onto arm C faster and more stably. With the baseline OFF, every positive reward inflates its arm, so early lucky pulls on a mediocre arm can derail learning. Toggle and compare from a fresh Reset.

One number can’t know every prompt

A single constant baseline has an obvious blind spot. Suppose your batch mixes an easy prompt (typical reward $+8$ ) with a hard one (typical reward $+2$ ). The batch mean sits around $+5$ , so every decent answer to the hard prompt looks like a failure, and every sloppy answer to the easy prompt looks like a triumph. The right “typical reward” to compare against depends on where you’re starting from.

What we want is a baseline that adapts to the situation: the reward we’d expect from this particular prompt. That state-dependent expectation has a name, the value function, and learning it is the job of a second model called the critic. That’s the next chapter.