← 返回 2026-05-28

通过最优系数校准实现强化学习中多令牌预测的联合训练 Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration

Zili Wang, Jiajun Chai, Lin Chen, Xiaohan Wang, Shiming Xiang, Guojun Yin 📅 2026-05-27 👍 6 2026-07-13 08:36
优化理论 多令牌预测 强化学习后训练 数学推理
arXiv HuggingFace PDF

提出自适应系数校准方法OCC，解决RL与MTP联合训练的性能下降问题

前置知识

Reinforcement Learning from Verifiable Rewards (RLVR)

RLVR是一种强化学习范式，用于在模型后训练阶段提升大语言模型的推理能力。它通过可验证的奖励信号（如数学题的正确性）来优化模型策略，代表性的算法包括PPO、GRPO、DAPO和GSPO。这些算法通过估计advantage（优势值）来对样本进行加权，高奖励样本被放大，低奖励样本被抑制，从而引导模型向产生更高质量输出的方向优化。

本文的研究背景正是RLVR在数学推理任务上的应用，理解RLVR的梯度特性（advantage零均值特性）是理解为什么Cross-Entropy Loss会失败的关键。

Multi-Token Prediction (MTP)

MTP是一种在预训练阶段广泛使用的技术，让模型同时预测未来的多个token。在主模型之后添加MTP头，训练时主模型预测下一个token，MTP头预测后续的k个token。MTP在预训练中已被证明能够提升样本效率、改善下游任务性能，并且可以作为投机解码的draft模型加速推理。代表性的采用MTP的模型包括DeepSeek-V3/V4、LongCat-Flash、StepFun-3.5-Flash、MiMo和Qwen3-Next等。

本文的核心问题是如何在RL后训练阶段保留MTP的益处。理解MTP的工作原理和训练方式是理解为什么联合训练会失败以及如何设计OCC方法的基础。

L-平滑性 (L-smoothness)

L-平滑性是优化理论中的一个标准假设，它要求目标函数的梯度是Lipschitz连续的。具体来说，对于L-平滑函数，存在常数L > 0使得对于所有参数θ和θ′，有‖∇J(θ) − ∇J(θ′)‖ ≤ L‖θ − θ′‖。这个假设允许我们使用二次函数来近似目标函数的局部行为，从而推导出单步改进的下界。在本文中，作者从L-平滑性出发，推导出了MTP对RL目标影响的分解公式。

L-平滑性是本文理论分析的起点，作者基于L-平滑性推导出了核心的公式（3）和（4），将MTP的影响分解为一阶相关项和二阶惩罚项。理解这个假设对于理解论文的理论贡献和局限非常重要。

研究动机

在强化学习后训练中联合训练Multi-Token Prediction（MTP）会导致严重的性能下降。这是一个普遍观察到的问题，包括veRL、slime等主流框架都明确报告了这种现象，并建议将MTP梯度从主模型detach。例如，论文Figure 1显示，在AIME24数据集上，直接使用Cross-Entropy Loss进行联合训练会使得准确率从45.3%下降到16.1%。即使使用与RL相同的Policy Loss，虽然早期表现超过detach基线（48.9 vs 48.3），但最终会下降到38.9%，低于detach的45.3%。这意味着现有方法无法安全地将MTP引入RL后训练，牺牲了多token监督可能带来的收益。

本文的目标是本文的目标是从优化角度深入理解为什么联合MTP-RL训练会失败，并设计一个能够真正实现有效联合训练的原理性方法。具体来说，作者希望：(1)建立一个理论框架来分析MTP对RL目标的影响机制；(2)解释为什么现有的三种训练策略（Detach、Cross-Entropy Loss、Policy Loss）会表现出不同的行为；(3)提出一个自适应的MTP系数校准策略，能够在整个训练过程中保持稳定的性能提升。

与已有工作不同的是，本文的独特切入角度是从优化理论的角度分析MTP-RL联合训练问题，而不是简单地尝试不同的损失函数或调参。作者发现MTP对RL目标的影响可以分解为一阶相关项和二阶惩罚项，这个分解统一了三种训练策略的解释。更重要的是，作者识别出一个关键现象：随着训练进展，梯度的对齐度会发生变化（从陡峭区域到平坦区域），导致固定的MTP系数无法追踪这种相变。本文的OCC方法是第一个基于理论分析的自适应系数校准策略，它使用对数概率代理在线追踪理论最优系数，计算开销可以忽略不计。

核心方法

OCC（Optimal Coefficient Calibration）的核心思路是动态调整MTP损失的系数λ，使得在每一步都能最大化MTP对RL目标的正面贡献。作者从L-平滑性出发，推导出了MTP对RL目标的单步影响下界：ΔMTP = ηλ(1−Lη)⟨gRL, gMTP⟩ − (Lη²λ²/2)‖gMTP‖²，其中第一项是一阶相关项（反映梯度方向的一致性），第二项是二阶惩罚项（反映MTP梯度引入的扰动）。当相关项大于惩罚项时，MTP能够提升RL性能；反之则会导致性能下降。OCC的目标就是自适应地调整λ，使得ΔMTP始终为正。

