MTP(多 token 预测)已经是主流大模型预训练的标配:论文点名 DeepSeek-V3/4、NVIDIA Nemotron 3、GLM-5、Kimi K3 都在用。它的好处一石二鸟——额外的预测头既能改善下游表现,推理时还能当投机解码的起草头。但 10 月 28 日挂上 arXiv 的单作者论文(MIT)指出:标准 MTP 训练目标里埋着一个低信噪比的梯度残差,模型越大,它越拖累学习。
占八成的低质梯度
论文把 MTP 的交叉熵梯度做了精确分解:一部分是 KL 自蒸馏项,把 MTP 头往模型自己的下一 token 预测分布上拉——这恰是投机解码想要的对齐信号;另一部分残差,本质是对数据的下一 token 交叉熵梯度,信噪比很低,却把有用的 KL 信号淹掉了。实测上,整个训练过程里残差的信噪比都比 KL 分量低约 4 倍,差距从 167M 的 4.1 倍拉大到 1.2B 的 4.6 倍;而流入共享主干网络的 MTP 梯度里,约 80% 由这个残差贡献。论文还做了单步探针:沿 KL 分量走一步,下一 token 损失下降;沿残差走一步,损失反而上升——它不只是噪声,还在跟主任务对着干。
修法:一行代码换干净梯度
直接的思路是最小化显式 KL,但那要物化词表大小的 logit 矩阵,内存吃不消,也跟主流预训练依赖的融合交叉熵内核不兼容。作者提出的 KL 梯度恢复技巧绕开了这一点:当两个分布共享线性输出头时,KL 的梯度可以由两条交叉熵梯度经直通式梯度路由精确组合出来。于是 Soft-MTP 只需一行代码改动:不加参数、不多算 logits、兼容融合内核。在 H200 上实测,吞吐与标准 MTP 只差 0.07%,FLOPs 与显存完全一致;官方仓库带了 float64 单元测试,验证恢复出的梯度与显式算 KL 机器精度等价。DeepSeek 式与 Meta 式 MTP 结构都支持。
账本:四档规模都占优,微调最亮眼
论文实测 Soft-MTP 在 167M、332M、689M、1.2B 四档规模(各训 31.5B FineWeb-Edu token)的每一档都优于标准 MTP:零样本平均准确率的增益随规模拉大,小模型约 0.1 分,689M 与 1.2B 约 0.5 分,并在 1.2B 反超纯 NTP 基线(平均 57.31,高于标准 MTP 的 56.80 与 NTP 的 56.75);投机解码验收长度五域平均提升 1.8–2.2%,加速提升 1.4–1.5%。最亮眼的是微调迁移:1.2B 分别在 21B token 的 Stack-Edu Python 与 4.2B token 的 OpenMathInstruct-2 上微调后,Soft-MTP 八格里赢下七格,HumanEval pass@64 领先 3.1 个百分点(43.90 对 40.85),且是唯一从不错破 NTP 底线的目标。对照显式方案 KL-MTP,Soft-MTP 多数指标持平或更好,而后者要多花 1.5% FLOPs、吞吐低 5.6%、显存多 21GB。
边界:单作者、自费算力、前沿规模未验证
冷水也要泼:这是单作者预印本,实验主要由作者本人在 Vast.ai 租 8 张 RTX PRO 6000 完成,验证上限是 1.2B 参数、31.5B token,离万亿级前沿训练还远;「免费午餐」是论文自己的说法,数学微调上 KL-MTP 仍占先(GSM8K 50.34 对 49.66)。论文第 7 节还补了一条背景:DeepSeek-V4.1-Flash 已经放弃 MTP、改用训练后得到的 DSpark 投机解码器,MTP 路线本身的前途尚在未定之天。但 KL 恢复技巧的适用面不止 MTP——只要教师与学生共享线性输出头,OPSD 这类训练后自蒸馏场景都能借它省掉昂贵的显式 KL。
对还在用 MTP 的团队,这是一次近乎零成本的工程升级;对所有人,它提醒的是:连「标配」都可能藏着没人算清的梯度级缺陷。四档规模的代码与 checkpoint 均已开源(GitHub)。