Self-Distillation for Multi-Token Prediction

📄 arXiv: 2603.23911v1 📥 PDF

作者: Guoliang Zhao, Ruobing Xie, An Wang, Shuaipeng Li, Huaibing Xie, Xingwu Sun

分类: cs.CL, cs.AI, cs.LG

发布日期: 2026-03-25

💡 一句话要点

提出MTP-D自蒸馏方法，提升LLM多Token预测的效率和接受率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 多Token预测 自蒸馏 推理加速 知识迁移

📋 核心要点

  1. 现有多Token预测（MTP）方法在大型语言模型推理加速中受限于MTP head接受率和多head联合训练的困难。
  2. 论文提出MTP-D自蒸馏方法，通过知识蒸馏提升MTP head的性能，同时保持主head性能，降低训练成本。
  3. 实验结果表明，MTP-D显著提升了MTP head的接受率和推理速度，并在多个基准测试中验证了其有效性。

📝 摘要（中文）

随着大型语言模型（LLMs）规模的扩大，推理效率成为关键瓶颈。多Token预测（MTP）通过并行预测多个未来token来加速LLM推理。然而，现有的MTP方法仍然面临两个挑战：MTP head的接受率有限，以及难以联合训练多个MTP head。因此，我们提出MTP-D，一种简单而有效的自蒸馏方法，以最小的额外训练成本，提高MTP head的接受率（+7.5%），同时最大限度地保持main-head的性能。我们还为MTP-D引入了一种循环扩展策略，从而能够有效且经济地扩展MTP head，并进一步显著加速推理，达到1-head MTP（+220.4%）。此外，我们通过在七个基准上的大量实验，系统地探索和验证了关于蒸馏策略和MTP潜在可扩展性的关键见解。这些结果表明，我们的MTP-D和循环扩展策略有效地提高了MTP-head的性能和推理效率，从而促进了MTP在LLM中的实际应用。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）推理效率低下的问题，特别是现有Multi-Token Prediction (MTP) 方法中MTP head接受率有限以及难以联合训练多个MTP head的痛点。现有的MTP方法虽然能够并行预测多个token，但由于head的性能问题，实际加速效果受限。

核心思路：论文的核心思路是利用自蒸馏技术，将主head（main-head）的知识迁移到MTP head，从而提升MTP head的性能和接受率。通过这种方式，MTP head能够更准确地预测多个token，从而提高整体的推理效率。自蒸馏方法降低了训练成本，同时保持了主head的性能。

技术框架：MTP-D的整体框架包括以下几个主要步骤：首先，训练一个标准的主head。然后，利用主head的输出作为“教师信号”，训练多个MTP head。MTP head的目标是模仿主head的输出，从而学习到主head的知识。为了进一步提高性能，论文还提出了一个循环扩展策略，允许逐步增加MTP head的数量，从而实现更大的推理加速。

关键创新：MTP-D的关键创新在于其简单而有效的自蒸馏方法，以及循环扩展策略。自蒸馏方法能够以最小的额外训练成本显著提升MTP head的性能，而循环扩展策略则允许灵活地调整MTP head的数量，以适应不同的推理需求。与现有方法相比，MTP-D不需要复杂的训练技巧或额外的模型结构。

关键设计：MTP-D的关键设计包括：1) 使用KL散度作为蒸馏损失函数，衡量MTP head和主head输出之间的差异。2) 设计了一个循环扩展策略，通过逐步增加MTP head的数量，实现更大的推理加速。3) 仔细调整了蒸馏过程中的超参数，以平衡MTP head的性能和主head的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MTP-D方法在提升MTP head接受率方面表现出色，相较于基线方法提升了7.5%。通过循环扩展策略，MTP-D实现了显著的推理加速，在1-head MTP设置下达到了220.4%的加速效果。实验结果在七个基准测试中得到了验证，证明了MTP-D的有效性和泛化能力。

🎯 应用场景

该研究成果可应用于各种需要加速大型语言模型推理的场景，例如在线对话系统、文本生成、机器翻译等。通过提高推理效率，可以降低计算成本，提升用户体验，并促进LLM在资源受限设备上的部署。未来，该方法有望进一步扩展到其他模型架构和任务中。

📄 摘要（原文）

As Large Language Models (LLMs) scale up, inference efficiency becomes a critical bottleneck. Multi-Token Prediction (MTP) could accelerate LLM inference by predicting multiple future tokens in parallel. However, existing MTP approaches still face two challenges: limited acceptance rates of MTP heads, and difficulties in jointly training multiple MTP heads. Therefore, we propose MTP-D, a simple yet effective self-distillation method with minimal additional training cost, which boosts MTP head acceptance rates (+7.5\%) while maximumly preserving main-head performance. We also introduce a looped extension strategy for MTP-D, enabling effective and economical MTP head extension and further significant inference speedup to 1-head MTP (+220.4\%). Moreover, we systematically explore and validate key insights on the distillation strategies and the potential scalability of MTP through extensive experiments on seven benchmarks. These results demonstrate that our MTP-D and looped extension strategy effectively enhance MTP-head performance and inference efficiency, facilitating the practical usage of MTP in LLMs.