The score function and baselines

How to take the gradient of an average over samples you can't differentiate through. The log-derivative trick behind every policy gradient, why baselines don't change it, and the reparameterisation alternative.

Key idea: The score function (log-derivative) estimator, control variates and baselines, leave-one-out, and the reparameterisation trick
Real world: A cook tweaks a recipe by tasting batches. They can't differentiate "how good does a random batch taste" with respect to the amount of salt. But they can note which batches tasted better than usual and lean the recipe towards whatever those batches had more of.
In an LLM: The reward of an answer can't be differentiated with respect to the model's weights: sampling a token is a discrete jump, and the checker is a black box. The policy gradient gets around this by differentiating the log-probability of the sampled answer instead, weighted by how much better than usual it scored.

The problem

RL for LLMs wants to raise the expected reward $J(\theta) = \mathbb{E}_{y \sim \pi_\theta}[r(y)]$ . Gradient descent needs $\nabla_\theta J$ . But:

The answer $y$ is sampled. A tiny change in the weights doesn't change which token comes out; it changes how likely each token is.
The reward $r(y)$ comes from a unit test, a maths checker or a separate reward model. There's no gradient through it.

So you can't backpropagate through the sample. You need a gradient of the distribution.

The log-derivative trick

The fix is one line of calculus. The gradient of a probability equals the probability times the gradient of its log:

\nabla_\theta \pi_\theta(y) = \pi_\theta(y)\, \nabla_\theta \log \pi_\theta(y)

Put that inside the expectation and the gradient becomes another expectation, which can be estimated by sampling:

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

$\nabla_\theta \log \pi_\theta(y)$ is called the score function. It's the direction in weight space that makes $y$ more likely. The estimate: sample answers, and push each one's log-probability up in proportion to its reward. This is REINFORCE, and every policy-gradient method for LLMs is a variant of it.

example · tasting batches of soup

Make 10 batches with slightly random salt. Batches that taste better than average had, on average, a bit more salt. Shift the recipe towards that. You never needed "taste" as a smooth function of salt; you only needed to know which random variations you made, and how they turned out.

example · an answer's score function

For a sampled answer,

\nabla \log \pi_\theta(y) = \sum_t \nabla \log \pi_\theta(y_t \mid y_{\lt t})

That's one ordinary backward pass through the tokens it contains, exactly like a supervised fine-tuning step on that answer. RL is "fine-tune on your own samples, weighted by reward".

Baselines are control variates

The score function averages to zero: summed over every possible answer, pushing all of them up in proportion to their probability changes nothing.

\mathbb{E}_{y \sim \pi_\theta}\big[\nabla_\theta \log \pi_\theta(y)\big] = \nabla_\theta \sum_y \pi_\theta(y) = \nabla_\theta 1 = 0

So you can subtract any constant $b$ from the reward without changing the expected gradient. It does change the variance, enormously. This is a control variate: subtract something with known mean (here zero) that is correlated with your noise.

example · pricing a house

To estimate a house's value, you could average the prices of nearby sales: very noisy, because houses differ. Or compare each sale to the area's known average and average the differences. Same target, far less noise, because the shared part has been subtracted out.

example · rewards from 0 to 1, or 10 to 11

If every answer gets reward 10 or 11 instead of 0 or 1, REINFORCE without a baseline pushes every sample up hard, and only the small difference between 10 and 11 carries information. With the group mean as baseline, the advantages are exactly what they were for rewards of 0 and 1, whatever the offset.

add c to every reward0.0

wrong Ar = 0.0+0.00

wrong Dr = 0.0+0.00

right Er = 1.0+1.00

right Cr = 1.0+1.00

wrong Ar = 0.0+0.00

wrong Br = 0.0+0.00

wrong Ar = 0.0+0.00

Push on each answer’s logit from this group (thin tick: the true gradient)

wrong A-0.10

wrong B-0.06

right C+0.09

wrong D-0.03

right E+0.11

baseline used0

agreement with the true direction (cosine)0.96

With rewards of 0 and 1 and no baseline, wrong answers get no push of their own and right ones get pushed up; everything else loses a little, because probabilities must add up to 1. It works, but adding a constant to every reward would wreck it. Try it.

One prompt, a group of 8 sampled answers. Each answer is pushed up (green) or down (red) in proportion to its advantage: reward minus baseline. The bottom panel adds up what the whole group does to each answer’s probability.

Which baseline?

