Soft-MTP：从噪声梯度中找回多Token预测的有效信号

导语

多Token预测（Multi-Token Prediction，MTP）已经成为语言模型预训练中的常见设计。它让模型不仅预测下一个Token，还尝试预测更远位置的Token，既可能提升下游任务表现，也能为推测解码提供更强的草稿预测能力。不过，MTP的收益并不只取决于增加几个预测头：梯度如何回传到共享主干，同样决定了训练信号的质量。

标准MTP的问题在哪里

这项研究指出，标准MTP训练中的交叉熵梯度并非完全由有用信号组成。研究者将其精确分解为两部分：一部分是KL散度形式的自蒸馏梯度，促使各个MTP头学习模型自身的下一Token预测分布；另一部分则是残差。前者有助于让多Token预测头与主模型分布保持一致，并改善推测解码中的候选接受效果。后者信噪比较低，却可能将噪声传回模型主干，同时干扰下一Token学习和多Token学习。

论文特别强调，这种干扰会随着模型规模扩大而更加突出。也就是说，简单地把标准MTP损失加入预训练目标，并不一定能充分利用MTP带来的额外监督。

Soft-MTP如何处理

研究者提出了“KL梯度恢复”技巧。其关键前提是，NTP和MTP使用共享的线性输出头。借助直通式梯度路由，可以组合现有的NTP与MTP交叉熵梯度，精确恢复KL散度对应的梯度，而不必直接计算需要词表大小logit矩阵的KL损失。

基于这一技巧，Soft-MTP只需对标准MTP进行一行代码级修改，就能移除低信噪比残差，仅反向传播KL自蒸馏信号。它不增加参数，也不引入可测量的额外计算开销，并且能够继续使用常见的融合交叉熵内核。论文称，该方法同时适用于DeepSeek式和Meta式MTP，并在167M至1.2B参数规模范围内优于标准MTP。

意义与影响

Soft-MTP的价值不在于重新设计MTP结构，而在于重新审视MTP训练中“哪些梯度值得保留”。如果额外预测头的主要作用是学习与主模型一致、同时适合推测解码的分布，那么直接提取KL自蒸馏信号可能比让所有梯度无差别回传更有效。

这一方法还具有较强的工程吸引力：它不依赖额外参数，不要求改变融合损失算子，也不需要构造庞大的中间分布矩阵。因此，对于已经部署MTP预训练流程的团队，Soft-MTP更像是一次低成本的训练目标修正，而不是全新架构迁移。需要注意的是，现有素材未提供完整实验表格，关于更大模型规模、更多预测头以及不同任务上的收益，仍需结合论文全文进一步评估。

来源：Hugging Face Daily Papers