The Optimal Token Baseline: Variance Reduction for Long-Horizon LLM-RL

Reinforcement Learning (RL) for Large Language Models (LLMs) often suffers from training collapse in long-horizon tasks due to exploding gradient variance. To mitigate this, a baseline is commonly introduced for advantage computation; however, traditional value models remain difficult to optimize, and standard group-based baselines overlook sequence heterogeneity. Although classic optimal baseline theory can achieve global variance reduction, it neglects token heterogeneity and requires prohibitive gradient-based computation. In this work, we derive the Optimal Token Baseline (OTB) from first principles, proving that gradient updates should be weighted inversely to their cumulative gradient norm. To ensure efficiency, we propose the Logit-Gradient Proxy that approximates the gradient norm using only forward-pass probabilities. Our method achieves training stability and matches the performance of large group sizes ($N=32$) with only $N=4$, reducing token consumption by over 65\% across single-turn and tool-integrated reasoning tasks.

- **Problem**: Reinforcement Learning (RL) has become the standard paradigm for aligning Large Language Models (LLMs) in complex reasoning and agentic tasks. However, as tasks grow in complexity and generation horizons expand, practitioners frequently observe **Training Collapse** in RL training, where the gradient norm suddenly surges, causing model performance to crater. - **Solution**: We project the global variance-minimization objective into a causal form to derive the **Optimal Token Baseline (OTB)**. The key to OTB lies in its use of accumulated gradient magnitude as the weights for each token. Since the gradient variance of a specific token is not an isolated value but rather an accumulation of stochasticity from the start of the sequence, OTB explicitly accounts for this accumulated noise to achieve effective variance reduction. - **Impact**: OTB ensures stable training and extreme sampling efficiency, allowing models to achieve high performance even with smaller group sizes ($N=4$) while significantly reducing the total token budget. Furthermore, they introduce the "Logit-Gradient Proxy," a computationally "free" proxy that estimates the gradient norm using only forward-pass probabilities without requiring any additional backward passes.