推理加速trick之：Multi-Token Prediction(MTP)

1. 问题来源：

DeepSeek-V3论文中提到Multi-Token Prediction，如下所示：
Architecture: Innovative Load Balancing Strategy and Training Objective

• On top of the efficient architecture of DeepSeek-V2, we pioneer an auxiliary-loss-free
  strategy for load balancing, which minimizes the performance degradation that arises
  from encouraging load balancing.
• We investigate a Multi-Token Prediction (MTP) objective and prove it beneficial to model
  performance. It can also be used for speculative decoding for inference acceleration

2. Multi-Token Prediction (MTP) 的含义

通常在自然语言处理（NLP）任务中，模型的训练目标是 单词级预测 或 单个 token 的预测，即根据上下文预测下一个单词或 token（如 GPT 模型的自回归生成目标）。这种目标称为 单一 token 预测。

Multi-Token Prediction (MTP) 是一种扩展目标，在训练时，模型不仅预测下一个单 token，而是同时预测多个未来的 token。例如：

• 在一个时间步上，模型预测未来 2 个、3 个或更多 token。
• 目标是通过一次前向计算得到多个 token 的概率分布，而不是逐个生成。

这种方法不仅能提升训练效率，还可以为推理提供新的加速手段。

---

3. MTP 提升模型性能的原理

MTP 对模型性能的提升来自以下几个方面：

(1) 提高上下文感知能力

• 传统单 token 预测：模型的目标是预测下一个单 token，训练过程中每次只关注当前上下文与下一个 token 的关系。
• MTP：模型需要预测多个未来 token，这使得它必须从更大范围的上下文中提取信息，从而提高上下文感知能力。
• 效果：模型对语言模式的学习更加全面，生成的句子更流畅且符合语义逻辑。

(2) 增加训练信号的密度

• 单 token 预测的每个时间步只提供一个训练信号（即目标 token 的误差），而 MTP 在每个时间步上提供多个 token 的训练信号。
• 效果：模型在相同的训练迭代下能更快地学到语言结构，提高训练效率。

(3) 改进序列建模的效率

• MTP 的目标更接近实际生成任务。传统的单 token 预测可能导致模型局限于逐步生成，而 MTP 能让模型更好地学习到多 token 的联合分布，提高建模性能。

---

4. Speculative Decoding（推测式解码）的含义

(1) 什么是 Speculative Decoding

Speculative Decoding 是一种用于加速推理的技巧，特别适用于 自回归模型（autoregressive models） 的生成任务。

• 传统自回归解码：生成文本时，模型逐个 token 地生成，每个 token 的生成都依赖于之前生成的结果。这种方法尽管准确，但非常慢，因为每次生成都需要进行完整的前向计算。
• Speculative Decoding：通过推测未来的多个 token，提前生成一部分候选结果，然后在这些候选结果中进行验证或修正，从而减少逐步生成的次数。

(2) MTP 与 Speculative Decoding 的关系

• MTP 的目标 是在一次前向计算中同时预测多个未来 token 的概率分布。
• 这种能力可以直接用于 Speculative Decoding：
  • 模型一次性生成多个 token 的候选结果，而不是逐个生成。
  • 通过某种机制（如更精确的校验模型或后处理步骤）验证和修正这些候选结果。
  • 如果候选 token 验证通过，可以直接使用，从而跳过若干步推理。

(3) Speculative Decoding 的加速原理

• 传统逐步生成方法需要对每个 token 做一次前向传播，推理时间与序列长度成正比。
• Speculative Decoding 通过一次性生成多个 token，可以减少前向传播的次数，从而显著加速推理。

---

5. MTP 和 Speculative Decoding 的结合优势

(1) MTP 为 Speculative Decoding 提供基础

• MTP 的多 token 预测能力恰好为 Speculative Decoding 提供了候选生成。
• 在一次前向传播中，模型可以同时生成多个 token 的候选，减少逐步生成的次数。

(2) 减少错误传播

• 在传统的逐步生成中，错误会随着生成序列的增长逐步积累。
• MTP 的多 token 预测可以通过一次性生成多个 token，并对其进行验证，从而减少错误传播。

(3) 提高推理速度

• Speculative Decoding 的本质是通过并行化和批量化的方式减少推理步骤，而 MTP 的能力直接支持这种并行化生成。

---

6. 案例：MTP 和 Speculative Decoding 的应用

以下是一个具体案例，说明 MTP 和 Speculative Decoding 的实际应用：

GPT 模型中的应用

假设我们有一个 GPT 模型用于生成长文本：

• 传统方法：GPT 基于单 token 预测，每次生成一个 token。生成长度为 𝐿 的文本需要进行 𝐿 次前向传播。
• 使用 MTP 和 Speculative Decoding：
  • 训练 GPT 时，引入 MTP 目标，让模型在每个时间步预测多个后续 token。
  • 在推理阶段，使用 Speculative Decoding：
    • 一次预测 5 个 token 候选。
    • 验证候选 token，如果通过，直接采用；如果不通过，只需重新生成部分 token。
    • 通过减少前向传播次数，加速推理。

效果

• 性能提升：通过 MTP，模型在训练阶段学到了多 token 的联合分布，因此生成的句子更加连贯。
• 推理加速：通过 Speculative Decoding，生成速度显著加快，尤其在长序列生成中效果更加显著。

---

7. 总结

这句话包含两个核心原理：

1. Multi-Token Prediction (MTP)：
   • 通过预测多个 token 提供更密集的训练信号，提高上下文建模能力，从而提升模型性能。
2. Speculative Decoding：
   • 在推理阶段利用 MTP 的能力，提前生成多个 token，并通过验证减少逐步生成的次数，从而加速推理。

两者的结合既能提升模型的训练效果，又能优化推理效率，是一种兼顾性能和速度的创新设计。这种方法特别适用于需要高效生成（如机器翻译、文本生成等）的任务。