核心创新点在于将MTP的系数λ从固定值改为自适应值，并且这个自适应值有理论依据。作者将ΔMTP视为λ的函数，发现它是一个关于λ的开口向下的抛物线：ΔMTP(λ) = η(1−Lη)c·λ − (Lη²v²/2)·λ²，其中c = ⟨gRL, gMTP⟩是梯度内积，v² = ‖gMTP‖²是MTP梯度的平方范数。这个抛物线的最大值在λ* = (1−Lη)c/(Lηv²)处取得。然而，在大规模分布式训练系统中计算完整的模型梯度非常昂贵（作者测量的开销是5.6倍）。作者使用对数概率变化作为梯度的代理：⟨g1, g2⟩ ∝ ⟨δ1, δ2⟩，‖g‖² ∝ ‖δ‖²，其中δ = logπθ − logπold是小步更新下的对数概率变化。这个代理的计算开销可以忽略不计，但与真实梯度量有很高的相关性（r = 0.78和r = 0.74）。

方法步骤详情

OCC方法包含三个主要步骤：(1)在每步训练中，首先进行标准的RL forward和MTP forward，得到logits和log probabilities；(2)计算两个在线统计量：对齐代理ĉt = ⟨δRL, δMTP⟩（估计梯度内积c）和方差代理ṽ²t = ‖δMTP‖²（估计MTP梯度的平方范数v²），其中δ = logπθ − logπold是当前策略和旧策略的log概率变化；(3)使用公式λt = λ+ · ĉt/(ṽ²t + ε)计算最优MTP系数，其中λ+是一个预定义的比例因子（吸收未知的光滑性前置因子），ε是防止除零的小常数。在backward阶段，使用计算出的λt来加权MTP损失，让梯度以适当的强度流入主模型。

技术新颖性

OCC的技术新颖性体现在多个方面。首先，它是第一个基于理论分析的自适应MTP系数校准方法，而不是启发式的调参。其次，它巧妙地避开了昂贵的全模型梯度计算，使用对数概率代理实现了高效的自适应调整（训练时间开销与detach基本相同，而全梯度计算需要5.6倍时间）。第三，OCC允许系数为负值，这意味着当MTP梯度与RL梯度方向相反时，系统会自动反转MTP梯度的方向，将其从有害信号转变为校正信号。论文的消融实验显示，允许负值的OCC-NoClip比裁剪到非负的OCC-Clip表现更好（平均准确率61.7 vs 60.3）。第四，OCC是算法无关的，在DAPO和GSPO上都表现良好，并且能够泛化到不同的模型规模（从7B到106B MoE）。

Evolution of the policy-aligned gain across training steps (Eq. (10)). Each curve is a downward-opening parabola in λ, coloured from early (dark purple) to late (bright yellow) steps. The vertical dashed line marks a fixed λ: at early steps, the intersection point of λ and the parabola lies above zero (positive gain), but as training progresses the parabola drifts so that the intersection point falls under zero (negative gain).
Figure 3: Evolution of the policy-aligned gain across training steps (Eq. (10)). Each curve is a downward-opening parabola in λ, coloured from early (dark purple) to late (bright yellow) steps. The vertical dashed line marks a fixed λ: at early steps, the intersection point of λ and the parabola lies above zero (positive gain), but as training progresses the parabola drifts so that the intersection point falls under zero (negative gain).
Curve of the absolute value of OCC coefficient λt.
Figure 4: Curve of the absolute value of OCC coefficient λt.
Relationship between the log-probability proxies and the true full-model gradient quantities. Each point corresponds to one training step. Left: the alignment proxy plotted against the full-model first-order correlation. Right: the variance proxy plotted against the full-model second-order penalty. The proxies exhibit highly positive correlation with the true gradient quantities (r is annotated above).
Figure 8: Relationship between the log-probability proxies and the true full-model gradient quantities. Each point corresponds to one training step. Left: the alignment proxy plotted against the full-model first-order correlation. Right: the variance proxy plotted against the full-model second-order penalty. The proxies exhibit highly positive correlation with the true gradient quantities (r is annotated above).

实验结果

