多Token预测还能更优？一行代码解锁Soft-MTP

多Token预测（Multi-Token Prediction, MTP）已经成为前沿预训练的标准配方，但一项最新研究指出，它的实现方式远未达到最优。研究者发现，标准MTP的梯度中隐藏着一个被噪声残差掩盖的KL自蒸馏信号——这个信号本可以提升推测解码（speculative decoding）的效果，而噪声残差却会损害下一Token预测（next-token prediction）的性能。通过一个名为Soft-MTP的简单修改，只需一行代码，就能将这两个信号解耦，让MTP在不同任务上获得普遍提升。

MTP的隐忧：梯度中的信号与噪声

MTP的核心思想是让模型在预测下一个Token的同时，也预测未来多个Token。这种做法在前沿预训练中已被广泛采用，因为它能提供更丰富的训练信号。然而，这项研究揭示了一个被忽视的问题：标准MTP的梯度并非纯粹的“有用信号”。

具体来说，梯度中包含了两个成分：一个是KL自蒸馏信号，它有助于模型学习更鲁棒的表示，从而改善推测解码；另一个是噪声残差，它会干扰下一Token预测的准确性。两者纠缠在一起，使得MTP在提升某些能力的同时，可能对基础语言建模造成损害。

作者将这一发现称为“标准MTP梯度隐藏了KL自蒸馏信号”。问题的关键在于，如何在不引入噪声的情况下，单独提取出有益的KL信号。

KL自蒸馏（KL Self-Distillation） 是知识蒸馏的一种变体：教师模型和学生模型共享同一套参数或结构，学生通过最小化与教师输出分布之间的KL散度来学习。在MTP场景中，模型对未来多个Token的预测分布可以视为一种“软标签”，用它来约束当前位置的下一Token预测分布，就形成了自蒸馏信号。这种信号之所以有益于推测解码，是因为它让模型输出的概率分布更加平滑、更接近真实的多步未来分布，从而在被草稿模型采样时更容易被目标模型接受。相反，噪声残差指的是MTP头引入的、与下一Token预测目标不一致的梯度成分，它并不携带关于下一个Token的额外信息，反而会拉偏共享主干网络的更新方向。

推测解码（Speculative Decoding） 是一种无损加速大模型推理的技术：先用一个轻量草稿模型（或模型自身的某个头）快速生成若干候选Token，再由目标模型并行验证这些候选，接受与目标分布一致的Token、拒绝并修正不一致的部分。其加速比直接取决于“接受率”——候选Token被目标模型接受的比例。MTP之所以能改善推测解码，正是因为它训练模型同时预测多个未来Token，使草稿分布更接近目标分布，从而提高接受率。理解这一点，才能明白为何保留MTP梯度中的KL自蒸馏信号、剔除噪声残差，会同时影响生成质量与推理效率。

一行代码的魔法：KL梯度恢复技巧

研究团队提出了一种优雅的解决方案：通过修改隐藏状态的残差连接，将KL梯度从噪声中“恢复”出来。具体操作仅需一行代码：

h_ntp += h_mtp.detach() - h_mtp

这行代码的作用是，在计算下一Token预测的隐藏状态时，加入MTP隐藏状态的“去梯度”版本与原始版本之差。这样，梯度在反向传播时只会流经KL自蒸馏部分，而噪声残差被有效阻断。作者将这一方法称为Soft-MTP，并认为它几乎是一种“免费的午餐”——不需要额外计算开销，也不改变模型结构，却能带来跨任务的性能提升。

detach() 与停止梯度（stop-gradient） 是PyTorch等深度学习框架中的基础操作：tensor.detach() 返回一个与原张量共享数据但脱离计算图的新张量，反向传播时梯度不会流经它。表达式 h_mtp.detach() - h_mtp 在数值上恒为零，因此前向传播时隐藏状态没有任何改变；但在反向传播时，由于减号右侧的 h_mtp 仍保留计算图，梯度会以取负的方式流经MTP分支。这种“前向为零、反向非零”的技巧常被称为梯度手术（gradient surgery）或零值残差注入，它让研究者能够在不改变模型输出的前提下，精确控制哪些梯度路径被保留、哪些被切断。

Soft-MTP：免费午餐的实证效果

根据论文描述，Soft-MTP在多种任务上均优于标准MTP。虽然原始素材未披露具体的基准测试数据，但作者强调“改进MTP across a range of tasks”，表明该方法具有较好的泛化性。

从原理上看，Soft-MTP的优势在于它让MTP的训练目标更加“纯净”：模型既能通过KL信号提升推测解码的接受率，又不会因噪声残差而牺牲下一Token预测的精度。这对于需要同时兼顾生成质量和推理效率的大模型来说，具有实际意义。

“免费午餐”在这里有明确的技术含义：Soft-MTP不增加任何前向或反向的计算量，不引入新参数，也不改变推理时的模型结构——所有改动只发生在训练阶段的梯度路径上。与之对比，许多提升MTP效果的方法需要增加预测头数量、扩大上下文窗口或引入额外的蒸馏教师模型，都会带来显存与算力开销。Soft-MTP的代价几乎为零，收益却体现在下一Token预测精度与推测解码接受率两个维度上，这正是作者称其为免费午餐的依据。不过需要注意，其效果仍依赖于原有MTP头的训练质量，若MTP头本身欠拟合，可分离出的KL信号也会有限。

开源资源与复现路径

研究团队已开源了代码、论文和HuggingFace模型集合，方便社区复现和进一步探索：

• 代码：https://github.com/JaydenTeoh/Soft-MTP
• 论文：https://arxiv.org/abs/2610.38331
• HuggingFace：https://huggingface.co/collections/JaydenTeoh/soft-multi-token-prediction

对于正在使用MTP进行预训练的研究者和工程师来说，这行代码改动成本极低，值得尝试。

对预训练实践的启示

这项工作的价值不仅在于提出一个技巧，更在于它提醒我们：即便是被广泛采纳的“标准配方”，也可能存在未被察觉的缺陷。MTP的梯度分解视角为后续研究提供了新的分析工具——通过分离不同训练信号，可以更精细地控制模型的学习过程。

未来，类似Soft-MTP的思路或许可以推广到其他多任务训练场景，帮助研究者在保持计算效率的同时，获得更干净、更可控的梯度信号。