Hugging Face Journal Club: How Score Centering Stabilizes Offline Reinforcement Learning

In reinforcement learning training for LLMs, mismatches between the inference engine and the training engine often lead to training instability. A paper recently discussed at the Hugging Face Journal Club mathematically analyzes the root cause of this mismatch and proposes a correction method called "Score Centering" that can significantly improve training stability.

Root Cause: Mismatch Between Inference and Training

The paper first defines the sampler (Q) and the trainer (P). Ideally, the two should be completely identical, but in reality, due to factors such as quantization and asynchronous updates, P and Q are not equal. This mismatch introduces a "drift" term, causing model weights to be erroneously updated even without a learning signal.

The speaker demonstrates how the drift term arises by deriving the policy gradient loss. When P≠Q, the expected gradient contains a covariance-related term, and the goal of the correction method is to eliminate this drift.

Intuitive Difference Between the Covariance Term and the Drift Term

In policy gradients, the signal we truly want to exploit is the covariance between "the change in log probability of an action" and "the return": the gradient direction is meaningful only when the policy increases the probability of high-return actions. The drift term is different—it is unrelated to the return and is produced solely by the systematic difference between the two distributions P and Q. It can be understood as a kind of "spurious gradient": even if all sampled returns are identical (no learning signal), this term remains nonzero and continuously pushes the weights in some direction. Under extreme quantization, the tail difference between P and Q is amplified, and the drift term may dominate the total gradient, manifesting as sudden loss divergence or policy collapse. The key observation of Score Centering is: the drift term originates from the nonzero expectation of the score function under Q, so subtracting this expectation exactly cancels it without affecting the covariance signal.

Why P and Q Become Inconsistent

Ideally, the inference engine and the training engine use the same set of weights and the same numerical precision, in which case P=Q. But in real systems there are multiple sources of inconsistency: the inference side often uses quantization (e.g., 8-bit weights, 4-bit KV cache) to reduce memory usage and latency, and quantization error shifts the output distribution; in asynchronous training, the weights used by the inference engine lag several steps behind the training engine, forming a "stale policy"; in addition, kernel implementations and batching methods across different frameworks also introduce subtle numerical differences. These factors combine to make the sampling distribution Q systematically deviate from the training distribution P. Score Centering does not attempt to eliminate these engineering differences, but instead mathematically compensates for their effect on the gradient, so it applies to all the above types of mismatch.

Mathematical Correction: Score Centering

The core idea is to center the score function, i.e., subtract its expected value from the raw score. The paper proves that with the corrected score, the drift term in the expected gradient is completely canceled, leaving only the covariance term (the learning signal).

Computing the correction term requires traversing the entire vocabulary, which is computationally expensive. To this end, the authors propose a top-k approximation: consider only the k tokens with the highest probabilities, and adjust the mass of the tail probabilities to ensure the accuracy of the approximation.

What Is the Score Function

In reinforcement learning and variational inference, the score function refers to the gradient of the log probability with respect to the parameters, i.e., ∇θ log πθ(a|s). It measures "how a small change in parameters would change the probability of an action being selected." The policy gradient theorem states that the gradient of the expected return can be written as the expectation of the product of the return and the score function, so the score is the bridge connecting sampling and parameter updates. When the sampling distribution Q is inconsistent with the training distribution P, we are actually computing the expectation of the score under Q, and it deviates from the true score under P—this is exactly the source of the drift term. Score Centering estimates the expectation of the score under Q and subtracts it from each sample's score, making the corrected score have zero mean under Q. In this way, any constant offset unrelated to the return can no longer contribute to the gradient, and only the covariance part that varies with the return remains.

The Principle of the top-k Approximation

Computing the exact expectation of the score requires summing over the entire vocabulary (typically tens of thousands to hundreds of thousands of tokens), which is too costly to do on every update. The top-k approximation exploits an empirical fact: the output distribution of language models is usually highly concentrated, with most of the probability mass falling on a few high-probability tokens. By computing the score and its weighted expectation only for these k tokens, and either redistributing the remaining tail probability mass proportionally or simply ignoring it, the full expectation can be approximated with very small error. In experiments, k=8 already covers most of the probability mass, and k=16 gives even higher coverage, while the computation drops by one to two orders of magnitude compared with the full vocabulary. This allows Score Centering to be embedded in the regular training loop without significantly slowing down throughput.

Practical Results: Stable Training and Efficient Inference

Experiments show that Score Centering remains stable under extreme quantization scenarios (e.g., 8-bit weights, 4-bit KV cache), while other methods collapse. In addition, the method is compatible with techniques such as truncated importance sampling and can be used in combination.

Community members also found that with the top-k approximation, k=8 covers most of the probability mass, k=16 gives higher coverage, and performance overhead can be reduced through optimization.

Conclusion

Score Centering provides a simple and effective stabilization solution for offline reinforcement learning, especially suitable for quantized inference and asynchronous training scenarios. The paper's mathematical derivation is rigorous and its experimental validation is thorough, and it is expected to become a default practice in RLHF training.

One-line Summary

Score Centering eliminates the drift term caused by inference-training mismatch by centering the score function, stabilizing LLM reinforcement learning training.

Paragraph Summary

The article discusses the training instability in LLM reinforcement learning caused by inconsistency between the inference engine and the training engine. The paper attributes the difference between the sampler Q and the trainer P to factors such as quantization and asynchronous updates, and proves that this mismatch introduces a return-independent drift term into the policy gradient, erroneously updating weights even without a learning signal. The authors propose the Score Centering method, which subtracts the expectation of the raw score under Q, mathematically canceling the drift term completely and retaining only the covariance learning signal. To address the high cost of full-vocabulary computation, the paper adopts a top-k approximation; experiments show that k=8 covers most of the probability mass. The method remains stable under extreme quantization such as 8-bit weights and 4-bit KV cache, and is compatible with truncated importance sampling.