Score Centering: Stabilize RL Training

October 26, 2026

How Thinking Machines’ hidden StableReinforce loss was reverse-engineered from its gradients, and a derivation from first principles of the method it turned out to be: score centering, from Marek & Ryabinin (2026).

References

Martin Marek and Max Ryabinin, Score Centering Stabilizes Off-Policy Reinforcement Learning, arXiv:2609.20807, 2026.
Martin Marek, “Tinker’s new StableReinforce objective is Reinforce + Score Centering. Here’s how I reverse-engineered it”, thread on X, October 24, 2026.

Part 1The reverse-engineering trick

A hidden loss function

On October 21, Thinking Machines added a new loss to Tinker, their post-training API, called StableReinforce. The pitch is that RL runs normally collapse when the policy that generates the data drifts away from the policy being trained, through async staleness or mismatched numerics, and that this loss keeps them stable. The docs say it “has led to extremely stable training, even under highly off-policy async RL and low precision numerics.” They do not say what the loss is.

Three days later, Martin Marek, first author of the score-centering paper, reverse-engineered the loss from the outside and posted the result in a thread on X. StableReinforce turns out to be REINFORCE plus score centering, the correction his paper introduces. The rest of Part 1 follows the thread: two tricks that recover the raw gradient of a hidden loss, and the clue they turned up, a single-token update that moves only eight logits. Part 2 goes back to the paper, builds score centering from scratch following its Section 4, and ends by showing why that clue gives the objective away.

Reverse-engineering a loss from its gradients

Tinker lets you send data, take an optimiser step and download the updated weights, but you never see the loss. Every loss leaves a fingerprint, though: the gradient it produces. If you can read off the raw gradient, you can write down a candidate loss, compute its gradient, and compare. That needs two things: the raw gradient from a weight update, and a place in the network where the gradient is easy to interpret.

Trick 1: make Adam behave like SGD

With plain SGD, the weight change is the gradient times a constant: $\Delta w = -\eta\, g$ . Tinker doesn’t offer SGD, only Adam, and Adam hides the gradient. On the first step its bias-corrected moments are exactly $\hat m = g$ and $\hat v = g^2$ , so the update is:

\Delta w \;=\; -\eta\, \frac{\hat m}{\sqrt{\hat v} + \epsilon} \;=\; -\eta\, \frac{g}{|g| + \epsilon}.

With the usual tiny $\epsilon$ (around $10^{-8}$ ) this is $-\eta\,\operatorname{sign}(g)$ . Every weight moves by the same amount and only the sign of the gradient survives, so the update tells you almost nothing about the loss.

The trick is to set $\epsilon = 10{,}000$ . Real gradients are tiny next to that, so $|g| + \epsilon \approx \epsilon$ and:

\Delta w \;\approx\; -\frac{\eta}{\epsilon}\, g \qquad\Longrightarrow\qquad g \;\approx\; -\frac{\epsilon}{\eta}\,\Delta w .

Adam has been turned into SGD with learning rate $\eta/\epsilon$ . Download the weights before and after, subtract, and rescale: that is the raw gradient. The relative error is about $|g|/\epsilon$ , which is negligible. You also need weight decay off, so that nothing else moves the weights.

Figure 1. Recovering a gradient from one Adam step, for eight weights with a made-up gradient. The panels show what you can infer as $-\Delta w/\eta$ , rescaled. With the default $\epsilon$ , every weight moves by the same amount and only the sign survives. With $\epsilon = 10{,}000$ , the update is proportional to $g$ .

Trick 2: a rank-1 LoRA exposes the logit gradient

Even with the raw gradient in hand, billions of numbers spread across every layer are hard to interpret. The cleanest place to look is the logits, because that is where a loss function acts directly. The idea is to train a rank-1 LoRA adapter on the unembedding, the matrix $W$ that turns the final hidden state $h$ into logits. A rank-1 adapter adds $B A h$ , with $A$ a $1 \times d$ row and $B$ a $V \times 1$ column:

z \;=\; W h \;+\; B\,\underbrace{(A h)}_{a\ \in\ \mathbb{R}} .

The term $a = A h$ is a single number. Write $g_z = \partial L / \partial z$ for the gradient with respect to the $V$ logits. By the chain rule:

\frac{\partial L}{\partial B} \;=\; g_z \cdot a .

The gradient on $B$ is the logit gradient multiplied by one scalar. For a single-token update you can read the whole logit gradient straight off $B$ , up to that scale factor. Which logits are non-zero, and their ratios, are exact. Combined with trick 1, a downloaded weight update now tells you exactly how the hidden loss pushes on every one of the vocabulary’s logits.

Figure 2. A rank-1 LoRA on the unembedding. The adapter’s path squeezes through a single number $a$ , so the gradient reaching $B$ (dashed) is the full logit gradient $\partial L/\partial z$ times one scalar.

