解决强化学习中多标记预测联合训练失效问题,提出自适应系数校准方法。
Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration

- 从优化角度分解多标记预测对强化学习的影响,揭示性能下降原因。
- 在六项数学推理基准上,新方法性能优于或等同于现有最优基线。
- 适合研究大模型推理能力提升与强化学习融合的学者参考。
基于可验证奖励的强化学习(RLVR)已成为提升大语言模型推理能力的标准范式,而多标记预测(MTP)则是预训练中的通用模块。将二者结合是自然思路,但当前强化学习实践因分离MTP梯度而导致联合训练性能下降。本文从优化视角重新审视这一失败:发现每步MTP对强化学习目标的影响可分解为一阶相关项与二阶扰动惩罚项。该分解统一了三种MTP训练范式——分离、交叉熵损失与策略损失,并解释其成败原因。进一步分析表明,尽管策略损失符合直觉,但性能仍下降,因相关项衰减而二次惩罚项持续存在。据此,本文提出最优系数校准(OCC),一种通过日志概率代理在线追踪最优系数的轻量级自适应方案。在六个竞赛级数学推理基准上,OCC始终达到或超越分离基线,显著提升联合MTP-RL训练表现。
原文摘要 · Abstract (English)
Reinforcement Learning from Verifiable Rewards (RLVR) has emerged as the standard paradigm for improving reasoning capability of large language models, while Multi-Token Prediction (MTP) has been a widely adopted module in pretraining. Combining them is a natural approach, yet current RL practices detach MTP gradients because joint training degrades the performance. We revisit this failure from an optimization perspective. We show that the per-step effect of MTP on the RL objective can be decomposed into two terms: a first-order correlation and a second-order perturbation penalty. This decomposition unifies three MTP training regimes: Detach, Cross-Entropy loss, and Policy loss, and explains why each succeeds or fails. Further analysis of policy loss reveals that, although it aligns with intuition, performance still degrades: the correlation term decays while the quadratic penalty persists. Guided by the analysis, we propose Optimal Coefficient Calibration (OCC), an adaptive scheme that tracks the optimal coefficient online via a log-probability proxy at negligible cost. Across six competition-level mathematical reasoning benchmarks, OCC consistently matches or exceeds the detach baseline, delivering improved joint MTP-RL training performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。