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.LGSep 26, 2026

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

Authors: Tianrun Yu, Kaixiang Zhao, Shangzhe Li, Yuxiao Yang, Porter Jenkins, Weitong Zhang, Taylor W. Killian

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 $\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.

Figures & tables

Figure 1: The inference and training engines assign different probabilities to the same tokens. CIS truncates the importance ratios with confidence-dependent caps, which are tighter for more confident tokens. The truncated ratios reweight the token-level GRPO loss ( Shao et al., 2024 ) .

Figure 2: (a) Per-logit perturbation

\delta_{j}

and (b) log-odds displacement

\varepsilon_{t}

, MoE vs. dense.

Figure 3:

\mathbb{E}[e^{\alpha\varepsilon_{t}}]

(left axis) and

\mathbb{E}[((1-p_{t})e^{\varepsilon_{t}})^{\alpha}]

(right axis) on the measured displacements of the trained MoE.

Figure 4: $\varepsilon_{t}$ against token uncertainty $1-p_{t}$ on the base MoE; the black line is the median.

Model	Method	GSM8K	MATH500	SVAMP	Minerva	Olympiad	Avg
Qwen1.5-MoE-A2.7B	Base model (no RL)	$55.50$	$21.20$	$59.00$	$4.04$	$4.90$	$28.93$
	No correction	$61.21_{\pm 1.17}$	$21.33_{\pm 1.40}$	$64.33_{\pm 1.73}$	$3.68_{\pm 1.27}$	$4.40_{\pm 0.45}$	$30.99_{\pm 0.22}$
	Exact ratio	$63.18_{\pm 2.97}$	$22.00_{\pm 1.22}$	$67.11_{\pm 3.75}$	$5.51_{\pm 0.64}$	$4.60_{\pm 0.15}$	$32.48_{\pm 1.21}$
	KPop †	$64.80_{\pm 1.29}$	$23.67_{\pm 1.29}$	$68.67_{\pm 3.71}$	$5.15_{\pm 1.47}$	$5.24_{\pm 0.67}$	$33.50_{\pm 1.04}$
	IcePop	$66.16_{\pm 0.96}$	$24.33_{\pm 1.03}$	$69.78_{\pm 1.02}$	$5.51_{\pm 0.64}$	$5.09_{\pm 1.11}$	$34.18_{\pm 0.55}$
	TIS	$61.03_{\pm 5.11}$	$20.73_{\pm 3.92}$	$66.78_{\pm 5.09}$	$4.54_{\pm 0.76}$	$3.91_{\pm 0.31}$	$31.40_{\pm 2.47}$

Table 1: Held-out accuracy (%) on five math benchmarks after RL on GSM8K. Mean

\pm

standard deviation over three seeds; bold marks the best mean per column within each model.

Figure 5: (a) Share of

\mathrm{Var}(k_{t})

on each side of

\log k_{t}=0

. (b) Truncation bias across confidence bins. (c) Overall bias relative to exact correction. (d) Added variance relative to no correction.

Variant	GSM8K	MATH500	SVAMP	Minerva	Olympiad	Avg
CIS ( $\lambda=2.3$ , $\kappa=5\times 10^{-3}$ )	$\mathbf{67.88}_{\pm 1.64}$	$\mathbf{24.53}_{\pm 1.03}$	$\mathbf{70.22}_{\pm 1.35}$	$\mathbf{5.88}_{\pm 0.64}$	$\mathbf{5.39}_{\pm 0.62}$	$\mathbf{34.78}_{\pm 0.25}$
Fixed ratio cap (TIS, $C=2$ )	$61.03_{\pm 5.11}$	$20.73_{\pm 3.92}$	$66.78_{\pm 5.09}$	$4.54_{\pm 0.76}$	$3.91_{\pm 0.31}$	$31.40_{\pm 2.47}$
Fixed ratio cap (TIS, $C=1.13$ )	$14.99_{\pm 17.04}$	$5.73_{\pm 5.06}$	$20.33_{\pm 20.53}$	$1.47_{\pm 1.91}$	$1.88_{\pm 1.86}$	$8.88_{\pm 9.24}$
Two-sided band ( $w_{t}\geq(1+p_{t})/2$ )	$5.31_{\pm 0.33}$	$4.40_{\pm 0.40}$	$15.67_{\pm 1.86}$	$1.47_{\pm 0.64}$	$3.21_{\pm 0.45}$	$6.01_{\pm 0.50}$
Raised floor ( $\kappa=0.1$ )	$64.42_{\pm 1.67}$	$22.60_{\pm 4.06}$	$64.44_{\pm 1.95}$	$4.90_{\pm 0.85}$	$4.50_{\pm 0.37}$	$32.17_{\pm 1.28}$
No floor ( $\kappa=0$ )	$55.62_{\pm 2.69}$	$20.87_{\pm 1.10}$	$59.78_{\pm 2.83}$	$5.88_{\pm 0.37}$	$2.97_{\pm 0.15}$	$29.02_{\pm 1.43}$

