[论文] 正确实现的多 Token 预测——Soft-MTP

多 Token 预测（MTP）现已成为语言模型预训练中的常见组成部分，因为它能提升下游性能，并通过投机解码实现更快的推理。我们发现了标准 MTP 训练中的一个缺陷：MTP 梯度的残差分量会同时干扰下一 token 学习和多 token 学习，而且随着规模增大，这种干扰会越来越严重。具体来说，MTP 交叉熵梯度可以精确分解为一个 KL 自蒸馏项和一个残差项；前者将每个 MTP 头拉向模型自身的下一 token 预测（NTP）分布，并改善投机解码，后者信噪比低，会掩埋有用的 KL 信号。我们引入了 KL 梯度恢复技巧：当两个分布共享一个线性输出头时，可以通过直通梯度路由组合它们的交叉熵梯度，从而精确恢复它们 KL 散度的梯度。与直接最小化 KL 不同，这一技巧不需要词表大小的 logit 矩阵，并且兼容标准的融合交叉熵内核。这使得 Soft 多 Token 预测（Soft-MTP）成为可能：只需对 MTP 做一行代码改动，就能去除噪声残差，并通过复用现有的 NTP 和 MTP 交叉熵计算，仅反向传播 KL 自蒸馏信号。Soft-MTP 同时支持 DeepSeek 风格和 Meta 风格的 MTP，不增加参数，也不产生可测量的开销。从 167M 到 1.2B 参数，Soft-MTP 在每个规模上都优于标准 MTP：更低的 NTP 困惑度、更快的投机解码和更高的零样本准确率；从 689M 起 MTP 困惑度更低；在 1.2B 上微调后数学和代码推理更强。在大多数指标上，Soft-MTP 也匹配甚至超过直接最小化 KL 散度。实际上，Soft-MTP 是 MTP 优化的免费午餐。arXiv：https://arxiv.org/abs/2610.38331 PDF：https://arxiv.org/pdf/2610.38331 GitHub：https://github.com/JaydenTeoh/Soft-MTP HuggingFace：https://huggingface.co/