The Most Intuitive Way to Understand Score Function Centering: Understand Trainer-Generator Mismatch in Async RL from the Score Function

You can always find my latest articles (also the Chinese version if you are a native Chinese speaker) in my personal blog webiste, which has better rendering and interactive experience and you can also leave comments there! It astonishes me that I need to pay for a pro subscription in order to post community article in HuggingFace. I probably will stop my subscription some day and you know where to find me ⬆️

I find that the explanations of score function centering available online are all more or less roundabout, and not direct enough. I want this article to give you the most direct explanation of score function centering, starting from intuition. This article was written entirely by hand, not by AI. Of course, if you are reading the English version, I am sorry, that one was translated by AI.

This article assumes you know how to use PPO or GRPO (what the loss looks like) and some simple probability theory (that really is enough). If you have not read it yet, please take ten minutes to read this article.

If any one of the following describes you, this article should suit you:

• You only know how to use GRPO, but you do not know the principle behind it
• You want to deepen your understanding of trainer-generator mismatch
• You work on async RL with long rollout training and are still not sure how to improve the stability of the algorithm
• You have no idea what score function centering is (you may not even know what the score here is) but you want to find out

Score Function Centering is not an established term. I know that every frontier lab has its own name for this technique, alternative token, covariance optimization and so on. Call it whatever you like. This article will also explain where these names come from.

TL;DR:

• After section 1 you will know that the score function is an important bridge for computing the policy gradient, and the most direct thing we can get out of auto-differentiation.
• After section 2 you will know the important property of the score function: it is already centered, so its expectation is 0. We need that expectation to stay 0 in order to replace reward with the so-called advantage in the policy gradient. You will also see how trainer-generator mismatch breaks that property. And you will see what breaking it actually costs: the mean advantage produces an extra gradient that pulls the current trainer toward the old generator, which is why we need to repair it.
• After section 3 you will know what score function centering is really doing, why it repairs that zero-expectation property, how we should use it in practice, and one direction for churning out papers in this area.

1. Let us start from the score function. What is a score? Why do we use it?

1.1 In GRPO/PPO, what gradient are we actually computing?

Below is a quick derivation of the policy gradient. If you are already quite familiar with this part, skip this subsection.

Suppose that at some state, that is, at some LLM prefix $xx$, the policy is:

$$p_{\theta }(a\mathrm{\mid }x)p\_\backslash theta(a|x)$$

For simplicity, let us drop $xx$ for now and write:

$$p_{\theta }(a)\mathrm{.}p\_\backslash theta(a).$$

Our goal is very natural:

"Make high reward actions easier for the model to sample."

So we can write the simplest possible objective:

$$J(\theta )={\mathbb{E}}_{a\sim p_{\theta }}[R(a)]\mathrm{.}J(\backslash theta)\; =\; \backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[R(a)].$$

What this says is: "Sample an action from the current model $p_{\theta }p\_\backslash theta$, compute the reward, and hope its expectation keeps going up."

The question is: how do we compute ${\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{a\sim p_{\theta }}[R(a)]\backslash nabla\_\backslash theta\; \backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[R(a)]$? There is a difficulty here: $\theta \backslash theta$ does not appear directly inside $R(a)R(a)$. It is hidden inside the sampling distribution $a\sim p_{\theta }a\backslash sim\; p\_\backslash theta$. In other words, we are not computing ${\mathrm{\nabla }}_{\theta }R(a)\backslash nabla\_\backslash theta\; R(a)$. In fact, the reward in RL is perfectly allowed to be non-differentiable. For example: 1 point for a correct answer to a maths problem and 0 otherwise, or a human expert handing out a score directly. None of that matters. What really depends on $\theta \backslash theta$ is $p_{\theta }(a)p\_\backslash theta(a)$.

Writing the expectation out:

$$J(\theta )=\sum _a{p}_{\theta }(a)R(a)\mathrm{.}J(\backslash theta)\; =\; \backslash sum\_a\; p\_\backslash theta(a)R(a).$$

Differentiating:

$${\mathrm{\nabla }}_{\theta }J(\theta )=\sum _{a}R(a){\mathrm{\nabla }}_{\theta }{p}_{\theta }(a)\mathrm{.}\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash sum\_a\; R(a)\backslash nabla\_\backslash theta\; p\_\backslash theta(a).$$

Then we use a very simple but extremely important identity:

$${\mathrm{\nabla }}_{\theta }{p}_{\theta }(a)=p_{\theta }(a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}\backslash nabla\_\backslash theta\; p\_\backslash theta(a)\; =\; p\_\backslash theta(a)\; \backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

Because:

$$\mathrm{\nabla }\log p=\frac{1}{p}\mathrm{\nabla }p\mathrm{.}\backslash nabla\; \backslash log\; p\; =\; \backslash frac\{1\}\{p\}\backslash nabla\; p.$$

So:

$${\mathrm{\nabla }}_{\theta }J(\theta )=\sum _a{p}_{\theta }(a)R(a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash sum\_a\; p\_\backslash theta(a)\; R(a)\; \backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

Writing it back as an expectation:

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{a\sim p_{\theta }}[R(a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)]\mathrm{.}\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}\; \backslash left[\; R(a)\; \backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a)\; \backslash right].$$

Here, at last, is the most critical object in the policy gradient:

$$s_{\theta }(a)={\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}s\_\backslash theta(a)\; =\; \backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

This thing is called the score function. Note that the score here is not a reward score. It is not "how many points did this answer get". It is a specific term from probability and statistics:

$$\text{score}={\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}\backslash text\{score\}\; =\; \backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

That is, the gradient, with respect to the model parameters, of the log probability of some sampled action.

1.2 Why do we study the score?

Because the core structure of the policy gradient really is just:

reward $\times \backslash times$ score

(I know you want to say that what we actually use is $\text{Advantage}\times \text{score}\backslash text\{Advantage\}\backslash times\backslash text\{score\}$, but hold that thought.)

Written the way we are more used to when training LLMs:

$$g=R(a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}g=R(a)\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

So you can read the two parts like this: $R(a)R(a)$ tells us "should this action be reinforced or weakened, and by how much?" while ${\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a)$ tells us "if I want to raise or lower the probability of this action, which direction should the parameters move in?"

This is also why I think the most important sentence for understanding the policy gradient is:

The score function is the bridge that turns "the gradient of an expectation over a distribution", which you have no way of computing, into "a gradient estimate you can compute by sampling".

More precisely:

distribution-level differentiation $\to \backslash rightarrow$ sample-level gradient estimation

What we originally wanted to compute was:

$${\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{a\sim p_{\theta }}[R(a)]\mathrm{.}\backslash nabla\_\backslash theta\; \backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[R(a)].$$

It looks like it requires differentiating through the whole action distribution.

But the score-function trick turns it into this: first sample an action

$$a\sim p_{\theta },a\backslash sim\; p\_\backslash theta,$$

then compute:

$$R(a){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}R(a)\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

Sample a few more times and average:

$${\mathrm{\nabla }}_{\theta }J\approx \frac{1}{N}\sum _{i=1}^{N}R(a_i){\mathrm{\nabla }}_{\theta }\log p_{\theta }(a_i)\mathrm{.}\backslash nabla\_\backslash theta\; J\backslash approx\backslash frac1N\backslash sum\_\{i=1\}^NR(a\_i)\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a\_i).$$