Table 2: Ablations of CIS on Qwen1.5-MoE-A2.7B. Held-out accuracy (%) under the protocol of Table 1 ; mean

\pm

standard deviation over three seeds.

Figure 6: (a) Threshold

\lambda

and (b) floor

\kappa

on Qwen1.5-MoE-A2.7B.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Complementary cumulative distributions of (a)

|\delta_{j}|

and (b)

|\varepsilon_{t}|

on the dense and MoE models.

	Dense	MoE
Paired entries	$1{,}521{,}672$	$1{,}457{,}477$
Median	$0.0000$	$0.0000$
Mean	$-0.0013$	$-0.0036$
$\|\mathrm{mean}\|/\mathrm{sd}$	$1.6\%$	$1.1\%$
$\mathrm{sd}(\delta_{j})$	$0.079$	$0.336$
$q_{99}(\|\delta_{j}\|)$	$0.250$	$1.144$

Appendix

Table 3: Per-logit perturbation

\delta_{j}

on the base checkpoints.

Figure 8: (a) Standard deviation of

\delta_{j}

and (b) share of entries with

|\delta_{j}|>0.5

against the number of layers with routing disagreement; dashed: dense model.

$N_{\mathrm{flip}}$	$0$	$1$	$2$ – $3$	$4$ – $6$	$\geq 7$
Positions (%)	$8.0$	$15.6$	$33.1$	$28.2$	$15.1$
$\mathrm{sd}(\delta_{j})$	$0.240$	$0.269$	$0.307$	$0.352$	$0.457$
$\Pr(\|\delta_{j}\|>0.5)$ (%)	$4.51$	$5.42$	$7.85$	$11.86$	$18.48$
$\Pr(\|\delta_{j}\|>1)$ (%)	$0.47$	$0.68$	$1.11$	$1.65$	$3.73$

Appendix

Table 4: Per-logit perturbation against the number of layers

N_{\mathrm{flip}}

with routing disagreement.

	$\mathrm{sd}(\delta_{j})$	$\Pr(\|\delta_{j}\|>0.5)$	$\Pr(\|\delta_{j}\|>1)$
Dense	$0.079$	$0.003\%$	$0.0001\%$
MoE, $N_{\mathrm{flip}}=0$	$0.240$	$4.51\%$	$0.47\%$
MoE, $N_{\mathrm{flip}}>0$	$0.343$	$10.35\%$	$1.62\%$

Appendix

Table 5: Per-logit perturbation on the dense model and on MoE positions with and without routing disagreement.

	Dense	MoE
Sampled tokens	$12{,}389{,}115$	$12{,}293{,}115$
Finite $\varepsilon_{t}$	$11{,}934{,}872$	$12{,}226{,}092$
Finite fraction	$96.3\%$	$99.5\%$
$\mathrm{sd}(\varepsilon_{t})$	$0.0997$	$0.2947$
$q_{95}(\|\varepsilon_{t}\|)$	$0.165$	$0.599$
$q_{99}(\|\varepsilon_{t}\|)$	$0.336$	$1.055$

Appendix

Table 6: Log-odds displacement

\varepsilon_{t}

in the static measurement on the base checkpoints.

Figure 9: (a) Median, 10–90% and 1–99% intervals of

\varepsilon_{t}

and (b) median absolute deviation of

\varepsilon_{t}

and

\log k_{t}

against

1-p_{t}

on the base MoE.