What the gradient showed

For a standard REINFORCE loss, every logit gets a non-zero gradient. (We derive why below: the softmax couples them all.) In StableReinforce’s single-token update, only eight logits had a non-zero gradient. Marek’s hypothesis was REINFORCE plus score centering, with the centering term estimated from eight tokens sampled from the inference engine. He implemented that as a custom Tinker loss, took a step, and found the recovered gradients matched StableReinforce’s “within numeric precision across all layers.”

Figure 3. The logit gradient for one token over a toy 40-token vocabulary, with advantage $+1$ . Standard REINFORCE moves every logit, because the softmax couples them. Score centering with 8 sampled tokens cancels the shared $-p$ part, leaving only the sampled token and the correction tokens. Illustrative numbers, not Marek’s measurements.

Two caveats. This is a reverse-engineering result, not something Thinking Machines has stated, though a match across every layer is strong evidence. And the Tinker docs also list an optional sampler–learner KL term (default coefficient 0.01) that the thread doesn’t discuss.

To see why “eight non-zero logits” gives the objective away, we need to understand what score centering is and what problem it solves. The rest of this post builds that from scratch.

Part 2What score centering actually does

This part follows Section 4 of the paper. It builds score centering from first principles, compares it with importance sampling, and ends by coming back to the eight logits from Part 1.

Two copies of the same model

RL post-training uses two copies of the model. The sampler $q$ is an inference engine (vLLM, SGLang) that generates rollouts. The trainer $p$ runs those finished rollouts through a training framework to compute gradients. Splitting the work is a large speed win, because generation goes one token at a time while scoring a finished sequence takes one parallel pass. (Appendix A goes through why the split exists.)

The two copies are meant to be identical, but in practice they aren’t: lower-precision inference, different kernels, different batch shapes, or weights that are a few steps stale. This gap is the training-inference mismatch, and Section 4 of the paper is about exactly what it breaks.

Notation, at one token position with a fixed prefix:

Symbol	Meaning
$p_v,\ q_v$	probability of vocabulary token $v$ under the trainer and the sampler (one number each)
$p,\ q$	the full next-token distributions at this position: vectors of all $V$ probabilities, $p = (p_1, \dots, p_V)$
$s_v = \nabla_\theta \log p_v$	the score of token $v$ : the direction that makes it more likely
$y_t$	the token the sampler actually picked
$R$	the reward or advantage
$\bar s = \sum_v q_v\, s_v$	the average score when tokens are drawn by the sampler

The vanilla policy-gradient update for one token is $R\, s_{y_t}$ : if the reward was good, make that token more likely.

On-policy, the average score is zero

If the trainer drew the tokens itself, the average score would be:

\begin{aligned} \mathbb{E}_{p}[s] &= \sum_v p_v \,\nabla \log p_v = \sum_v p_v \cdot \frac{\nabla p_v}{p_v} = \sum_v \nabla p_v = \nabla \Big(\sum_v p_v\Big) = \nabla(1) = 0 . \end{aligned}

Probabilities sum to one, so raising one token’s probability must lower the others, and on average the pushes cancel. A useful consequence: if every rollout gets the same reward, the expected update is $R \cdot 0 = 0$ . A constant reward carries no information, and on-policy you correctly learn nothing from it.

Under mismatch, the update splits into drift and signal

Now the tokens come from $q$ but the scores are computed through $p$ . The cancellation above needed both to be the same distribution, so in general $\bar s = \sum_v q_v \nabla \log p_v \neq 0$ . Using the covariance identity $\mathbb{E}[AB] = \mathbb{E}[A]\,\mathbb{E}[B] + \operatorname{Cov}(A,B)$ with $A = R$ and $B = s_{y_t}$ , the expected update at a prefix splits into two parts (the paper’s Eq. 3):

\underbrace{\mathbb{E}_q[R\, s_{y_t}]}_{\text{update}} \;=\; \underbrace{\mathbb{E}_q[R]\;\bar s}_{\text{drift}} \;+\; \underbrace{\operatorname{Cov}_q(R,\, s_{y_t})}_{\text{signal}} . \tag{3}

The signal asks whether the tokens that led to high reward differ from the ones that led to low reward. That is the learning we want. The drift is the mean reward times the average score. It doesn’t know which token earned which reward, so it is a pure artefact of the mismatch, and on-policy it vanishes.

A two-token example

Take a model with a two-token vocabulary $\{A, B\}$ whose only parameters are the two logits $z_A, z_B$ . This example is mine, not the paper’s. The softmax gives $p_A = e^{z_A}/(e^{z_A} + e^{z_B})$ , so $\log p_A = z_A - \log(e^{z_A} + e^{z_B})$ . Differentiating with respect to each logit:

\frac{\partial \log p_A}{\partial z_A} = 1 - p_A, \qquad \frac{\partial \log p_A}{\partial z_B} = -p_B .

The same pattern holds for any vocabulary size $V$ . Differentiating $\log p_v$ with respect to logit $z_u$ gives $1 - p_u$ when $u = v$ and $-p_u$ otherwise. Written as a vector over all $V$ logits:

s_v = \mathbf{e}_v - p ,

where $\mathbf{e}_v$ is the one-hot vector for token $v$ : length $V$ , with a 1 in $v$ ‘s slot and 0 everywhere else. In the two-token model, $\mathbf{e}_A = (1, 0)$ and $\mathbf{e}_B = (0, 1)$ . So the score says: push the sampled token’s logit up by 1, and push every logit down by its current probability. Note that this is dense. Even though only $A$ was sampled, every logit moves. That is the “standard REINFORCE puts a gradient on all logits” from the thread.

Let the trainer be $p = (0.5, 0.5)$ and the sampler, slightly off from quantisation, $q = (0.6, 0.4)$ . Then $s_A = (0.5, -0.5)$ (raise $A$ , lower $B$ ) and $s_B = (-0.5, 0.5)$ . The average score depends on who draws the tokens:

Tokens drawn by	Average score	Result
trainer, $p = (0.5, 0.5)$	$0.5\,(0.5,-0.5) + 0.5\,(-0.5,0.5)$	$(0,\ 0)$
sampler, $q = (0.6, 0.4)$	$0.6\,(0.5,-0.5) + 0.4\,(-0.5,0.5)$	$(0.1,\,-0.1)$

This matches the shortcut $\bar s = q - p$ . With a constant reward of $+1$ there should be no update. Instead the logits move toward $A$ . The pushes for $A$ and $B$ are the same size, because the trainer computes them believing the odds are 50/50. But the sampler decides how often each push happens, and it picks $A$ 60% of the time. With learning rate 1, the logits go to $(0.1, -0.1)$ and $p_A$ goes from 0.5 to about 0.55, halfway toward the sampler. The model learned nothing about the task. It just copied the sampler.

Figure 4. The two-token example with a constant reward of $+1$ . Each bar is one token’s push on $A$ ‘s logit (±0.5) times how often that token is drawn. When the trainer draws, the pushes cancel. When the sampler draws, $A$ ‘s push wins, and the net update (◆) is the drift, $+0.10$ .

Drift is the trainer copying the sampler, and the sampler keeps moving

So what is the drift actually doing? It turns out to be something familiar. If you wanted to train $p$ to imitate $q$ , you would use the cross-entropy loss $\mathrm{CE}(q,p) = -\sum_v q_v \log p_v$ . Its gradient is $-\sum_v q_v \nabla \log p_v$ , which is exactly $-\bar s$ . In other words, $\bar s$ points the same way as an SFT step on the sampler’s outputs. Under mismatch, every policy-gradient step quietly includes a little “make the trainer look more like the sampler,” scaled by the average reward. Nobody asked for it, and it has nothing to do with the task.

On its own that wouldn’t be a disaster. Copying a fixed teacher is just distillation: you get a bit closer every step and eventually match it. The trouble is that the sampler isn’t a fixed teacher. It’s a slightly distorted copy of the trainer, and after every step it is refreshed with the trainer’s new weights, distortion and all:

Every time the trainer catches up, the sampler is refreshed from it and the distortion is added back, so the target stays one step ahead.

Follow it round once. The trainer takes a step toward the sampler. Its new weights are copied into the sampler, which adds its usual distortion on top. The trainer then steps toward that, and so on. The target never stays still long enough to be reached, so instead of settling down, the error builds up step after step. This is what the paper means when it says drift accumulates over training.

It also clears up something that looks like a paradox. SFT on text written by a completely different model, which is a huge mismatch, works fine. Yet online RL can fall apart over tiny numerical differences. What matters isn’t how far apart the two models are. It’s whether the teacher stays put or follows the student around. The same idea explains why all-positive ( $+1/0$ ) rewards were the least stable online setting in the paper: the drift is scaled by the average reward, and with no negative rewards to balance it out, that average is as large as it gets.

Solution: score centering

The fix is to make the average score zero again, as it is on-policy, by subtracting it (Eq. 4):

\tilde s_{y_t} = s_{y_t} - \bar s, \qquad \mathbb{E}_q[\tilde s] = \bar s - \bar s = 0 . \tag{4}

Both averages use the same distribution $q$ , so the cancellation is exact even off-policy. The expected update becomes (Eq. 5):