This is the core computational mechanism of REINFORCE and of policy gradient. And although GRPO adds group-relative advantage, importance ratio, clipping and so on, the basic object inside it that actually produces the policy gradient still comes down to:

$${\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

At the same time, you will notice that this is the most direct thing a single backward hands you when you actually use GRPO.

2. The score function and trainer-generator mismatch

2.1 What is the beautiful property of the score function, and why does it matter?

Let us now talk about the advantage. We all know that in GRPO/PPO we do not use the reward directly, we use the advantage (why?). The advantage is the reward minus a constant:

$$A(a)=R(a)-b,A(a)=R(a)-b,$$

Here $bb$ can be a value function, it can be a group mean, and it can be any baseline that depends only on the current state and not on the specific action. In GRPO, for instance, the thing we are most familiar with is sampling many responses for the same prompt and then using some group-relative reward as the advantage. Never mind exactly how the normalization is done, in essence it is all doing one thing: subtracting from the raw reward a part that everyone shares, leaving only whether this action is better or worse than average. But there is a very basic question here that is also very easy to overlook:

What gives us the right to subtract this baseline at all?

The original policy gradient is:

$$g={\mathbb{E}}_{a\sim p_{\theta }}[R(a)s_{\theta }(a)]\mathrm{.}g=\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[R(a)s\_\backslash theta(a)].$$

If we swap the reward for the advantage:

$$A(a)=R(a)-b,A(a)=R(a)-b,$$

then the new gradient is:

$$g^{\mathrm{\prime }}={\mathbb{E}}_{a\sim p_{\theta }}[(R(a)-b)s_{\theta }(a)]\mathrm{.}g\text{'}=\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[(R(a)-b)s\_\backslash theta(a)].$$

Expanding it:

$${\mathbb{E}}_{a\sim p_{\theta }}[R(a)s_{\theta }(a)]-b{\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[R(a)s\_\backslash theta(a)]\; -b\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)].$$

The first term is our original policy gradient. In other words, if we want the expectation of the policy gradient to be completely unchanged after subtracting the baseline, the second term has to vanish:

$$b{\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]=0.b\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]\; =\; 0.$$

And $bb$ is clearly not necessarily 0, so what we really need is:

$${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]=0.\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]\; =\; 0.$$

This is a very important property of the score function inside the policy gradient. You can ignore for now why the property holds and simply remember what it means in practice:

As long as the expectation of the score is 0, we can subtract from the reward any baseline that does not depend on the action, without changing the expected policy gradient.

This is why we can write the reward as an advantage, and why reward centering, value baselines and group-relative baselines all work inside PPO and GRPO. Some papers explain this starting from covariance, writing something like:

$${\mathbb{E}}_{a\sim p_{\theta }}[R(a)s_{\theta }(a)]=\mathrm{Cov}(R(a),s(a))+\mathbb{E}[R(a)]\mathbb{E}[s(a)],\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[R(a)s\_\backslash theta(a)]\; =\; \backslash operatorname\{Cov\}(R(a),s(a))\; +\; \backslash mathbb\; E[R(a)]\backslash mathbb\; E[s(a)],$$

and then, because $\mathbb{E}[s]=0\backslash mathbb\; E[s]=0$, only the covariance is left. There is of course nothing wrong with writing it this way, but personally I feel it makes something very simple sound roundabout. What we actually care about is one sentence: I want to subtract a baseline from the reward, and I do not want that baseline to produce a gradient of its own. That naturally requires $\mathbb{E}[s]=0\backslash mathbb\; E[s]=0$. It really is that simple.

And this property fits intuition very well. Suppose the rewards of two actions are:

$$R(A)=101,R(B)=99.R(A)=101,\backslash qquad\; R(B)=99.$$

and:

$$R^{\mathrm{\prime }}(A)=2,R^{\mathrm{\prime }}(B)=0.R\text{'}(A)=2,\backslash qquad\; R\text{'}(B)=0.$$

From the point of view of policy optimization, these two sets of rewards carry exactly the same information: $AA$ is 2 points better than $BB$ in both. We clearly do not want the expected update direction of the model to change simply because we added 99 to every reward. Adding or subtracting the same constant from all rewards never changes the expected policy gradient. This is also why a baseline is fundamentally a very natural operation and not merely a mathematical trick.

Now the question is:

Why does the expectation of the score happen to equal 0?

To answer that, we finally need to look at what the score really is.

2.2 What exactly is the score function?

Next we do some calculus on the score function, the kind anyone who has taken a calculus course can do. We defined earlier:

$$s_{\theta }(a)={\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}s\_\backslash theta(a)=\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

Suppose the current model outputs, at some prefix, the logits of all tokens in the vocabulary:

$$z_{\theta }(1),z_{\theta }(2),\dots ,z_{\theta }(V)\mathrm{.}z\_\backslash theta(1),z\_\backslash theta(2),\backslash ldots,z\_\backslash theta(V).$$

We assume as usual that the output distribution is given by Softmax:

$$p_{\theta }(a)=\frac{e^{z_{\theta }(a)}}{\sum _v{e}^{z_{\theta }(v)}}\mathrm{.}p\_\backslash theta(a)\; =\; \backslash frac\{e^\{z\_\backslash theta(a)\}\}\{\backslash sum\_v\; e^\{z\_\backslash theta(v)\}\}.$$

So:

$$s_{\theta }(a)={\mathrm{\nabla }}_{\theta }{z}_{\theta }(a)-{\mathrm{\nabla }}_{\theta }\log \sum _v{e}^{z_{\theta }(v)}\mathrm{.}s\_\backslash theta(a)\; =\; \backslash nabla\_\backslash theta\; z\_\backslash theta(a)\; -\; \backslash nabla\_\backslash theta\; \backslash log\backslash sum\_v\; e^\{z\_\backslash theta(v)\}.$$

Let us give the first term above a nicer name, the logit Jacobian:

$$J_{\theta }(v)={\mathrm{\nabla }}_{\theta }{z}_{\theta }(v)\mathrm{.}J\_\backslash theta(v)=\backslash nabla\_\backslash theta\; z\_\backslash theta(v).$$

You can read this $J_{\theta }(v)J\_\backslash theta(v)$ as: if I change the current trainer's model parameters slightly, in which direction does the raw logit of token $vv$ move.

Here is a concrete example. Suppose the model parameters change from $\theta \backslash theta$ to:

$$\theta +\mathrm{\Delta }\theta ,\backslash theta+\backslash Delta\backslash theta,$$

then to first order:

$$z_{\theta +\mathrm{\Delta }\theta }(v)\approx z_{\theta }(v)+J_{\theta }(v{)}^T\mathrm{\Delta }\theta \mathrm{.}z\_\{\backslash theta+\backslash Delta\backslash theta\}(v)\; \backslash approx\; z\_\backslash theta(v)\; +\; J\_\backslash theta(v)^T\backslash Delta\backslash theta.$$

So $J_{\theta }(v)J\_\backslash theta(v)$ is essentially a "logit sensitivity direction" for token $vv$ in parameter space. It tells us which way the parameters have to move to push this token's logit up or pull it down.

Next we deal with the second term:

$${\mathrm{\nabla }}_{\theta }\log \sum _v{e}^{z_{\theta }(v)}\backslash nabla\_\backslash theta\; \backslash log\backslash sum\_v\; e^\{z\_\backslash theta(v)\}$$

This needs the chain rule twice. First differentiate the outer log, which puts the sum inside the log into the denominator:

$${\mathrm{\nabla }}_{\theta }\log \sum _v{e}^{z_{\theta }(v)}=\frac{1}{\sum _u{e}^{z_{\theta }(u)}}{\mathrm{\nabla }}_{\theta }\sum _v{e}^{z_{\theta }(v)}\mathrm{.}\backslash nabla\_\backslash theta\backslash log\backslash sum\_v\; e^\{z\_\backslash theta(v)\}=\backslash frac\{1\}\{\backslash sum\_u\; e^\{z\_\backslash theta(u)\}\}\backslash nabla\_\backslash theta\backslash sum\_v\; e^\{z\_\backslash theta(v)\}.$$

Then move the gradient inside the sum and use the chain rule once more on every $e^{z_{\theta }(v)}e^\{z\_\backslash theta(v)\}$:

$${\mathrm{\nabla }}_{\theta }\sum _v{e}^{z_{\theta }(v)}=\sum _v{\mathrm{\nabla }}_{\theta }{e}^{z_{\theta }(v)}=\sum _v{e}^{z_{\theta }(v)}{\mathrm{\nabla }}_{\theta }{z}_{\theta }(v)\mathrm{.}\backslash nabla\_\backslash theta\backslash sum\_v\; e^\{z\_\backslash theta(v)\}=\backslash sum\_v\backslash nabla\_\backslash theta\; e^\{z\_\backslash theta(v)\}=\backslash sum\_v\; e^\{z\_\backslash theta(v)\}\backslash nabla\_\backslash theta\; z\_\backslash theta(v).$$

Putting the two steps together gives:

$$\frac{\sum _v{e}^{z_{\theta }(v)}{\mathrm{\nabla }}_{\theta }{z}_{\theta }(v)}{\sum _u{e}^{z_{\theta }(u)}}\mathrm{.}\backslash frac\{\; \backslash sum\_v\; \backslash color\{orange\}\{e^\{z\_\backslash theta(v)\}\}\; \backslash color\{blue\}\{\backslash nabla\_\backslash theta\; z\_\backslash theta(v)\}\; \}\{\; \backslash color\{orange\}\{\backslash sum\_u\; e^\{z\_\backslash theta(u)\}\}\; \}.$$

And it happens that:

$$p_{\theta }(v)=\frac{e^{z_{\theta }(v)}}{\sum _u{e}^{z_{\theta }(u)}},J_{\theta }(v)={\mathrm{\nabla }}_{\theta }{z}_{\theta }(v)\mathrm{.}p\_\backslash theta(v)\; =\; \backslash color\{orange\}\{\backslash frac\{e^\{z\_\backslash theta(v)\}\}\{\backslash sum\_u\; e^\{z\_\backslash theta(u)\}\}\},\; \backslash quad\; J\_\backslash theta(v)=\; \backslash color\{blue\}\{\backslash nabla\_\backslash theta\; z\_\backslash theta(v)\}.$$

So the second term can be written as:

$$\sum _v{p}_{\theta }(v)J_{\theta }(v)\mathrm{.}\backslash sum\_v\; p\_\backslash theta(v)J\_\backslash theta(v).$$

That is:

$${\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

Therefore:

$$s_{\theta }(a)=J_{\theta }(a)-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}s\_\backslash theta(a)\; =\; J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

This expression is especially important, because it tells us that inside the score function, Softmax is already centering the logit Jacobian by itself. $J_{\theta }(a)J\_\backslash theta(a)$ is the logit gradient of the action we just sampled, while ${\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]$ is the average logit gradient over all actions, weighted by the current policy $p_{\theta }p\_\backslash theta$.

So the score really is the logit Jacobian of the sampled action minus the average logit Jacobian under the current policy.

With that expression in hand, the fact that the expectation of the score is 0 barely needs proving. Because the action itself is sampled from:

$$a\sim p_{\theta }a\backslash sim\; p\_\backslash theta$$

we have:

$${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]={\mathbb{E}}_{a\sim p_{\theta }}[J_{\theta }(a)-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]=\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}\backslash left[J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]\backslash right].$$

Splitting the expectation:

$$={\mathbb{E}}_{a\sim p_{\theta }}[J_{\theta }(a)]-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}=\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(a)]-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

Here $aa$ and $vv$ are just two dummy variables with different names. They both come from exactly the same distribution $p_{\theta }p\_\backslash theta$, so the two terms are the same thing:

$${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]={\mathbb{E}}_{p_{\theta }}[J_{\theta }]-{\mathbb{E}}_{p_{\theta }}[J_{\theta }]=0.\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]=\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]-\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]=0.$$

So now we can put the whole story together: the Softmax score is centered by construction, therefore:

$${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]=0,\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]\; =\; 0,$$

and this zero-mean property in turn guarantees that we can use reward minus baseline in place of the raw reward in the policy gradient, without changing the expected policy gradient.

Up to here everything is lovely. But notice that the whole derivation above quietly uses one assumption we have not yet stated carefully:

The distribution the expectation is computed over, and the distribution the action is generated from, are both $p_{\theta }p\_\backslash theta$.

In other words, we assume:

the generator that produces the rollout is the very trainer that is now running backward.

The story only really begins once that assumption fails.

2.3 What exactly does trainer-generator mismatch break?

It must break ${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]=0\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]\; =\; 0$, the zero-mean property of the score function, otherwise all that discussion above would have been for nothing.

But the most important question here is not "it is broken", it is: which step exactly changed?

Let us first separate the trainer and the generator clearly. Suppose the trainer policy that is currently running backward and about to update the parameters is:

$$p_{\theta }(a)\mathrm{.}p\_\backslash theta(a).$$

And the generator that actually produced the rollout is a slightly older policy:

$$q_{\varphi }(a)\mathrm{.}q\_\backslash phi(a).$$

I deliberately used different parameters $\theta \backslash theta$ and $\varphi \backslash phi$ here. You can read $\theta \backslash theta$ as the trainer's current parameters, and $\varphi \backslash phi$ as the version the generator took away to do the rollout. In synchronous training we approximately have:

$$q_{\varphi }\approx p_{\theta },q\_\backslash phi\; \backslash approx\; p\_\backslash theta,$$

and we can even simply treat it as:

$$q_{\varphi }=p_{\theta }\mathrm{.}q\_\backslash phi=p\_\backslash theta.$$

But in async RL, especially when rollouts are long, there are many generators, and the trainer keeps updating, the two distributions naturally start to drift apart:

$$q_{\varphi }\mathrm{\ne }{p}_{\theta }\mathrm{.}q\_\backslash phi\backslash neq\; p\_\backslash theta.$$

Now recall the formula we just used to prove the zero-mean property:

$${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)]={\mathbb{E}}_{a\sim p_{\theta }}[J_{\theta }(a)]-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]=0.\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)]=\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(a)]-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]=0.$$

Why is it 0? Essentially because both expectations are taken over the same $p_{\theta }p\_\backslash theta$. In the first expectation, the action $aa$ is sampled from $p_{\theta }p\_\backslash theta$. In the second, Softmax itself uses $p_{\theta }p\_\backslash theta$ to center all the logit Jacobians. So the two terms match exactly.

But now the generator has become $q_{\varphi }q\_\backslash phi$. The action that actually reaches the trainer no longer satisfies:

$$a\sim p_{\theta },a\backslash sim\; p\_\backslash theta,$$

but instead:

