Match the Distribution, Not the Compute: Post-Training Multi-Token Prediction Heads

📄 arXiv: 2610.00888v1 📥 PDF

作者: Prachi Badarayani, Aidan Jay, Chenghui Zhou, Dayquan Julienne, Yuan Gao, Tianwei Chen, George Zerveas, Ishmam Zabir, Xiren Zhou, Chris Quirk, Xia Song

分类: cs.LG, cs.AI

发布日期: 2026-10-01

💡 一句话要点

提出后训练多标记预测头以提升推理效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 多标记预测 后训练 推理效率 自然语言处理 模型优化

📋 核心要点

  1. 现有的多标记预测方法需要在预训练阶段与基础模型联合训练，导致训练成本高且灵活性不足。
  2. 论文提出了一种后训练方案，通过在冻结模型上进行轻量级的训练，显著降低了所需的训练标记数量。
  3. 实验结果显示，后训练方案在多个基准测试中实现了与联合训练模型相当或更好的性能提升，速度提升幅度达到12%至16%。

📝 摘要（中文）

多标记预测（MTP）通过在每次前向传播中生成多个下一个标记来提高自回归生成的吞吐量，同时通过验证步骤确保基础模型的标记分布得以保留。本文探讨在冻结推理模型上进行轻量级后训练是否足以达到与联合训练模型相同的吞吐量提升。研究结果表明，使用约25亿个标记的简单交叉熵后训练方案，能够在数学、编码和知识基准上达到或超过联合训练模型的预期速度提升。此外，提出了一种链感知的草稿标记验证规则放宽方法和自适应控制器，进一步优化了推理过程。

🔬 方法详解

问题定义：本文旨在解决现有多标记预测方法在预训练阶段需要大量标记的高成本问题，同时探讨冻结模型的后训练是否能达到相似的性能。

核心思路：提出了一种后训练方案，通过在冻结的推理模型上进行轻量级的训练，利用少量标记达到与联合训练相似的性能。

技术框架：整体架构包括冻结的推理模型和后训练阶段，后训练使用简单的交叉熵损失函数，并引入链感知的草稿标记验证规则和自适应控制器。

关键创新：最重要的创新在于提出了链感知的草稿标记验证规则的放宽方法，使得标记分布可以有一定的偏离，从而提升了速度而不损失准确性。

关键设计：后训练过程中使用了约25亿个标记，相比于联合训练减少了1000到10000倍的标记需求。同时，自适应控制器根据推理时的需求动态选择参与的多标记预测头数量。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果表明，在冻结的Qwen3-8B模型上，使用后训练方案的速度提升达到了预期的12%至16%，并且在数学、编码和知识基准测试中表现优异，显示出与联合训练模型相当的性能，同时显著减少了训练所需的标记数量。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的文本生成、对话系统和编程辅助工具等。通过提升推理效率，能够在实际应用中实现更快的响应时间和更高的用户体验，具有重要的实际价值和未来影响。

📄 摘要（原文）

Multi-token prediction (MTP) improves the throughput of autoregressive generation by enabling the language model to draft multiple next tokens per forward pass, while a verification step over draft tokens ensures that token distribution of the backbone is preserved. Every open MTP-family release (MiMo-7B, DeepSeek-V3, Qwen3) trains its heads jointly with the backbone over the full pretraining run of tens of trillions of tokens, thus setting the drafter quality at pretraining time. We ask whether a lightweight post-training pass on target-generated chain-of-thought is enough to reach the same expected throughput speedup on a frozen reasoning model, and study how a serving-time system built on such a checkpoint can be optimized. We present three findings. 1) On a frozen Qwen3-8B with $K{=}3$ chained MTP heads, we show that a post-training recipe with plain cross-entropy on $\approx!2.5$B tokens reaches or exceeds the expected speedup of jointly trained MiMo-7B on math, coding and knowledge benchmarks. Our post-training recipe utilizes $10^3$-$10^4\times$ less MTP-training tokens as compared with joint pre-training of MiMO-7B MTP baseline. 2) We propose a chain-aware relaxation of draft token verification rule that allows a bounded drift from backbone language model token distribution. We show that this relaxation lifts expected speedups by $+12$ to $+16\%$ per benchmark while preserving task accuracy. 3) We propose an adaptive controller that dynamically chooses the number of MTP heads to be engaged at inference time and demonstrate recovery of upto $11$--$14\%$ loss in speedup using fixed maximum MTP draft length.