投机解码 Speculative Decoding 和多token预测 MTP-平芜编程栈

模型采样

投机解码是一种采样方式，先来看模型是如何采样的，模型的最后一层输出是一个长度等于词典大小的向量，表示下一个token是词典中每个词的概率。

采样就是根据这个概率向量，选出来下一个token，也就是不会改变模型的推理架构。最简单的采样方式就是谈心采样，每一步都选向量中概率最大的token。也

投机解码是想要优化这个采样过程，思想是，原来的解码方法每次只能解码一个token，是严格串行的，一个自然的想法就是能不能并行这个过程，一次预测多个token。

投机解码的做法是，让一个小模型预测草稿，连续调用这个小模型预测多步，一次性输出多个token，然后让大模型并行验证每一个预测是不是对的，如果对的则接受，不对则根据策略来接受或修改。

这里的巧妙之处在于，小模型预测多次是串行的，但是小模型参数量很少，比大模型推理要快很多。大模型只负责检验，检验由于是在已有的基础上做的，通过类似训练时的因果掩码，可以同时预测多个前缀的下一个token，也就是检验阶段可以并行。整体下来，时间是比大模型串行推理多个token要快的。

找到大模型结果和小模型结果对比，第一个验证失败的token，这个位置用大模型的预测结果替代，小模型前面的预测结果接受，后面的丢掉。

验证最简单的就是对比大模型预测和小模型预测的最大概率token是不是同一个，但这有点过于苛刻了。实际预测采取的是一个更灵活的方法。

小模型连续预测了多个位置，大模型并行检验，也预测了相同数量的位置。此时，对于每个位置，我们手上有两组概率：

大模型会从第一个草稿 Tokenx 1 x_1x1开始逐个验证

如果大模型觉得小模型猜得很好（p ( x 1 ) ≥ q ( x 1 ) p(x_1) \ge q(x_1)p(x1)≥q(x1)）也就是说，大模型自己算出来的概率，比小模型还要高（或者相等）。直接接受小模型的预测，继续看下一个位置。
如果大模型觉得小模型在瞎猜（p ( x 1 ) < q ( x 1 ) p(x_1) < q(x_1)p(x1)<q(x1)）这意味着小模型太自信了，但大模型觉得这个词没那么靠谱。这时候，大模型不会一棒子打死，而是会计算一个接受概率：接受概率 = p ( x 1 ) q ( x 1 ) \text{接受概率} = \frac{p(x_1)}{q(x_1)}接受概率=q(x1)p(x1)
大模型会采样一个0到1的随机数，看是否坐落在[ 0 , p ( x 1 ) q ( x 1 ) ] [0,\frac{p(x_1)}{q(x_1)}][0,q(x1)p(x1)]内如果在则接受小模型的预测，继续看下一个；如果不接受，后面的预测不用看了，都作废。
一旦一个位置的预测被大模型拒绝了，这个位置之后的丢掉，这个位置会使用大模型的预测结果作为替代

推理时选择的小模型也有讲究。

一种策略是，选择一个独立的预测模型，这样的优点是预测模型可以很小，推理快，部署需要的资源少，且可以对这个小模型做单独训练。缺点是，小模型和大模型可能预测出的概率分布不一样，很容易被拒绝，被拒绝太多，就相当于小模型预测一个，大模型预测一个，没有加速效果了。
另一种策略是，直接使用验证模型本身进行预测，但对验证模型做一些裁剪，比如下图，4层注意力，只运行2层就输出了，降低推理时间，这被称为自投机解码。这样的优点是，本质上是相同模型的小参数版本，预测和验证时分布更接近，预测更容易被接受，且只需要维护一个模型，更简单。缺点是验证模型裁剪后推理速度可能还是不如小模型快，加速效果不明显，且不能对预测模型做独立训练

另一个和投机解码有点类似的概念是MTP，多token预测。和投机解码一样也可以一次预测多个token，但区别在于，投机解码本身是一种后处理，不懂模型结构，是拿到概率后的操作，但MTP是在模型层面就设计成一次输出多个token

最早MTP由meta提出，思路是在注意力层结束后，设置多个并行的分类头，同时输出多个token的概率分布，分别作为接下来第一个，第二个，第三个……第k个token，这样的缺点是并行的头，不好训练

MTP第一次在工业生产中被大规模使用，是DeepSekk V3，此时为了使模型更好训练，把并行预测多个token改成串行预测了，如下图，进行多个注意力层，最后几个注意力层，每层预测一个新token。这样虽然牺牲了一些并行度，但梯度爆炸减少了，更好训练。

注意不要搞混，串行MTP看起来也是运行一步输出一个token，和非MTP很像。但仔细看的话，会发现输出多个token都是模型一次前向里做的，相当于把原本最后几层注意力拿出来，每层接上一个分类头。而非MTP架构，一次前向传播只输出一个token，还是不一样的。