Revisiting the Training-Inference Gap in LLM Reinforcement Learning: Origins and Remedies

Research 0 views

CIS fixes LLM RL training-inference mismatch with confidence-aware importance sampling, boosting MoE math reasoning accuracy on benchmarks.

Key Takeaways

  • Calibrated Importance Sampling (CIS) is proposed to reduce the training–inference mismatch in LLM reinforcement learning, where identical model weights can still produce different token probabilities due to inference and training engine differences.
  • Standard fixed-ratio importance sampling corrections can mistakenly treat engine-induced probability mismatches as signs of model uncertainty, potentially clipping or downweighting valid policy updates.
  • CIS reformulates the mismatch in log-odds space and uses confidence-aware truncation: it applies stricter corrections to high-confidence tokens while preserving more flexibility for low-confidence tokens.
  • The method is reported to achieve the highest average accuracy among evaluated baselines on three Mixture-of-Experts models across five mathematical reasoning benchmarks, with code and theoretical/empirical support provided.
  • The work is currently at the research-release stage, so its key next validation is whether independent evaluations reproduce the results and whether the approach generalizes beyond mathematical reasoning tasks.

Meta description: CIS addresses training-inference mismatch in LLM reinforcement learning, raising accuracy across three MoE models on five mathematical benchmarks.

Tags: LLM Reinforcement Learning, Training-Inference Mismatch, Calibrated Importance Sampling, MoE Models, Mathematical Reasoning Benchmarks

Featured image: Abstract technology style visual showing sampling paths, model nodes, and probability curves; no real people, no company logos, no text.

Calibrated Importance Sampling (CIS) has been introduced to address training–inference mismatch in LLM reinforcement learning, with reported accuracy results across three MoE models and five mathematical reasoning benchmarks. The issue matters because a model can have identical weights in two systems while still assigning different token probabilities, and those differences can distort updates if the correction rule assumes the probability gap means confidence.

Background: When identical weights do not produce identical token probabilities

In large-scale language model training, reinforcement learning often separates two roles: an inference engine generates candidate tokens or sequences, while a training engine computes gradients and applies updates. These engines may use different kernels, batching strategies, precision formats, or scheduling paths. Even when the checkpoint is the same, the resulting token probability can differ.

Importance sampling is a common correction for such distribution mismatch: it reweights samples generated under one policy for use under another policy. The CIS announcement identifies a subtle problem in simpler corrections. A fixed importance-ratio threshold treats every large ratio in the same way. But a large ratio can come from two different sources: the model genuinely being uncertain, or the training and inference engines disagreeing. When the threshold conflates mismatch with token confidence, it can clip or downweight updates for reasons unrelated to the policy’s actual uncertainty.

Development process: From fixed ratios to confidence-aware truncation

CIS changes the correction target. Instead of applying one blunt cap to probability ratios, it characterizes the mismatch in log-odds space. Log-odds, a standard way to express probabilities as logarithmic margins, can make small probability differences more visible when tokens are near certainty.

From that representation, CIS derives a confidence-aware truncation rule. The rule applies tighter caps to high-confidence tokens, where mismatch-induced changes can look disproportionately large, and allows less truncation bias for low-confidence tokens, where the model’s uncertainty should not be over-penalized.

In practical terms, the method tries to separate two questions: how much should a sample be corrected, and how much should it be clipped because it is unreliable? The reported design answers both with a calibration step rather than a single threshold.

Latest progress: Reported results on three MoE models and five mathematical reasoning benchmarks

The material reports that CIS achieves the highest average accuracy among evaluated baselines on each of three MoE models across five mathematical reasoning benchmarks. MoE, or mixture-of-experts, models are large language architectures that route different parts of computation to specialized expert subnetworks; they are often used to increase capacity without activating every parameter for every token.

Mathematical reasoning benchmarks provide a useful test bed because they reward exact problem solving, making them sensitive to unstable or noisy policy updates. The CIS announcement also states that the claims are supported by theoretical analysis and empirical diagnostics, and it links to code at https://github.com/kzhk5/CIS-RL.

Current stage: Research code available, accuracy results reported

The introduction is at the research-release stage. It provides a method, reported benchmark results, and code, but the material does not describe deployment across production systems or public dates beyond the announcement itself.

For readers tracking LLM post-training, the next signal will be whether independent evaluations reproduce the highest average accuracy among evaluated baselines on the same three MoE models and five mathematical reasoning benchmarks, and whether the log-odds calibration rule behaves consistently outside math-focused tasks.

Tags: LLM Reinforcement LearningTraining-Inference MismatchCalibrated Importance SamplingMoE ModelsMathematical Reasoning Benchmarks

Comments

No comments yet. Be the first to comment.

Leave a comment