Multi-Token Prediction Done Right

Soft-MTP 通过从现有 cross-entropy 计算中精确恢复 KL self-distillation 梯度，从标准 multi-token prediction 梯度中去除带噪声的 residual，在不增加成本的情况下提升 NTP、speculative decoding 和 zero-shot accuracy。

标准 MTP 训练的 cross-entropy 梯度包含一个 residual 分量，会干扰 next-token 和 multi-token 学习，并随规模增大而噪声加剧。有用的 KL self-distillation 信号被这个低 SNR residual 掩盖。直接最小化 KL 代价高，因为它需要 vocabulary-sized logit 矩阵，且与 fused cross-entropy kernels 不兼容。

论文将 hard MTP 梯度精确分解为 KL self-distillation 项和 residual 项。它引入 KL gradient-recovery trick：将 stop-gradient 路由后的 MTP hidden states 加到 NTP hidden states 上，增强后的 NTP loss 在 MTP hidden state 处注入负的 residual 梯度。将其与 MTP cross-entropy loss 相加后，在 MTP head 处恰好只剩 KL 梯度。Soft-MTP 应用这种路由，使 output head 仍接收 cross-entropy 梯度，而 MTP heads 和 shared trunk 只接收恢复出的 KL 梯度。它支持 DeepSeek- 和 Meta-style MTP，且不增加参数。

• 在 1.2B 规模下，与标准 MTP 相比，Soft-MTP 降低了 MTP 分布相对于模型自身 NTP verifier 的 teacher-forced KL divergence 和 TV distance。 §6.2
• 在 167M 到 1.2B 参数范围内，Soft-MTP 在每个规模下的 self-speculative decoding 表格中都赢得所有 metric。 §6.2
• Soft-MTP 的训练成本与标准 MTP 相当：matmul FLOPs、MFU 和 fused cross-entropy peak memory 相同，throughput 差异在 0.07% 以内。 §6.3
• 在每个规模下，residual 在 shared trunk 中承载约 80% 的 MTP 梯度，尽管其 signal-to-noise ratio 显著低于 KL 分量。 §6.1
• 在 332M 规模、K=3、lambda_MTP=1 和 7.3B FineWeb-Edu tokens 下使用 Meta-style MTP heads 时，Soft-MTP 比标准 MTP 获得更低的 NTP 和 MTP loss。 §C

• 如论文在其 bias-variance 观点中所述，Soft-MTP 偏向 NTP 分布而非数据分布；这对 speculative decoding 有利，但也是明确说明的 tradeoff。
• 实验覆盖 167M 到 1.2B 参数；论文在 §7 讨论了 frontier-scale 的相关性，但未报告 frontier-scale 实验，因此该 regime 在实证上仍待探索。
• Meta-style MTP 仅在 Appendix C 中的 332M 规模上测试，因此实验中对 DeepSeek-style MTP 之外的架构覆盖有限。
• 论文仅从 689M 起报告更低的 MTP perplexity，因此在更小的 167M 和 332M 规模上，MTP-loss 收益尚未确立。

• hard MTP 梯度精确分解为 KL self-distillation 分量和 residual 分量：g_hard = g_KL + g_res。 §4
• g_KL 分量等于从 stop-gradient 对齐的 NTP 分布到 MTP 分布的前向 KL divergence 的梯度。 §4.1
• KL gradient-recovery trick 用路由后的 MTP hidden states 增强 NTP hidden states，使增强后的 NTP loss 在 MTP hidden state 处注入负的 g_res；与 MTP loss 相加后恰好只剩 g_KL。 §5.1
• Soft-MTP 将 cross-entropy 梯度送入 shared output head，但将恢复出的 KL 梯度送入 MTP heads 和 shared trunk，从而避免 vocabulary-sized logit 矩阵。 §5.2
• residual 包含 NTP prediction error 加 one-hot sampling noise；Soft-MTP 以偏向 NTP 分布的 bias 为代价去除 sampling variance。 §5.3
• Soft-MTP 同时支持 DeepSeek- 和 Meta-style MTP，不增加参数，且没有可测量的 overhead。 §Abstract
• Soft-MTP 在 FLOPs、MFU 和 fused cross-entropy peak memory 上与标准 MTP 相当，并在所有报告的成本轴上比 direct KL-MTP 更高效。 §6.3
• 在 167M 到 1.2B 的每个测试规模上，Soft-MTP 都改善 NTP perplexity、zero-shot accuracy 和 self-speculative decoding。 §6.2