多词预测的梯度里混着噪声：只保留自蒸馏项，1.2B模型零样本准确率反超单词预测
核心概要
多词预测（MTP，让模型一次预测多个未来词元）已是语言模型预训练的常见组件，但标准MTP的交叉熵梯度里混有一项低信噪比的残差，随规模增大而更嘈杂。Soft-MTP把这一残差从梯度中剔除，只回传其中隐含的KL自蒸馏信号，做法是复用已有的交叉熵梯度、不额外生成词表大小的logit矩阵，代码改动仅一行。在167M到1.2B参数的LLaMA式模型上，Soft-MTP在每个规模都优于标准MTP：下一词困惑度更低、自推测解码更快、零样本准确率更高；从689M起多词困惑度也更低，1.2B微调后数学与代码推理更强。
深度剖析
多词预测训练中一直混着一项对学习没有帮助的梯度成分，现在可以把它单独去掉，而保留其中真正有用的部分。 标准做法是用交叉熵对数据训练MTP头，前沿模型也如此；此前若想让MTP头对齐模型自身的下一词分布，需要显式最小化KL散度，必须同时生成教师和学生的全词表分布。 在共享主干上，残差项虽然信噪比低，却占到达主干的MTP梯度的80%；沿残差方向走一步会持续抬高下一词损失，而沿KL分量走一步则降低它。
去掉残差后，多词预测在语言建模、推测解码和下游任务上同时变好，且模型越大优势越明显。 标准MTP在小规模上并不比单词预测提升多选题准确率，这与已有观察一致。 在167M到1.2B四个规模上，Soft-MTP的下一词困惑度、五个领域的推测解码接受长度与加速比、零样本平均准确率均优于标准MTP；1.2B时零样本平均准确率为57.31，高于标准MTP的56.80和单词预测的56.75。
实现方式是在隐藏状态层面把KL梯度从已有计算中还原出来，不新增参数，也不增加可测开销。 直接最小化KL需要物化词表大小的logit矩阵，与预训练常用的融合交叉熵核不兼容。 1.2B规模下，Soft-MTP的矩阵乘FLOPs为每词元10.28 GFLOP、融合交叉熵峰值内存107.0 GB，与标准MTP相同，吞吐在1%以内；KL-MTP则为10.43 GFLOP、128.1 GB。
启示与展望
实验在LLaMA式模型上进行，规模从167M到1.2B参数，使用DeepSeek式顺序MTP结构并加三个MTP头，主实验在FineWeb-Edu上训练约31.5B词元；Meta式并行MTP的结果在332M规模、7.3B词元条件下给出。作者还讨论了该方法与数据并行、ZeRO、专家并行和流水线并行的兼容性，并指出序列切分是唯一需要额外通信的维度。
作者把Soft-MTP描述为以偏差换方差：去掉残差后MTP头偏向模型自身的下一词分布而非数据分布，这一取舍随规模增大而更有利，但偏差本身仍在。梯度信噪比与梯度占比的分析是在冻结检查点上离线测量的，作者说明这只是刻画主干与MTP头交界处的行为，实际训练更新还取决于优化器。多词预测在预训练中的长期地位，作者也明确表示尚不确定。