论文
匹配分布,而不是计算:后训练多标记预测头
Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads
摘要
多词元预测 (MTP) 通过使语言模型能够在每次前向传递中起草多个下一个词元来提高自回归生成的吞吐量,同时对草案词元的验证步骤可确保保留主干的词元分布。每个开放 MTP 系列版本(MiMo-7B、DeepSeek-V3、Qwen3)都会在数十万亿词元的完整预训练运行中与主干网络联合训练其头部,从而在预训练时设置起草者质量。我们询问目标生成的思想链上的轻量级后训练传递是否足以在冻结推理模型上达到相同的预期吞吐量加速,并研究如何优化建立在此类检查点上的服务时间系统。我们提出三项发现。 1) 在具有 $K{=}3$ 链式 MTP 头的冻结 Qwen3-8B 上,我们表明在 $\approx\!2.5$B 词元上具有简单交叉熵的后训练配方达到或超过了联合训练的 MiMo-7B 在数学、编码和知识基准方面的预期加速。与 MiMO-7B MTP 基线的联合预训练相比,我们的后训练配方使用了 $10^3$-$10^4\times$ 更少的 MTP 训练词元。 2)我们提出了一种链感知放宽词元验证规则草案,允许主干语言模型词元分布有界漂移。我们表明,这种放松将每个基准的预期加速从 $+12$ 提升到 $+16\%$,同时保持任务准确性。 3) 我们提出了一种自适应控制器,可以动态选择在推理时参与的 MTP 磁头数量,并使用固定的最大 MTP 草案长度来演示恢复高达 $11$--$14\%$ 的加速损失。