精读:把 MTP 梯度拆成自蒸馏与噪声残差——Soft-MTP 的一行改动、80% 的噪声占比与那份免费午餐的账单
标准多 token 预测的交叉熵梯度可精确分解为把 MTP 头拉向模型自身 NTP 分布的 KL 自蒸馏项与一个低信噪比、随规模恶化的残差项;Soft-MTP 借直通式梯度路由把残差从 MTP 头与主干的更新里精确减掉,一行代码、零参数、零开销,167M–1.2B 全尺度投机解码全指标领先,1.2B 微调后数学代码 7/8 格领先标准 MTP——但账单上仍有三笔没付清的账:1.2B 域内困惑度仍输纯 NTP、332M 档有输格、数学两项输给显式 KL。
1 论文概览
一句话总结:标准多 token 预测(MTP)的交叉熵梯度可以被一行代数精确拆成两半——把 MTP 头拉向模型自身下一 token 预测(NTP)分布的 KL 自蒸馏项 ,与一个低信噪比、随规模恶化的残差项 ;Soft-MTP 借直通式梯度路由把残差从 MTP 头与共享主干的更新里精确减掉,一行代码、零新增参数、吞吐开销 0.07%,在 167M–1.2B 上对标准 MTP 全尺度改进。
核心问题:MTP 已进入 DeepSeek-V3/V4、Nemotron 3、GLM-5、Kimi K3 的预训练配方,用途有二:对冲下一 token 监督的短视、给投机解码当草稿器。可 MTP 头自己该用什么目标训练,没人拆开看过。投机解码里,MTP 起草的 token 只有与模型自己的 NTP 分布一致才会被接受——顺着这个机制,理想目标就是把 MTP 头直接对齐 NTP 分布的自蒸馏;但显式 KL 要物化教师与学生两份词表大小的 logit 矩阵,与主流 fused cross-entropy kernel 不兼容,前沿模型至今仍在用对数据的普通交叉熵硬训 MTP 头。论文问三件事:标准 MTP 梯度里究竟装了什么;那个残差是不是噪声;KL 信号能不能不付显式 KL 的代价、免费拿到手。
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.
(Abstract)
主要贡献:
- 精确分解(式 2)。,恒等式、无近似;前一半恰是 (式 3),后一半恰是对齐位置上的 NTP 交叉熵梯度。标准 MTP 头从训练第一天起就在被隐式自蒸馏,只是信号被残差盖着。
- KL gradient–recovery trick(式 4–7)。两个分布共享线性输出头时,它们的 KL 梯度可以从各自交叉熵梯度的差里精确恢复,全程不物化任何 logit——直通估计器的管道被反过来用:不是拿有偏梯度近似离散操作,而是拿零值前向/非零反向的修正项精确合成另一个梯度。
- Soft-MTP。给 NTP 隐状态加一项前向恒为零的 ,增广 NTP 损失在每个 MTP 隐状态处注入 ,与原 MTP 损失相加后,MTP 头与主干只吃 ;输出头 照旧吃对数据的交叉熵(表 1 的梯度路由表)。
- 全谱实证。167M–1.2B、FineWeb-Edu 31.5B token(最大模型约 1.3× Chinchilla)、DeepSeek 式 K=3 顺序头、Muon 优化器:残差 SNR 比 KL 项低约 4 倍(167M 4.1× → 1.2B 4.6×)却占进入主干 MTP 梯度的约 80%;Soft-MTP 在 NTP/MTP 困惑度、投机解码、zero-shot、数学代码微调上全尺度领先, 与 扫描全档 Pareto 支配,与显式 KL-MTP 打平或反超且省 21GB 峰值显存。
2 方法详解
分解的代数与两个前提。k 号 MTP 头在位置 t 的预测 监督 ;把「预测同一个 token 的 NTP 分布」定义为对齐分布 。插入 即得式 2。两个前提缺一不可:输出头 由 NTP 与 MTP 共享(两份梯度才能在同一空间里相加减);MTP 与 NTP 由同一段数据监督( 对 的交叉熵梯度才有「NTP 梯度」的身份)。这两条恰是 DeepSeek 式与 Meta 式 MTP 的公共结构,分解对两种架构都成立(附录 C 另报了 Meta 式并行头的结果)。
KL gradient–recovery 的机关。式 4 给每个 NTP 隐状态加上 :前向数值恒为零(stop-gradient 项减去自身),NTP 优化完全不变;反向时 sg 项截断、非 sg 项活着,增广 NTP 损失(式 5)于是对每个 注入 (式 6)。与普通 MTP 损失的 相加,恰好剩 (式 7)。改动落在 NTP 分支的一行张量运算上,每个头照旧走自己的 fused CE kernel,不新增任何词表维张量——论文附录 A 给出含 加权与各头目标数不齐重标定的完整 PyTorch 实现。
梯度路由作为设计空间。表 1 把三类目标画成一张路由表:
| 目标 | 输出头 | MTP 头与共享主干 |
|---|---|---|
| 标准 MTP | CE 梯度 | CE 梯度 |
| KL-MTP | KL 梯度 | KL 梯度 |
| Soft-MTP | CE 梯度 | KL 梯度 |
中间格不是折中审美。输出头若吃 KL,teacher 分布 本身就是这颗头读出来的,KL 有一条捷径可走:把读出对 MTP/NTP 隐状态的差异方向脱敏,损失照样降:
This creates a shortcut because the KL can be reduced not only by improving the MTP hidden states, but also by making the shared readout less sensitive to the directions in which the MTP and NTP hidden states differ.
(§6.3)
交叉熵没有这条捷径(目标不依赖 )。图 6 用 logit 映射 (unembedding 与最终 RMSNorm 增益之积)的有效秩量了塌缩程度:NTP 基线 496,标准 MTP 塌到 131,KL-MTP 321,Soft-MTP 437。路由表还允许做差分归因——Soft-MTP 与 KL-MTP 只差输出头吃什么(437 对 321,CE 在头拿回 116),与标准 MTP 只差主干吃什么(437 对 131,去掉残差拿回 306):塌缩的主犯是残差经主干回传改变了表征,从犯才是头吃 KL。
偏差-方差读法(式 8)。残差再拆一层:,前项是 NTP 误差(偏差),后项是 one-hot 采样噪声(方差、零均值)。两者的规模行为相反:
The bias term is the NTP teacher’s own prediction error, which keeps shrinking as the base model scales and/or trains longer. The variance term is the irreducible noise of sampling a one-hot target, which no amount of scale or training removes.
(§5.3)
Soft-MTP 用一个随规模收缩的偏差换掉一个不可约的方差,这笔交易随规模单调变好——§6.1 的 SNR 曲线把这个预言直接验证了。
与基线方法的区别。与显式自蒸馏一族(NextLat 逐格算 KL、top-k logit 近似、argmax 蒸馏、缓存 teacher 梯度但仍要物化 teacher 全分布且冻结 student 头)相比,Soft-MTP 不算 KL、不近似、不冻结任何东西,代价是只有 MTP 头与主干拿到精确 KL 梯度、输出头仍吃 CE——表 6 的对照显示这个「不完整」反而更好。与最小化 TV 距离的草稿器训练(Samarin 等)相比,论文的辩护是 KL 在 TV 不光滑处光滑、对欠预测 token 保持非零梯度,接受率经 Pinsker 上界间接受益。
3 实验结果
残差的体检报告。SNR 按批期望方向的平方范数对总方差之比定义(McCandlish 口径的无偏估计,256 个独立批、每批 4×2048 token):
| 诊断量 | KL 项 | 残差 |
|---|---|---|
| SNR | 高(基准) | 低约 4 倍,167M→1.2B 从 4.1× 恶化到 4.6× |
| 占进主干 MTP 梯度的份额 | 约 20% | 约 80% |
| 单位步代理对 NTP 损失的影响 | 下降 | 上升 |
The residual is much noisier (lower SNR) than the KL component, yet accounts for roughly 80% of the MTP gradient.
(§6.1 小结)
单位步代理实验里,沿 对主干隐状态走一步,NTP 损失全程下降;沿 走一步,全程上升——残差不只是噪声,还在与 NTP 更新方向反着对冲。作者标明这只是主干-头界面的刻画,实际更新还取决于优化器;但 Soft-MTP 训练中 NTP 损失更低、且差距随规模拉大,为这条机制注了脚。
投机解码(表 2):唯一一张「每一格每一尺度」都赢的表。1.2B 上 MTP 分布对自家 NTP 验证器的 KL 从 0.384 降到 0.364、TV 从 0.260 到 0.253;五个域平均,接受长度 +1.8–2.2%、加速比 +1.4–1.5%。幅度是百分比级——投机解码端的白拿收益是「零成本换 2%」,不是数量级改善。
下游 zero-shot(表 3):平均准确率四档全赢(167M 46.91 对 46.84;332M 50.46 对 50.38;689M 55.01 对 54.47;1.2B 57.31 对 56.80),前两档 +0.1 级、后两档 +0.5 级,增益随规模长的形态与偏差-方差论证一致;1.2B 处 Soft-MTP 的 WikiText(17.49)与 LAMBADA(14.07)困惑度、平均准确率都是三个目标里最好的。账本背面要翻:332M 档 LAMBADA 困惑度 28.57 比 MTP 的 27.00 差了近 1.6 个点,PIQA(68.72 对 69.59)与 HellaSwag(45.58 对 45.92)也有输格;167M/332M 两档上标准 MTP 本就不及纯 NTP(与 Gloeckle 一致),Soft-MTP 在这两档只是补齐落差、没有反超 NTP。
微调迁移(表 4):1.2B 预训练后丢弃 MTP 头、仅用 NTP 损失微调(代码 Stack-Edu Python 子集 21B token、数学 OpenMathInstruct-2 4.2B token)。8 格里 Soft-MTP 领先 7 格,是唯一没有任何一格跌破 NTP 基线的目标:HumanEval pass@64 43.90 对 40.85(+3.05)、MBPP pass@64 61.00 对 59.80、GSM8K 49.66 对 48.60、MATH-500 25.00 对 23.20。代码端优势随采样预算 k 增大而扩大——预训练目标留下的表征红利在多次采样下被放大。
超参与成本。 与 的扫描(332M、7.3B token)上,Soft-MTP 在 NTP 损失对 MTP 损失、NTP 损失对验证器 KL 两张图上全档 Pareto 支配标准 MTP——「只是变相调小了 MTP 权重」的质疑被这两张图挡住。成本表(1.2B、单卡 H200):Soft-MTP 与标准 MTP 的 FLOPs(10.28 GFLOP/token)、MFU(51.7%)、fused-CE 峰值显存(107.0GB)逐项相同,吞吐差 0.07%(49,748 对 49,783 tok/s);显式 KL-MTP 每一项都最差——多 1.5% FLOPs、吞吐低 5.6%、峰值显存多 21GB,且吃不到 fused kernel 的红利。
与显式 KL 的对照(表 6),全文最值得多看的一眼。Soft-MTP 与 KL-MTP 的主干梯度完全相同,只差输出头吃什么:困惑度与 zero-shot Soft-MTP 平或赢,代码 6 格全赢;数学两格 KL-MTP 反超(GSM8K 50.34 对 49.66、MATH-500 26.00 对 25.00)。0.7–1.0 分的数学差额完全归因于「输出头吃 CE 还是吃 KL」这一个变量——把头拟合到数据不全是好处,长尾推理 token 上的 one-hot 信号与 KL 的平滑信号各有所长。论文用有效秩解释了代码端,数学端这格逆差没有给出解释。
4 个人分析
亮点。
- 恒等式先于方法。式 2 只是一行加减法,但它重写了 MTP 的自我理解:标准 MTP 头一直在被隐式自蒸馏,信号被约 80% 的噪声残差压着。精确分解不依赖假设、不依赖规模、不依赖架构细节(顺序/并行头通吃),这类「先拆账、再动手」的手艺与本专栏 FORK 期把 rubric 奖励拆成召回率加精确率是同一路数——修法方向由分解唯一确定。
- 工程约束长出了技巧。显式 KL 的问题从来不是算不准而是算不起:词表级 logit 物化与 fused kernel 不兼容,在无张量并行的训练配置(DeepSeek-V3 式)里更贵。KL gradient–recovery 把加减法挪进隐状态空间,每个头保住自己的 fused CE——这个解法的形状是被真实约束逼出来的,不是审美偏好。
- 直通估计器的反转用法。STE 十来年都被当作「用有偏梯度近似不可导离散操作」的妥协工具;这里前向恒零/反向非零的同一根管道被用来精确合成另一个梯度。工具箱里多了一种借法。
- 诊断三件套可移植:梯度按批 SNR 分解、单位步代理、readout 有效秩监控。三件都不依赖 MTP 语境,任何「辅助损失到底在帮忙还是在添乱」的问题都能用这套流程审一遍。
局限性。
- 规模与种子。最大 1.2B、31.5B token(约 1.3× Chinchilla),单种子。核心主张「交易随规模越来越划算」恰恰是在未测区间外推:1.2B 处 FineWeb-Edu 域内困惑度 10.741 仍高于纯 NTP 的 10.572,MTP 对主干的困惑度税只被减半、没有退掉,更大规模上这笔账怎么走没有数据。
- 「全指标全胜」的口径。只有表 2 成立。表 3 有输格(332M 的 LAMBADA/PIQA/HellaSwag),表 6 数学两格输给 KL-MTP。摘要里 improves over standard MTP 的强度,逐表核对后要打八折。
- 「残差=噪声」的叙事有一个未决的对立面。Gloeckle 一系对 MTP 的辩护是对冲 NTP 的短视,而承担「看远」的那部分监督,恰是残差里对真实未来 token 的硬交叉熵。Soft-MTP 把 MTP 头推向纯草稿器、把学新东西的通道收敛到 KL 对齐;1.2B 的微调增益支持这个取舍,但「残差有害」与「MTP 反短视靠残差」两种读法在本文的尺度与数据量内无法裁决。数学格输给 KL-MTP 这件事也提示:残差里经输出头的 one-hot 信号,在长尾推理 token 上另有一功。
- 诊断口径。SNR 用 4×2048 token 的微批估计(训练批是 0.5M token),约 80% 的梯度份额是范数口径而非学习贡献口径;单位步代理只在主干-头界面成立,作者自己标注了这层限定。
- 单作者、自费租卡(8×RTX PRO 6000,Vast.ai,致谢里写明)——透明度值得肯定,复算资源与多种子稳健性自然受限。
对大模型领域的影响。一行改动、零参数、零开销、逐表不亏,这是采纳摩擦最低的一类贡献,已有 MTP 配方的预训练管线没有理由不试。更宏观的一层:「MTP 还值不值得留在预训练里」正处在十字路口——
DeepSeek-V4.1-Flash recently dropped MTP in favor of DSpark, a speculative decoder trained after pre-training.
(§7)
这篇论文的分解把问题重述为「MTP 头的训练信号要不要洗干净」:洗过之后,MTP 以更低的干扰留在预训练里。KL gradient–recovery 的外延也比 MTP 大:凡 teacher 与 student 共享线性头的 forward KL 都适用,论文点名了 Nemotron 3 Ultra 的 MTP Boosting、on-policy 自蒸馏(OPSD)与 off-policy RL 的 score centering——score centering 本专栏 9 月精读过(Marek & Ryabinin,2609.20807),那条「把漂移项精确减掉」的思路与本文同族,现在多了一件把 KL 计算摊销掉的新工具。
5 实践启示
可以借鉴的技术。
- KL gradient–recovery trick 直接搬。任何「两个分布共享输出头、要它们的 forward KL 梯度」的场景——自蒸馏辅助损失、草稿器对齐、RL 里的蒸馏正则——都能用 CE 梯度相减替代显式 KL 或 top-k 近似,免去词表维物化。附录 A 的实现逐字可抄。
- 梯度路由表当设计语言。「损失在哪个参数子集上走哪个梯度」是比「换损失函数」更细的一层设计空间;评估任何辅助目标前先画它的路由表,Soft-MTP 的全部巧思就压缩在表 1 的三个格子里。
- 辅助损失的体检流程:批间 SNR 分解(信号是噪声的几倍)→ 单位步代理(这一项把主任务损失往哪边推)→ 有效秩监控(readout 有没有被辅助目标带塌)。三步成本低于一次完整消融,给出的是机制而不只有效果。
- 把「哪个分量承重」写进训练日志。残差占 80% 梯度范数却只添噪声——不拆开看,这类隐性浪费在 loss 曲线上完全不可见。
潜在应用场景。
- 存量 MTP 配方升级:DeepSeek 式顺序 MTP 头上加一行张量修正,投机解码接受率与 zero-shot 拿白捡增量,不动并行化方案(§7 逐项过了一遍:数据/专家/流水线并行零额外通信——MTP 模块本就与输出头同 stage;序列切分只需边界处交换至多 K 份隐状态梯度)。
- RL 后训练的 rollout 吞吐:rollout 生成是 RL 后训练的主要瓶颈,接受长度 +2% 直接折算成采样吞吐,对 RLVR 系训练是真金白银。
- 数据受限的预训练:论文引 NanoGPT Slowrun 的证据,数据受限时 MTP 提样本效率;Soft-MTP 让这笔效率以更干净的信号入账,小模型与重复数据配方里优先级更高。
- 1B 级自研基座:小尺度上 MTP 对 zero-shot 近乎白给(表 3 前两档),Soft-MTP 把「近乎」补成「确实」,配合 3.2% 参数量的 MTP 头配置,是小模型的低成本增强件。
参考文献
- Teoh, J., Multi-Token Prediction Done Right, arXiv:2610.38331(本文,全部公式与表格出处:式 1–8、表 1–6、图 1–6、§1/§5.1/§5.3/§6.1/§6.3/§7、附录 A/C)。
- Gloeckle, F., Idrissi, B. Y., Rozière, B., Lopez-Paz, D., Synnaeve, G., Better & Faster Large Language Models via Multi-token Prediction, ICML 2024, arXiv:2404.19737(Meta 式并行 MTP 与反短视论证)。
- DeepSeek-AI, DeepSeek-V3 Technical Report, arXiv:2412.19437(顺序 MTP 模块,被多家前沿模型沿用)。
- Stern, M., Shazeer, N., Uszkoreit, J., Blockwise Parallel Decoding for Deep Autoregressive Models, NeurIPS 2018(多 token 预测源头)。
- Leviathan, Y., Kalman, M., Matias, Y., Fast Inference from Transformers via Speculative Decoding, ICML 2023, arXiv:2211.17192(投机解码与接受规则)。
- Teoh, J., Tomar, M., Ahn, K., Hu, E. S., Pearce, T., Sharma, P., Krishnamurthy, A., Islam, R., Lamb, A., Langford, J., Next-Latent Prediction Transformers Learn Compact World Models, arXiv:2511.05963(作者前作,显式 KL 的来源与摊销初衷)。
- Marek, M., Ryabinin, M., Score Centering Stabilizes Off-policy Reinforcement Learning, arXiv:2609.20807(KL gradient–recovery 的候选受益方,本专栏 2026-09-20 精读)。
- Bengio, Y., Léonard, N., Courville, A., Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation, arXiv:1308.3432;van den Oord, A. et al., Neural Discrete Representation Learning, NeurIPS 2017(直通估计器两处源头)。
- McCandlish, S., Kaplan, J., Amodei, D., et al., An Empirical Model of Large-Batch Training, arXiv:1812.06162(梯度噪声尺度与 SNR 无偏估计口径)。
- Penedo, G. et al., The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale, arXiv:2406.17557(训练语料)。
相关精读
-
本篇把本专栏连续五十六期的主线再挪一格:此前五十五期回答的是「挑哪些 token、挑哪些轨迹、信用按哪套口径记账、信号被消费时旧了多少、数据分布离基座多远、整条响应有没有资格参与这次更新、策略容量会不会被优化动力学耗光、教师分布里哪些位置允许比较、事后经验除了改行为还能不能改行动之前的那次预测、参考模型的 KL 该从哪条边进组相对更新、教师的意见从哪条通道进 RL、off-policy 校正的记账区间有多长、裁剪区间那两个端点由什么决定、干预落在流水线的哪一环、逐步信用从哪个模型里来、一个已经报出来的增益该怎么拆账、这些信号最终以什么形状落到参数上、结局空间该分成几档、同一笔优势在响应内部怎么分、教师交出来的那笔位移里哪些前缀有资格进损失、这笔信用在更新里当斜率还是当终点、一次迭代该对齐到任务的哪一格、一份外部语料要进模型,是作为检索出来的文本进上下文,还是作为预测出来的隐状态进解码器、模型内部被拆出来的那批坐标能不能当作可以相加的行为解释单元、思考预算拧大之后,模型换掉的是对一个句子的读法,还是只换掉了答对的比例、一组 token 压成一个向量时,这个压缩动作本身该不该由输入自己解出来、一份原生多模态预训练预算该在三个模态之间怎么切成参数、token 与重复次数、生成过程里没有左右顺序时 token 级信用该记在什么量上、当奖励本身就是一份由打分器算出来的清单覆盖率时,优化器在最大化哪一样,以及给这份奖励补上一个精确率项之后,模型缩短的究竟是废话还是内容、当参数量被钉死、深度靠同一组权重反复复用来买时,多出来的计算到底在买什么」;这一期问的是MTP 的交叉熵梯度里,到底有多少在教模型它还不会的事、有多少只是在低信噪比地复读它已经会的下一 token 分布——把后者单独留下来,要不要付钱——上一期退到最上游看计算的组织方式,这一期回到损失与梯度本身做一次精确分解:先把信号拆到不能再拆,再决定留哪一半。
-
精读:三明治包络、输入注入与状态组织的三轴审判——Looped Models Done Right 与那笔藏在全胜账本背面的 5.4 分