$p_{t}$ range	Tokens	Med. $\varepsilon$	MAD $(\varepsilon)$	$q_{05}(\varepsilon)$	$q_{95}(\varepsilon)$	MAD $(\log k)$
$p_{t}<0.5$	$2{,}824{,}021$	$-0.0059$	$0.156$	$-0.375$	$0.259$	$1.15{\times}10^{-1}$
$0.5$ – $0.9$	$2{,}337{,}554$	$-0.0011$	$0.175$	$-0.376$	$0.355$	$3.96{\times}10^{-2}$
$0.9$ – $0.99$	$1{,}951{,}864$	$-0.0005$	$0.186$	$-0.453$	$0.420$	$6.43{\times}10^{-3}$
$0.99$ – $0.999$	$1{,}549{,}915$	$-0.0019$	$0.203$	$-0.492$	$0.472$	$6.78{\times}10^{-4}$
$p_{t}>0.999$	$3{,}562{,}738$	$0.0000$	$0.237$	$-0.504$	$0.501$	$9.90{\times}10^{-6}$

Appendix

Table 7: Conditional statistics of

\varepsilon_{t}

and

\log k_{t}

on the base MoE.

Method	Weight on the token loss	Hyperparameters
No correction	$1$ (decoupled loss disabled)	—
Exact ratio	$k_{t}$	clamp to $[10^{-6},10^{6}]$
TIS	$\min(k_{t},C)$	$C=2$
IcePop	$k_{t}\,\mathbf{1}\{k_{t}\in[\alpha,\beta]\}$	$[\alpha,\beta]=[0.5,5]$
KPop †	$k_{t}\,\mathbf{1}\{\max(\mathrm{KL}_{\mathrm{B}}(p\\|q),\mathrm{KL}_{\mathrm{B}}(q\\|p))\leq\phi\}$ , float64	$\phi=2$
Seq-TIS	$\min(K,C)$ for every token of the response	$C=2$

Appendix

Table 8: Operators and hyperparameters of the arms in Table 1 .

K=\prod_{t}k_{t}

is the product over a response and

\varphi_{t}=\max(1-p_{t},\kappa)

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

Appendix

Table 9: Evaluation benchmarks.

Panel of Figure 10 — Figure 10: Training dynamics on Qwen1.5-MoE-A2.7B (mean over seeds, band spanning the seeds, five-step smoothing). (a) Reward. (b) Gradient norm. (c) Entropy. (d) Response length. (e) Fraction of tokens acted on, counting whole responses for Seq-TIS and Seq-MIS. (f) Largest $|\log k_{t}|$ in the batch.

Method	Reward	Grad. norm	Peak grad.	Entropy	Length	Acted on (%)	$\max_{t}\|\log k_{t}\|$
No correction	$0.732_{\pm 0.004}$	$0.75_{\pm 0.02}$	$2.5_{\pm 0.4}$	$0.220_{\pm 0.005}$	$322_{\pm 2}$	$0.00_{\pm 0.00}$	$9.5_{\pm 0.5}$
Exact ratio	$0.829_{\pm 0.012}$	$3.58_{\pm 0.75}$	$69.5_{\pm 41.4}$	$0.102_{\pm 0.013}$	$227_{\pm 4}$	$0.06_{\pm 0.03}$	$25.9_{\pm 3.4}$
KPop †	$0.847_{\pm 0.016}$	$3.09_{\pm 0.76}$	$29.2_{\pm 23.5}$	$0.093_{\pm 0.033}$	$245_{\pm 29}$	$1.33_{\pm 0.05}$	$27.3_{\pm 1.7}$
IcePop	$0.859_{\pm 0.005}$	$1.12_{\pm 0.04}$	$2.8_{\pm 1.5}$	$0.113_{\pm 0.035}$	$224_{\pm 7}$	$3.64_{\pm 0.40}$	$24.0_{\pm 2.3}$
TIS ( $C=2$ )	$0.861_{\pm 0.014}$	$0.97_{\pm 0.11}$	$3.6_{\pm 1.6}$	$0.064_{\pm 0.004}$	$226_{\pm 9}$	$0.76_{\pm 0.07}$	$37.7_{\pm 4.4}$
TIS ( $C=1.13$ )	$0.780_{\pm 0.065}$	$0.91_{\pm 0.08}$	$8.5_{\pm 11.7}$	$0.200_{\pm 0.142}$	$224_{\pm 23}$	$4.32_{\pm 0.64}$	$29.7_{\pm 4.4}$

