PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.32444v1cs.LGcs.AI

TL;DR

In RLVR, separate inference and training engines can assign different probabilities to the same tokens, creating a training–inference mismatch. CIS characterizes that discrepancy and truncates it with a single threshold that maps to confidence-dependent importance-ratio caps. Across three MoE models, CIS achieves the highest five-benchmark average among evaluated baselines.

  • Problem

    Separate inference and training engines assign different probabilities to the same tokens, making nominally on-policy RLVR off-policy and particularly challenging for MoE training.

    Source: §1 · p. 1

  • Method

    CIS truncates a confidence-independent discrepancy component at one threshold, producing importance-ratio caps that tighten as token confidence increases.

    Source: §1 · p. 2

  • Results

    CIS achieves the highest five-benchmark average among evaluated baselines on all three MoE models.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    CIS bounds the second-moment term that can destabilize exact correction, at the cost of bias controlled by the truncated excess.

    Source: §1 · p. 2, §6 · p. 10

    The stated theoretical bounds assume bounded token contributions and independence across tokens, with an effective sample-size adjustment for correlated contributions.

    Source: §3.2 · p. 5, §3.3 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Separate inference and training engines assign different probabilities to the same tokens, making nominally on-policy RLVR off-policy and motivating better mismatch correction. CIS truncates a confidence-independent discrepancy component, yielding confidence-aware ratio caps; experiments report the strongest five-benchmark average among evaluated baselines on three MoE models.

1 INTRODUCTION

  • Different engine probabilities make nominally on-policy training off-policy, and sharp token-level disagreements can make exact importance correction highly variable.The mismatch is particularly severe for MoE models, where it can cause training instability or collapse.Source: §1 · p. 1
  • The importance ratio combines engine discrepancy with token confidence, so fixed thresholds can concentrate truncation bias on low-confidence tokens.Existing methods truncate or mask ratios or related quantities using heuristic rules.Source: §1 · p. 2
  • CIS applies one threshold to the discrepancy component, yielding importance-ratio caps that tighten as token confidence increases.The discrepancy distribution varies only slightly with token confidence, so discrepancies of the same size are treated alike across confidence levels.Source: §1 · p. 2
  • Large-scale analyses find that the discrepancy component varies little with token confidence, while MoE routing disagreement amplifies its heavy tail.CIS is motivated by characterizing the mismatch itself rather than thresholding the mixed importance ratio.Source: §1 · p. 2
  • CIS achieves the highest five-benchmark average among evaluated baselines on three MoE models.Diagnostic analyses find that CIS places less truncation bias on low-confidence tokens than fixed-ratio capping and attains lower overall bias.Source: §1 · p. 3

2 PRELIMINARIES

Separate inference and training engines can assign different probabilities to the same token, making rollouts nominally on-policy but off-policy for training. Exact ratio correction is unbiased but can have uncontrolled variance, so existing methods truncate or mask ratios and accept bias.

2.1 TRAINING–INFERENCE MISMATCH

  • Separate engines can assign different probabilities to the same token, making nominally on-policy rollouts off-policy with respect to the training engine.For MoE models, a small numerical perturbation can change discrete expert-routing decisions.Source: §2.1 · p. 3
  • The mismatch ratio corrects the next-token distribution given the observed prefix, while the update ratio is handled by standard PPO clipping.The mismatch ratio does not reweight the distribution of the prefix itself.Source: §2.1 · p. 3

2.2 CORRECTING THE MISMATCH

  • Exact mismatch-ratio correction is unbiased, but sharp engine disagreements can produce very large ratios and uncontrolled gradient variance.Truncating or masking the ratio trades some bias for lower variance.Source: §2.2 · p. 3
  • Existing corrections cap or mask token- or sequence-level ratios, change the objective, or reduce engine discrepancy through precision or routing replay.The paper motivates deriving correction thresholds from the mismatch statistics rather than applying rules directly to the ratio.Source: §2.2 · p. 3

3 CALIBRATED IMPORTANCE SAMPLING

The section characterizes engine mismatch as an additive log-odds displacement and explains why exact importance correction is unstable. It then introduces CIS, which truncates positive displacement with a single threshold, bounding the second moment while trading off bias.

3.1 FROM LOGIT PERTURBATION TO LOG-ODDS DISPLACEMENT

  • The engines’ mismatch can be written as an additive log-odds displacement ε_t determined by their per-logit perturbation; MoE displacements have a heavy tail.For a given token confidence, the displacement factor e^ε_t moves the ratio more on low-confidence tokens than on high-confidence tokens.Source: §3.1 · p. 4