$$a\sim q_{\varphi }\mathrm{.}a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}.$$

Note that the trainer has not changed the way it computes the score. What we run backward is still the current trainer's:

$$s_{\theta }(a)=J_{\theta }(a)-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}s\_\backslash theta(a)=J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

The $J_{\theta }J\_\backslash theta$ here still comes from the current trainer, and the Softmax centering still uses the current trainer distribution $p_{\theta }p\_\backslash theta$. Only one thing has actually changed: which distribution the action on the outside was sampled from.

So what we should really be looking at now is no longer:

$${\mathbb{E}}_{a\sim p_{\theta }}[s_{\theta }(a)],\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[s\_\backslash theta(a)],$$

but:

$${\mathbb{E}}_{a\sim q_{\varphi }}[s_{\theta }(a)]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[s\_\backslash theta(a)].$$

Expanding the score:

$${\mathbb{E}}_{a\sim q_{\varphi }}[s_{\theta }(a)]={\mathbb{E}}_{a\sim q_{\varphi }}[J_{\theta }(a)-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[s\_\backslash theta(a)]=\backslash mathbb\; E\_\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}\backslash left[J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]\backslash right].$$

Splitting the expectation:

$$={\mathbb{E}}_{a\sim q_{\varphi }}[J_{\theta }(a)]-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}=\backslash mathbb\; E\_\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[J\_\backslash theta(a)]-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

Now the problem is obvious. Originally we had:

$${\mathbb{E}}_{a\sim p_{\theta }}[J_{\theta }(a)]-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)],\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(a)]-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)],$$

where both expectations are over the same distribution, so it equals 0. Now it has become:

$${\mathbb{E}}_{a\sim q_{\varphi }}[J_{\theta }(a)]-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[J\_\backslash theta(a)]-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

The left side averages under the generator distribution $q_{\varphi }q\_\backslash phi$, while the right side still averages under the trainer distribution $p_{\theta }p\_\backslash theta$. As long as:

$$q_{\varphi }\mathrm{\ne }{p}_{\theta },q\_\backslash phi\backslash neq\; p\_\backslash theta,$$

we generally can no longer expect:

$${\mathbb{E}}_{a\sim q_{\varphi }}[J_{\theta }(a)]={\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(a)]=\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)].$$

So in the end:

$${\mathbb{E}}_{a\sim q_{\varphi }}[s_{\theta }(a)]={\mathbb{E}}_{q_{\varphi }}[J_{\theta }]-{\mathbb{E}}_{p_{\theta }}[J_{\theta }]\mathrm{\ne }0.\backslash mathbb\; E\_\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[s\_\backslash theta(a)]=\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]-\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]\backslash neq\; 0.$$

This is the most direct effect of trainer-generator mismatch on the score function.

You will find there is nothing mysterious about this at all. The score did not suddenly go bad, the Softmax formula did not change, and the $J_{\theta }J\_\backslash theta$ the trainer computes did not change either. The only thing that actually changed is that the distribution the action is sampled from went from $p_{\theta }p\_\backslash theta$ to $q_{\varphi }q\_\backslash phi$.

And the zero-mean property happens to depend on "the distribution the action is sampled from" and "the distribution the centering inside the score uses" being the same distribution. Those two things are now different:

$$\underset{\text{generator sampling distribution}}{\underbrace{a\sim q_{\varphi }}},\underset{\text{trainer centering distribution}}{\underbrace{{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]}}\mathrm{.}\backslash underbrace\{a\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}\_\{\backslash text\{generator\; sampling\; distribution\}\},\backslash qquad\; \backslash underbrace\{\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]\}\_\{\backslash text\{trainer\; centering\; distribution\}\}.$$

So the old:

$${\mathbb{E}}_{p_{\theta }}[J_{\theta }]-{\mathbb{E}}_{p_{\theta }}[J_{\theta }]=0\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]-\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]=0$$

has become:

$${\mathbb{E}}_{q_{\varphi }}[J_{\theta }]-{\mathbb{E}}_{p_{\theta }}[J_{\theta }]\mathrm{\ne }0.\backslash mathbb\; E\_\{\backslash color\{red\}\{q\_\backslash phi\}\}[J\_\backslash theta]-\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]\backslash neq\; 0.$$

At this point we have in fact explained the most critical layer of trainer-generator mismatch:

In synchronous RL, the generator and the trainer are the same distribution, so the standard score is zero-mean by nature. In async RL, the action comes from the old generator $q_{\varphi }q\_\backslash phi$ while the score is still centered by the current trainer $p_{\theta }p\_\backslash theta$, and so the zero-mean property is broken.

2.4 What does this non-zero mean actually cause in practice?

At this point you may feel that ${\mathbb{E}}_{q_{\varphi }}[s_{\theta }]\mathrm{\ne }0\backslash mathbb\; E\_\{q\_\backslash phi\}[s\_\backslash theta]\backslash neq\; 0$ still sounds a bit abstract. We said above that it breaks baseline invariance, but what kind of gradient does it actually produce in training?

This becomes very intuitive if we look at it directly in logit space. For a Softmax, the gradient of the log probability of some action $aa$ with respect to the whole logit vector $zz$ is:

$${\mathrm{\nabla }}_z\log p_{\theta }(a)=e_a-p_{\theta },\backslash nabla\_z\backslash log\; p\_\backslash theta(a)=e\_a-p\_\backslash theta,$$

where $e_{a}e\_a$ is the one-hot vector of action $aa$. In synchronous training $a\sim p_{\theta }a\backslash sim\; p\_\backslash theta$, so taking the expectation of it gives:

$${\mathbb{E}}_{a\sim p_{\theta }}[{\mathrm{\nabla }}_z\log p_{\theta }(a)]=p_{\theta }-p_{\theta }=0.\backslash mathbb\; E\_\{a\backslash sim\; p\_\backslash theta\}[\backslash nabla\_z\backslash log\; p\_\backslash theta(a)]=p\_\backslash theta-p\_\backslash theta=0.$$

But now the action is sampled by the old generator $q_{\varphi }q\_\backslash phi$:

$${\mathbb{E}}_{a\sim q_{\varphi }}[{\mathrm{\nabla }}_z\log p_{\theta }(a)]=q_{\varphi }-p_{\theta }\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash nabla\_z\backslash log\; p\_\backslash theta(a)]=q\_\backslash phi-p\_\backslash theta.$$

This formula makes the intuition behind it easier to see than ${\mathbb{E}}_q[J]-{\mathbb{E}}_p[J]\backslash mathbb\; E\_q[J]-\backslash mathbb\; E\_p[J]$ did: if some token is more common in the old generator than in the current trainer, that is $q_{\varphi }(v)>p_{\theta }(v)q\_\backslash phi(v)\backslash gt\; p\_\backslash theta(v)$, this direction pushes its logit up. The other way around, if $q_{\varphi }(v)<p_{\theta }(v)q\_\backslash phi(v)\backslash lt\; p\_\backslash theta(v)$, this direction pulls it down.

Now let us go back to the baseline we have been discussing all along. Suppose you add a constant $cc$ that has nothing to do with the action to every reward or advantage. In a normal on-policy setting this constant should not produce any expected gradient. But after trainer-generator mismatch it produces an extra:

$$\mathrm{\Delta }{g}_z=c{\mathbb{E}}_{a\sim q_{\varphi }}[{\mathrm{\nabla }}_z\log p_{\theta }(a)]=c(q_{\varphi }-p_{\theta })\mathrm{.}\backslash Delta\; g\_z=c\backslash ,\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash nabla\_z\backslash log\; p\_\backslash theta(a)]=c(q\_\backslash phi-p\_\backslash theta).$$

If $c>0c\backslash gt\; 0$, what this gradient does is very direct: it pulls the current trainer $p_{\theta }p\_\backslash theta$ toward the distribution of the old generator $q_{\varphi }q\_\backslash phi$.

The other way around, if $c<0c\backslash lt\; 0$, the trainer is pushed away from the generator. That is a bad thing to have, because the $cc$ here is playing the role of the baseline, and in theory a baseline is free to be positive or negative. So we have effectively added a perturbation to the distance between the trainer and the generator, and the direction and the size of that perturbation are decided entirely by the value of the baseline.

There is another way to look at this direction. Suppose we really do take the samples of the old generator and maximize the log likelihood of the current trainer:

$${\mathbb{E}}_{a\sim q_{\varphi }}[\log p_{\theta }(a)],\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash log\; p\_\backslash theta(a)],$$

its gradient in logit space is exactly:

$${\mathrm{\nabla }}_z{\mathbb{E}}_{a\sim q_{\varphi }}[\log p_{\theta }(a)]=q_{\varphi }-p_{\theta }\mathrm{.}\backslash nabla\_z\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash log\; p\_\backslash theta(a)]=q\_\backslash phi-p\_\backslash theta.$$

In other words:

$$q_{\varphi }-p_{\theta }=-{\mathrm{\nabla }}_z{D}_{\mathrm{K}\mathrm{L}}(q_{\varphi }\mathrm{\parallel }{p}_{\theta })\mathrm{.}q\_\backslash phi-p\_\backslash theta=-\backslash nabla\_zD\_\{\backslash mathrm\{KL\}\}(q\_\backslash phi\backslash |p\_\backslash theta).$$

So when the mean reward or the mean advantage is positive, that thing which was supposed to be no more than "the part that everyone shares" quietly turns into a gradient whose direction is exactly the direction of doing behavior cloning on the old generator. This is quite strange: the average says nothing about which action is better or worse than which, so what gives it the right to decide whether we move toward the old generator, and by how much?

The same thing shows up in the parameter-space notation we used earlier. Split the advantage into:

$$A(a)=(A(a)-{\bar{A}}_q)+{\bar{A}}_q,{\bar{A}}_q={\mathbb{E}}_{q_{\varphi }}[A]\mathrm{.}A(a)=(A(a)-\backslash bar\; A\_q)+\backslash bar\; A\_q,\backslash qquad\; \backslash bar\; A\_q=\backslash mathbb\; E\_\{q\_\backslash phi\}[A].$$

Then:

$${\mathbb{E}}_{q_{\varphi }}[A(a)s_{\theta }(a)]={\mathbb{E}}_{q_{\varphi }}[(A(a)-{\bar{A}}_q)s_{\theta }(a)]+{\bar{A}}_q{\mathbb{E}}_{q_{\varphi }}[s_{\theta }(a)]\mathrm{.}\backslash mathbb\; E\_\{q\_\backslash phi\}[A(a)s\_\backslash theta(a)]=\backslash mathbb\; E\_\{q\_\backslash phi\}[(A(a)-\backslash bar\; A\_q)s\_\backslash theta(a)]+\backslash bar\; A\_q\backslash ,\backslash mathbb\; E\_\{q\_\backslash phi\}[s\_\backslash theta(a)].$$

And the second term is:

$${\bar{A}}_q({\mathbb{E}}_{q_{\varphi }}[J_{\theta }]-{\mathbb{E}}_{p_{\theta }}[J_{\theta }])\mathrm{.}\backslash bar\; A\_q\backslash left(\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]-\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]\backslash right).$$

This is one of the most practical consequences of breaking the zero-mean property: the term that really describes how good one token is relative to another is the first one, but the mean advantage additionally multiplies a trainer-generator mismatch direction.

Looking at Score Function Centering now, the motivation is more direct. We do not only want a mathematical expectation to look like 0 again, we also want to take this old-policy cloning direction, the one quietly controlled by "the average reward level", out of the reward gradient. How far apart the trainer and the generator should be allowed to drift is a separate question, and ${\bar{A}}_q\backslash bar\; A\_q$ should not get to decide it for us in passing here.

3. What is Score Function Centering actually doing?

The answer is almost shouting at you by now. Score Function Centering simply wants the broken zero-mean property of the score function to hold again, so that we can go back to happily using a baseline.

We just saw that after trainer-generator mismatch, the real problem becomes:

$${\mathbb{E}}_{a\sim q_{\varphi }}[s_{\theta }(a)]={\mathbb{E}}_{q_{\varphi }}[J_{\theta }]-{\mathbb{E}}_{p_{\theta }}[J_{\theta }]\mathrm{\ne }0.\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[s\_\backslash theta(a)]=\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]-\backslash mathbb\; E\_\{p\_\backslash theta\}[J\_\backslash theta]\backslash neq\; 0.$$

Why is this not 0? The reason is very simple: the action was sampled from the generator's $q_{\varphi }q\_\backslash phi$, but the "center" subtracted inside the score is still the trainer's $p_{\theta }p\_\backslash theta$.

To correct that, given that the action on the outside already comes from $a\sim q_{\varphi }a\backslash sim\; q\_\backslash phi$, the most direct approach is of course: stop centering with $p_{\theta }p\_\backslash theta$, and center with $q_{\varphi }q\_\backslash phi$, the distribution that actually generated these actions.

That is, take the original standard score:

$$s_{\theta }(a)=J_{\theta }(a)-{\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]s\_\backslash theta(a)=J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]$$

and change it to:

$${\tilde{s}}_{\theta ,\varphi }(a)=J_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]\mathrm{.}\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)=J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[J\_\backslash theta(v)].$$

This is the core step of Score Function Centering.

Note that there are two distributions and model parameters here, and they should not be mixed up. $J_{\theta }(v)J\_\backslash theta(v)$ is still:

$$J_{\theta }(v)={\mathrm{\nabla }}_{\theta }{z}_{\theta }(v),J\_\backslash theta(v)=\backslash nabla\_\backslash theta\; z\_\backslash theta(v),$$

so it is still the logit Jacobian of the current trainer. We are not going back to the old generator to compute gradients, and we are not updating $\varphi \backslash phi$. The only thing that changes is the probability distribution used for centering:

$${\mathbb{E}}_{v\sim p_{\theta }}[J_{\theta }(v)]⟶{\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; p\_\backslash theta\}[J\_\backslash theta(v)]\backslash quad\backslash longrightarrow\backslash quad\backslash mathbb\; E\_\{v\backslash sim\; \backslash color\{red\}\{q\_\backslash phi\}\}[J\_\backslash theta(v)].$$

You can remember it as this: the generator decides "how to average", the trainer decides "which parameter direction to update".

With that single change, the zero-mean property comes straight back. Because the real action was from:

$$a\sim q_{\varphi },a\backslash sim\; q\_\backslash phi,$$

we get:

$${\mathbb{E}}_{a\sim q_{\varphi }}[{\tilde{s}}_{\theta ,\varphi }(a)]={\mathbb{E}}_{a\sim q_{\varphi }}[J_{\theta }(a)]-{\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)]=\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(a)]-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(v)].$$ $aa$ and $vv$ are still two dummy variables, and this time they finally come from the same distribution again:

$${\mathbb{E}}_{a\sim q_{\varphi }}[{\tilde{s}}_{\theta ,\varphi }(a)]={\mathbb{E}}_{q_{\varphi }}[J_{\theta }]-{\mathbb{E}}_{q_{\varphi }}[J_{\theta }]=0.\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)]=\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]-\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]=0.$$

So we recover:

$${\mathbb{E}}_{a\sim q_{\varphi }}[(R(a)-b){\tilde{s}}_{\theta ,\varphi }(a)]={\mathbb{E}}_{a\sim q_{\varphi }}[R(a){\tilde{s}}_{\theta ,\varphi }(a)]\mathrm{.}\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[(R(a)-b)\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)]=\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[R(a)\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)].$$

In other words, under the distribution $q_{\varphi }q\_\backslash phi$ that actually produced the rollout, the baseline once again stops producing an extra expected gradient of its own.

This is the most essential intuition of Score Function Centering: align the centering distribution of the score with the distribution that actually generates the actions.

3.1 But there is a big problem in practice: how do we get ${\mathbb{E}}_{q_{\varphi }}[J_{\theta }]\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]$?

If you are really training an LLM and you use the method above, you immediately run into a very practical problem:

$${\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]=\sum _{v\in \mathcal{V}}{q}_{\varphi }(v)J_{\theta }(v)\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(v)]=\backslash sum\_\{v\backslash in\backslash mathcal\; V\}q\_\backslash phi(v)J\_\backslash theta(v).$$

Suppose the vocabulary has 128K tokens. Then this expectation appears to require knowing the generator's probability for all 128K tokens at this prefix:

$$q_{\varphi }(1),q_{\varphi }(2),\dots ,q_{\varphi }(V)\mathrm{.}q\_\backslash phi(1),q\_\backslash phi(2),\backslash ldots,q\_\backslash phi(V).$$

The problem is that this information only exists at the moment the generator produces the token, while the trainer that actually runs backward may be on another machine, in another pipeline stage, or even dozens of train steps later.

Of course you could say: then I will just store the generator's entire vocabulary distribution and send it to the trainer, problem solved.

In theory you can. In engineering you cannot.

Because for every token position in the rollout you would have to store an extra distribution of length vocabulary size. Suppose the vocab is 128K and a rollout has thousands or even tens of thousands of tokens, then multiply by batch size, by the number of generators, and by the number of parallel rollouts. That communication and memory overhead very quickly becomes completely unacceptable.

So the practical problem becomes:

We want the expectation under $q_{\varphi }q\_\backslash phi$, but we cannot move the complete $q_{\varphi }q\_\backslash phi$ to the trainer. What do we do?

The answer is Monte Carlo again: since the expectation is an average over $q_{\varphi }q\_\backslash phi$ in the first place, let us sample a few more tokens from $q_{\varphi }q\_\backslash phi$ on the generator side.

3.2 Alternative tokens: do not give the trainer the whole distribution, a few samples are enough

Suppose the token actually chosen by the generator and entered into the rollout is:

$$a\sim q_{\varphi }\mathrm{.}a\backslash sim\; q\_\backslash phi.$$

Besides this primary action, at exactly the same prefix we additionally sample $KK$ tokens from the generator distribution:

$$v_1,v_2,\dots ,v_K\sim q_{\varphi }\mathrm{.}v\_1,v\_2,\backslash ldots,v\_K\backslash sim\; q\_\backslash phi.$$

These tokens are never actually appended to the trajectory, and they do not change the subsequent generation. They only tell the trainer:

"This is roughly what my distribution looked like at this prefix."

Different places give these tokens different names, alternative tokens, alter tokens, negative tokens, auxiliary samples, and so on. They are all doing the same thing in essence: using a small number of samples to describe the generator's local action distribution at the time.

If these $v_{k}v\_k$ really are sampled IID from $q_{\varphi }q\_\backslash phi$, then:

$${\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]\approx \frac{1}{K}\sum _{k=1}^K{J}_{\theta }(v_k)\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(v)]\backslash approx\backslash frac\{1\}\{K\}\backslash sum\_\{k=1\}^KJ\_\backslash theta(v\_k).$$

So the SFC score can be approximated as:

$${\widehat{s}}_{\theta ,\varphi }(a)=J_{\theta }(a)-\frac{1}{K}\sum _{k=1}^K{J}_{\theta }(v_k)\mathrm{.}\backslash hat\; s\_\{\backslash theta,\backslash phi\}(a)=J\_\backslash theta(a)-\backslash frac\{1\}\{K\}\backslash sum\_\{k=1\}^KJ\_\backslash theta(v\_k).$$

And this Monte Carlo estimator is unbiased for the behavior-centered score we actually want:

$${\mathbb{E}}_{v_1,\dots ,v_K\sim q_{\varphi }}[\frac{1}{K}\sum _{k=1}^K{J}_{\theta }(v_k)]={\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]\mathrm{.}\backslash mathbb\; E\_\{v\_1,\backslash ldots,v\_K\backslash sim\; q\_\backslash phi\}\backslash left[\backslash frac\{1\}\{K\}\backslash sum\_\{k=1\}^KJ\_\backslash theta(v\_k)\backslash right]=\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(v)].$$

So the generator does not have to hand the whole $q_{\varphi }q\_\backslash phi$ to the trainer. It only needs to attach a small amount of extra information at each token position, such as:

$$[a,v_1,v_2,\dots ,v_K]\mathrm{.}[a,v\_1,v\_2,\backslash ldots,v\_K].$$

More realistically, the log probability of these tokens under the generator is stored alongside them:

$$[\log q_{\varphi }(a),\log q_{\varphi }(v_1),\dots ,\log q_{\varphi }(v_K)]\mathrm{.}[\backslash log\; q\_\backslash phi(a),\backslash log\; q\_\backslash phi(v\_1),\backslash ldots,\backslash log\; q\_\backslash phi(v\_K)].$$

This information then travels with the rollout all the way from the generator to the trainer. So you can picture an ordinary RL rollout token as originally carrying only:

$$\text{token id}+\text{reward/advantage}+\log q_{\varphi }(a),\backslash text\{token\; id\}+\backslash text\{reward/advantage\}+\backslash log\; q\_\backslash phi(a),$$

while after SFC it also carries a very small local policy sketch:

$$\{v_k,\log q_{\varphi }(v_k){\}}_{k=1}^K\mathrm{.}\backslash \{v\_k,\backslash log\; q\_\backslash phi(v\_k)\backslash \}\_\{k=1\}^K.$$

This local policy sketch tells the trainer: besides the action it actually chose, which other tokens the generator considered possible at the time.

Clever as you are, you will have noticed that if you want to churn out a few conference papers, you can swap in a few different sampling methods for this generator distribution, see which beats which, and there are several papers right there (LoL).

3.3 Hold on, do we really have to compute $J_{\theta }(v)J\_\backslash theta(v)$ explicitly in the trainer?

At this point your head may be starting to hurt, because:

$$J_{\theta }(v)={\mathrm{\nabla }}_{\theta }{z}_{\theta }(v)J\_\backslash theta(v)=\backslash nabla\_\backslash theta\; z\_\backslash theta(v)$$

and $\theta \backslash theta$ is every parameter of the whole LLM. For a 70B model, $J_{\theta }(v)J\_\backslash theta(v)$ is itself a vector with tens of billions of dimensions.