Baseline	Biased?	Notes
None	no	Very noisy; sensitive to reward offset
Learned value $V(s)$ (PPO's critic)	no	Needs a second network
Group mean, including the sample itself	slightly	Shrinks the gradient by $(G-1)/G$ ; GRPO
Leave-one-out (RLOO)	no	Mean of the other $G - 1$ samples
Group mean ÷ group std (GRPO)	yes	Rescales each prompt's gradient; Dr. GRPO removes the std

A baseline must not depend on the sample it's applied to, or the zero-mean argument breaks. Leave-one-out exists precisely to satisfy that.

Try it

add c to every reward0.0group size G8

no baselinegroup mean, including itselfleave-one-out (RLOO)true expected reward V

no baseline0.161bias 0%

group mean, including itself0.122bias 13%

leave-one-out (RLOO)0.142bias 0%

true expected reward V0.129bias 0%

All four point the same way on average, apart from the group mean that includes the answer itself: it shrinks the gradient by (G − 1)/G, a 13% bias at G = 8. Leave-one-out is unbiased and almost as quiet as knowing the true value V. Now add a constant to every reward.

A prompt with five kinds of answer, two of them correct (reward 1). Each point is the noise of the gradient estimate from one group, measured over 800 groups. The dots mark the group size you chose.

The other way: reparameterisation

If the randomness can be moved outside the model, the gradient can go straight through. Instead of sampling $z \sim \mathcal N(\mu_\theta, \sigma_\theta^2)$ , sample noise $\epsilon \sim \mathcal N(0, 1)$ and compute $z = \mu_\theta + \sigma_\theta \epsilon$ . Now $z$ is a differentiable function of $\theta$ . This reparameterisation trick gives much lower-variance gradients and powers VAEs and diffusion models.

It doesn't work for picking tokens: a discrete choice has no smooth path. The Gumbel-softmax relaxation fakes one by replacing the hard choice with a soft, temperature-controlled mixture, which is useful inside a network but not for training a model to write text, where the score function remains the standard.

Where this shows up

REINFORCE, PPO, GRPO, RLOO are all score-function estimators with different baselines: policy gradient, GRPO, fixes.
Advantage normalisation is a choice of baseline and scale: GRPO fixes.
DPO avoids the estimator by solving for the optimum in closed form: tilted distributions.
Gumbel-max sampling, the hard cousin of Gumbel-softmax: sampling text.

going deeper

The estimator

\nabla_\theta\, \mathbb{E}_{y \sim \pi_\theta}[r(y)] = \mathbb{E}_{y \sim \pi_\theta}\big[(r(y) - b)\, \nabla_\theta \log \pi_\theta(y)\big] \approx \frac1G \sum_{i=1}^{G} (r_i - b_i)\, \nabla_\theta \log \pi_\theta(y_i)

For a softmax over logits $\theta$ ,

\frac{\partial \log \pi(a)}{\partial \theta_j} = \mathbb{1}[j = a] - \pi_j

The widget computes this per sample and measures the total variance across 800 groups.

Baselines in papers

\text{RLOO:}\quad \hat A_i = r_i - \frac{1}{G-1} \sum_{j \ne i} r_j = \frac{G}{G-1}\big(r_i - \bar r\big) \qquad \text{GRPO:}\quad \hat A_i = \frac{r_i - \operatorname{mean}(\mathbf r)}{\operatorname{std}(\mathbf r)}

RLOO is the group-mean advantage rescaled by $G/(G-1)$ , which undoes the shrinkage. GRPO's division by the std weights prompts unequally: prompts with nearly all-same rewards have tiny std and get boosted. Dr. GRPO drops it.

The variance-minimising constant baseline is not the mean reward but a reward average weighted by gradient size:

b^* = \frac{\mathbb{E}\big[r\, \lVert \nabla \log \pi \rVert^2\big]}{\mathbb{E}\big[\lVert \nabla \log \pi \rVert^2\big]}

In practice the mean is close enough.

Reparameterisation and Gumbel-softmax

\nabla_\theta\, \mathbb{E}_{\epsilon \sim p(\epsilon)}\big[f(g_\theta(\epsilon))\big] = \mathbb{E}_{\epsilon}\big[\nabla_\theta f(g_\theta(\epsilon))\big] \qquad \text{Gumbel-softmax: } y_k = \frac{\exp\big((\log \pi_k + g_k)/\tau\big)}{\sum_j \exp\big((\log \pi_j + g_j)/\tau\big)}

with $g_k$ i.i.d. Gumbel noise. As $\tau \to 0$ , $y$ becomes the one-hot Gumbel-max sample from $\pi$ .

Papers

Williams, Simple statistical gradient-following algorithms for connectionist reinforcement learning (1992): REINFORCE
Monte Carlo Gradient Estimation in Machine Learning (Mohamed et al., 2019): survey of score-function, reparameterisation and control variates
Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback (Ahmadian et al., 2024): RLOO for LLMs
Understanding R1-Zero-Like Training (Liu et al., 2025): Dr. GRPO, the std and length biases
Auto-Encoding Variational Bayes (Kingma & Welling, 2013): reparameterisation
Categorical Reparameterization with Gumbel-Softmax (Jang et al., 2016) and The Concrete Distribution (Maddison et al., 2016)

The problem

The log-derivative trick

Baselines are control variates

Which baseline?

Try it

The other way: reparameterisation

Where this shows up

The estimator

Baselines in papers

Reparameterisation and Gumbel-softmax

Papers

Related