MTP（Multi-Token Prediction）：提升大模型性能的利器

一、引言

在自然语言处理领域，大型语言模型（LLMs）已经取得了显著的成果。然而，这些模型在训练和推理过程中面临着一些挑战，其中最主要的问题之一是生成效率低下。传统的单token预测方法（Next-Token Prediction）在训练和推理阶段都需要逐个生成token，这不仅耗时，而且难以学习长距离的依赖关系。为了解决这些问题，MTP（Multi-Token Prediction）技术应运而生。

二、为什么要做 MTP

1. 背景

当前主流的大模型都是基于解码器（decoder-based）的结构，在训练和推理阶段，对于一个序列的生成过程，都是token-by-token的。每次生成一个token时，都需要频繁与访存交互，加载KV-Cache，再通过多层网络做完整的前向计算。这种访存密集型的任务通常会因为访存效率形成训练或推理的瓶颈。

2. MTP方法的作用

MTP的核心思想是通过解码阶段的优化，将1-token的生成转变为multi-token的生成，从而提升训练和推理的性能。具体来说：

• 训练阶段：一次生成多个后续token，可以一次学习多个位置的label，进而有效提升样本的利用效率，提升训练速度。
• 推理阶段：通过一次生成多个token，实现成倍的推理加速来提升推理性能。

三、MTP的核心原理与公式推导

1. 核心原理

MTP方法的核心是在训练阶段让模型一次性预测多个未来的token，而不是像传统方法那样只预测下一个token。这样可以迫使模型学习更长的token依赖关系，从而更好地理解上下文，避免陷入局部决策的学习模式。

2. 公式推导

在传统的单token预测方法中，模型的训练目标是最大化下一个token的预测概率，其损失函数可以表示为：

$$L_1=-\sum _t\log P_{\theta }(x_{t+1}|x_{1:t})$$

其中，$P_{\theta }$ 是模型的参数，$x_{t+1}$ 是下一个token，$x_{1:t}$ 是当前的token序列。

而在MTP方法中，模型需要一次性预测接下来的n个token，其损失函数可以表示为：

$$L_n=-\sum _t\log P_{\theta }(x_{t+n:t+1}|x_{1:t})$$

这里，$x_{t+n:t+1}$ 表示从 $t+1$ 到 $t+n$ 的token序列。

通过对多个token的预测，模型可以同时学习到多个位置的label，从而提高样本的利用效率，加速模型的收敛。

四、MTP的实际案例

1. Meta的MTP

Meta在2024年4月发表了一篇关于MTP的论文，提出了一种简单的多token预测架构。该架构在训练时间和内存使用上没有额外开销，实验证明在大规模模型上是有效的，平均可以解决大约15%以上的编程问题。

• 方法实现：Meta的MTP方法使用一个共享的Transformer主网络，上面接入多个并行预测头，针对输入token分别预测后续的多个token。
• 优势：通过预测多个token，模型能够捕捉更长距离的依赖关系，提高训练效率和模型性能。

2. DeepSeek的MTP

DeepSeek-V3也采用了MTP技术，通过引入多token预测目标来增强模型的性能。

• 实现过程：DeepSeek-V3使用多个顺序模块来预测多个未来token。每个模块包括共享的嵌入层、共享的输出头、一个Transformer块和一个投影矩阵。
• 训练目标：对于每个预测深度，计算交叉熵损失，并对所有深度的损失进行平均，得到总体MTP损失。
• 优势：MTP目标使得训练信号更加密集，从而可能提高模型在数据利用上的效率，增强预测能力。

五、总结

MTP（Multi-Token Prediction）技术通过在训练阶段一次性预测多个token，显著提升了大模型的训练效率和推理性能。它不仅能够提高样本的利用效率，加速模型的收敛，还能够增强模型对长距离依赖关系的学习能力。Meta和DeepSeek的实践案例充分证明了MTP的有效性和实用性，为大模型的优化提供了新的思路和方法。