\begin{aligned} \mathbb{E}_q[R\, \tilde s_{y_t}] &= \mathbb{E}_q[R]\;\underbrace{\mathbb{E}_q[\tilde s_{y_t}]}_{=\,0} \;+\; \operatorname{Cov}_q(R,\, \tilde s_{y_t}) \\ &= \operatorname{Cov}_q(R,\, s_{y_t} - \bar s) \;=\; \operatorname{Cov}_q(R,\, s_{y_t}) . \end{aligned} \tag{5}

The last step holds because $\bar s$ depends only on the prefix, not on the sampled token, so it is a constant as far as the covariance is concerned. The drift is gone and the signal is untouched. The one remaining difference from a true on-policy update is that the covariance is measured under $q$ rather than $p$ . When the two agree, $\bar s = 0$ and score centering does nothing.

In the toy model, $\tilde s_A = (0.4, -0.4)$ and $\tilde s_B = (-0.6, 0.6)$ , and under the sampler they average to $0.6 \cdot 0.4 + 0.4 \cdot (-0.6) = 0$ .

The one-hot vectors from the two-token section give an even simpler form. Averaging one-hot vectors over the sampler’s choices just gives back the sampler’s probabilities, $\sum_v q_v \mathbf{e}_v = q$ . (In the toy model, $0.6\,(1, 0) + 0.4\,(0, 1) = (0.6, 0.4)$ .) So the average score is

\bar s = \sum_v q_v (\mathbf{e}_v - p) = \sum_v q_v \mathbf{e}_v - p \sum_v q_v = q - p ,

and the centered score of the sampled token $y$ is

\tilde s_{y} = (\mathbf{e}_y - p) - (q - p) = \mathbf{e}_y - q .

The trainer’s $p$ cancels out completely. Check it on the toy model: $\mathbf{e}_A - q = (1, 0) - (0.6, 0.4) = (0.4, -0.4)$ , the same $\tilde s_A$ as above. In words: push the sampled token up by 1, and push every token down by how often the sampler picks it. Hold on to this formula; it is the key to the eight logits.

Try it

The two-token model from above. Move the sliders to set the trainer and the sampler, and see the expected update on $A$ ‘s logit for each method ( $B$ ‘s is always the negative). The dashed line is the on-policy target, what you would get with no mismatch at all. Hover a row for details. The importance-sampling rows, and the TIS cap slider, are explained in the next section; they’re worth coming back to after reading it.

Trainer p_A:

Sampler q_A:

TIS cap c:

Reward

Presets:

Method ← toward Btoward A → Update Gap

pushes toward A pushes toward B on-policy target gap = update − target

With a constant reward, vanilla policy gradient shows the full drift $q_A - p_A$ , while score centering sits on zero however far apart the sampler and trainer are. With a token-dependent reward, score centering always points the right way but lands short of the target by an amount that grows with the mismatch. That is the “covariance under $q$ ” caveat from Eq. 5.

The importance-sampling rows preview the next section. Exact importance sampling also sits on zero, and so does TIS until its cap binds: with $p_A = 0.5$ and a cap of 2, that happens once $q_A > 0.75$ , after which drift leaks back in and TIS + score centering removes it again. One thing the playground can’t show is variance. These are expected values, so exact importance sampling always looks perfect here, even though variance is exactly what sinks it in practice.

Doesn’t importance sampling already fix this?

Exact importance sampling does. The idea is to reweight each sampled token by how much more or less likely the trainer is to produce it than the sampler, using the importance ratio

r_t = \frac{p(y_t)}{q(y_t)} .

The per-token loss multiplies the usual policy-gradient loss by this ratio, held fixed with a stop-gradient:

L_{\text{IS}} = -\,\mathrm{sg}[r_t]\; R\, \log p(y_t) \qquad\Longrightarrow\qquad -\nabla L_{\text{IS}} = r_t\, R\, s_{y_t} .

Tokens the sampler over-picks ( $q > p$ ) get a weight below 1, and tokens it under-picks get a weight above 1. Under a constant reward the average update is

\sum_v q_v \cdot \frac{p_v}{q_v} \cdot s_v = \sum_v p_v\, s_v = 0,

so the $q$ ‘s cancel and the drift is gone. In the toy model the ratios are $r_A = 0.5/0.6 = 0.83$ and $r_B = 0.5/0.4 = 1.25$ , and they cancel the drift exactly. The problem is variance. For a token the sampler is much less likely to produce than the trainer ( $q \ll p$ ), the ratio $p/q$ can be enormous, even if both probabilities are small: $p = 0.01$ and $q = 0.0001$ give $r = 100$ . That token is rarely sampled, but when it is, it can dominate a step. In the paper’s experiments the raw version (“naive IS”) collapses under staleness and ends below 5% at the heaviest 30B quantisation.

So nobody uses it raw. The most common fix is truncated importance sampling (TIS), which caps the weight at a constant $C$ :

w_t = \min(r_t,\ C), \qquad L_{\text{TIS}} = -\,\mathrm{sg}[w_t]\; R\, \log p(y_t),