Appendix

Table 10: Training statistics on Qwen1.5-MoE-A2.7B, mean

\pm

standard deviation over three seeds: third-epoch mean (median for the gradient norm and

\max_{t}|\log k_{t}|

), and peak gradient norm and fraction acted on over the whole run.

$\lambda$	GSM8K	MATH500	SVAMP	Minerva	Olympiad	Avg
$0$	$14.71_{\pm 13.47}$	$6.80_{\pm 5.67}$	$19.33_{\pm 17.17}$	$0.98_{\pm 1.39}$	$1.93_{\pm 1.12}$	$8.75_{\pm 7.44}$
$0.45$	$64.82_{\pm 1.88}$	$23.33_{\pm 2.44}$	$70.33_{\pm 2.08}$	$7.11_{\pm 0.21}$	$3.66_{\pm 0.17}$	$33.85_{\pm 0.27}$
$0.89$	$66.59_{\pm 1.95}$	$22.87_{\pm 1.75}$	$70.22_{\pm 1.90}$	$6.25_{\pm 1.33}$	$4.45_{\pm 0.83}$	$34.08_{\pm 0.19}$
$1.6$	$65.68_{\pm 0.46}$	$23.33_{\pm 1.63}$	$68.78_{\pm 2.17}$	$6.74_{\pm 0.56}$	$4.90_{\pm 0.59}$	$33.89_{\pm 0.46}$
$2.3$ (default)	$67.88_{\pm 1.64}$	$24.53_{\pm 1.03}$	$70.22_{\pm 1.35}$	$5.88_{\pm 0.64}$	$5.39_{\pm 0.62}$	$34.78_{\pm 0.25}$
$3.1$	$65.07_{\pm 0.70}$	$23.67_{\pm 0.92}$	$67.89_{\pm 2.50}$	$6.13_{\pm 1.49}$	$4.01_{\pm 0.90}$	$33.35_{\pm 0.66}$

Appendix

Table 11: Accuracy (%) of CIS on Qwen1.5-MoE-A2.7B for different thresholds

\lambda

with

\kappa=5\times 10^{-3}

. Mean

\pm

standard deviation over three seeds.

$\kappa$	GSM8K	MATH500	SVAMP	Minerva	Olympiad	Avg
$0$	$55.62_{\pm 2.69}$	$20.87_{\pm 1.10}$	$59.78_{\pm 2.83}$	$5.88_{\pm 0.37}$	$2.97_{\pm 0.15}$	$29.02_{\pm 1.43}$
$10^{-3}$	$64.01_{\pm 0.65}$	$24.00_{\pm 0.20}$	$68.67_{\pm 0.67}$	$6.86_{\pm 0.21}$	$3.51_{\pm 0.09}$	$33.41_{\pm 0.36}$
$5\times 10^{-3}$ (default)	$67.88_{\pm 1.64}$	$24.53_{\pm 1.03}$	$70.22_{\pm 1.35}$	$5.88_{\pm 0.64}$	$5.39_{\pm 0.62}$	$34.78_{\pm 0.25}$
$2\times 10^{-2}$	$65.93_{\pm 2.31}$	$24.73_{\pm 0.95}$	$70.56_{\pm 2.46}$	$7.11_{\pm 0.21}$	$3.61_{\pm 0.09}$	$34.39_{\pm 1.20}$
$0.1$	$64.42_{\pm 1.67}$	$22.60_{\pm 4.06}$	$64.44_{\pm 1.95}$	$4.90_{\pm 0.85}$	$4.50_{\pm 0.37}$	$32.17_{\pm 1.28}$

Appendix

Table 12: Accuracy (%) of CIS on Qwen1.5-MoE-A2.7B for different floors

\kappa

with

\lambda=2.3

. Mean

\pm

standard deviation over three seeds.

