多Token预测的正确打开方式

Hugging Face Daily Papers 论文
multi-token-prediction language-models self-distillation speculative-decoding training-optimization

摘要

Soft-MTP 提出一种 KL 梯度恢复技巧，去除多token预测训练中的噪声残差梯度，在无额外开销下改进困惑度、投机解码与推理能力。

多Token预测（Multi-Token Prediction, MTP）已成为语言模型预训练中的常见组件，因为它能够提升下游性能，并通过投机解码实现更快的推理速度。我们发现了标准MTP训练中的一个缺陷：MTP梯度中的一个残差分量会同时干扰下一个token的学习和多token预测学习，且随着模型规模增大干扰愈发严重。具体来说，MTP的交叉熵梯度可以精确分解为两部分：一项KL自蒸馏项，将每个MTP头拉向模型自身的下一个token预测（NTP）分布，从而改善投机解码；以及一个信噪比较低的残差项，它会淹没有用的KL信号。我们提出了KL梯度恢复技巧：当两个分布共享同一个线性输出头时，可以通过直通梯度路由（straight-through gradient routing）将它们的交叉熵梯度组合，精确恢复出它们之间KL散度的梯度。与直接最小化KL散度不同，该技巧无需词表大小的logit矩阵，且与标准的融合交叉熵内核完全兼容。这使得Soft多Token预测（Soft-MTP）成为可能——仅需一行代码修改即可在MTP中去除噪声残差，通过复用已有的NTP和MTP交叉熵计算，仅反向传播KL自蒸馏信号。Soft-MTP同时支持DeepSeek风格和Meta风格的MTP，不增加任何参数，且无可测量的额外开销。从1.67亿到12亿参数规模，Soft-MTP在所有规模上均优于标准MTP，表现为更低的NTP困惑度、更快的投机解码和更高的零样本准确率；从6.89亿参数起，MTP困惑度也更低；在12亿参数下微调后，数学和代码推理能力也更强。Soft-MTP在大多数指标上与直接最小化KL散度持平甚至超越。可以说，Soft-MTP是MTP优化的一顿免费午餐。我们的代码已开源：https://github.com/JaydenTeoh/Soft-MTP。
查看原文
查看缓存全文

缓存时间: 2026/10/30 11:09

论文页面 - 正确实现多Token预测

来源: https://huggingface.co/papers/2610.38331

摘要

Soft-MTP 通过直通 KL 恢复技巧去除了多Token预测中带有噪声的梯度残差，从而在不增加额外开销的情况下提升了困惑度、投机解码和推理能力。

