Score Centering Stabilizes Off-policy Reinforcement Learning

Paper: 2609.20807 Authors: Martin Marek, Max Ryabinin Categories: cs.LG

The Gap

Reinforcement learning for large language models (PPO, GRPO, REINFORCE) is bottlenecked by rollout throughput. Generating rollouts from the full-precision training model in native PyTorch is painfully slow. Modern RL pipelines decouple training from inference: high-throughput serving systems (such as vLLM or TensorRT-LLM) generate rollouts using aggressive optimizations—FP8 or INT4 quantization, custom fused kernels, and asynchronous weight updates.

This creates Training-Inference Mismatch (TIM): the rollout distribution $\pi_{\text{rollout}}$ deviates subtly but persistently from the training policy $\pi_\theta$.

Conventionally, practitioner folklore attributes RL collapse under TIM to high variance or out-of-distribution token probabilities, attempting to patch it with importance sampling (IS) ratios $\frac{\pi_\theta(a|s)}{\pi_{\text{rollout}}(a|s)}$. But IS ratios explode under quantization discrepancies and require storing per-token logits from the inference engine, adding memory and network bottlenecks.

   THE TRAINING-INFERENCE MISMATCH (TIM) TRAP

   Rollout Engine (vLLM / FP8 Quant / Async)
                     |  Generates trajectories
                     v
             Divergence / Mismatch
                     |
                     v
   Training Engine (BF16 PyTorch / Gradient Accumulation)
                     |
   +-------------------------------------------------+
   | Conventional View: Mismatch = High Variance     |
   | Common Fix: Importance Sampling Ratios          |
   | Consequence: Ratio explosion under quantization |
   +-------------------------------------------------+
                     |
                     v
   PAPER'S CORE INSIGHT: Mismatch = Persistent Systematic DRIFT
   Discrepancies don't average out; they accumulate step-over-step
                     |
                     v
   SOLUTION: Additive Score Centering Correction
   Cancels the mean drift directly without unstable ratios

The Increment

One sentence: By identifying persistent gradient drift—rather than random variance—as the root cause of RL failure under training-inference mismatch, the authors derive an additive score-centering term that stabilizes training across 0.6B to 30B parameter scales while bypassing the ratio-explosion problems of importance sampling.

Core Mechanism

Under standard policy gradients, the gradient estimator is: $g = \mathbb{E}_{\tau \sim \pi_{\text{rollout}}} [R(\tau) \nabla_\theta \log \pi_\theta(\tau)]$

When $\pi_{\text{rollout}} \neq \pi_\theta$, taking the expectation under the rollout distribution introduces a persistent non-zero bias: $\mathbb{E}_{\tau \sim \pi_{\text{rollout}}} [\nabla_\theta \log \pi_\theta(\tau)] \neq 0$

Under on-policy conditions ($\pi_{\text{rollout}} = \pi_\theta$), the expected score function $\mathbb{E}[\nabla \log \pi] = 0$ holds identically. But under TIM, this expectation is non-zero and points consistently in the direction of the inference engine’s systematic distortions (e.g. rounding bias in FP8 quantization or stale parameter differences). As training progresses, this error acts like a constant non-zero force pushing parameters off the optimization cliff.

Score Centering subtracts this expected score vector directly: $\tilde{g} = \sum_{t} \left( R_t \nabla_\theta \log \pi_\theta(a_t|s_t) - \bar{R} \cdot \mu_{\text{score}} \right)$

where $\mu_{\text{score}}$ tracks the empirical mean of the score function under rollout samples.

   SCORE CENTERING VS IMPORTANCE SAMPLING

   Raw Policy Gradient under TIM:
     Gradient = Signal + [Persistent Systematic DRIFT] -> Training Diverges

   Importance Sampling Fix:
     Gradient = Signal * [Ratio w] -> w explodes when rollout != train

   Score Centering (This Paper):
     Gradient = Signal + Drift - [Estimated Drift Center] -> Zero-Mean Robust

The structural metaphor is a rowing boat steering across a river with a constant cross-current.

• The target destination is policy improvement.
• Random turbulence is gradient variance (normal SGD noise).
• The cross-current is TIM drift caused by quantized rollouts: every stroke pushes the boat subtly toward the bank. If you ignore the current, even if your strokes are powerful, you end up miles downstream on the rocks.
• Importance sampling tries to look through high-powered magnifying glasses at every tiny wave, multiplying each stroke’s angle by an erratic compensation factor that frequently violently jerks the oars.
• Score centering simply measures the steady current once and offsets the rudder by that exact angle: a calm, additive counter-torque that cancels the drift while keeping every stroke smooth.

Key Concepts

• Training-Inference Mismatch (TIM): The systematic difference between the execution graph and arithmetic precision used to generate trajectory tokens versus the graph used to compute backward passes.
• Persistent Drift vs. Noise: Zero-mean noise in stochastic gradients averages out over batches. Systematic bias from numerical precision or parameter staleness has a non-zero mean and compounds monotonically across optimization steps.
• Additive vs. Multiplicative Correction: Importance sampling uses multiplicative weights ($w = p/q$) which suffer from unbounded variance when support differs. Score centering is additive ($g - c$), guaranteeing stable magnitude regardless of distribution divergence.

Framework Shift

Before (Multiplicative Ratio Compensation):
  Quantized Rollout -> Compute per-token logits -> Form ratio p_train / p_rollout
  -> Ratios explode under FP8/INT4 -> Optimization collapses unless clipped hard

After (Additive Score Centering):
  Quantized Rollout -> Compute standard score vector -> Subtract estimated drift center
  -> Additive correction preserves gradient scale -> Stable training from 0.6B to 30B

From “correcting rollout probabilities by multiplying per-token likelihood ratios,” the core shift is centering the gradient score vector additively to cancel systematic engine drift.

Expert Assessment

Problem choice: Critical for real-world scaling. Modern post-training clusters burn hundreds of thousands of GPU hours running vLLM rollouts paired with Megatron/FSDP training. TIM is one of the dirtiest open secrets in LLM post-training, frequently causing silent training divergence or forcing teams onto expensive full-precision rollout setups.

Method maturity: Grounded and mathematically clean. Moving from multiplicative ratios to additive centering of the score function is intuitive yet theoretically sound. The property that score centering composes gracefully with importance sampling when handling stale checkpoints makes it directly applicable to asynchronous distributed RL.

Experimental integrity: Strong breadth of scale. Experiments span small models (0.6B) up to 30B parameters. The evaluation tests both quantization-induced TIM (FP8/INT8) and staleness-induced TIM (delayed actor updates), clearly charting where score centering alone suffices and where hybrid composition with IS provides the best stability.

Writing quality: Transparent and precise. The distinction between variance-induced failure and drift-induced failure is articulated crisply with both theoretical derivations and empirical trajectories.

Verdict: strong accept — A principled, low-overhead fix for a pervasive systems bottleneck in large-scale LLM reinforcement learning.

Takeaways

• If your RL run becomes unstable when switching to quantized rollout engines, do not assume you need smaller learning rates or tighter PPO clipping; suspect systematic TIM drift.
• Implement score centering by tracking running score means $\mathbb{E}[\nabla \log \pi]$ during batch evaluation.
• For asynchronous RL with stale rollout actors, compose score centering with moderate importance sampling rather than relying on importance sampling alone.