with $C = 2$ in the paper’s experiments. Other methods limit the ratio in other ways. Masked importance sampling (MIS) sets $w_t = 0$ whenever $r_t$ falls outside a range ( $[0.5, 5]$ in the paper), and PPO-style objectives clip $r_t$ to $[1-\varepsilon,\ 1+\varepsilon]$ . All of them trade some correctness for lower variance. Once the weights are altered, the $q$ ‘s no longer cancel and drift leaks back in.

Take a badly skewed sampler, $q = (0.9, 0.1)$ , with $p = (0.5, 0.5)$ . The ratios are $r_A = 0.5/0.9 = 0.56$ and $r_B = 0.5/0.1 = 5$ . TIS caps $r_B$ at 2, so the average push on $A$ ‘s logit under a constant reward is

\underbrace{0.9 \times 0.56 \times 0.5}_{A\ \text{sampled}} \;+\; \underbrace{0.1 \times 2 \times (-0.5)}_{B\ \text{sampled}} \;=\; 0.25 - 0.10 \;=\; 0.15 .

Without the cap, B’s term would be $0.1 \times 5 \times (-0.5) = -0.25$ and the total would be exactly zero.

The table below does the same calculation for each method. Weights (A, B) are the weights $w_A$ and $w_B$ that the method multiplies the update by when $A$ or $B$ is sampled. The average push on $A$ ‘s logit adds up each token’s contribution: how often the sampler draws it ( $q_v$ ), times its weight, times its push on $A$ ‘s logit ( $+0.5$ for $A$ , $-0.5$ for $B$ ), with $R = 1$ :

\text{average push} = q_A\, w_A \,(+0.5) \;+\; q_B\, w_B\,(-0.5) = 0.9\, w_A\,(0.5) - 0.1\, w_B\,(0.5).

Method	Weights (A, B)	Calculation	Average push on A
Vanilla policy gradient	1, 1	$0.45 - 0.05$	0.40
Exact importance sampling	0.56, 5	$0.25 - 0.25$	0
TIS, cap at 2	0.56, 2	$0.25 - 0.10$	0.15
TIS + score centering	0.56, 2	$0.15 - 0.15$	0

In the last row, score centering subtracts the average weighted score, which is exactly the 0.15 that TIS leaves behind.

Clipping cuts the drift from 0.40 to 0.15, but not to zero. The two corrections also work differently. Importance sampling is multiplicative, random, and unbounded. Score centering is additive and fixed once you know the prefix. Because they are independent, they compose: apply clipped weights $w_v$ , then subtract the average weighted score $\sum_v q_v w_v s_v$ (Appendix A.2 of the paper).

Method	Drift	Variance
Vanilla policy gradient	full	low
Exact importance sampling	none	very high
Clipped or masked IS (TIS, MIS, PPO-style)	partial	moderate
Score centering	none	low
Clipped IS + score centering	none	moderate

So what does the paper add?

Two things. The first is a diagnosis. Importance sampling fixes several things at once, so its success never told us which problem it was fixing. Score centering removes only the drift, with no ratios, no clipping and no masking, and that alone is enough to stabilise training under severe quantisation. That identifies drift as the cause of the instability.

The second is a method that holds up when the mismatch gets bad. On Qwen3-30B-A3B with math RL, under three levels of sampler quantisation:

Sampler	Outcome
FP8 weights and activations, FP8 KV cache	even uncorrected policy gradient trains stably (58%); nothing to separate
FP8 + FP4 KV cache	policy gradient collapses, MIS collapses late; score centering 52%, TIS 51%
INT8 + INT4 KV cache	score centering 30%, TIS 12%, every other method below 5%

The gap grows as the mismatch gets worse. Under heavy staleness (the sampler refreshed only every 64 steps), TIS + score centering beat score centering alone. The authors attribute this to the Eq. 5 caveat: when the trainer moves far from the sampler, measuring the signal under $q$ starts to cost something, and importance sampling helps. Their recommendation is to compose the two whenever the trainer can move far between syncs.

The settings were made severe on purpose. Under mild mismatch the best methods overlapped within the authors’ compute budget, and they argue that mild mismatch causes the same problem over longer runs. The experiments cover Qwen3 models on Countdown and math only, with some single-seed 30B runs.

Making it cheap: top-k and a tail model

Score centering needs $\bar s$ as an actual number at every position. As a loss, with $\mathrm{sg}$ meaning stop-gradient:

L = -R\Big(\log p_{y_t} - \sum_v \mathrm{sg}[q_v]\,\log p_v\Big) \quad\Longrightarrow\quad -\nabla L = R\,(s_{y_t} - \bar s) = R\,\tilde s_{y_t}.