Do we really have to materialize a huge Jacobian for every alternative token and then compute:

$$J_{\theta }(a)-\frac{1}{K}\sum _k{J}_{\theta }(v_k)?J\_\backslash theta(a)-\backslash frac\{1\}\{K\}\backslash sum\_kJ\_\backslash theta(v\_k)?$$

Of course not. Here $J_{\theta }J\_\backslash theta$ is mainly a mathematical object that helps us understand the algorithm. When we actually write the code, all we need is to construct a scalar surrogate and let autograd differentiate it for us. Go back to the gradient SFC wants:

$$J_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]\mathrm{.}J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(v)].$$

Because:

$$J_{\theta }(v)={\mathrm{\nabla }}_{\theta }{z}_{\theta }(v),J\_\backslash theta(v)=\backslash nabla\_\backslash theta\; z\_\backslash theta(v),$$

we have:

$$J_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[J_{\theta }(v)]={\mathrm{\nabla }}_{\theta }(z_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[z_{\theta }(v)])\mathrm{.}J\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[J\_\backslash theta(v)]=\backslash nabla\_\backslash theta\backslash left(z\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[z\_\backslash theta(v)]\backslash right).$$

So the scalar we actually need to construct is very simple:

$$c_{\theta ,\varphi }(a)=z_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[z_{\theta }(v)]\mathrm{.}c\_\{\backslash theta,\backslash phi\}(a)=z\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[z\_\backslash theta(v)].$$

With the $KK$ alternative samples from before:

$${\widehat{c}}_{\theta ,\varphi }(a)=z_{\theta }(a)-\frac{1}{K}\sum _{k=1}^K{z}_{\theta }(v_k)\mathrm{.}\backslash hat\; c\_\{\backslash theta,\backslash phi\}(a)=z\_\backslash theta(a)-\backslash frac\{1\}\{K\}\backslash sum\_\{k=1\}^Kz\_\backslash theta(v\_k).$$

Then multiply by the advantage, just like an ordinary policy gradient:

$$L_{\mathrm{S}\mathrm{F}\mathrm{C}}=-A(a){\widehat{c}}_{\theta ,\varphi }(a)\mathrm{.}L\_\{\backslash mathrm\{SFC\}\}=-A(a)\backslash hat\; c\_\{\backslash theta,\backslash phi\}(a).$$

Run one ordinary backward on it:

$${\mathrm{\nabla }}_{\theta }{L}_{\mathrm{S}\mathrm{F}\mathrm{C}}=-A(a)(J_{\theta }(a)-\frac{1}{K}\sum _{k=1}^K{J}_{\theta }(v_k))\mathrm{.}\backslash nabla\_\backslash theta\; L\_\{\backslash mathrm\{SFC\}\}=-A(a)\backslash left(J\_\backslash theta(a)-\backslash frac\{1\}\{K\}\backslash sum\_\{k=1\}^KJ\_\backslash theta(v\_k)\backslash right).$$

This is exactly the SFC gradient we wanted, and real code works the same way. The rough logic is:

1. the generator tells the trainer which token is the primary one and which are the alternative ones;
2. the trainer runs a normal forward of the current model;
3. gather the logits of those tokens from the current trainer's logits;
4. compute primary_logit - mean(alternative_logits);
5. multiply by the advantage;
6. backward.

3.4 Why do we not even need a full log-softmax here?

All along we have been talking about the score:

$${\mathrm{\nabla }}_{\theta }\log p_{\theta }(a),\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a),$$

so why does the SFC implementation suddenly work directly on logits, $z_{\theta }(a)-{\mathbb{E}}_{q_{\varphi }}[z_{\theta }(v)]z\_\backslash theta(a)-\backslash mathbb\; E\_\{q\_\backslash phi\}[z\_\backslash theta(v)]$?

The reason is that the Softmax normalization term cancels exactly. We can start from this quantity:

$$\log p_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[\log p_{\theta }(v)]\mathrm{.}\backslash log\; p\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[\backslash log\; p\_\backslash theta(v)].$$

Because:

$$\log p_{\theta }(v)=z_{\theta }(v)-\log Z_{\theta },\backslash log\; p\_\backslash theta(v)=z\_\backslash theta(v)-\backslash log\; Z\_\backslash theta,$$

where:

$$Z_{\theta }=\sum _u{e}^{z_{\theta }(u)},Z\_\backslash theta=\backslash sum\_u\; e^\{z\_\backslash theta(u)\},$$

we get:

$$\log p_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[\log p_{\theta }(v)]=z_{\theta }(a)-\log Z_{\theta }-{\mathbb{E}}_{v\sim q_{\varphi }}[z_{\theta }(v)-\log Z_{\theta }]\mathrm{.}\backslash log\; p\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[\backslash log\; p\_\backslash theta(v)]=z\_\backslash theta(a)-\backslash log\; Z\_\backslash theta-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[z\_\backslash theta(v)-\backslash log\; Z\_\backslash theta].$$

For a fixed prefix, $\log Z_{\theta }\backslash log\; Z\_\backslash theta$ has nothing to do with which $vv$ was sampled, so:

$${\mathbb{E}}_{v\sim q_{\varphi }}[\log Z_{\theta }]=\log Z_{\theta }\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[\backslash log\; Z\_\backslash theta]=\backslash log\; Z\_\backslash theta.$$

The two normalization terms therefore cancel outright:

$$\log p_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[\log p_{\theta }(v)]=z_{\theta }(a)-{\mathbb{E}}_{v\sim q_{\varphi }}[z_{\theta }(v)]\mathrm{.}\backslash log\; p\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[\backslash log\; p\_\backslash theta(v)]=z\_\backslash theta(a)-\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[z\_\backslash theta(v)].$$

This is very nice from an engineering point of view: for this SFC loss path, we only need the logits of the selected tokens, and we do not need to materialize the log-softmax over the whole vocabulary for its sake.

To be a little more rigorous: this does not mean the LLM forward never needs full vocabulary logits again. It only means that the SFC loss and gradient path itself can depend on a small number of selected token logits.

3.5 Purely random sampling is not good enough: why do people use head plus tail in practice?

If we simply sample at random from $q_{\varphi }q\_\backslash phi$:

$$v_1,\dots ,v_K\sim q_{\varphi },v\_1,\backslash ldots,v\_K\backslash sim\; q\_\backslash phi,$$

we do of course already get an unbiased estimator. But the next-token distribution of an LLM is often extremely uneven: a small number of top tokens can take up the vast majority of the probability mass, and behind them is a very long tail.

Suppose that at some prefix:

$$q_{\varphi }(A)=0.40,q_{\varphi }(B)=0.25,q_{\varphi }(C)=0.15,q\_\backslash phi(A)=0.40,\backslash qquad\; q\_\backslash phi(B)=0.25,\backslash qquad\; q\_\backslash phi(C)=0.15,$$

the first three tokens alone already account for 80% of the probability mass.

If we still treat every token as an ordinary Monte Carlo sample here, it is a bit wasteful: $A\mathrm{/}B\mathrm{/}CA/B/C$ clearly matter almost every time, and yet we are still leaving it to chance whether we happen to draw them.

So a very natural approach is: compute the high-probability head exactly, and use Monte Carlo sampling only for the long tail.

Suppose the generator picks the top-\(H\) tokens to form the head set:

$$\mathcal{H}=\mathrm{TopH}(q_{\varphi }),\backslash mathcal\; H=\backslash operatorname\{TopH\}(q\_\backslash phi),$$

whose total probability mass is:

$$\alpha =\sum _{h\in \mathcal{H}}{q}_{\varphi }(h)\mathrm{.}\backslash alpha=\backslash sum\_\{h\backslash in\backslash mathcal\; H\}q\_\backslash phi(h).$$

Then any expectation can be split as:

$${\mathbb{E}}_{v\sim q_{\varphi }}[f(v)]=\sum _{h\in \mathcal{H}}{q}_{\varphi }(h)f(h)+(1-\alpha ){\mathbb{E}}_{v\sim q_{\varphi }(\cdot \mid v\mathrm{\notin }\mathcal{H})}[f(v)]\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[f(v)]=\backslash sum\_\{h\backslash in\backslash mathcal\; H\}q\_\backslash phi(h)f(h)+(1-\backslash alpha)\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi(\backslash cdot\backslash mid\; v\backslash notin\backslash mathcal\; H)\}[f(v)].$$

The first term is computed exactly, and for the second we sample $MM$ tokens from the conditional tail distribution:

$$t_1,\dots ,t_M\sim q_{\varphi }(\cdot \mid v\mathrm{\notin }\mathcal{H})\mathrm{.}t\_1,\backslash ldots,t\_M\backslash sim\; q\_\backslash phi(\backslash cdot\backslash mid\; v\backslash notin\backslash mathcal\; H).$$

So:

$${\mathbb{E}}_{v\sim q_{\varphi }}[f(v)]\approx \sum _{h\in \mathcal{H}}{q}_{\varphi }(h)f(h)+(1-\alpha )\frac{1}{M}\sum _{m=1}^{M}f(t_m)\mathrm{.}\backslash mathbb\; E\_\{v\backslash sim\; q\_\backslash phi\}[f(v)]\backslash approx\backslash sum\_\{h\backslash in\backslash mathcal\; H\}q\_\backslash phi(h)f(h)+(1-\backslash alpha)\backslash frac\{1\}\{M\}\backslash sum\_\{m=1\}^Mf(t\_m).$$

In SFC, substitute:

$$f(v)=J_{\theta }(v)f(v)=J\_\backslash theta(v)$$

and you are done.

You can understand this method with one plain sentence:

I want all of the few most important tokens; for the long tail, which has tiny probabilities but an enormous number of entries, I use a few samples to stand in for it.

This head/tail sampling is nothing new in itself, it is essentially stratified Monte Carlo. But for a distribution like an LLM's, with a large vocabulary and probability highly concentrated, it is very practical.

3.6 So what does SFC actually require in a real system?

At this point we can walk through the whole procedure again in very concrete terms.

At some prefix $x_{t}x\_t$, the generator has its own behavior policy:

$$q_{\varphi }(\cdot \mathrm{\mid }{x}_t)\mathrm{.}q\_\backslash phi(\backslash cdot|x\_t).$$

It samples, as usual, one token that really enters the trajectory:

$$a_t\sim q_{\varphi }\mathrm{.}a\_t\backslash sim\; q\_\backslash phi.$$

At the same time it keeps a small number of extra tokens that describe the current distribution. The simplest option is to sample $KK$ alternative tokens directly from $q_{\varphi }q\_\backslash phi$. More realistically, it can keep the top-\(H\) head tokens and sample a few more from the conditional tail.

The generator then passes these token ids, their $q_{\varphi }q\_\backslash phi$ probabilities or log probabilities, and its own model version at the time, along with the rollout, to the trainer.

Once the trainer receives this data, it runs a normal forward with the current parameters $\theta \backslash theta$ and then gathers only:

$$z_{\theta }(a_t),z_{\theta }(v_{t,1}),\dots ,z_{\theta }(v_{t,K})\mathrm{.}z\_\backslash theta(a\_t),\backslash qquad\; z\_\backslash theta(v\_\{t,1\}),\backslash ldots,z\_\backslash theta(v\_\{t,K\}).$$

It uses these selected logits to build a sparse estimator of:

$${\mathbb{E}}_{q_{\varphi }}[z_{\theta }(v)]\backslash mathbb\; E\_\{q\_\backslash phi\}[z\_\backslash theta(v)]$$

and finally forms:

$$z_{\theta }(a_t)-{\widehat{\mathbb{E}}}_{q_{\varphi }}[z_{\theta }(v)]\mathrm{.}z\_\backslash theta(a\_t)-\backslash widehat\{\backslash mathbb\; E\}\_\{q\_\backslash phi\}[z\_\backslash theta(v)].$$

Multiply by the advantage:

$$L_t=-A_t(z_{\theta }(a_t)-{\widehat{\mathbb{E}}}_{q_{\varphi }}[z_{\theta }(v)]),L\_t=-A\_t\backslash left(z\_\backslash theta(a\_t)-\backslash widehat\{\backslash mathbb\; E\}\_\{q\_\backslash phi\}[z\_\backslash theta(v)]\backslash right),$$

and then run a normal backward.

Autograd will naturally give:

$$-A_t(J_{\theta }(a_t)-{\widehat{\mathbb{E}}}_{q_{\varphi }}[J_{\theta }(v)])\mathrm{.}-A\_t\backslash left(J\_\backslash theta(a\_t)-\backslash widehat\{\backslash mathbb\; E\}\_\{q\_\backslash phi\}[J\_\backslash theta(v)]\backslash right).$$

So if I had to sum up Score Function Centering in one sentence: the generator no longer tells the trainer only "which token I finally chose", it also brings back a little information about "what my local distribution looked like at the time"; the trainer then uses these sparse samples to turn a score that was centered on $p_{\theta }p\_\backslash theta$ into one centered on $q_{\varphi }q\_\backslash phi$, the distribution that actually generated the data. The problem it solves is:

We want to compute ${\mathbb{E}}_{q_{\varphi }}[J_{\theta }]\backslash mathbb\; E\_\{q\_\backslash phi\}[J\_\backslash theta]$, but we cannot move the complete $q_{\varphi }q\_\backslash phi$ from the generator to the trainer.

So all we can do is find a way to keep, on the generator side, a sparse approximation that is cheap enough and still informative enough.

One last boundary is worth stressing. What SFC restores is:

$${\mathbb{E}}_{a\sim q_{\varphi }}[{\tilde{s}}_{\theta ,\varphi }(a)]=0,\backslash mathbb\; E\_\{a\backslash sim\; q\_\backslash phi\}[\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)]=0,$$

that is, the zero-mean and baseline-invariance property under the behavior distribution. It does not mean that:

$${\tilde{s}}_{\theta ,\varphi }(a)={\mathrm{\nabla }}_{\theta }\log p_{\theta }(a)\mathrm{.}\backslash tilde\; s\_\{\backslash theta,\backslash phi\}(a)=\backslash nabla\_\backslash theta\backslash log\; p\_\backslash theta(a).$$

As long as:

$$q_{\varphi }\mathrm{\ne }{p}_{\theta },q\_\backslash phi\backslash neq\; p\_\backslash theta,$$

these two are generally not the same quantity. So SFC is not "we fixed the off-policy bias". More precisely, what it does is:

given that trainer-generator mismatch already exists, rebuild a centered gradient signal that is aligned with the real sampling distribution $q_{\varphi }q\_\backslash phi$, so that we at least get back the two very important structures of a zero-mean score and baseline invariance.