Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

The paper reframes RL instability in LLM fine-tuning as a dynamic optimization problem driven by gradient noise and training-inference mismatch, rather than a static numerical artifact. It introduces a reactive learning rate scheduler that decays the learning rate when the average response length surges, with a decay period approximately 1.8× the surge duration, to suppress gradient noise and keep the mismatch at bay. Empirical results on Qwen3 base models show improved stability and higher peak validation performance compared to constant LR and IS patches, across BF16/FP16 settings. This work provides a practical, scalable approach to stabilizing RL training for large language models and advances understanding of how optimization dynamics interact with numerical drift across heterogeneous engines.

Abstract

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.

Beyond Precision: Training-Inference Mismatch is an Optimization Problem and Simple LR Scheduling Fixes It

TL;DR

The paper reframes RL instability in LLM fine-tuning as a dynamic optimization problem driven by gradient noise and training-inference mismatch, rather than a static numerical artifact. It introduces a reactive learning rate scheduler that decays the learning rate when the average response length surges, with a decay period approximately 1.8× the surge duration, to suppress gradient noise and keep the mismatch at bay. Empirical results on Qwen3 base models show improved stability and higher peak validation performance compared to constant LR and IS patches, across BF16/FP16 settings. This work provides a practical, scalable approach to stabilizing RL training for large language models and advances understanding of how optimization dynamics interact with numerical drift across heterogeneous engines.

Abstract

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.
Paper Structure (17 sections, 4 theorems, 22 equations, 13 figures, 1 table, 1 algorithm)

This paper contains 17 sections, 4 theorems, 22 equations, 13 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Training-Inference Mismatch and Importance Sampling Correction
  4. Related Work
  5. Methodology
  6. Why LR Scheduler
  7. Traditional Scheduler is Not Feasible
  8. Our Method: Decay by Response Length
  9. Experiments
  10. Our Scheduler Stabilizes Training and Reduces Mismatch
  11. Criteria of Choosing Decay Period
  12. Ablation Experiments on Importance Sampling
  13. Ablations on FP16 Precision
  14. Conclusion
  15. Brief discussion on LR decay and gradient noise
  16. ...and 2 more sections

Key Result

Theorem 3.1

Let $\nabla_{\theta} \mathcal{J}(\theta)$ and $\nabla_{\theta} \mathcal{J}_{actual}(\theta)$ be the ideal gradient for backpropagation and the real calculated gradient defined in eq:ideal graident and eq:biased gradient, respectively. Then, under mild assumptions, the following inequality holds: We provide detailed proof in Appendix Appendix:proof.

Figures (13)

  • Figure 1: Results of Qwen3-4B-Base on filtered DAPO training dataset. The vertical dotted line indicates the point (300 step) at which the methods collapse.
  • Figure 2: Training stability comparison with/without IS patches.
  • Figure 3: Smoothed gradient norm. The vertical dotted line is at 300 step.
  • Figure 4: LR:1e-6 vs 1e-7. It’s worth noting that the blue line isn’t simply 10x slower than the green line. The green line keeps rapid and homogeneous growing and then rapidly collapses; the blue line, in contrast, has a sudden surge around 1.5k steps, which exhibits heterogeneity in the growing process.
  • Figure 5: Experiment with partial training dataset. The collapse of orange line with 25% dataset is even slower than using full dataset. The green line with 2.5% dataset still grows for 300 steps.
  • ...and 8 more figures

Theorems & Definitions (7)

  • Theorem 3.1: informal
  • Theorem B.1
  • Lemma B.4: Conversion to token-level objective
  • proof
  • Lemma B.5: State Occupancy Drift
  • proof
  • proof