🤖 machine learning

Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration

本文提出最优系数校准（OCC），这是一种基于梯度交互优化分析得出的自适应方案，旨在实现多令牌预测与可验证奖励强化学习的有效联合训练，从而克服以往的性能退化问题，并在数学推理基准测试中取得更优结果。

原作者： Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

发布于 2026-05-28

📖 1 分钟阅读☕ 轻松阅读

原作者： Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用简单语言和创意类比对这篇论文的解读。

全局概览：两位教练，一位运动员

想象你正在训练一个巨大的、超级聪明的机器人（大型语言模型）来解决高难度的数学问题。

主教练（RL）： 这位教练教导机器人如何获胜。他们查看机器人的答案，对正确的解法竖起“大拇指”，对错误的答案竖起“大拇指向下”。机器人学会重复有效的做法，停止无效的做法。这被称为强化学习（RL）。
助理教练（MTP）： 这位教练是一个特殊模块，试图预测机器人接下来要说的几个词，而不仅仅是下一个词。这被称为多词元预测（MTP）。这就像一位帮助机器人提前规划句子的教练。

问题所在：
此前，研究人员尝试让两位教练同时训练机器人。但事情出了差错。机器人的表现变得更差，甚至比只有一位主教练时还要差。

因此，大多数团队决定在“夺冠”训练阶段“解雇”助理教练。他们让助理教练在一旁观看，但不允许助理教练的建议改变机器人的大脑。他们“断开”了助理教练。

核心问题：
本文的作者问道：为什么拥有两位教练会伤害机器人？我们能否修复它，让两位教练协同工作而不破坏训练？

诊断：“推与拉”的问题

作者审视了训练背后的数学原理并找到了答案。他们意识到，当助理教练试图提供建议时，它会在机器人的大脑上产生两种相互对抗的力量：

有益的推力（相关性）： 有时，助理教练的建议与主教练的方向一致。这是好事！就像两个人朝同一个方向推车。
恼人的晃动（惩罚）： 有时，助理教练的建议只是“噪声”。它将机器人的大脑推向随机的、无助于获胜的方向。这就像有人在你想开车时摇晃车子。

旧方法为何失败：

方法 A（断开）： 他们关闭了助理教练的推力。安全，但错失了“有益推力”。
方法 B（交叉熵）： 他们让助理教练施加推力，但助理教练试图教导机器人平等地预测每一个词，包括错误的词。这产生了大量的“晃动”，而“推力”很少。机器人感到困惑，表现变差。
方法 C（策略损失）： 他们告诉助理教练使用与主教练相同的“获胜”规则。起初，这效果极佳！“推力”很强。但随着机器人变得更强，两位教练开始产生细微分歧。“晃动”保持不变，但“有益推力”却变弱了。最终，“晃动”占据了上风，机器人的表现崩溃了。

类比：
想象你正在走钢丝（RL 训练）。

断开是独自行走。安全，但缓慢。
旧联合训练是有一个朋友牵着你的手，但他喝醉了。他把你从绳子上拉下来。
“起落”问题是有一个清醒的朋友在开始时帮助你，但当你疲惫时，他们开始把你往错误的方向拉，导致你跌落。

解决方案：“智能调节器”（OCC）

作者意识到问题不在于助理教练不好，而在于他们声音的音量被锁定在一个固定的水平。

训练早期： 助理教练非常有用。我们应该将他们的音量调高。
训练晚期： 助理教练开始变得令人分心。我们应该将他们的音量调低。

他们创建了一个名为**最优系数校准（OCC）**的新系统。

工作原理（创意隐喻）：
将 OCC 想象成训练过程的智能恒温器。

它不是将温度（训练权重）设定为一个固定数字，而是不断检查房间状况。
它使用一个“代理”（一种快速、廉价的传感器）来检查助理教练当前是在帮忙还是在添乱。
如果助理教练朝正确的方向推动，恒温器将音量调高。
如果助理教练开始漂移并制造噪声，恒温器将音量调低。

这一过程自动、逐步地进行，无需停下来计算会拖慢整体速度的复杂数学公式。

结果：圆满结局

作者在六项高难度的数学竞赛（如 AIME 和奥林匹克竞赛）上测试了这种新的“智能调节器”。

它击败了旧方法： 使用“智能调节器”（OCC）训练的机器人，比仅使用主教练（断开）训练的机器人解决了更多问题。
它修复了崩溃： 与起初强劲随后失败的“策略损失”方法不同，OCC 方法从第一步到最后一步都保持强劲。
它速度很快： “智能调节器”没有拖慢训练速度。它和旧的“断开”方法一样快，因为它使用了一个巧妙的捷径来检查音量，而不是进行繁重的计算。

总结

这篇论文证明，你可以用一位主教练和一位助理教练来训练机器人，但你不能让他们保持在固定设置下。你需要一个动态系统，该系统能倾听训练过程，并实时调整助理教练的影响力。当他们这样做时，机器人学得更好、更快，达到了前所未有的水平。

技术摘要：通过最优系数校准实现多 Token 预测与强化学习的联合训练

问题陈述

基于可验证奖励的强化学习（RLVR）已成为增强大语言模型（LLM）推理能力的标准范式。与此同时，多 Token 预测（MTP）作为一种在预训练期间训练以预测多个未来 Token 的辅助目标，已在样本效率方面展现出显著优势，并可作为推测性解码的草稿模型。