实验结果在多个竞争级数学推理基准上验证了OCC的有效性。在MiMo-7B-RL使用DAPO的设置下，OCC在所有六个基准上都达到或超过了detach基线：AIME24（45.9 vs 45.3）、AIME25（46.7 vs 36.7，+10.0）、AMC（83.1 vs 81.9）、MATH（91.0 vs 90.0）、Minerva（43.4 vs 39.7）、OlympiadBench（60.3 vs 60.1），平均准确率61.7 vs 58.9（+2.8）。使用GSPO算法的结果类似：OCC平均60.1 vs detach的57.8。在更大的GLM-4.5-Air（106B-A12B MoE）模型上，OCC也取得了稳定的改进：平均67.6 vs 65.9（+1.7）。Cross-Entropy Loss在所有设置中都表现最差，平均准确率仅47.7（vs detach的58.9），最大差距出现在AIME24（16.1 vs 45.3，-29.2）。Policy Loss呈现出先升后降的趋势：在训练第30步时准确率为48.9（超过detach的48.3），但最终下降到38.9（低于detach的45.3）。OCC的训练曲线显示它能够有效捕捉训练过程中的相变，动态调整系数并保持性能提升。消融实验验证了对数概率代理的保真度：ĉt与真实梯度内积c的相关系数为0.784，ṽ²t与真实二阶惩罚v²的相关系数为0.740。OCC的训练开销与detach基本相同（8.07s vs 8.11s），远低于全梯度计算的45.06s（5.6倍慢）。

Main results on mathematical reasoning benchmarks across models and algorithms. We report avg@32 accuracy (%). Best results are in bold; second best are underlined.
Table 1: Main results on mathematical reasoning benchmarks across models and algorithms. We report avg@32 accuracy (%). Best results are in bold; second best are underlined.
Comparison of policy loss with λ ∈{0.1, 0.2, 0.5, 1.0} against Detach and OCC.
Table 2: Comparison of policy loss with λ ∈{0.1, 0.2, 0.5, 1.0} against Detach and OCC.
Sensitivity analysis of λ+. We sweep λ+ ∈{0.1, 0.2, 0.5, 1.0, 2.0} and compare the performance across various benchmarks. λ+ = 1.0 yields the best average performance, indicating it as the preferred choice based on overall results. Best in bold.
Table 3: Sensitivity analysis of λ+. We sweep λ+ ∈{0.1, 0.2, 0.5, 1.0, 2.0} and compare the performance across various benchmarks. λ+ = 1.0 yields the best average performance, indicating it as the preferred choice based on overall results. Best in bold.
Left: Accuracy on AIME24. CE loss shows great degradation, while policy loss initially surpasses but degrades as training progresses. In contrast, OCC sustains gains throughout training and consistently outperforms. Middle & Right: Illustrations of the MTP training regimes. Detach training stops the gradient from MTP into the main model. Joint MTP-RL training allows this gradient to flow back into the main model, enabling joint updates.
Figure 1: Left: Accuracy on AIME24. CE loss shows great degradation, while policy loss initially surpasses but degrades as training progresses. In contrast, OCC sustains gains throughout training and consistently outperforms. Middle & Right: Illustrations of the MTP training regimes. Detach training stops the gradient from MTP into the main model. Joint MTP-RL training allows this gradient to flow back into the main model, enabling joint updates.
查看结构化数据
任务指标本文基线提升
AIME24 avg@32 accuracy (%) 45.9 Detach (45.3) +0.6
AIME25 avg@32 accuracy (%) 46.7 Detach (36.7) +10.0
AMC avg@32 accuracy (%) 83.1 Detach (81.9) +1.2
MATH avg@32 accuracy (%) 91.0 Detach (90.0) +1.0
Minerva avg@32 accuracy (%) 43.4 Detach (39.7) +3.7
OlympiadBench avg@32 accuracy (%) 60.3 Detach (60.1) +0.2

局限与改进

作者承认了几个局限性。首先，理论分析依赖于L-平滑性假设，这在非凸的深度学习景观中可能不完全成立，因此推导的界限提供的是定性指导而非定量预测。其次，对数概率代理是一阶近似，在大学习率或激进策略更新时可能丧失保真度。第三，实验主要集中在数学推理的可验证奖励场景（RLVR），泛化到其他RL后训练设置（如RLHF使用学习到的奖励模型或开放生成任务）尚待验证。第四，预定义的比例因子λ+吸收了一个未知的光滑性前置因子，虽然消融实验显示在较大范围内都具有鲁棒性，但最优选择可能因训练配置而异。

独立分析的弱点

独立分析的第一个弱点是OCC的性能提升幅度相对有限（+2.8平均准确率）。在AIME24和OlympiadBench等基准上，提升甚至不到1个百分点。这可能意味着在这些任务上，MTP本身能提供的额外监督信息有限。改进方向是探索如何增强MTP head的容量或设计更好的MTP损失函数，使其能提供更有价值的信号。第二个弱点是OCC依赖于对数概率代理的保真度，如果代理失效（例如在训练初期策略变化剧烈时），性能可能不稳定。改进方向是探索更鲁棒的代理机制或自适应的代理加权策略。第三个弱点是实验只验证了数学推理任务，在其他任务上（如代码生成、对话优化）的效果未知。改进方向是进行更广泛的任务验证，并根据任务特点调整OCC的超参数。

