Source-linked AI summary
Beyond Precision: Training-Inference Mismatch is an Optimization Problem and Simple LR Scheduling Fixes It
Yaxiang Zhang, Yingru Li, Jiacai Liu, Jiawei Xu, Ziniu Li, Qian Liu, Haoyuan Li
TL;DR
RL training instability is associated with training-inference mismatch and can persist despite importance-sampling corrections. This paper analyzes the mismatch as an optimization-coupled dynamic failure and introduces response-length-triggered learning-rate decay, which consistently stabilizes training across model scales.
Problem
RL training for LLMs is fragile, while importance-sampling corrections can still suffer training collapse during extended runs.
Method
The paper analyzes mismatch dynamics and proposes a reactive LR scheduler that triggers decay when average response length surges.
Results
The method consistently stabilizes training, suppresses engine-level distribution mismatch, and achieves superior peak performance compared with standard importance sampling and constant learning-rate strategies.
Takeaways & Limitations
Training-inference mismatch is a dynamic optimization failure, and response-length surges provide an early-warning signal for adaptive stabilization.
Takeaways & Limitations
Importance-sampling corrections can introduce hyperparameter tuning, gradient-estimation bias, and only extend the stable training window in some experiments.
Abstract
from arXiv · showhide
Reinforcement Learning (RL) for training Large Language Models is notoriously unstable. While recent studies attribute this to "training inference mismatch stemming" from inconsistent hybrid engines, standard remedies, such as Importance Sampling, might fail during extended training runs. In this work, we analyze this instability through the lens of optimization, demonstrating that gradient noise and training-inference mismatch escalate in tandem as training progresses. Meanwhile, we find that the mismatch can be effectively suppressed by shrinking the update size. Taken together, we deduce that the mismatch is not merely a static numerical discrepancy, but a dynamic failure coupled with the model's optimization. Based on this insight, we propose a simple yet effective solution: a specialized Learning Rate (LR) scheduler. Instead of pre-defined decay schedule in traditional LR scheduler, our method dynamically triggers LR decay based on response length, which we identify as a reliable early-warning signal for impending instability. Empirical evidence suggests that by reducing the learning rate as gradient noise rises, we can consistently stabilize RL training and keep the training-inference mismatch at a safe level.
1 Introduction
RL training instability is linked to training-inference mismatch, which can worsen during optimization rather than behaving as static noise. The paper proposes adaptive learning-rate decay triggered by response-length surges to stabilize training.
- RL training can abruptly collapse, reducing rewards, validation accuracy, sample efficiency, and computational economy.
- Training-inference mismatch arises from different rollout and training engines whose finite-precision operation sequences introduce optimization noise.
- Importance Sampling corrections may require masking or clipping and can precipitate catastrophic collapse during extended training runs.
- The mismatch is a dynamic failure that escalates with gradient noise and the model’s optimization trajectory.
- Epoch-based LR schedules are poorly aligned with gradient SNR and optimization dynamics, causing suboptimal or premature decay.
- The proposed scheduler decays LR when average response length surges, halving it every decay period until a preset minimum; a period near 1.8× the surge consistently stabilizes training.
2 Background and Related Work
This section explains how heterogeneous rollout and training engines create training-inference mismatch and off-policy bias in RL for LLMs. It reviews importance-sampling corrections and motivates optimization-aware scheduling because existing corrections can add bias, require tuning, or fail during extended training.
- Training-Inference Mismatch: RL policy gradients can be evaluated on rollout trajectories from µθ rather than the target training policy πθ, introducing off-policy bias.The mismatch arises because modern systems use different engines for rollout generation and parameter updates, whose finite-precision calculation orders differ.
- Importance Sampling Correction: Importance Sampling corrects for differing sampling and target distributions, with token-level and sequence-level variants defining different response units.Token-level IS applies ratios autoregressively, whereas sequence-level IS treats the whole response as y.
- Importance Sampling Correction: Vanilla Importance Sampling is impractical for long responses because extreme probability ratios produce high variance.Truncation and masking techniques trade variance against bias, including TIS and MIS.
- Importance Sampling Correction: Clipping and masking introduce additional tuning requirements and can reintroduce bias into gradient estimation, undermining theoretical soundness.The clipping hyperparameter C requires extra tuning, while experiments indicate these corrections may only extend the stable training window before collapse.
- Beyond Importance Sampling: Engineering workarounds such as changing numerical precision remain insufficient because autoregressive inference and parallel training use fundamentally different floating-point computations.The section therefore motivates a solution grounded in optimization dynamics rather than only computational alignment.
- Learning Rate Scheduling: Predefined LR schedules are poorly matched to RL because training duration is unpredictable and they are blind to model state and system-level discrepancies.The proposed direction is a reactive scheduler triggered by early-warning signals of numerical instability rather than static time-based decay.
3 Methodology
The analysis links training-inference mismatch to optimization dynamics: gradient noise rises during training, and reducing update size suppresses instability. This motivates a reactive learning-rate scheduler triggered by response-length surges rather than fixed training time.
- Why LR Scheduler: Around 300 steps, the onset of distribution mismatch coincides with validation-accuracy collapse, while most Importance Sampling patches only prolong stability.Only token-level TIS prevents collapse in this experiment; other methods eventually fail, sometimes more severely than baseline.
- Why LR Scheduler: The escalating gradient norm suggests increasingly noise-dominated gradient estimation, motivating a learning-rate reduction from 1e-6 to 1e-7.The paper states that lowering η reduces the effects of gradient-noise bias and variance.
- Why LR Scheduler: Shrinking the update size suppresses mismatch, indicating that numerical discrepancy depends on model weight location and optimization dynamics rather than acting as static noise.The authors connect later instability to malignant optimization dynamics and potentially sharper, higher-curvature regions.
- Traditional Scheduler is Not Feasible: Fixed LR schedules can decay too early during benign training or too late after numerical instability has begun, making their timing poorly matched to RL dynamics.The paper argues that epoch count is an unreliable proxy for gradient SNR and optimization state.
- Our Method: Decay by Response Length: The proposed scheduler preserves the initial LR until instability signals appear, then halves it every decay period until reaching a floor of 10% of η0.The algorithm takes η0, η∞, Tdecay, and total training steps as inputs and outputs ηt for each step.
- Our Method: Decay by Response Length: A response-length surge, in which average output length can triple within a few steps, serves as the scheduler’s early-warning signal.The proposed decay period is approximately 1.8× the surge interval, balancing learning speed and stability.
4 Experiments
Experiments across Qwen3 model scales show that response-length-triggered learning-rate decay stabilizes RL training and suppresses training-inference mismatch. The scheduler remains effective alongside importance sampling and under FP16 precision.
- The evaluation uses Qwen3 Base models with filtered DAPO data and VeRL for RL training.The initial learning rate is 1e-6 and the training batch size is 64.
- 4.1 Our Scheduler Stabilizes Training and Reduces Mismatch: The length-decay scheduler improves peak validation performance while stabilizing gradient norms and keeping log ppl abs diff within a safe regime.It also preserves higher policy entropy, helping avoid premature mode collapse.
- 4.2 Criteria of Choosing Decay Period: For Qwen3-4B-Base, decay periods of 125 and 204 stabilize training, while period 250 only extends the stable window before eventual collapse.The response-length surge occurs near step 110, and period 204 performs slightly better than 125.
- 4.2 Criteria of Choosing Decay Period: For Qwen3-8B-Base, the epoch-aligned period 204 reduces AIME24 accuracy, whereas period 160 performs best and scales with surge onset by approximately 1.8×.The response-length surge occurs earlier, at step 90.
- 4.3 Compatibility with Importance Sampling: The scheduler stabilizes training when token-level MIS fails, further improves validation accuracy with TIS, and dampens mismatch-indicator spikes.These results support scheduling as complementary to importance-sampling corrections.
- 4.4 Experiment on FP16: Under FP16, training still exhibits collapse without the scheduler, while the same length-triggered schedule successfully stabilizes training.This supports the claim that the instability is not resolved merely by changing numerical precision.
5 Conclusion
The paper characterizes training-inference mismatch as a dynamic optimization failure coupled to gradient noise and model evolution. It introduces a response-length-triggered scheduler that stabilizes training, suppresses mismatch, and improves peak performance across model scales.
- Training-inference mismatch is presented as a dynamic optimization failure coupled with gradient noise and the model’s optimization trajectory.
- The proposed scheduler uses a surge in average response length as an early-warning signal for impending instability.
- Across model scales, the method stabilizes training, suppresses engine-level distribution mismatch, and achieves superior peak performance over importance sampling and constant learning-rate strategies.
A Brief discussion on LR decay and gradient noise
The optimization analysis explains why reducing the learning rate can resist gradient noise. Progress and bias effects scale roughly linearly with the step size, while the noise penalty decreases quadratically.
- The smoothness-based analysis separates the update effect into true progress, bias effect, and noise penalty terms.
- The stochastic gradient estimate is decomposed into the true gradient, bias, and zero-mean noise with variance Var.
- For small η, true progress and bias effects scale roughly linearly with η, while the noise penalty decreases quadratically.
- Therefore, learning-rate decay slows progress but reduces the influence of gradient noise faster than it reduces progress.
B Proof of Theorem 1
The appendix proves a theorem under bounded local numerical drift and bounded score-function assumptions. Supporting lemmas connect sequence-level gradients to token-level objectives and bound the accumulation of state-occupancy drift.
- Theorem B.1 assumes binary rewards together with assumptions on local numerical drift and bounded score functions.
- Lemma B.4 converts the sequence-level policy-gradient objective into an equivalent token-level form.
- Lemma B.5 bounds the L1 distance between policy and rollout state-occupancy distributions as accumulating linearly over time.
- The proof proceeds using the state-transition distribution and the two auxiliary lemmas.
C Hyperparameters for reproducibility
This section presents key configurations from the main experiment to support reproducibility.
- The main experiment includes key configurations intended to support reproducibility.The passage introduces these configurations specifically for reproducibility purposes.
- The section focuses on experiment configuration information rather than reported results.
- Table 1 lists the key configurations used in the experiment.