3.2 WHY EXACT CORRECTION IS UNSTABLE

  • At α = 2, the weighted moment reaches 22.9, and the exact-correction bound is more than twenty times the uncorrected bound.The second-moment term dominates the bound, so exact correction’s variance is driven by the upper tail of ε_t.Source: §3.2 · p. 5

3.3 CORRECTION IN THE DISPLACEMENT COORDINATE

  • 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.The reported Spearman rank correlation between |ε_t| and log(1 − p_t) is −0.028.Source: §3.3 · p. 5
  • CIS applies the same upper displacement threshold across confidence levels, which maps to a ratio cap that tightens as token confidence increases.The confidence-dependent cap follows from mapping the displacement cutoff back to ratio space.Source: §3.3 · pp. 5–6
  • CIS replaces the unbounded second-moment term with a bound of (1 + λ)^2, at the cost of bias from truncated excess.λ controls the trade-off; as λ approaches infinity, CIS recovers exact correction.Source: §3.3 · p. 6
  • The implementation floors token uncertainty at κ = 5 × 10^-3 because a narrower cap can act on log-probability rounding error and lower accuracy.Without the floor, the reported average accuracy drops from 34.78 to 29.02.Source: §3.3 · p. 6

4 EXPERIMENTS

Across three MoE models, CIS achieves the highest five-benchmark average among the evaluated methods. The experiments examine its performance, truncation choices, bias–variance behavior, and hyperparameter sensitivity.

4.1 SETUP

  • The shared evaluation trains three MoE models on GSM8K and measures greedy pass@1 on five mathematical reasoning benchmarks, reporting mean and standard deviation over three seeds.Baselines cover token-level, sequence-level, and numerical approaches; methods share the training recipe and differ in mismatch handling.Source: §4.1 · pp. 6–7, §4 · p. 6

4.2 RQ1: DOES CIS IMPROVE HELD-OUT PERFORMANCE?

  • 34.78 is CIS’s five-benchmark average on Qwen1.5-MoE-A2.7B, the highest among evaluated methods; CIS also leads on DeepSeek-V2-Lite (37.16) and Qwen3-30B-A3B (69.88).On Qwen1.5-MoE-A2.7B, the uncorrected average is 30.99; individual transfer-benchmark margins over IcePop are not interpreted as statistically separated.Source: Table 1 · p. 7, §4.2 · p. 7

4.3 RQ2: WHY DOES CIS WORK, AND WHICH DESIGN CHOICES MATTER?

  • 99.96% of Var(kt) comes from tokens with kt > 1, versus 0.04% from kt < 1; lifting small displacements to e^εt ≥1/2 lowers average accuracy from 34.78 to 6.01.CIS therefore truncates only the upper side and leaves the bounded lower side unchanged.Source: §4.3 · p. 8, Figure 5 · p. 8
  • Below pt = 0.3, CIS’s truncation bias is less than half of TIS’s, while CIS takes on more bias at high confidence.CIS thresholds the displacement factor rather than the final ratio, making the induced ratio cap adapt to token confidence.Source: §4.3 · p. 8, Figure 5 · p. 8
  • 0.0064 versus 0.0078 is the overall bias for CIS versus TIS; added variance is 0.0081 for CIS versus 0.0132 for Exact Ratio.These estimates compare bias with exact correction and added variance with no correction.Source: Figure 5 · p. 8, §4.3 · p. 9
  • 31.40 with the fixed TIS cap, 6.01 with two-sided truncation, and 29.02 without the numerical floor are below CIS’s 34.78 average.A fixed cap C = 1.13, despite truncating a similar share of tokens as CIS, gives an average of 8.88.Source: Table 2 · p. 8, §4.3 · p. 9

4.4 RQ3: HOW SENSITIVE IS CIS TO ITS HYPERPARAMETERS?

  • 34.78 is the highest average at λ = 2.3; λ = 0 gives 8.75, while κ = 5 × 10^-3 and κ = 2 × 10^-2 give 34.78 and 34.39, respectively.Removing the floor gives 29.02, while raising it to κ = 0.1 gives 32.17.Source: Figure 6 · p. 8, §4.4 · p. 9

5 RELATED WORK

Related methods address engine mismatch through numerical interventions or importance-weight truncation and masking. CIS instead derives its ratio cap from log-odds displacement.