未来方向

作者提出和基于成果可延伸的未来研究方向包括：(1)将OCC扩展到RLHF等其他RL后训练场景，验证其在使用学习到的奖励模型时的有效性；(2)探索如何更好地选择或调整λ+这个预定义比例因子，可能根据训练阶段的自动调整；(3)研究OCC在更大规模模型（如百亿参数以上）上的效果和计算效率；(4)分析OCC在不同奖励机制下的行为特性，特别是在奖励信号噪声较大或延迟较高的情况下；(5)探索将OCC的思想应用到其他联合训练场景，例如多任务学习、多模态训练等；(6)研究如何将OCC与其他训练技术（如课程学习、对抗训练）结合，进一步提升模型性能。

复现评估

论文的复现性相对较好。作者提供了开源代码（https://github.com/MarkXCloud/RL-MTP），使用了公开可用的训练数据（DAPO-Math-17k，Apache-2.0许可）和评估基准。实验基于veRL框架实现，使用Megatron-LM作为训练后端，SGLang作为rollout引擎。训练配置（200步、batch size 128、学习率1e-6、temperature 1.0、top-p 0.7）都详细报告。硬件资源是128张NVIDIA H20 GPU。然而，复现成本较高：仅训练一次就需要128张H20 GPU，这对于大多数研究团队来说是不可及的。此外，论文没有提供checkpoint或中间结果，使得结果比对更加困难。代码的详细程度和文档质量未知，需要实际测试才能评估复现的难度。总体而言，有资源的团队应该能够复现主要结果，但独立验证的成本很高。

论文图表

Figure 2: Illustration of the training dynamics of coefficient components. (a) Curves of |c| of full model, CE loss, and alignment proxy. While the correlation of CE loss (green curve) remains near zero, the full model correlation of policy loss declines steadily. The proxy |ˆct| effectively captures the drift trends. (b) Curves of v2 of full model and variance proxy. ˆv2t also captures v2 and stabilizes around a similar magnitude in the later stage. (c) Curves of first-order correlation term and second-order penalty term. The first-order correlation term decays while the second-order penalty stays persistent. Their crossing flips ∆MTP from positive gain to negative gain.

这张图由三个子图组成，展示了训练过程中MTP系数各组成部分的动态变化。子图(a)展示了梯度内积绝对值|c|的曲线，CE loss的相关性保持在零附近，Policy loss的相关性稳步下降，而OCC的对齐代理ĉt有效捕捉了这种漂移趋势。子图(b)展示了梯度平方范数v²的曲线，OCC的方差代理ṽ²t能够捕捉v²并在后期稳定在相似的量级。子图(c)展示了一阶相关项和二阶惩罚项的曲线，一阶相关项随训练衰减，而二阶惩罚项持续存在，两者的交叉点使得ΔMTP从正收益翻转为负收益。

这张图是理解OCC理论基础的关键，它将论文的核心数学公式可视化。子图(c)特别重要，因为它展示了为什么Policy Loss会失败：一阶相关项衰减但二阶惩罚项持续存在，导致相变。子图(a)和(b)验证了对数概率代理的有效性，这是OCC能够高效运行的前提。

Figure 5: Wall-clock training time of OCC, Detach and Full-model.

这张图对比了Detach、OCC和全模型梯度计算的每步训练时间。OCC的平均步时间为8.07秒，与Detach的8.11秒基本相同，而全模型反向传播需要45.06秒，导致5.6倍的慢速。

这张图证明了OCC的效率优势。它可以在不增加显著开销的情况下实现自适应MTP系数校准，避免了计算完整模型梯度的 prohibitive 成本。这使得OCC适用于大规模分布式RL训练。

Figure 6: Training reward of CE Loss, Detach, Policy Loss and OCC.

这张图展示了CE Loss、Detach、Policy Loss和OCC四种策略的训练奖励曲线。CE Loss的奖励明显低于其他策略，Detach和Policy Loss在训练后期都落后于OCC。OCC达到了最高的奖励，展示了其改善联合MTP-RL训练的能力。

这张图提供了Table 1准确率结果的训练过程视角。它验证了OCC在整个训练过程中都保持优势，而不仅仅是最终评估点。这也再次证明了CE Loss的问题和Policy Loss的不稳定性。

Figure 7: Training reward of the Policy Loss with different coefficient.

这张图展示了使用不同系数λ（0.1、0.2、0.5、1.0）的Policy Loss的训练奖励曲线。所有固定系数的Policy Loss在训练后期都落后于OCC。

这张图补充了Table 2的分析，从训练过程的角度进一步验证了没有任何固定系数能够匹配OCC的性能。它强调了OCC的自适应调整带来的真正优势。