Rethinking Training-Inference Mismatch in LLM Reinforcement Learning: Where It Arises and How to Correct It
Organizations: Brigham Young University · University of North Carolina at Chapel Hill
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 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.
Figures & tables
| Model | Method | GSM8K | MATH500 | SVAMP | Minerva | Olympiad | Avg |
|---|---|---|---|---|---|---|---|
| Qwen1.5-MoE-A2.7B | Base model (no RL) | ||||||
| No correction | |||||||
| Exact ratio | |||||||
| KPop † | |||||||
| IcePop | |||||||
| TIS |
| Variant | GSM8K | MATH500 | SVAMP | Minerva | Olympiad | Avg |
|---|---|---|---|---|---|---|
| CIS ( , ) | ||||||
| Fixed ratio cap (TIS, ) | ||||||
| Fixed ratio cap (TIS, ) | ||||||
| Two-sided band ( ) | ||||||
| Raised floor ( ) | ||||||
| No floor ( ) |
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
| Dense | MoE | |
|---|---|---|
| Paired entries | ||
| Median | ||
| Mean | ||
| – | – | ||||
|---|---|---|---|---|---|
| Positions (%) | |||||
| (%) | |||||
| (%) |
| Dense | |||
|---|---|---|---|
| MoE, | |||
| MoE, |
| Dense | MoE | |
|---|---|---|
| Sampled tokens | ||
| Finite | ||
| Finite fraction | ||
| range | Tokens | Med. | MAD | MAD | ||
|---|---|---|---|---|---|---|
| – | ||||||
| – | ||||||
| – | ||||||
| Method | Weight on the token loss | Hyperparameters |
|---|---|---|
| No correction | (decoupled loss disabled) | — |
| Exact ratio | clamp to | |
| TIS | ||
| IcePop | ||
| KPop † | , float64 | |
| Seq-TIS | for every token of the response |
| Benchmark | Domain | Split | Scoring | Items |
|---|---|---|---|---|
| GSM8K ( Cobbe et al., 2021 ) | math | test | math-verify | 1,319 |
| MATH-500 ( Hendrycks et al., 2021 ; Lightman et al., 2024 ) | math | test subset | math-verify | 500 |
| SVAMP ( Patel et al., 2021 ) | math | test | math-verify | 300 |
| Minerva-Math ( Lewkowycz et al., 2022 ) | math | test | math-verify | 272 |
| OlympiadBench ( He et al., 2024 ) | math | English, text-only, open-ended | math-verify | 674 |
| MMLU-STEM ( Hendrycks et al., 2020 ) | knowledge | test, 19 STEM subjects | option letter | 3,153 |
| Method | Reward | Grad. norm | Peak grad. | Entropy | Length | Acted on (%) | |
|---|---|---|---|---|---|---|---|
| No correction | |||||||
| Exact ratio | |||||||
| KPop † | |||||||
| IcePop | |||||||
| TIS ( ) | |||||||
| TIS ( ) |
| GSM8K | MATH500 | SVAMP | Minerva | Olympiad | Avg | |
|---|---|---|---|---|---|---|
| (default) | ||||||
| GSM8K | MATH500 | SVAMP | Minerva | Olympiad | Avg | |
|---|---|---|---|---|---|---|
| (default) | ||||||
| Method | MMLU | ARC-C | OBQA | HumanEval | MBPP | Science | Code | All |
|---|---|---|---|---|---|---|---|---|
| Base model (no RL) | ||||||||
| No correction | ||||||||
| Exact ratio | ||||||||
| KPop † | ||||||||
| IcePop | ||||||||
| TIS |
| Method | Rollout | Rescore | Train step | Advantage | Total | vs. no corr. | Length |
|---|---|---|---|---|---|---|---|
| No correction | — | ||||||
| Exact ratio | |||||||
| KPop † | |||||||
| IcePop | |||||||
| TIS | |||||||
| Seq-TIS |
Reinforcement Learning
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
See all 1,742 papers