Deepseek技术原理2：最详细图解模型结构MTP和思考

浪漫的数据分析

于 2025-03-14 09:00:00 发布

阅读量1.7k

点赞数 20

文章标签：人工智能深度学习机器学习

本文链接：https://blog.youkuaiyun.com/weixin_43290383/article/details/146245802

版权

作者： 罗辑

Meta GPU批次并行预测（batch）

DeepSeek改进版MTP

改进分析

DeepSeek GPU批次并行预测（batch）

DeepSeek MTP优缺点

什么时候做MTP（when）

个人的思考和改进（my suggestion）

参考文献

摘要：火热的DeepSeek模型结构的改进上，除了MLA、MOE，还有个MTP（Multi-Token Prediction）--多词预测，简单说是：让模型在训练时，一次性预测多个未来词（token），而不是仅仅预测下一个词（token）。这样计算loss时，除了和next token计算loss外，还可以和未来的token标签进行多个loss的计算，效果上可以加速模型收敛。

DeepSeek V3原始论文中的示意图，把几件事情杂糅在一起，所以看着理解起来有点费劲，别急，我会详细的剖析，拆分成多个环节来讲，结合图示讲解。做一次史上最详细讲解。

为什么要做MTP（why）

当前主流的大模型(LLM)都是decoder模型结构，传统训练方法是一个词一个词的做文本续写，叫token-by-token。实际操作过程中会面临两个瓶颈。

访存密集型瓶颈

每次在生成一个token的时候，都要频繁从显存中读取缓存的向量（加载KV-Cache），频繁跟访存交互，再通过多层神经网络做完整的前向计算。GPU擅长大规模并行计算，但显存访存效率形成训练和推理的瓶颈。

训练和推理效率瓶颈

训练时，每次在只用下一个token的预测结果，来计算loss更新模型，效率不高，同时，推理是也是token-by-token一个个的预测，效率也低，也跟人类做文本续写方式不太一致。虽然我们没法下笔时全篇都考虑清楚了，但是当前词语的下下一个词，多半提前已经在心里构思好了，不需要等到写完一个词语再停下来想一个词。

什么是MTP（what）

MTP(Multi-Token Prediction)实际上就是将大模型原始的1-token的生成，转变成multi-token的生成，从而提升训练和推理的性能。具体来说，在训练阶段，一次生成多个后续token，可以一次学习多个后续位置上的label，进而有效提升样本的利用效率，提升训练速度；在推理阶段通过一次生成多个后续token，实现成倍的推理加速来提升推理性能。

用个示意图可以表示为：

通过【1】预测后续的【2345】，假如【2345】都是准的话，下一个预测的词就是【6】了，这样就提高了3倍预测效率。

说的深刻一点。传统方法的问题（预测下一个token）：

训练阶段：token-by-token生成，是一种局部感知的训练方法，难以学习长距离的依赖关系。

推理阶段：逐个token生成，推理速度较慢

MTP方法（一次预测多个token）：

训练阶段：通过预测后续多步的token，迫使模型学到更长的token依赖关系，从而更好理解上下文，避免陷入局部决策的学习模式。同时一次预测多个后续token，可大大提高样本的利用效率，相当于一次预估可生成后续多个<predict, label>样本，能收集到多个loss来更新模型，有助于模型加速收敛。
推理阶段：并行预估多个token，可提升推理速度。

怎么样实现MTP（how）

原始版MTP

MTP最早可以追溯到Google在18年发表在NIPS上的工作，18年是Transformer诞生年，还没有大模型，但那时提出MTP的原理更现在的MTP相差不大。

直观的示意图如下：

可以看到，模型在做词语接龙时，把前文进行了编码，然后并行预测了后面几个词。

原始论文的示意图如下：

原始论文示意为用“I saw a dog ride”去同时预测后续的“in”“the”“bus”。

Meta改进版

原始版MTP有个什么问题呢？因为那时候当前LLM的decoder架构还不受到重视，因此meta结合当前LLM的架构，重新设计了更符合大模型的MTP。

结合前面那个中文预测的示意图，可以理解为：