The trainer already has every $\log p_v$ , because its softmax runs over the whole vocabulary. The problem is $q$ . Storing the sampler’s full 152K-entry distribution for every token would take about 20 TB per batch. So the sampler logs only its top $k = 128$ probabilities (the head $H$ ), and the tail $T$ is filled in with the trainer’s own probabilities, rescaled to the right mass (Eq. 7):

\hat q_v = \begin{cases} q_v & v \in H \\ \rho\, p_v & v \in T \end{cases} \qquad \rho = \frac{1 - \sum_{v\in H} q_v}{1 - \sum_{v\in H} p_v}. \tag{7}

Why model the tail at all, instead of just using the top 128? Because the head alone doesn’t sum to one. Treating the missing mass as zero, or renormalising the head to cover it, both misstate the sampler’s distribution, so the subtracted term no longer has exactly zero mean and a little drift survives. Filling the tail with the trainer’s probabilities, which are already computed, only assumes the two engines agree on the tail’s shape up to its total mass, and Eq. 8 shows it costs nothing extra.

Since $\sum_v p_v s_v = 0$ , the tail sum equals minus the head sum, $\sum_T p_v s_v = -\sum_H p_v s_v$ , and the whole expectation collapses onto the head (Eq. 8):

\mathbb{E}_{\hat q}[s] = \sum_{H} q_v s_v + \rho \sum_{T} p_v s_v = \sum_{v\in H} (q_v - \rho\, p_v)\, s_v . \tag{8}

That gives the paper’s implementation loss (Eq. 6), a sum over just 128 tokens per position:

L = -R\Big(\log p_{y_t} - \sum_{v\in H} \mathrm{sg}\big[q_v - \rho\, p_v\big]\,\log p_v\Big). \tag{6}

def score_centering_loss(train_logp, samp_logp, topk_ids, sampled_token, advantage):
    train_head_logp = train_logp[topk_ids]
    tail_mass_ratio = (1 - jnp.exp(samp_logp).sum()) / (1 - jnp.exp(train_head_logp).sum())   # ρ
    head_prob_residual = jnp.exp(samp_logp) - tail_mass_ratio * jnp.exp(train_head_logp)      # q_v − ρ p_v
    logp_correction = (stop_gradient(head_prob_residual) * train_head_logp).sum()
    return -advantage * (train_logp[sampled_token] - logp_correction)                         # Eq. 6

The reported overhead is under 1% of wall-clock time, and $k = 32$ already matches full-vocabulary centering in every setting they tested.

Back to the eight logits

We can now decode the thread. In logit space, the centered score is always $\mathbf{e}_y$ minus some estimate $\hat q$ of the sampler’s distribution. What differs between implementations is which estimate you use.

The paper’s top-k version. Expanding Eq. 8 with $s_v = \mathbf{e}_v - p$ , and using $\sum_H (q_v - \rho p_v) = 1 - \rho$ :

\tilde s_y = \mathbf{e}_y - \sum_{v \in H} (q_v - \rho\, p_v)\,\mathbf{e}_v - \rho\, p \;=\; \mathbf{e}_y - \hat q .

The $-\rho\, p$ term touches every logit, so this gradient is dense. It is deterministic, and slightly biased wherever the tail model is wrong.

StableReinforce’s sampled version. Draw $K = 8$ tokens $v_1, \dots, v_K$ from the sampler and estimate the average score by their mean:

\bar s \approx \frac{1}{K}\sum_{j=1}^{K} s_{v_j} = \frac{1}{K}\sum_{j} \mathbf{e}_{v_j} - p \quad\Longrightarrow\quad \tilde s_y = (\mathbf{e}_y - p) - \Big(\frac{1}{K}\sum_j \mathbf{e}_{v_j} - p\Big) = \mathbf{e}_y - \frac{1}{K}\sum_{j=1}^{K}\mathbf{e}_{v_j}.

The dense $-p$ part is shared by every token’s score, so it cancels exactly. What remains is non-zero only on the sampled token and the correction tokens: a handful of logits out of 152K. In Marek’s single-token test that came to eight. Since $\mathbb{E}[\mathbf{e}_{v_j}] = q$ , the estimate is unbiased; it averages to $\mathbf{e}_y - q$ .

Figure 5. Three versions of the centered logit gradient for the same toy token as Figure 3. The paper’s top-k version (here $k = 8$ to match the sample budget; the paper uses 128) is dense and nearly identical to the exact one. The 8-sample version is sparse and noisy, but equal to the exact one on average.

	Paper: top-k + tail	StableReinforce: 8 sampled tokens
Estimate of $q$	top-128 probabilities, trainer fills the tail	empirical distribution of 8 draws
Bias	slight (tail model)	none
Variance	none (deterministic)	some, but bounded: each term is a one-hot vector
Logit gradient	dense	sparse (sampled token + correction tokens)
What the sampler logs	128 logprobs per position	8 token ids per position

