通过最优系数校准在强化学习中联合训练多标记预测
机器学习
2026-05-28 v1
摘要
从可验证奖励的强化学习（RLVR）已成为提高大语言模型推理能力的标准范式，而多标记预测（Multi-Token Prediction, MTP）则是预训练中广泛采用的模块。将两者联合训练是自然的做法，但当前的 RL 实践因联合训练导致性能下降而分离 MTP 梯度。我们从优化视角重新审视这一失败。我们表明，MTP 对 RL 目标的每一步效应可分解为两个术语：一个一阶相关性和一个二阶扰动惩罚。这一分解统一了三个 MTP 训练范式：Detach、交叉熵损失和策略损失，并解释了每种方法为何成功或失败。进一步的政策损失分析揭示，尽管其与直觉一致，性能仍下降：相关项衰减，而二次惩罚持续存在。基于我们的分析，我们提出了最优系数校准（Optimal Coefficient Calibration, OCC），一种通过 log 概率代理在线跟踪最优系数的自适应方案，开销极低。在六个竞赛级别的数学推理基准测试中，OCC 始终匹配或超过分离基准，实现了改进的联合 MTP-RL 训练性能。
引用
@article{arxiv.2605.28184,
title = {Joint Training of Multi-Token Prediction in Reinforcement Learning via Optimal Coefficient Calibration},
author = {Zili Wang and Jiajun Chai and Lin Chen and Xiaohan Wang and Shiming Xiang and Guojun Yin},
journal= {arXiv preprint arXiv:2605.28184},
year = {2026}
}