Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Taming the Tail: Stable LLM Reinforcement Learning via Dynamic Vocabulary Pruning

Yingru Li, Jiawei Xu, Jiacai Liu, Yuxuan Tong, Ziniu Li, Tianle Cai, Ge Zhang, Qian Liu, Baoxiang Wang

TL;DR

This work tackles training instability in RL for large language models caused by a mismatch between fast inference and high-precision training. By redesigning the RL objective to operate over a dynamically pruned safe vocabulary using min-P filtering, the authors avoid tail-induced gradient bias rather than patching after the fact. They derive and validate an asymmetric vulnerability bound $|\Delta_a|\le 2\epsilon_{\max}(1-p_a)$, show that tail tokens drive instability, and propose a constrained gradient estimator with a provably bounded optimization bias. Empirically, Dynamic Vocabulary Pruning yields stable training and improved mathematical reasoning performance, offering a principled route to reliable RL for LLMs in high-throughput settings.

Abstract

Reinforcement learning for large language models (LLMs) faces a fundamental tension: high-throughput inference engines and numerically-precise training systems produce different probability distributions from the same parameters, creating a training-inference mismatch. We prove this mismatch has an asymmetric effect: the bound on log-probability mismatch scales as $(1-p)$ where $p$ is the token probability. For high-probability tokens, this bound vanishes, contributing negligibly to sequence-level mismatch. For low-probability tokens in the tail, the bound remains large, and moreover, when sampled, these tokens exhibit systematically biased mismatches that accumulate over sequences, destabilizing gradient estimation. Rather than applying post-hoc corrections, we propose constraining the RL objective to a dynamically-pruned ``safe'' vocabulary that excludes the extreme tail. By pruning such tokens, we trade large, systematically biased mismatches for a small, bounded optimization bias. Empirically, our method achieves stable training; theoretically, we bound the optimization bias introduced by vocabulary pruning.

Taming the Tail: Stable LLM Reinforcement Learning via Dynamic Vocabulary Pruning

TL;DR

This work tackles training instability in RL for large language models caused by a mismatch between fast inference and high-precision training. By redesigning the RL objective to operate over a dynamically pruned safe vocabulary using min-P filtering, the authors avoid tail-induced gradient bias rather than patching after the fact. They derive and validate an asymmetric vulnerability bound , show that tail tokens drive instability, and propose a constrained gradient estimator with a provably bounded optimization bias. Empirically, Dynamic Vocabulary Pruning yields stable training and improved mathematical reasoning performance, offering a principled route to reliable RL for LLMs in high-throughput settings.

Abstract

Reinforcement learning for large language models (LLMs) faces a fundamental tension: high-throughput inference engines and numerically-precise training systems produce different probability distributions from the same parameters, creating a training-inference mismatch. We prove this mismatch has an asymmetric effect: the bound on log-probability mismatch scales as where is the token probability. For high-probability tokens, this bound vanishes, contributing negligibly to sequence-level mismatch. For low-probability tokens in the tail, the bound remains large, and moreover, when sampled, these tokens exhibit systematically biased mismatches that accumulate over sequences, destabilizing gradient estimation. Rather than applying post-hoc corrections, we propose constraining the RL objective to a dynamically-pruned ``safe'' vocabulary that excludes the extreme tail. By pruning such tokens, we trade large, systematically biased mismatches for a small, bounded optimization bias. Empirically, our method achieves stable training; theoretically, we bound the optimization bias introduced by vocabulary pruning.
Paper Structure (29 sections, 6 theorems, 20 equations, 1 figure)

This paper contains 29 sections, 6 theorems, 20 equations, 1 figure.

Table of Contents

  1. Introduction
  2. The Speed-Stability Dilemma
  3. Our Approach: Objective Redesign over Reactive Patching
  4. Contributions
  5. Preliminaries
  6. Diagnosis: The Fundamental Instability
  7. The Training-Inference Mismatch
  8. Modeling the Mismatch: Logit Perturbations
  9. Asymmetric Vulnerability
  10. Summary
  11. Solution: Constrained Optimization via Vocabulary Pruning
  12. Dynamic Vocabulary Pruning via Min-P Filtering
  13. Constrained Policies and Objective
  14. The Stable Gradient Estimator
  15. Why It Works: Avoiding Systematic Bias
  16. ...and 14 more sections

Key Result

Theorem 3.1

Let $g = \mathbb{E}_{y \sim \pi^{\text{train}}_\theta}[\nabla_\theta \log \pi^{\text{train}}_\theta(y|x) \cdot R(x,y)]$ be the ideal gradient and $g' = \mathbb{E}_{y \sim \pi^{\text{infer}}_\theta}[\nabla_\theta \log \pi^{\text{train}}_\theta(y|x) \cdot R(x,y)]$ the practical gradient. The bias $b = where $\Delta_y = \log \pi^{\text{train}}_\theta(y|x) - \log \pi^{\text{infer}}_\theta(y|x)$ is the

Figures (1)

  • Figure 1: Training-Inference PPL Gap (Left) and AIME25 scores (Right) across training steps.

Theorems & Definitions (17)

  • Definition 2.1: LLM Generation MDP
  • Theorem 3.1: Gradient Bias from Mismatch
  • Proposition 3.2: Asymmetric Vulnerability
  • Proposition 3.3: Signature of Failure
  • Definition 4.1: Min-P Safe Action Sets
  • Definition 4.2: Min-P Constrained Policies
  • Theorem 4.3: Constrained Gradient Estimator
  • Remark 4.4: Support Condition
  • Remark 4.5: Fixed Safe Sets in Gradient Computation
  • proof
  • ...and 7 more