Rethinking Training-Inference Mismatch in LLM Reinforcement Learning: Where It Arises and How to Correct It

AI summary

Overview

Research area: Reinforcement learning for large language models (LLM RL post-training), specifically reinforcement learning with verifiable rewards (RLVR), and the numerical discrepancy between the inference engine that generates rollouts and the training engine that computes gradients.

Technical level: Intermediate. The paper assumes familiarity with policy-gradient methods (PPO-style clipping, GRPO), importance sampling, and mixture-of-experts (MoE) architectures, though its central idea is explained through a simple ratio-displacement identity.

Scope: The paper characterizes training–inference mismatch as an additive log-odds displacement, proposes a confidence-aware truncation called Calibrated Importance Sampling (CIS), and evaluates it on three mixture-of-experts models across five mathematical reasoning benchmarks.

What This Paper Is About

Modern LLM reinforcement learning pipelines split work between an inference engine (e.g., vLLM, SGLang) that samples rollouts and a separate training engine (e.g., FSDP, Megatron) that computes gradients. Even though both load identical parameters, differences in kernel implementations, reduction orders, and numerical precision mean they assign different probabilities to the same token, turning nominally on-policy training into off-policy training. The paper's goal is to explain the statistical structure of this mismatch and to design a correction that controls the resulting variance without dumping all of its truncation bias onto low-confidence tokens, as fixed-threshold methods do.

Key Contributions

1. A decomposition of the mismatch ratio. The authors prove (Lemma 3.1) that the mismatch ratio factors exactly into a token-confidence term and a discrepancy term, the log-odds displacement ε_t. Large-scale measurements show the distribution of ε_t barely varies with token confidence, and that routing disagreement between the two engines substantially amplifies its heavy tail on MoE models.

2. Calibrated Importance Sampling (CIS). CIS truncates the displacement factor rather than the ratio itself, applying a single constant threshold exp(ε_t) ≤ 1 + λ at every confidence level. Mapped back to ratio space this yields the cap min{k_t, 1 + λ(1 − p_t)}, which tightens automatically as token confidence increases.

3. A theoretical bound. The paper shows that the unbounded second-moment term governing the error of exact correction is replaced by a term bounded by (1 + λ)², at the cost of a bias controlled by the truncated excess. Any finite cap admits a similar bound; the authors argue CIS's advantage lies in where the bias falls, not in the bound itself.

4. Empirical validation on MoE models. Across three mixture-of-experts models and five mathematical reasoning benchmarks, CIS achieves the highest five-benchmark average on all three models among the evaluated baselines, together with diagnostics showing lower overall truncation bias than a fixed ratio cap.

Main Findings

• CIS leads on all three models. On Qwen1.5-MoE-A2.7B the five-benchmark average improves from 30.99 without correction to 34.78 with CIS, compared with 32.48 (Exact Ratio), 33.50 (KPop†), 34.18 (IcePop), and 31.40 (TIS). CIS reaches averages of 37.16 on DeepSeek-V2-Lite and 69.88 on Qwen3-30B-A3B, the highest mean in both blocks. On the training-domain GSM8K benchmark CIS reaches 67.88 on Qwen1.5-MoE-A2.7B.

• The variance problem is one-sided. On the sampled tokens of the trained MoE at step 86, tokens with k_t > 1 account for 99.96% of Var(k_t), whereas tokens with k_t < 1 contribute only 0.04%. This motivates truncating only large positive displacements.

• Exact correction is dramatically high-variance. The raw moment E[e^(αε_t)] is 7.07 at α = 1 but reaches 9.2 × 10⁶ at α = 2. The weighted moment E[((1 − p_t)e^(ε_t))^α], the quantity entering the bound, rises from 0.16 at α = 1 to 22.9 at α = 2.

• Displacement is nearly invariant to confidence. Across six orders of magnitude of token uncertainty, the median displacement stays at zero and its spread changes by less than a factor of two, with a Spearman rank correlation of only −0.028.

• CIS shifts truncation bias away from low-confidence tokens. Below p_t = 0.3, the bias induced by CIS is less than half that of TIS, while CIS takes on more bias at high confidence. More than half of the weight CIS removes comes from tokens with p_t < 0.9, and for p_t < 0.565 the CIS cap is looser than a fixed cap of C = 2.

• Lower overall bias and variance. CIS has overall bias 0.0064 versus 0.0078 for TIS, while its added variance, 0.0081, remains well below that of Exact Ratio at 0.0132.