多Token预测（https://huggingface.co/papers?q=Multi-Token%20Prediction)（MTP）如今已成为语言模型预训练的常见组成部分，因为它能够提升下游性能并通过投机解码（https://huggingface.co/papers?q=speculative%20decoding) 实现更快的推理。我们发现了标准 MTP 训练中的一个缺陷：MTP 梯度中的一个残差分量会同时干扰下一个Token和多Token学习，且随着规模增大干扰愈发严重。具体而言，MTP 的交叉熵梯度（https://huggingface.co/papers?q=cross-entropy%20gradient) 可以被精确分解为两个部分：一个是 KL 自蒸馏（https://huggingface.co/papers?q=KL%20self-distillation) 项，它将每个 MTP 头拉向模型自身的下一个Token预测（https://huggingface.co/papers?q=next-token%20prediction)（NTP）分布并改善投机解码（https://huggingface.co/papers?q=speculative%20decoding)；另一个是信噪比极低的残差，它会掩盖有用的 KL 信号。我们提出了 KL 梯度恢复技巧：当两个分布共享一个线性输出头时，可以通过直通梯度路由（https://huggingface.co/papers?q=straight-through%20gradient%20routing) 组合它们的交叉熵梯度（https://huggingface.co/papers?q=cross-entropy%20gradient) 来精确恢复其 KL 散度的梯度。与直接最小化 KL 不同，该技巧无需词表大小的 logit 矩阵，并且与标准的融合交叉熵内核兼容。由此实现了 Soft 多Token预测（https://huggingface.co/papers?q=Multi-Token%20Prediction)（Soft-MTP（https://huggingface.co/papers?q=Soft-MTP)），只需对 MTP 进行一行代码的修改即可移除噪声残差，通过复用现有的 NTP 和 MTP 交叉熵计算来仅反向传播 KL 自蒸馏（https://huggingface.co/papers?q=KL%20self-distillation) 信号。Soft-MTP（https://huggingface.co/papers?q=Soft-MTP) 同时支持 DeepSeek 和 Meta 风格的 MTP，不增加任何参数，且无可测量的额外开销。从 167M 到 1.2B 参数量，Soft-MTP（https://huggingface.co/papers?q=Soft-MTP) 在每个规模下均优于标准 MTP，具有更低的 NTP 困惑度、更快的投机解码（https://huggingface.co/papers?q=speculative%20decoding) 和更高的零样本准确率，从 689M 起具有更低的 MTP 困惑度，并在 1.2B 微调后具有更强的数学和代码推理能力。Soft-MTP（https://huggingface.co/papers?q=Soft-MTP) 在大多数指标上也达到了甚至超越了直接 KL 散度最小化的效果。Soft-MTP（https://huggingface.co/papers?q=Soft-MTP) 实际上是 MTP 优化的一顿免费午餐。我们的代码可在 https://github.com/JaydenTeoh/Soft-MTP (https://huggingface.co/papers?q=Soft-MTP) 获取。

查看 arXiv 页面 (https://arxiv.org/abs/2610.38331) · 查看 PDF (https://arxiv.org/pdf/2610.38331) · GitHub1 (https://github.com/JaydenTeoh/Soft-MTP) · 添加到收藏 (https://huggingface.co/login?next=%2Fpapers%2F2610.38331)

在你的 agent 中获取这篇论文：

hf papers read 2610.38331

还没有最新版 CLI？curl -LsSf https://hf.co/cli/install.sh | bash

引用此论文的模型 12

JaydenTeoh/ntp-167m-model 文本生成 • 0.2B • 更新于约 4 小时前 • 1 (https://huggingface.co/JaydenTeoh/ntp-167m-model)

JaydenTeoh/hard-mtp-167m-model 文本生成 • 0.2B • 更新于约 4 小时前 • 1 (https://huggingface.co/JaydenTeoh/hard-mtp-167m-model)

JaydenTeoh/soft-mtp-167m-model 文本生成 • 0.2B • 更新于约 4 小时前 (https://huggingface.co/JaydenTeoh/soft-mtp-167m-model)

JaydenTeoh/hard-mtp-332m-model 文本生成 • 0.4B • 更新于约 4 小时前 (https://huggingface.co/JaydenTeoh/hard-mtp-332m-model)

浏览 12 个引用此论文的模型 (https://huggingface.co/models?other=arxiv:2610.38331)

引用此论文的数据集 0

没有数据集链接此论文

在数据集的 README.md 中引用 arxiv.org/abs/2610.38331 以从此页面链接。

引用此论文的 Spaces 0

没有 Space 链接此论文

在 Space 的 README.md 中引用 arxiv.org/abs/2610.38331 以从此页面链接。

包含此论文的合集 1

相似文章

AdaMTP: An Adaptive Training Paradigm for Multi-Token Prediction

arXiv cs.CL

This paper introduces AdaMTP, an adaptive training paradigm for multi-token prediction that dynamically aligns prediction horizons with sequence predictability using entropy-based segmentation, consistently outperforming standard MTP on math, code, and general benchmarks across three LLM backbones.

多令牌残差预测

arXiv cs.LG

引入多令牌残差预测（MRP），这是一个用于扩散语言模型的轻量级模块，能够在单次主干前向传播中实现依赖感知的多令牌去噪，实现高达1.42倍的无损加速。