Method	MMLU	ARC-C	OBQA	HumanEval	MBPP	Science	Code	All
Base model (no RL)	$47.95$	$72.61$	$66.60$	$44.51$	$33.40$	$55.88$	$36.14$	$53.49$
No correction	$49.96_{\pm 0.36}$	$72.90_{\pm 0.65}$	$69.67_{\pm 0.42}$	$48.37_{\pm 1.86}$	$33.67_{\pm 1.40}$	$57.58_{\pm 0.23}$	$37.30_{\pm 1.52}$	$55.12_{\pm 0.34}$
Exact ratio	$50.49_{\pm 2.21}$	$67.49_{\pm 5.67}$	$69.20_{\pm 2.11}$	$50.00_{\pm 1.22}$	$33.27_{\pm 1.21}$	$56.56_{\pm 3.00}$	$37.40_{\pm 1.11}$	$54.24_{\pm 2.56}$
KPop †	$51.33_{\pm 0.23}$	$71.87_{\pm 1.08}$	$71.07_{\pm 1.21}$	$50.20_{\pm 3.57}$	$33.33_{\pm 1.89}$	$58.36_{\pm 0.33}$	$37.50_{\pm 2.19}$	$55.84_{\pm 0.25}$
IcePop	$51.79_{\pm 0.40}$	$72.81_{\pm 1.44}$	$71.40_{\pm 2.42}$	$52.64_{\pm 0.35}$	$35.40_{\pm 1.64}$	$58.93_{\pm 0.10}$	$39.66_{\pm 1.30}$	$56.60_{\pm 0.09}$
TIS	$48.76_{\pm 1.81}$	$69.77_{\pm 6.27}$	$70.87_{\pm 1.79}$	$48.17_{\pm 1.22}$	$33.47_{\pm 1.60}$	$56.15_{\pm 2.85}$	$37.10_{\pm 1.51}$	$53.85_{\pm 2.32}$

Appendix

Table 13: Out-of-domain accuracy (%) on Qwen1.5-MoE-A2.7B after RL on GSM8K. Mean

\pm

standard deviation over three seeds; Science, Code, and All are item-weighted over the knowledge, code, and all benchmarks.

Panel of Figure 11 — Figure 11: Compute time per step on Qwen1.5-MoE-A2.7B, excluding checkpointing and evaluation. (a) Time by phase; dashed: no correction. (b) Mean response length.

Method	Rollout	Rescore	Train step	Advantage	Total	vs. no corr.	Length
No correction	$0.44_{\pm 0.02}$	$0.00_{\pm 0.00}$	$29.41_{\pm 0.08}$	$0.82_{\pm 0.02}$	$30.66_{\pm 0.09}$	—	$321_{\pm 1}$
Exact ratio	$0.48_{\pm 0.03}$	$6.50_{\pm 0.08}$	$23.26_{\pm 0.40}$	$0.72_{\pm 0.01}$	$30.97_{\pm 0.48}$	$+1.0\%$	$252_{\pm 3}$
KPop †	$0.45_{\pm 0.01}$	$6.64_{\pm 0.58}$	$23.90_{\pm 2.08}$	$0.74_{\pm 0.03}$	$31.73_{\pm 2.68}$	$+3.5\%$	$260_{\pm 24}$
IcePop	$0.49_{\pm 0.03}$	$6.51_{\pm 0.23}$	$23.12_{\pm 0.67}$	$0.77_{\pm 0.06}$	$30.88_{\pm 0.93}$	$+0.7\%$	$250_{\pm 8}$
TIS	$0.48_{\pm 0.03}$	$6.40_{\pm 0.13}$	$22.89_{\pm 0.56}$	$0.74_{\pm 0.01}$	$30.51_{\pm 0.69}$	$-0.5\%$	$249_{\pm 6}$
Seq-TIS	$0.49_{\pm 0.02}$	$7.39_{\pm 0.04}$	$26.44_{\pm 0.24}$	$0.72_{\pm 0.01}$	$35.03_{\pm 0.28}$	$+14.3\%$	$289_{\pm 2}$

Appendix

Table 14: Compute time per optimizer step (seconds) on eight A100 GPUs, mean

\pm

standard deviation over three seeds. Rescore is the training-engine forward pass that produces the training-side log-probabilities.

Panel of Figure 12 — Figure 12: Tokens kept (blue) and truncated or masked (red) in the run of each operator (seed two), early, middle, and late in training; panel titles give the share acted on. (a) CIS. (b) IcePop. (c) TIS. (d) KPop † .

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

May 14, 2026

Diagnosing Training Inference Mismatch in LLM Reinforcement Learning via a Zero-Mismatch Reference