直观而言，在后期训练中联合优化 MTP 和 RL 目标，应能使模型在获得推理能力的同时受益于多 Token 监督。然而，实证证据表明，联合 MTP-RL 训练往往会导致主模型性能显著下降。因此，当前的实践（例如在 veRL 和 slime 等框架中）通常在 RL 训练期间将 MTP 梯度与主模型解耦。这种解耦迫使 MTP 参数仅通过独立的后期微调进行适应，从而牺牲了来自同步多 Token 监督的潜在收益。核心问题在于：为什么联合 MTP-RL 训练会失败，以及能否设计一种原则性策略以实现有效的联合训练？

方法论

理论分析

作者从优化角度切入，分析 MTP 对 RL 目标 $J(\theta)$ 的每一步影响。假设目标函数为 $L$ -平滑，他们推导出了每一步改进的下界。这种分解揭示了 MTP 的影响由两个相互竞争的项支配：

一阶相关性项：由 RL 梯度（ $g_{RL}$ ）与 MTP 梯度（ $g_{MTP}$ ）之间的方向一致性（内积）决定。
二阶扰动惩罚：一个非正项，反映了由 MTP 梯度幅度引入的扰动。

仅当相关性项超过惩罚项时，性能才会提升。

现有训练机制分析

利用这一分解，本文统一并解释了三种现有的训练机制：

解耦（Detach）：针对主模型， $g_{MTP} = 0$ 。 $\Delta_{MTP} = 0$ 。这是安全的，但放弃了潜在收益。
交叉熵（CE）损失：相关性项在期望上接近于零，因为 RL 会对高奖励样本进行加权，而 CE 则平等对待所有样本。因此，负向的二阶惩罚占主导地位，导致性能下降。
策略损失（Policy Loss）：对 MTP 使用相同的 RL 目标公式。在初期，梯度是对齐的，相关性项占主导地位，从而产生收益。然而，随着训练进行，模型进入损失景观的平坦区域，相关性衰减，而梯度范数（及其导致的惩罚）依然存在。这导致性能曲线呈现“先升后降”的趋势，固定系数无法跟踪这种相变。

proposed 解决方案：最优系数校准（OCC）

为了解决相变问题，作者提出了最优系数校准（OCC），这是一种在线跟踪最优 MTP 系数（ $\lambda$ ）的自适应方案。

闭式最优解：分析得出了每一步最优 $\lambda^*$ 的闭式表达式，其正比于梯度相关性（ $c$ ）与梯度范数平方（ $v^2$ ）之比。
对数概率代理：在大规模分布式系统中，计算用于 $c$ $c$ 和 $v^2$ $v^{2}$ 的全模型梯度代价过高。OCC 利用从微小步长更新下的对数概率变化（ $\delta \approx \log \pi_\theta - \log \pi_{old}$ $δ \approx lo g π_{θ} - lo g π_{o l d}$ ）导出的对数概率代理。
- 对齐代理（ $\hat{c}_t$ ）：估计内积。
- 方差代理（ $\hat{v}^2_t$ ）：估计范数平方。
自适应更新：系数更新为 $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$ 。这使得 OCC 能够动态调整 MTP 的权重，在整个训练过程中保持正向收益，而无需计算全模型梯度。

主要贡献

理论统一：本文提供了理论分析，统一了解耦、交叉熵和策略损失机制，基于一阶相关性与二阶扰动之间的相互作用解释了它们的成功或失败。
OCC 算法：提出了最优系数校准，这是一种利用对数概率代理在线校准 MTP 系数的自适应过程。该方法计算开销极小，且避免了全模型梯度计算。
实证验证：在多个模型（MiMo-7B、GLM-4.5-Air）和算法（DAPO、GSPO）上，针对竞赛级数学推理基准进行了广泛实验，结果表明 OCC 始终匹配或优于解耦基线。

实验结果

作者在六个基准测试中评估了 OCC：AIME24、AIME25、AMC、MATH、Minerva 和 OlympiadBench。

性能：在所有基准测试和模型规模（7B 稠密模型和 106B MoE 模型）上，OCC 始终优于解耦基线和固定系数策略损失。
- 在 MiMo-7B 配合 DAPO 的实验中，OCC 实现了 61.7% 的平均准确率，超过了解耦（58.9%）和策略损失（57.4%）。
- 在 GLM-4.5-Air 上，OCC 实现了 67.6% 的平均准确率，而解耦为 65.9%。
训练动态：策略损失显示出先升后降的趋势（在后期阶段性能下降），而 OCC 在整个训练过程中保持了稳定的提升。CE 损失始终表现不如解耦。
效率：OCC 引入了可忽略的开销。OCC 每步的墙钟训练时间为 8.07 秒，与解耦（8.11 秒）相当，而全模型反向传播需要 45.06 秒（慢了 5.6 倍）。
消融实验：敏感性分析表明，预定义的缩放因子 $\lambda_+$ 在一系列数值范围内具有鲁棒性，其中 $\lambda_+ = 1.0$ 产生了最佳平均性能。

意义与主张

本文主张，联合 MTP-RL 训练的失败并非固有属性，而是源于静态系数与训练期间梯度对齐的动态漂移之间的不匹配。通过提供优化景观的理论分解并提出 OCC，作者证明了 MTP 可以被安全地重新引入到 RL 后期训练中。

其意义在于使模型能够同时通过 RL 获得推理能力，并利用多 Token 表示，而无需牺牲性能或承担过高的计算成本。这项工作表明，自适应校准对于非凸深度学习景观中的联合训练至关重要，为当前梯度解耦的行业标准提供了一种原则性的替代方案。

局限性

作者指出，他们的理论分析依赖于 $L$ -平滑假设，这在非凸景观中可能不完全成立，因此提供的是定性而非定量的指导。对数概率代理是一阶近似，在大学习率下可能会失去保真度。此外，实验主要集中在具有可验证奖励的数学推理上；将其推广到具有学习奖励模型的 RLHF 或开放式生成等场景仍有待验证。