The Optimal Token Baseline: Variance Reduction for Long-Horizon LLM-RL
This addresses training instability for researchers and practitioners in LLM-RL, offering an incremental improvement over existing baselines.
The paper tackled the problem of training collapse in long-horizon LLM-RL due to gradient variance by deriving the Optimal Token Baseline, which weights gradient updates inversely to their cumulative gradient norm, achieving training stability and matching performance with large group sizes using only N=4, reducing token consumption by over 65%.
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.