Modern LLM RL systems separate rollout generation from policy optimization. These two stages are expected to produce token probabilities that match exactly. However, implementation differences can make them assign different values to the same sequence under the same model weights, inducing Training-Inference Mismatch (TIM). TIM is difficult to inspect because it is entangled with off-policy drift and common stabilization mechanisms. In this work, we isolate TIM in a zero-mismatch diagnostic setting (VeXact), and show that small token-level numerical disagreements can independently cause training collapse. We further show that TIM changes the effective optimization problem, and identify a set of remedies that could mitigate TIM. Our results suggest that TIM is not benign numerical noise, but a systems-level perturbation that should be treated as a first-order factor in analyzing LLM RL stability.

Tianle Zhong, Neiwen Ling, Yifan Pi +5

ByteDance · The University of Virginia

Reinforcement Learning LLM Inference RL Stability RL Post-Training

Ask PDF

May 13, 2026

AIS: Adaptive Importance Sampling for Quantized RL

Reinforcement learning (RL) for large language models (LLMs) is dominated by the cost of rollout generation, which has motivated the use of low-precision rollouts (e.g., FP8) paired with a BF16 trainer to improve throughput and reduce memory pressure. This introduces a rollout-training mismatch that biases the policy gradient and can cause training to collapse outright on reasoning benchmarks. We show that the mismatch is non-stationary and acts as a double-edged sword: early in training it provides a stochastic exploration bonus, exposing the gradient to trajectories the trainer would otherwise under-sample, but the same perturbation transitions into a destabilizing source of bias as the policy concentrates. To solve this, we propose Adaptive Importance Sampling (AIS), a correction framework that adjusts the strength of its intervention on a per-batch basis. AIS combines three real-time diagnostics, namely weight reliability, divergence severity, and variance amplification, into a single mixing coefficient that interpolates between the uncorrected and fully importance-weighted gradients, suppressing the destabilizing component of the mismatch while preserving its exploratory benefit. We integrate AIS into GRPO and evaluate it on the diffusion-based LLaDA-8B-Instruct and the autoregressive Qwen3-8B and Qwen3.5-9B across mathematical reasoning and planning benchmarks. AIS matches the BF16 baseline on most tasks while retaining the 1.5 to 2.76x rollout speedup of FP8.

Jiajun Zhou, Wei Shao, Lingchao Zheng +2

Huawei · The University of Hong Kong

RL for Language Models Reinforcement Learning Importance Sampling for RL RL for Language Model Reasoning

Ask PDF

Jul 6, 2026

Turning Off-Policy Tokens On-Policy: A Plug-in Approach for Improving LLM Alignment

Reinforcement learning (RL) post-training for large language models (LLMs) follows a efficient paradigm of "rollout then update", which inevitably results in off-policy training data. To resolve this, Importance sampling (IS) is proposed, while the token-level ratios compound over long sequences, causing severe variance exploded. A natural idea is "transferring" these off-policy token into on-policy token, so that the importance scores for correction are unnecessary. Following this idea, we propose Selective Importance Sampling (SIS), which is inspired by rejection sampling. Concretely, SIS implements by viewing off-policy model as proposal distribution, and implement a token-level rejection test: accepted tokens are viewed as on-policy, so that receive unit importance score, while rejected tokens retain the standard IS correction. Our proposed SIS is theoretically proved reducing the gap between token-level and sequence-level off-policy gradient estimators. The SIS acts as a plug-in that only modifies the importance ratio in the policy loss, adding negligible wall-clock overhead, and can be combine with a vast vary of RL post-training algorithms. Experiments on dense and MoE LLMs across math and agent benchmarks show that SIS consistently improves all objectives, while providing substantially stronger robustness under off-policy data.

Yu Li, Xiuyu Li, Mingyang Yi +4

Renmin University of China

LLM Alignment Importance Sampling RL Post-Training Rejection Sampling Off-Policy RL Importance Sampling for RL

Ask PDF

Abstract

Reinforcement Learning

Explore similar work

Diagnosing Training Inference Mismatch in LLM Reinforcement Learning via a Zero-Mismatch Reference

AIS: Adaptive Importance Sampling for Quantized RL

Turning Off-Policy Tokens On-Policy: A Plug-in Approach for Improving LLM Alignment