MIPU changes what LLM RL is actually optimizing: instead of improving the training-engine policy (which may not match the deployed inference-engine policy), it proposes candidate updates weighted against the sampler, then rejects any synchronized update whose measured Post-update inference gap proxy (T̂_post) goes negative. Under FP8-quantized rollout, this keeps training from collapsing where Group Relative Policy Optimization (GRPO) and other fixes eventually degrade.
You’re running RL post-training on a reasoning model. For efficiency, rollouts come from vLLM in FP8, but gradients are computed in a training framework at higher precision. Same weights, different probabilities per token. Your loss curve looks fine for a while, then accuracy on your eval set falls off a cliff. Existing fixes (TIS (Truncated Importance Sampling), token filtering via MIS (Masked Importance Sampling), learning-rate decay) all try to make the training-side update look more like the inference distribution. This paper argues that’s fixing the wrong thing. The number the training loop optimizes and the number your users see are two different policies, and no amount of importance-weighting on the training side is guaranteed to move the deployed one in the right direction.
Call the training-engine policy π and the inference-engine policy μ. They share weights but assign different probabilities. Standard RL asks: did this gradient step improve π? MIPU asks the harder question: after we push π back into the inference engine to get a new μ, did that improve?
The authors write the inference-side improvement as three terms: the pre-update gap between π and μ, the training-side update itself, and the post-update gap between the new π and new μ. Step 1 handles the first two by weighting the Group Relative Policy Optimization (GRPO) loss with a truncated ratio π_k/μ_k (bounded to prevent variance blow-up from rare tokens), keeping the standard clip only on the current update. Step 2 handles the third term: after syncing weights, sample a small validation batch from the new μ, estimate whether μ is under-weighting the good responses relative to π, and if the signal is clearly negative, roll back both the model and the optimizer state.
for step in range(num_steps):
save_checkpoint(theta, opt_state, mu)
batch = sample_prompts(); rollouts = mu.generate(batch)
w = clamp(pi_k / mu_k, max=2.0) # truncated sampler ref
loss = w * clipped_grpo(rollouts, advantages)
theta = update(theta, loss); mu = sync(theta)
T_post = -mean(rho * advantage_on_val_from_mu_new)
if T_post < -c: # inference-gap check
restore_checkpoint()
The validation-side estimate uses a length-normalized log-ratio between the new π and new μ instead of a raw sequence ratio, which they note is too high-variance to be usable.
The prevailing view treats training-inference mismatch as a numerical bug to reduce, via FP16 rollouts, learning-rate decay, or ratio corrections on the training loss. This paper reframes it as an objective bug. A gradient step that provably improves the training-engine policy carries no guarantee about the inference-engine policy that actually gets deployed, so acceptance of an update should be gated by an inference-side check, not just a training-side loss going down. The load-bearing evidence isn’t the headline benchmark score; it’s the random-rollback ablation showing that rejecting the same fraction of updates without the inference-gap signal still collapses.
The most telling result is the random-rollback control. Step 2 rejects roughly 53.5% of candidate updates on Qwen3-4B under FP8; a random baseline set to reject 67% (more conservative) still collapses after a transient peak, while Step 2’s signal-gated rejection keeps a stable trajectory. Simply throwing away updates doesn’t help. Throwing away the right updates does.
•
On FP8-quantized rollout with Qwen3-4B, MIPU hits 66.71% average pass@1 across five math benchmarks; on Qwen3-1.7B, 53.97%. Baselines (Group Relative Policy Optimization (GRPO), MIS (Masked Importance Sampling), LR-decay) reach competitive intermediate peaks but degrade under continued training. MIPU stays stable to the end.
•
Ablation: Step 2 alone prevents collapse but can’t improve a bad proposal, so it just holds the previous policy. Step 1 alone produces better candidates but accepts everything, so mismatch fluctuations still accumulate. Both are needed.
•
The gap proxy T̂_post is not noise: on the smaller 1.7B model the training-inference KL is larger and T̂_post oscillates more strongly, matching the intuition that quantization hurts smaller models more.
•
Sensitivity: setting the tolerance too strict causes persistent rollback after ~280 steps, i.e., learning stalls. A dynamic tolerance annealed over the first 100 steps works best.
Reach for this when you’re running RLVR-style post-training with a split rollout/training stack, especially if you’re using low-precision inference for throughput. The concrete change is small: keep your Group Relative Policy Optimization (GRPO) loop, add a truncated π/μ weight on each token, and after every sync run a short validation rollout to compute the acceptance proxy. If it’s clearly negative, roll back the optimizer state, not just the weights. The rollback cost is one extra validation batch per step plus occasional wasted gradient work.
Code and configs are on the project page. Training uses the ROLL framework (Apache-2.0) with vLLM rollouts. Datasets used (DAPO-Math-17k, DeepMath-103k, MATH-500, AIME24, AMC23, Minerva, OlympiadBench) are all publicly available under permissive licenses. Experiments are limited to Qwen3-1.7B and Qwen3-4B; the authors flag that they haven’t tested larger scales.
Optimize the policy you’ll deploy, not the one your trainer sees. When your inference and training stacks disagree on probabilities, a training-side loss going down is a proxy, not the goal. Cheap validation rollouts after each sync are worth their cost when the alternative is a silent collapse hours into a run.
•
Only tested at 1.7B and 4B under one specific mismatch source (FP8 rollout). Whether the acceptance signal remains informative at 70B+ or under other mismatch sources (kernel differences, speculative decoding, different attention backends) is unknown.
•
The rollback needs a preserved optimizer state and the previous inference-engine weights. If your infrastructure streams weights or reuses optimizer memory aggressively, adopting Step 2 has a nontrivial engineering cost.
•
The tolerance parameter c is hand-annealed and per-model. The paper’s own sensitivity analysis shows stricter isn’t safer, so this is a knob that needs tuning, not a drop-in default.