5 RELATED WORK

  • FP16 reduces numerical error [Qi et al., 2025], while routing replay aligns expert selections [Ma et al., 2025]; precision does not enforce route agreement, and replay requires routing traces.Adaptive learning-rate decay stabilizes optimization [Zhang et al., 2026b], while RSPO addresses routing drift across policy updates [Zhang et al., 2026a].Source: §5 · p. 9
  • TIS caps token ratios, IcePop masks ratios outside fixed bounds [Zhao et al., 2025; Ling Team & Inclusion AI, 2025], and KPop uses binary-KL masking.Sequence-level alternatives truncate or optimize response ratios, and CISPO clips current-to-old policy weights [Zheng et al., 2025b; Chen et al., 2025].Source: §5 · p. 9
  • CIS derives a training–inference ratio cap from a common threshold on log-odds displacement, so the cap tightens with token confidence while retaining nonzero weight for truncated tokens.Unlike fixed ratio thresholds, this cap accounts for token-confidence heterogeneity.Source: §5 · p. 9

6 CONCLUSION

CIS uses a log-odds displacement characterization to truncate mismatch and achieves the highest five-benchmark average across three MoE models. The accompanying measurements examine how the heavy tail relates to architecture and routing disagreement.

6 CONCLUSION

  • CIS bounds the second moment that makes exact correction unstable, at the cost of bias controlled by truncated excess, and achieves the highest five-benchmark average on three MoE models.The conclusion also reports less low-confidence truncation bias than TIS and lower held-out accuracy when the lower side is truncated.Source: §6 · p. 10

A EMPIRICAL ANATOMY OF THE TRAINING–INFERENCE MISMATCH

  • The appendix measurements examine dense-versus-MoE perturbation tails, routing disagreement, propagation from per-logit perturbation δ to displacement εt, confidence dependence, and the second moment relevant to importance-sampling variance.These measurements support empirical claims used in the paper’s mismatch characterization.Source: A EMPIRICAL ANATOMY OF THE… · pp. 14–15

A.1 MEASUREMENT PROTOCOL

  • The protocol compares Qwen1.5-MoE-A2.7B-Chat with a dense Qwen1.5-14B-Chat control using identical parameters across bf16 training and inference engines.It samples 6,000 positions and records paired top-256 logits; the mean top-256 intersection rates are 0.949 on MoE and 0.991 on the dense model.Source: §A.1 · p. 15

A.2 DENSE VERSUS MOE: A HEAVY TAIL IN THE PER-LOGIT PERTURBATION

  • 9.87% of MoE paired logit entries exceed |δj| = 0.5, versus 0.003% on the dense control; the MoE standard deviation is 0.336 versus 0.079.The 99.9th percentile is 2.062 on MoE and 0.312 on dense, indicating a pronounced MoE tail.Source: Figure 7 · p. 15, §A.2 · pp. 15–16

A.5 THE HEAVY TAIL PROPAGATES FROM δ TO ε

  • 2.044 versus 0.693 is the 99.9th percentile of |εt| on MoE versus dense; the maximum is 19.38 versus 2.13, so the heavy tail persists after mapping δ to εt.The dense CCDF’s high-confidence steps reflect discrete probability numerics, not a continuous heavy tail.Source: Figure 7 · p. 15, §A.5 · pp. 17–18

A.3 ROUTING DISAGREEMENT AND THE MOE TAIL

  • From no routing disagreement to disagreement in at least seven layers, sd(δj) rises from 0.240 to 0.457 and Pr(|δj| > 0.5) rises from 4.51% to 18.48%.The association also holds under a stricter hard-flip definition based on gate margin.Source: Figure 8 · p. 16, §A.3 · pp. 16–17

A.4 CONDITIONAL ROUTING CONTROL

  • Even at Nflip = 0, MoE perturbation standard deviation is about three times the dense-control value; the measurements cannot isolate earlier-prefix effects because only 0.25% of tokens have fully route-agreeing prefixes.The authors interpret routing disagreement as a major tail amplifier, not its sole source; causal decomposition would require replaying identical routes.Source: §A.4 · p. 17, Table 5 · p. 17

A.6 THE DISPLACEMENT IS APPROXIMATELY INVARIANT TO TOKEN CONFIDENCE

The appendix tests whether log-odds displacement varies with token confidence and relates this characterization to the second-moment behavior of exact correction. It also records identities used in the proofs and the variance bound under correlated token contributions.

