arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2602.07078cs.LGcs.AI
收藏

最优Token基线:长程LLM-RL的方差缩减

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

National University of Singapore(新加坡国立大学);Nanyang Technological University(南洋理工大学);Fudan University(复旦大学);Hong Kong University of Science(香港科学大学);The Chinese University of Hong Kong, Shenzhen(香港中文大学(深圳));Vector Institute(向量研究所)

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

AI总结:

针对长程LLM-RL训练中梯度方差爆炸问题,从第一性原理推导最优Token基线,提出仅用前向概率近似梯度范数的Logit-Gradient Proxy,以N=4分组达到N=32性能,减少65%以上Token消耗。

摘要:

大型语言模型(LLM)的强化学习(RL)在长程任务中常因梯度方差爆炸导致训练崩溃。为缓解此问题,通常引入基线进行优势计算;然而,传统价值模型难以优化,且标准分组基线忽略了序列异质性。尽管经典最优基线理论可实现全局方差缩减,但它忽略了Token异质性,且需要代价高昂的基于梯度的计算。在本工作中,我们从第一性原理推导出最优Token基线(OTB),证明梯度更新应与其累积梯度范数成反比加权。为确保效率,我们提出Logit-Gradient Proxy,仅使用前向概率近似梯度范数。我们的方法实现了训练稳定性,并在单轮和工具集成推理任务中,以仅N=4的分组大小达到了大分组(N=32)的性能,减少了超过65%的Token消耗。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