• Ablations confirm the design. Ablations on Qwen1.5-MoE-A2.7B: a fixed ratio cap with C = 2 gives 31.40 average, C = 1.13 gives 8.88, a two-sided band gives 6.01 (versus CIS's 34.78), a raised floor κ = 0.1 gives 32.17, and no floor (κ = 0) gives 29.02.

• The confidence floor is necessary. Without flooring 1 − p_t, the truncation begins acting on rounding error rather than on the mismatch, and average accuracy drops from 34.78 to 29.02.

• Every positive threshold helps. The paper reports that every positive value of λ outperforms the uncorrected run; the reported configuration is λ = 2.3 with κ = 5 × 10⁻³.

• Out-of-domain results are not statistically separated. On three knowledge and two code benchmarks that the RL stage never uses, CIS has the highest pooled accuracy, but its margin over IcePop is within the seed spread. Similarly, its margins over IcePop on several transfer benchmarks are smaller than the across-seed variation.

Methodology in Plain English

The authors begin by measuring, on large-scale rollouts, the probabilities that the two engines assign to the same sampled token. They observe that in MoE models the gating network's top-k selection is a discrete decision, so a small numerical difference can change which experts are activated, and this propagates through all subsequent layers, producing a heavy-tailed distribution of per-logit perturbations. Rounding error alone accounts for the body of that distribution; routing disagreement is a major amplifier of the tail.

They then derive an exact identity expressing the mismatch ratio in terms of two interpretable quantities: the token's own probability p_t and the log-odds displacement ε_t, which is simply the difference between the log-odds the two engines assign to the sampled token. Because ε_t is determined only by the perturbation, it can be computed directly from the log-probabilities the two engines record.

Since the important ratio is a deterministic increasing function of ε_t, and since only large positive displacements can make the second moment explode, the authors truncate only the upper side: they cap exp(ε_t) at 1 + λ. Mapped back to ratio space, this becomes a per-token cap that tightens with confidence. In practice the method needs only the difference of the two log-probabilities, which stays finite even when a probability rounds to one in bf16, and it can be implemented as one elementwise pass with no extra forward or backward computation. One practical wrinkle is that when 1 − p_t gets very small, the cap approaches the storage resolution of the log-probabilities, so the authors floor 1 − p_t at κ = 5 × 10⁻³.

Evaluation uses three MoE models (Qwen1.5-MoE-A2.7B-Chat, DeepSeek-V2-Lite, Qwen3-30B-A3B), all trained on GSM8K and evaluated with greedy pass@1 on GSM8K, MATH500, SVAMP, Minerva-Math, and OlympiadBench. All methods share the same training recipe and code path, differing only in how they handle the mismatch, and results are reported as mean and standard deviation over three training seeds.

Why This Matters

Impact on research. The paper reframes training–inference mismatch from a per-token hyperparameter tuning problem into a statistical characterization problem. It shows that fixed thresholds on the importance ratio conflate two different quantities — the actual engine disagreement and the token's confidence — and that separating them yields a correction with demonstrably different bias placement. The ratio-displacement identity gives the community a common coordinate system for comparing existing methods (TIS, IcePop, KPop, Seq-TIS, Seq-MIS, GSPO) and for reasoning about why MoE models are especially vulnerable to mismatch.

Real-world applications:

• Post-training and fine-tuning pipelines for open-weight LLMs where rollout generation and gradient computation are served by separate frameworks.
• Reasoning-focused model deployment in domains such as mathematics and code, where verifiable rewards are readily available and RLVR is the dominant post-training method.
• Serving infrastructure for mixture-of-experts models, where routing decisions make numerical divergence between engines more consequential than for dense models.
• Any distributed or mixed-precision training setup where the numerics of sampling and of gradient computation diverge.

Industry relevance. The method is cheap: it adds one elementwise pass and requires no extra forward or backward computation, which matters for teams already paying the cost of decoupled rollout and training workers. Its requirement that the inference-side log-probability be the raw value recorded at the sampling step, and that the weight be detached, are concrete implementation constraints that practitioners adopting it will need to respect. Because mismatch can cause training instability or even collapse in MoE models, a correction that stabilizes training without harming throughput is directly relevant to anyone running large-scale RL post-training.

Future Directions

• Characterizing the source of the heavy tail more fully. The paper attributes the body of the perturbation distribution to rounding error and identifies routing disagreement as an amplifier of the tail, but the relative contributions are analyzed only through the diagnostic experiments in the appendices. A more complete decomposition could inform whether routing replay or precision changes should accompany CIS.

• Tying λ to the statistics of the mismatch. λ is treated as a hyperparameter, and the paper shows every positive threshold beats the uncorrected run. Whether an optimal λ can be computed from the observed displacement distribution rather than tuned is left open.

• Correlated token contributions. The theoretical bounds assume independent token contributions, with a note that n can be replaced by n_eff when contributions are correlated. The effective sample size and its behavior under correlation is not developed.

• Whether the same characterization transfers beyond math reasoning. Evaluation covers five mathematical reasoning benchmarks, three knowledge benchmarks, and two code benchmarks, with out-of-domain margins within seed spread. Broader domains and non-MoE architectures remain to be tested.

Target Audience

Researchers and engineers working on LLM reinforcement learning and post-training infrastructure, particularly those using decoupled rollout/training architectures or mixture-of-experts models. It is also relevant to readers interested in the numerics of distributed training and in the theory of importance sampling under bias–variance trade-offs. Readers without some background in policy-gradient methods and importance sampling will need to work through the preliminaries carefully before the theoretical results are accessible.

Authors’ abstract

We study training-inference mismatch in reinforcement learning with verifiable rewards (RLVR) for large language models, where rollouts are sampled by an inference engine while gradients are computed by a training engine, and the two engines assign different probabilities to the same tokens. To account for this discrepancy in policy updates, we introduce calibrated importance sampling (CIS). CIS is motivated by an empirically supported logit-displacement characterization that expresses the mismatch as an additive displacement $\varepsilon_t$ in log-odds, determined by the per-logit perturbation before the softmax, whose distribution is approximately invariant to token confidence. This characterization motivates a confidence-aware truncation: large positive displacements are truncated at a single constant threshold, which maps back to an importance-ratio cap that tightens as token confidence increases. Theoretically, we show that CIS replaces the unbounded second moment that governs the error of exact importance sampling with a term bounded by a constant, at the cost of a bias controlled by the truncated excess. In evaluation across three mixture-of-experts models and five mathematical reasoning benchmarks, CIS achieves the highest five-benchmark average on all three models among the evaluated baselines. Diagnostic analyses show that CIS places less truncation bias on low-confidence tokens than truncated importance sampling, while upward clipping of small importance weights reduces held-out accuracy.