A.6 THE DISPLACEMENT IS APPROXIMATELY INVARIANT TO TOKEN CONFIDENCE

  • Across confidence bins, εt's median stayed within 0.006 of zero and its median absolute deviation varied 1.8-fold, versus a 3.2 × 10^5-fold change for log kt.The displacement distribution is not exactly invariant, and becomes moderately wider toward the highest-confidence region.Source: Figure 9 · p. 18, §A.6 · p. 18

A.7 THE SECOND MOMENT ON THE TRAINED MODEL

  • On the trained MoE, E[e^2εt] = 9.2 × 10^6, and the second moment is determined by the upper tail of εt.The measurement scores 510,915 sampled tokens from 4,000 GSM8K rollouts.Source: §A.7 · p. 19

B.1 PRELIMINARIES AND KEY LEMMAS

  • Under correlated token contributions, the variance bounds remain valid with n replaced by the effective sample size n_eff.Independence is used to remove off-diagonal covariance terms.Source: §B.1 · p. 22

B.2 PROOF OF THEOREM 3.2

  • Exact correction has zero truncation bias, but its variance depends on the untruncated second-moment term E[(1 − pt)^2e^2εt].The bound requires no structural assumption on εt beyond the moments it explicitly contains.Source: §B.2 · p. 23

B.3 PROOF OF THEOREM 3.3

The proof bounds CIS's second-moment contribution while controlling the bias introduced by truncation. The numerical floor preserves this variance control, and the implementation applies operators through a shared token-level branch.

B.3 PROOF OF THEOREM 3.3

  • CIS bounds the second-moment contribution by (1 + λ)^2 regardless of the positive tail of εt, while its bias is controlled by the truncated excess.The resulting trade-off is finite upper-tail bias for a uniformly bounded second-moment contribution.Source: §B.3 · p. 24

B.4 NUMERICAL FLOOR IN ALGORITHM 1

  • The numerical floor raises the cap only at extreme confidence, removes no more exact weight than the ideal operator, and keeps the applied weight bounded by 1 + λ.It prevents the confidence-dependent cap from becoming narrower than stored log-probability resolution.Source: §B.4 · pp. 24–25

C.1 IMPLEMENTATION AND HARDWARE

  • Each arm used 87 optimizer steps on one eight-A100 node; operators shared a token-level branch and differed only in token-loss weight.CIS applies min{kt, 1 + λφt} as a weight without a gradient.Source: §C.1 · p. 25

C.2 DATASETS AND EVALUATION PROTOCOL

This section specifies the training and evaluation data, task prompts, and mismatch diagnostics. Final-checkpoint results use held-out evaluation, while the static mismatch measurement compares sampled-token probabilities across engines.

C.2 DATASETS AND EVALUATION PROTOCOL

  • Training uses 7,473 GSM8K problems, and reported results come from the final checkpoint; held-out evaluation uses greedy decoding and benchmark-specific scoring protocols.The GSM8K test split is logged once per epoch but is not used to select a checkpoint.Source: §C.2 · pp. 25–26

C.4 DIAGNOSTICS

  • The static mismatch set contains 2,500 DAPO-Math-17k problems; rescoring the same sequence twice gives zero difference, and mean kt is 1.0000.The measurement uses teacher-forced rescoring of sampled tokens with the same parameter snapshot.Source: §C.2 · p. 25, §C.4 · p. 27

C.3 PROMPTS

  • Prompts request boxed math answers, option-letter responses for multiple-choice tasks, or complete Python functions for code benchmarks.The math evaluation prompt asks for step-by-step reasoning.Source: §C.3 · pp. 26–27

C.4 DIAGNOSTICS

  • Training logs record rollout-side and training-side token log-probabilities and rollout-version staleness every ten optimizer calls.The logged training-side probability is taken at the start of the update.Source: §C.4 · p. 27

C.5 PREREGISTERED EVALUATION CONTRACT

The evaluation contract fixes comparison criteria and selected settings before the corresponding runs. It uses final-checkpoint held-out accuracy as the arbiter and applies the same recipe without per-arm tuning.

C.5 PREREGISTERED EVALUATION CONTRACT

  • Held-out accuracy at the final checkpoint is the arbiter of every comparison; training reward is reported but does not select methods.Source: §C.5 · p. 27
  • The floor κ = 5 × 10^-3 is set from log-probability storage resolution before training, and the slope grid is λ ∈ {1.6, 2.3, 3.1}.Source: §C.5 · pp. 27–28
  • All arms use the same recipe without per-arm tuning, and the main comparison and Table 2 ablations use three seeds.Source: §C.5 · p. 28

D.1 TRAINING DYNAMICS

