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

Yingru Li, Jiawei Xu, Ziniu Li, Jiacai Liu, Wei Liu, Yuxuan Tong, Longtao Zheng, Zhenghai Xue, Yaxiang Zhang, Tianle Cai, Ge Zhang, Qian Liu, Baoxiang Wang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:71553-71575, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-li26hy, title = {The Optimal Token Baseline: Variance Reduction for Long-Horizon {LLM}-{RL}}, author = {Li, Yingru and Xu, Jiawei and Li, Ziniu and Liu, Jiacai and Liu, Wei and Tong, Yuxuan and Zheng, Longtao and Xue, Zhenghai and Zhang, Yaxiang and Cai, Tianle and Zhang, Ge and Liu, Qian and Wang, Baoxiang}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {71553--71575}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/li26hy/li26hy.pdf}, url = {https://proceedings.mlr.press/v306/li26hy.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T The Optimal Token Baseline: Variance Reduction for Long-Horizon LLM-RL %A Yingru Li %A Jiawei Xu %A Ziniu Li %A Jiacai Liu %A Wei Liu %A Yuxuan Tong %A Longtao Zheng %A Zhenghai Xue %A Yaxiang Zhang %A Tianle Cai %A Ge Zhang %A Qian Liu %A Baoxiang Wang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-li26hy %I PMLR %P 71553--71575 %U https://proceedings.mlr.press/v306/li26hy.html %V 306 %X 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.
APA
Li, Y., Xu, J., Li, Z., Liu, J., Liu, W., Tong, Y., Zheng, L., Xue, Z., Zhang, Y., Cai, T., Zhang, G., Liu, Q. & Wang, B.. (2026). The Optimal Token Baseline: Variance Reduction for Long-Horizon LLM-RL. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:71553-71575 Available from https://proceedings.mlr.press/v306/li26hy.html.

Related Material