除了将模型结构更具体化成transformer的block外，还保留了在训练时的前后依赖关系，业界称为“因果关系”。

Meta GPU批次并行预测（batch）

大模型所谓的“因果关系”具体过程如下：

第一步：当输入第一个词时，同时预测后续4个词，即为图中的粉红色[2,3,4,5]。示意图如下：

第二步：当输入第一、二个词时，同时预测后续4个词，即为图中的粉红色[3,4,5,6]。示意图如下：

也就是在预测[3,4,5,6]时，输入的是[1,2]，而不是单单[2]，但实际操作时，因为第一步预测时，已经输入了[1]，而且[1]已经在GPU中缓存起来了，所以实际操作时，仅再补充输入[2]就行了，所以有些示意图就简单画的输入[2]，我们要理解它是有用到[1]的。

同理，

当输入第一、二、三个词[1,2,3]时，同时预测后续4个词，即为图中的粉红色[4,5,6,7]。示意图如下：

因为GPU在训练时，是支持批量同时训练的，多条数据组成一个batch批次，同时并行计算，因此，可以把上面几个训练画在一张图里，就变成了如下原始论文形式：

如果没有上面的铺垫，就很难看懂原始论文中的图示意思，包括DeepSeek论文中的MTP示意图。

DeepSeek改进版MTP

Meta版MTP有什么问题没呢？那就是除了next token（下一个词）比较准之外，后面的第二、第三个词，不太准。这个也很好理解。比如我们用“春天到了，又”去同时预测后续的“到了”“万物”“复苏”的三个词。我们知道，两个词语越接近，就关联性越强，也就越容易预测，反之则越远越难预测，用“春天到了，又”预测“复苏”难度太大。所以第二、第三个词基本上没法预测准，因为缺失太多中间词语，语义信息有缺失。这会导致预测第二、第三个词的神经网络没法训练到收敛。

但Meta版MTP也提到另一个重要的点，就是虽然第二、第三个词预测不太准，但训练时，它因为和shared编码网络有神经元连接，在训练阶段，可以通过梯度回传，更新shared编码网络参数，达到加速shared编码网络训练的过程。而这正是DeepSeek想要的，加快训练过程，从而节省训练的时间成本，也能节省电力开销。

添加的第二个词预测头（head2）的loss产生的梯度传递如下图所示。

奔着“加快训练过程”这个目的去的看DeepSeek的模型改进，就非常能理解他的做法。首先，目的是为了得到一个能预测next token的完整LLM大模型，我们称之为主模型。第二，通过一定的方式，使用第二、第三个词预测来加快主模型的训练速度。

完整的主模型长啥样呢？如下图。一个标准的且普通得不能再普通的大模型结构，通过用词语t1，去预测生成T2，和真实标签t2比对，计算交叉熵损失函数，然后反向梯度回传，更新整个模型。

那如果是领导要求你改一下模型结构，去辅助预测后续词语token，还能反向更新整个模型参数，你会从那里搭线出去？

改进分析

我们知道，图中那个output head接收的transformer block处理后的向量，专门预测t2的，不可能强制让他变成预测t3，也就是不能从图中output head之后接额外的网络结构。那如果从L个（L=32）transformer block中间引线出来，比如第30个transformer block中间引线出来，预测后续词语token，那第31、第32个transformer block因为没有连接后续词语token的预测，无法梯度回传到第31、第32个transformer block，他两个block的参数就无法得到更新。

自然而然，我们能想到的，就是从最后一个block搭线出去。如下图：

这样后续词语token的预测loss，能梯度回传到所有transformer block，就能最大程度覆盖主模型的所有神经元，除了output head专门用来预测t2的结构外。

搞定一个搭线问题后，我们解决第二个问题，应该引入什么样的网络结构来预测后续token？我们最初分析了，本身后续第二、第三个词基本上没法预测准，我们就简化处理，同时参照了meta MTP那种因果关系，实际上用的是[t1,t2]去预测[t3].最简单的方式，就是再用一个transformer block专门去预测[t3]，请注意，用1个，而不是L=32个。

因此，设计的示意图就如下：