On Qwen1.5-MoE-A2.7B, token-level ratio-bounding methods reach higher training reward with lower peak gradient norms than exact correction and KPop†. The CIS hyperparameter sweeps favor the default threshold and a floor near the log-probability storage resolution.

D.1 TRAINING DYNAMICS

  • 0.859–0.867 reward with peak gradient norm below four for TIS, IcePop, and CIS; exact ratio and KPop† reach lower rewards, with peak norms of 70 and 29.The uncorrected and fp16 runs stay near reward 0.74; Seq-MIS removes 95% of responses.Source: Figure 10 · p. 28, Table 10 · p. 28, §D.1 · p. 28
  • 4%–5% of tokens are acted on by CIS after the first third of training, following an initial rise.The reported fraction then remains in that range.Source: Figure 10 · p. 28, §D.1 · p. 29

D.2 HYPERPARAMETER SWEEPS

  • 34.78 is the highest average accuracy for the default λ = 2.3; λ = 0 yields 8.75, while every positive threshold beats the uncorrected run at 30.99.Other positive thresholds range from 32.53 to 34.08.Source: Table 11 · p. 29, §D.2 · p. 29
  • 34.78 and 34.39 are reached at floors κ = 5 × 10^-3 and κ = 2 × 10^-2, respectively; removing the floor lowers the average to 29.02.A floor of 10^-3 loses 1.4 points, and κ = 0.1 gives 32.17.Source: Table 12 · p. 29, §D.2 · p. 29

D.3 OUT-OF-DOMAIN EVALUATION

On out-of-domain knowledge and code benchmarks, every trained arm finishes above the untrained model, though CIS is not separated from IcePop. The appendix also reports compute costs and differences in where the correction operators act.

D.3 OUT-OF-DOMAIN EVALUATION

  • 56.79 pooled accuracy is highest for CIS across 5,489 knowledge and code items, but its margin over IcePop at 56.60 is smaller than the standard deviations.Every trained arm ends above the untrained model.Source: Table 13 · p. 29, §D.3 · pp. 29–30

D.4 COMPUTATIONAL COST

  • 3.2% more compute time per step is reported for CIS than the uncorrected run; Seq-TIS and Seq-MIS cost 14% and 18% more, respectively.The sequence-level methods also generate longer responses, confounding direct step-time comparisons.Source: Table 14 · p. 30, Figure 11 · p. 30, §D.4 · p. 30

D.5 WHERE EACH OPERATOR ACTS

  • CIS truncates above the diagonal, while IcePop mainly masks below it; TIS acts mostly at low inference-side probability, and KPop† masks near the plane’s edges.Early in training, every operator acts on less than 1% of tokens; the operators act on different regions as the policy sharpens.Source: §D.5 · pp. 30, 32, Figure 12 · p. 31
  • 83% of CIS-truncated tokens have training probability at least 0.9, versus 25% for TIS, but more than half of CIS’s removed weight comes from tokens below 0.9.By removed weight, CIS’s truncation is not concentrated on high-confidence tokens.Source: §D.5 · p. 32
  • 0.45% is the KPop mask rate after double-precision recomputation, compared with the framework-reported 22% intervention rate attributed largely to probability saturation.The corrected variant is reported as KPop†.Source: §D.5 · p. 32

D.6 LIMITATIONS

The evaluation trains only mixture-of-experts models using bf16 and specific engines; the theory and token-level correction also have stated scope limits.

D.6 LIMITATIONS

  • Training covers only mixture-of-experts models; dense-model correction is expected to matter less, and other kernels or numerical formats may require retuning λ.The analysis does not separately model policy lag in asynchronous training; stale tokens are under 1% of tokens in these runs.Source: §D.6 · p. 32
  • The threshold and floor are swept on Qwen1.5-MoE-A2.7B and fixed for the other two models, while baselines use their reference defaults.Comparable baseline sweeps and combinations with mismatch-reduction methods are left to future work.Source: §D.6 · p. 32
  • The bounds assume bounded per-token gradients and independent token contributions; with correlated contributions, n is replaced by n_eff.The theory explains why truncation is needed but does not rank CIS against other finite caps.Source: §D.6 · p. 32
  • Token-level correction addresses each token’s conditional distribution given its prefix, not the prefix distribution itself.Sequence-level methods address the prefix, but Seq-MIS rejects 95% of responses and Seq-TIS caps 12%–17% in these runs.Source: §D.6 · p. 32

Related papers

Loading 2609.32444v1…