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

Rethinking Training-Inference Mismatch in LLM Reinforcement Learning. Where It Arises and How to Correct It. By Tianrun Yu and six others, from Brigham Young University and the University of North Carolina at Chapel Hill. Published on the twenty-sixth of September. Its code repository is public and real: fifteen source files.

In reinforcement learning with verifiable rewards for large language models, two different engines are involved. An inference engine samples the rollouts. A training engine computes the gradients. The two do not agree exactly: they assign different probabilities to the same tokens. The model being trained is, in a small but real sense, not the model that produced the data it is being trained on. Policy updates have to account for that discrepancy somehow. Exact importance sampling is one answer, reweighting each token by the ratio between the two engines, but its error is governed by a second moment that is unbounded. The title asks two things, where the mismatch arises and how to correct it, and the paper answers them in that order.

The paper's first step is to characterize where the mismatch comes from, and the characterization is empirically supported. The mismatch shows up as an additive displacement in log-odds, fixed by the perturbation each logit receives before the softmax. The key observation is that the distribution of that displacement is roughly the same whatever the token's confidence. That motivates calibrated importance sampling, or CIS.

CIS applies a confidence-aware truncation. Large positive displacements are cut off at a single constant threshold. Translated back into importance ratios, one fixed threshold on displacement becomes a cap on the ratio that tightens as the token's confidence rises. A confident token is allowed less room for its ratio to stray than an uncertain one, even though the threshold on displacement never changes. On the theory side, the authors show that CIS swaps the unbounded second moment of exact importance sampling for a term bounded by a constant. The price is a bias, and that bias is controlled by how much gets truncated.

The evaluation covers three mixture-of-experts models and five mathematical reasoning benchmarks. On all three models, CIS gets the highest average across the five benchmarks among the baselines evaluated. Diagnostic analyses add two details. CIS puts less truncation bias on low-confidence tokens than truncated importance sampling does. And clipping small importance weights upward, rather than leaving them alone, lowers held-out accuracy.

Much of the work here is done before any correction is proposed. Once the mismatch is described as a displacement in log-odds that does not care about confidence, a single threshold is enough, and the confidence-dependent cap on the ratio falls out of the translation rather than being designed in separately.