Estimating $\bar s$ by sampling might look like it reintroduces the variance problem. It does add noise, but of a very different kind from importance sampling’s. Every term in the estimate is a one-hot vector, so nothing can blow up the way $p/q$ does on a rare token. Thinking Machines evidently decided that trade was worth it.

Summary

Mismatch adds a drift term. Under training-inference mismatch, the policy gradient contains $\mathbb{E}_q[R]\,\bar s$ , which ignores which token earned which reward.
Drift is distillation toward a moving teacher. The trainer is pulled toward a biased copy of itself that is re-synced every step, so the error compounds instead of converging.
Score centering removes it exactly. Subtracting the average score $\bar s$ cancels drift at every prefix and leaves the signal intact, measured under $q$ instead of $p$ .
Importance sampling fixes it only in theory. Exact ratios cancel drift but are too noisy, and the clipped versions everyone uses let some of it back. The two corrections compose, and the combination is best under heavy staleness.
It is cheap. The expected score is a sum over the next token’s distribution, computable from the top-128 sampler logprobs (the paper) or estimated from 8 sampled tokens (StableReinforce, as reverse-engineered).

The one-sentence version: when the sampler and trainer disagree, the policy gradient secretly includes an “imitate the sampler” term that compounds every step, and subtracting the average score removes exactly that term while keeping all the real learning signal.

Appendix A: why there is a sampler and a trainer at all

The short answer is speed. Generating is sequential, because each token depends on the previous one, so a 1,000-token answer needs 1,000 forward passes in a row. Scoring a finished sequence is parallel: all tokens are known, so one forward pass with a causal mask computes $\log p(y_t \mid y_{\lt t})$ at every position, and one backward pass gives all the gradients.

Figure 6. Why the work is split. Generating needs one forward pass per token, each waiting for the last. Scoring the finished sequence is a single parallel pass, which is also where the gradients come from.

	Sampler (inference engine)	Trainer (training engine)
Examples	vLLM, SGLang	FSDP, Megatron, JAX trainers
Job	write rollouts, token by token	score finished rollouts and compute gradients
Built for	throughput: KV cache, continuous batching, fused kernels, often lower precision	backprop: keeps activations, holds optimiser state, shards the model
Outputs	tokens, plus the logprobs it used ( $q$ )	$\log p$ with an autograd graph

You can’t skip the trainer’s forward pass and reuse the sampler’s numbers. Gradients need the trainer’s autograd graph, and the sampler’s numbers are slightly wrong anyway, which is the whole problem. Splitting the work buys a large speed-up, because generation usually dominates RL wall-clock time. The price is that the model that wrote the tokens isn’t exactly the model being trained. The paper’s position is that it’s better to keep the speed and correct the drift in the loss.

Appendix B: why is there a sum over the vocabulary?

A natural objection to everything above: when you actually train, the loss is computed on the one token the sampler picked, A or B, never on the whole vocabulary. So where does a sum over A and B ever happen?

The short answer is that in vanilla training it never does, at least not explicitly. The sum is shorthand for what you get when you add up lots of single-token updates. (One small correction to the setup: RL rollouts are sampled with some randomness, not decoded greedily. With greedy decoding the sampler would pick A every time. Sampling is what makes A come up 60% of the time and B 40%.)

The sum happens by adding up updates

Go back to the two-token model: trainer $p = (0.5, 0.5)$ , sampler $q = (0.6, 0.4)$ , constant reward $+1$ . Each sampled token gives one ordinary single-token update: $s_A = (+0.5, -0.5)$ if A was picked and $s_B = (-0.5, +0.5)$ if B was. Note that both logits move even though only one token was sampled. That’s the softmax: raising one token’s probability means lowering the others.

A gradient step doesn’t use one token, though. It averages over every token in the batch. Suppose this prefix comes up in ten rollouts:

Draw	1	2	3	4	5	6	7	8	9	10
Token	A	A	B	A	B	A	A	B	A	B
Push on $z_A$	+0.5	+0.5	−0.5	+0.5	−0.5	+0.5	+0.5	−0.5	+0.5	−0.5

Adding them up gives $6 \times 0.5 - 4 \times 0.5 = +1.0$ , or $+0.1$ per token. That is exactly $\sum_v q_v s_v = 0.6 \times 0.5 + 0.4 \times (-0.5)$ . Nobody computed a sum over the vocabulary: it came out of adding ten single-token updates, with A appearing more often because the sampler picks it more often. If the trainer had done the sampling you’d expect five of each, and the pushes would cancel to zero.

In real RL the same prefix never repeats, but that doesn’t matter. Every position shares the same weights, so the small same-direction pushes from thousands of different positions add up in the same way, which is also why SFT works on data where no prefix is ever seen twice.

The one place the sum is actually computed