当然有个小细节就是，transformer block只接受一个token的向量，我们刚才的设计草稿，让transformer block输入了两个token的向量（一个是红色箭头，一个是t2），因此需要合并一下，如何合并，就是添加一个线性层，将两个向量压缩融合成一个向量，在送进transformer block。2个向量融合成1个向量如下图：

因此，我们添加的辅助预测后续词语的网络结构就设计出炉了。

图中红框1就是向量融合，红框2就是1个标准的transformer block，输出的output就是预测t3、t4等后续词语。

同时为了进一步压缩模型参数量，我们用相同的embedding层，用来把文字词语编码成向量，也就是embedding layer是共享的。同时，output head也可以共享。只要output head输入的向量不同，他就可以用来预测t3，t4。

因此，我们把他级联起来，可以得到如下一个图。

再次提醒，新加的辅助模块用来预测[t3,t4]的transformer block是1个，而不是L=32个。

DeepSeek GPU批次并行预测（batch）

参照3.2.1中meta MTP的并行“因果关系”训练的具体过程，可以把一个batch中的多条记录，画在一张图上，就成了如下所示：

美化一下，调整一下排版，就成了DeepSeek论文中的样子。

这样分析下来，仿佛我如果在DeepSeek工作的话，我也能提出这种结构上的优化的错觉，没有任何难度，哈哈哈哈，当事后诸葛亮就是爽。

DeepSeek MTP优缺点

缺点：根据之前的分析，MTP原始的优势：加速推理，在DeepSeek这种模型结构下不适用，从模型结构上也可以看出，预测[t3]时，[t2]仅仅经过了1个transformer block，肯定对信息的提取程度是不够的。所以理论上这套模型架构就无法准确预测[t3,t4,t5,t6],所以论文里面也提到，训练完模型做推理时，基本上还是预测下一个词，丢掉了预测后续多个词的辅助模型结构，也就是仅仅保留了main model。反过来说，如果[t3,t4,t5,t6]能一次性预测准，按照中国人的调性，牛逼得吹得老高了，而他们没有提这点，而是用的“directly discard”，肯定也是因为不准。
看看论文中的表述：MTP in Inference. Our MTP strategy mainly aims to improve the performance of the main model, so during inference, we can directly discard the MTP modules and the main model can function independently and normally.
优点1：meta MTP有提到通过后续词语的loss，辅助主模型的梯度更新，这个优点在DeepSeek模型架构上，得到最大程度的保留。首先后续词语的loss能通过梯度回传，覆盖main model的所有transformer block，能辅助他们加快收敛。其次output head也是共享的，理论上也能接收到后续词语的loss影响，而且embedding layer也是共享的，所以整个main model的所有部件，都能受到后续词语的loss影响，最大程度的加快了main model的收敛。

优点2：因为main model的所有部件，都参与到了后续多个词语的训练，理论上这样的模型，多少也有点远程的规划能力，就是说第一个词语时，大概也有点知道第3,4,5个词了，所以也会稍微比只知道第二个词的模型聪明一点，这点DeepSeek做了实验。也有结果对比。会稍微好点，但不多。

什么时候做MTP（when）

经过刚才的详细分析，知道DeepSeek的MTP预测后续多个词语不准，但是可以加速模型收敛，所以在模型急需收敛的预训练阶段，加入这项技术是最合适的，在后训练以及微调时，因为语料也不是很大了，模型急需收敛的矛盾不突出，另外它的第二个优点：提升模型长程推理能力，这个效果并不太显著，属于锦上添花。所以结论就是：在预训练阶段做MTP最合适。

个人的思考和改进（my suggestion）

DeepSeek的MTP有优点但也称不上完美，还可以进一步改进，个人认为可以从两个方面。第一，我们知道，两个单词之间的距离越远，相关性越差，预测也就越不准。那在计算loss时，需要设置一个衰减因子来调节后续词语loss的权重。DeepSeek这里简单设置成一个平均值λ，个人认为这里的衰减因子应该设置成一个随距离呈指数衰减的形式更合理，更接近实际文本之间的相关性，参照GAE对未来奖励的估计那种的设计。