arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2610.38331cs.LG
收藏

正确实现的多令牌预测

Multi-Token Prediction Done Right

MIT(麻省理工学院)

Jayden Teoh

AI总结:

本文针对多令牌预测(MTP)训练的缺陷,提出KL梯度恢复技巧,实现Soft-MTP,仅需一行代码修改,无额外参数与开销,在各规模下均优于标准MTP,提升了语言模型的性能。

摘要:

多令牌预测(Multi-Token Prediction, MTP)现已成为语言模型预训练中的常用组件,它可提升下游性能并通过推测解码实现更快推理。我们发现标准MTP训练存在一个缺陷:MTP梯度的残差分量会同时干扰下一个令牌(next-token)和多令牌学习,且这种干扰会随模型规模增大而加剧。具体而言,MTP交叉熵梯度可精确分解为两部分:一是KL自蒸馏项,它会将每个MTP头拉向模型自身的下一个令牌预测(next-token prediction, NTP)分布,从而改进推测解码;二是残差项,其信噪比低,会掩盖有用的KL信号。我们提出了KL梯度恢复技巧:当两个分布共享一个线性输出头时,通过直通梯度路由组合它们的交叉熵梯度,可精确恢复其KL散度的梯度。与直接最小化KL散度不同,该技巧无需词汇量大小的logit矩阵,且兼容标准的融合交叉熵内核。这使得Soft多令牌预测(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散度最小化相当甚至更优,它实际上是MTP优化的免费午餐。我们的代码可在该https URL获取。

英文摘要:

Multi-Token Prediction (MTP) is now a common ingredient in language-model pre-training, as it improves downstream performance and enables faster inference through speculative decoding. We identify a flaw in standard MTP training, in which a residual component of the MTP gradient interferes with both next-token and multi-token learning, increasingly so with scale. Specifically, the MTP cross-entropy gradient can be decomposed exactly into a KL self-distillation term that pulls each MTP head toward the model's own next-token prediction (NTP) distribution and improves speculative decoding, and a residual that has a low signal-to-noise ratio and buries the useful KL signal. We introduce the KL gradient-recovery trick: when two distributions share a linear output head, the gradient of their KL divergence can be recovered exactly by combining their cross-entropy gradients through straight-through gradient routing. Unlike minimizing the KL directly, this trick requires no vocabulary-sized logit matrices and is compatible with standard fused cross-entropy kernels. This enables Soft Multi-Token Prediction (Soft-MTP), a one-line code change to MTP that removes the noisy residual and backpropagates only the KL self-distillation signal by reusing the existing NTP and MTP cross-entropy computations. Soft-MTP supports both DeepSeek- and Meta-style MTP, adds no parameters, and incurs no measurable overhead. From 167M to 1.2B parameters, Soft-MTP improves over standard MTP with lower NTP perplexity, faster speculative decoding and higher zero-shot accuracy at every scale, lower MTP perplexity from 689M onward, and stronger math and code reasoning after fine-tuning at 1.2B. Soft-MTP also matches, if not exceeds, direct KL divergence minimization on most metrics. Soft-MTP is, in effect, a free lunch for MTP optimization. Our code can be found at https://github.com/JaydenTeoh/Soft-MTP.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