Score centering subtracts $\bar s = (0.1, -0.1)$ from every single-token update. Each update becomes $(+0.4, -0.4)$ when A is picked and $(-0.6, +0.6)$ when B is, and the same ten draws now net to $6 \times 0.4 - 4 \times 0.6 = 0$ , as they should. To subtract $\bar s$ , though, you have to know it as a number, which means knowing the sampler’s probabilities, not just the token it picked. That is the only point where a sum over the vocabulary is evaluated explicitly.

Here is what that looks like at a realistic position. The prefix is The answer is, the sampler picked 42, and the top of the menu has these made-up numbers:

Token	Sampler $q_v$	Trainer $p_v$	$c_v = q_v - \rho\, p_v$
42 (picked)	0.70	0.64	+0.06
41	0.12	0.15	−0.03
forty	0.05	0.07	−0.02
newline	0.03	0.04	−0.01
rest of vocabulary	0.10	0.10	tails match, $\rho = 1$

GRPO’s loss uses one line, $-A \log p(\text{“42”})$ . The score-centered loss uses that line plus a few others at the same position:

L = -A\Big(0.94 \log p(\text{“42”}) + 0.03 \log p(\text{“41”}) + 0.02 \log p(\text{“forty”}) + 0.01 \log p(\text{“newline”})\Big).

The picked token gets slightly less credit because the sampler over-picks it, and the under-picked tokens get a little back. If the two engines agreed, every $c_v$ would be zero and this would be the GRPO loss exactly.

In summary

Where	What happens to the sum over the vocabulary
Vanilla training	Never computed. It appears implicitly when many single-token updates are added up, each token weighted by how often the sampler picks it.
The analysis in the paper	Written out explicitly, to describe what that adding-up produces on average. This is how the drift is found.
The score-centering correction	Actually computed, because the offset $\bar s$ has to be known before it can be subtracted.

The computed sum is cheap. The trainer’s softmax already sums over all 152K logits to produce $\log p(\text{“42”})$ , so reading off a few more $\log p_v$ at the same position costs almost nothing. The sampler only has to log its top- $k$ logprobs, or, in StableReinforce’s case, a few sampled token ids.

Appendix C: how the loss relates to GRPO and other methods

The paper compares every correction on top of one shared objective, REINFORCE with group-centered rewards (their Eq. 13):

J_{\text{base}}(\theta) = \frac{\sum_i A_i \sum_{t=1}^{L_i} \log p_\theta(y_{i,t} \mid y_{i,\lt t})}{\sum_i L_i}, \qquad A_i = R_i - \bar R_{\text{group}} .

Here $i$ indexes rollouts, $L_i$ is a rollout’s length, and $\bar R_{\text{group}}$ is the mean reward of the 8 completions for the same prompt. Against that base, each method only changes two things in the per-token loss: a weight $w_t$ on the sampled token, and a centering term $c_v$ .

L_t = -A\Big(\mathrm{sg}[w_t]\,\log p(y_t) \;-\; \sum_{v\in H}\mathrm{sg}[c_v]\,\log p_v\Big)

Method	Weight $w_t$	Centering $c_v$	Drift
REINFORCE (the base)	$1$	$0$	full
Naive IS	$r_t = p(y_t)/q(y_t)$	$0$	none, but huge variance
TIS	$\min(r_t, 2)$	$0$	partial
MIS	$r_t$ if $0.5 \le r_t \le 5$ , else $0$	$0$	partial
PPO / GRPO clip	$r_t$ , or $0$ where the clip binds	$0$	partial
Score centering	$1$	$q_v - \rho\, p_v$	none
TIS + score centering	$\min(r_t, 2)$	$q_v w_v - \alpha\, p_v$	none

For the combined row, $w_v = \min(p_v/q_v, 2)$ and $\alpha = \rho \min(1/\rho, 2)$ . The centering term is the average of the weighted score instead of the plain one.

PPO and GRPO belong in the importance-sampling rows. Their clipped objective contains $r_t \hat A$ , and since $\nabla r = r\, \nabla \log p$ , its gradient is $\hat A\, r_t\, \nabla \log p(y_t)$ : policy gradient with importance weight $r_t$ , set to zero wherever the clip binds. That makes them clipped importance sampling, with partial drift. In the paper, PPO and DAPO survived staleness but collapsed under quantisation and weight noise, consistent with a clip designed for ratios that come from policy updates rather than numerical error.

So score centering isn’t a new RL algorithm. It is one extra term in the per-token loss of whatever you already run, it costs a sum over a few logprobs at positions you are already training on, and it disappears when the sampler and trainer agree.

The paper and thread are recent. The equations here were transcribed from the arXiv PDF and the X thread, so check them against the originals before relying on exact notation. The toy examples, the “42” menu and the playground are my own illustrations, not from the paper.