Rethinking Training-Inference Mismatch in LLM Reinforcement Learning: Where It Arises and How to Correct It
Tianrun Yu, Kaixiang Zhao, Shangzhe Li, Yuxiao Yang, Porter Jenkins, Weitong Zhang, Taylor W. Killian
cs.LG, cs.AI
2026-09-26
CIS truncates log-odds displacement into a confidence-tightening IS cap and posts the highest five-benchmark average on three MoEs (34.78 vs 30.99 uncorrected on Qwen1.5-MoE).
RL with verifiable rewards now splits rollout and gradients across two engines. An inference stack (vLLM, SGLang) samples tokens; a training stack (FSDP, Megatron) computes gradients. Same weights, different kernels, reduction order, and numeric precision, so the two engines disagree on the probability of the same token. Nominally on-policy GRPO is already off-policy. Mixture-of-experts models make this worse: top-k routing is a discrete choice, and a small perturbation can flip experts and destabilize training.
The usual fix is importance sampling. Each token loss is reweighted by k = ptrain / pinfer. Occasional sharp disagreements blow k up, and the gradient variance with it. TIS caps k at a constant, IcePop masks tokens outside a fixed band, KPop masks on binary KL, and other lines truncate at sequence level or switch the pipeline to FP16. All of those rules act on k, or on quantities computed from the two probabilities. k mixes two things: the actual engine discrepancy, and how confident the token already is. The same discrepancy throws k far from one on low-confidence tokens and barely moves it on high-confidence ones. Fixed thresholds therefore land mostly on low-confidence tokens, and so does the truncation bias.
The mismatch has an exact identity: k = p + (1-p) exp(ε), where ε is the log-odds gap between the two engines, fixed by the per-logit perturbation before softmax. On a dense control, that perturbation sits near zero, consistent with bf16 rounding. On MoE the body is still near zero, but the tail is heavy. Across 24 layers the two engines disagree on the top-4 expert set in 3.69 layers on average; only 8.0% of positions agree everywhere. As disagreeing layers go from 0 to at least 7, the share of |δ| > 0.5 rises from 4.51% to 18.48%. Even with no routing disagreement at the current step, MoE perturbation width is still about 3× the dense control.
ε itself barely tracks token confidence. Across six orders of magnitude of 1-p, the median stays at zero, the spread changes by less than 2×, and Spearman correlation is -0.028. Variance is one-sided: after RL, tokens with k>1 carry 99.96% of Var(k). Exact correction has E[e^{2ε}] = 9.2×10^6; the weighted second moment that enters the error bound is 22.9, versus 1 with no correction.
CIS truncates only the upper tail. It caps exp(ε) at 1+λ, which maps to min(k, 1+λ(1-p)). More confident tokens get a tighter ratio cap. A floor κ=5×10^{-3} on 1-p stops the cap from shrinking below log-probability storage resolution; without it, average accuracy falls from 34.78 to 29.02. Default λ=2.3. The extra work is one elementwise pass; the weight is detached, and the inference log-probability must be the raw value recorded at sample time. The unbounded second-moment term of exact IS is replaced by something at most (1+λ)^2. The bias is the truncated excess.
Three MoE models, one GSM8K recipe (7,473 problems, three epochs, 87 steps, no KL), greedy pass@1 on five math benchmarks, three seeds. CIS has the highest five-benchmark mean on every model.
| Method | Qwen1.5-MoE | DeepSeek-V2-Lite | Qwen3-30B-A3B |
| No correction | 30.99 | 35.83 | 68.26 |
| Exact ratio | 32.48 | 35.47 | 68.72 |
| IcePop | 34.18 | 36.73 | 68.99 |
| TIS (C=2) | 31.40 | 36.02 | 69.14 |
| GSPO | 28.90 | 36.47 | 69.23 |
| CIS | 34.78 | 37.16 | 69.88 |
On Qwen1.5-MoE, CIS reaches 67.88 on GSM8K against 61.21 with no correction. FP16 averages 31.02, essentially the uncorrected run. Gaps versus IcePop on several transfer columns sit inside seed noise; the paper does not treat those columns as separated. Qwen3-30B already scores 95.22 on GSM8K before RL, so the training task has little headroom.
Ablations pin the design. Adding a lower clip drops the average from 34.78 to 6.01, below the 28.93 base model. A fixed cap C=1.13 that truncates a similar share of tokens (4.3% vs 3.8%) falls to 8.88. Removing the floor yields 29.02. λ=0 collapses to 8.75. Every positive λ beats the uncorrected 30.99; default 2.3 is best, other positive values sit between 32.53 and 34.08. Training reward for CIS is 0.867 with peak gradient norm 3.0, against 69.5 for the exact ratio. CIS acts on about 3.8% of tokens. Overall truncation bias is 0.0064 versus 0.0078 for TIS; added variance is 0.0081 versus 0.0132 for exact correction.
On MoE RLVR, training-inference mismatch is a gradient-stability problem, not leftover rounding. CIS does not change precision or replay routes. It only rescales the token loss and drops into the same token-level branch as TIS and IcePop. For teams already truncating IS, this is a mechanistic, incremental swap: clip a similar 4% of tokens, but clip in displacement space so the bias does not pile onto low-confidence tokens. It targets residual mismatch in the gradient. FP16 and routing replay still address other mechanisms. The 0.65-point edge on Qwen3-30B over the strongest baseline is small; the 3.8-point jump on the 2.7B-active model is the result that matters for adoption.
On several transfer benchmarks the IcePop gap is inside seed standard deviation. Out-of-domain knowledge and code pooled accuracy is also highest for CIS, and that margin is inside seed spread too. The routing study is observational, not an intervention, so prefix-level route drift cannot be isolated. ε is only approximately confidence-invariant and widens slightly at high confidence.
The experimental slice is narrow: 87 steps of GSM8K-only RL, all three models are MoE, dense models appear only in mismatch measurements, and routing replay is absent from the main table. k corrects the next-token distribution given the observed prefix, not the prefix law itself. The theorems assume independent token contributions; correlation replaces n with neff. λ near 0 collapses, and positive λ still swings the average by up to 2.3 points. Whether 2.3 transfers across frameworks and precisions is untested.