Score Centering targets RL gradient bias from training-inference mismatch

• Trainer-inference mismatch can move a policy even without reward differences, according to nrehiew_'s explanation.
• Score Centering subtracts that drift using an additive correction, approximated with top-128 sampler probabilities in nrehiew_'s implementation notes.
• Score Centering alone still collapsed in Modal's mixed-hardware coding run, highlighted by SemiAnalysis_; combining it with masked importance sampling finished stably.

The paper makes the failure beautifully blunt: give every rollout +1 and the model can still drift. The reference code includes PyTorch and JAX losses, and reports that saving just 32 sampler probabilities matched full-vocabulary centering in the authors' experiments.

Together AI's Martin Marek and Max Ryabinin analyze two distributions in their September paper: the rollout sampler q and the trainer p. Different kernels or numerical precision can make them disagree even with identical weights.

At a fixed prefix, the “score” means the gradient of a token's log-probability: s = ∇θ log p(token | prefix). The expected policy-gradient update splits into:

• Drift: E_q[R] × E_q[s], which depends on reward only through its mean.
• Signal: Cov_q(R, s), which captures the relationship between token choice and reward.

On-policy, E_p[s] = 0, so the drift vanishes. With positive expected reward under mismatch, drift fits the trainer toward the sampler, creating the moving-teacher effect described by nrehiew_.

Refreshing the sampler from the trainer feeds that bias into subsequent updates. Centering advantages across a prompt's rollouts does not eliminate it: the expected advantage at an individual prefix can remain positive or negative, according to the paper's prefix-level analysis.

Score Centering replaces each token score with s - E_q[s], restoring zero mean under the sampler. The practical implementation approximates that expectation without logging the sampler's entire vocabulary.

The released implementation has three mechanics:

1. Sampler logging: retain the top-128 token IDs and log-probabilities at each position, plus the sampled token's log-probability even when it falls outside that head.
2. Tail-mass ratio: for saved token set K, compute rho = (1 - Σ_K q_v) / (1 - Σ_K p_v), with small numerical floors in the code.
3. Tail approximation: model unrecorded sampler probabilities as q̂_v = rho × p_v, preserving the sampler's remaining probability mass.

For pure centering, the PyTorch loss reduces to the following, where head_logp and token_logp are trainer log-probabilities, and p_head and q_head are probabilities over the saved token IDs:

The stop-gradient freezes the correction's coefficients while leaving the trainer log-probabilities differentiable. The implementation uses float32 log-probabilities and leaves the saved top-k head unnormalized.

Exact score centering leaves Cov_q(R, s), while the on-policy update uses Cov_p(R, s), as the derivation makes explicit. Removing drift therefore leaves a difference in how tokens are weighted.

Importance sampling reweights updates toward p, but large probability ratios increase variance. Clipping or masking those ratios, the alternatives discussed in nrehiew_'s thread, bounds their contribution while reintroducing drift.

Centering can also operate on the weighted score, including the drift introduced by clipping or masking. The paper reports different results by mismatch source:

• Sampler quantization: standalone centering matched or outperformed importance-sampling baselines on Qwen3-0.6B Countdown and Qwen3-30B-A3B-Base INTELLECT-2 math experiments.
• Stale rollouts: centering combined with importance sampling outperformed pure importance-sampling baselines.

On October 13, Modal reported fully asynchronous coding-agent training of Qwen3.6-35B-A3B, with contexts up to 256K tokens. A Megatron-LM trainer ran in BF16 on B200 GPUs, while SGLang samplers used six accelerator types:

• B200 and B300: NVFP4 weights.
• H100 and H200: FP8 and BF16 weights.
• A100 and RTX PRO 6000: BF16 weights.

Every sampler used an FP8 KV cache, and training imposed no staleness limit. Only Score Centering plus masked importance sampling, with or without top-p sampling, trained stably to the end; each correction used alone ultimately finished at or below the base model.

With the full combination, SWE-bench Pro V2 pass@4 rose from 89% to about 94%. The top-p variant replayed the sampler's candidate mask in the trainer, giving both policies the same token support.

Mixed samplers matched the homogeneous B200/BF16 run's performance despite an average sampler-trainer KL 4.7 times larger. Their sampling cost per update was about 15% lower for most of training.