图1：原始自回归推理（为了便于直观理解，假设每个字都为单个token）。大模型在收到用户提问后会做出单个字的预测

图3：多头美杜莎的top-1推理（假设每个字都为单个token）。大模型在收到用户提问后会做出多个字的预测，这里大模型会一次性处理多字。具体为主预测头会预测下1个字，第一个和第二个美杜莎头分别会预测第2个和第3个字。更多的美杜莎头也以此类推。美杜莎头为独立模块，可加在现有预训练/微调好的基础模型中（例如Vicuna-13B）

投机采样虽然提升了自回归大语言模型的推理速度，然而在某些场景下，小模型的选择是件棘手的事。此外，对于现有的框架来说，如何同时部署大模型和小模型也是一种挑战。因此，今年9月，学者们提出了一种名为多头美杜莎（Medusa: Multiple Decoding Heads）[4] 的方法来解决这些问题。多头美杜莎的理念与投机采样类似，都是通过增加额外的计算量来减少数据搬运的代价，从而提升整体的推理速度。它们的区别在于，多头美杜莎利用了多个预测头（language model heads）来进行多字预测（图3），这些额外的预测头被称为美杜莎头。这种方法基于一个更早的工作 [5]，该工作研究了自回归模型如何使用多个预测头进行多字预测。然而，与之不同的是，多头美杜莎可以应用于已经训练好的模型中，而无需改变原有模型的结构，只需在现有的模型中加入美杜莎头即可。其训练过程是独立的，将原有模型的参数都固定，不进行训练，只针对额外加入的美杜莎头进行参数训练。这样做大大减少了计算量，并且训练的收敛难度也大幅降低。

每一个美杜莎头的结构都为单层或多层ResBlock接上Final Linear层（隐空间转换到字典维度的线性层）。每层ResBlock的结构为x+SiLU(Linear(x))，这里x为隐空间输入，SiLU为激活函数，Linear为线性层。在训练时，只需要使用原始网络的最终隐空间输出，而不需要经过Final linear层（主预测头）的logit。这是因为我们并不想改变原始网络的参数。每多一个美杜莎头则会预测更多一个字。例如，主预测头预测第1个字（不参与多头美杜莎训练，这里只是为了举例），第一个美杜莎头预测第2个字，第二个美杜莎头预测第3个字，并以此类推。在训练中，原始模型只参与前向传播过程而不会进行反向传播更新参数，外加上美杜莎头通常只使用2-3个，并且每个美杜莎头只有单层ResBlock，所以整体多头美杜莎训练是非常轻量级的，需要的硬件资源和时间要比原始模型少很多。

在实际推理场景中，如果直接使用top-1（贪婪）策略来进行推理会很容易掉进局部最优概率组合。因此，为了提高推理效果，多头美杜莎使用了top-k的方式来进行推理。这里举个实际的例子，假设使用两个美杜莎头。主预测头使用贪婪解码，第一个美杜莎头使用top-7解码，第二个美杜莎头使用top-6解码，这样会有42*3=126个tokens需要输入给模型判断。这里，多头美杜莎在top-k预测中做了一个创新改进，用来提升性能，具体为这126个tokens内部其实是一个树结构，里面含有重复的tokens组合，因此我们可以将其转换成树状结构：1+7+7*6=50个tokens输入给模型（图4）。这么做的确减少了模型的计算量，不过这里需要在序列混合的地方进行调整，去除树状结构中不相关tokens之间的关联，也就是注意力模块和位置编码。对此，原文提出树状注意力掩码（tree attention mask）[6]，用于掩盖不可能的tokens组合（图5）。

图4：树状索引和原始组合索引，树状组合可以减少模型前向计算量。可在模型输出后转换回原始组合结构来做最后的预测

图5：多头美杜莎的树状注意力机制，掩码用来屏蔽分支与分支之间字的关联。图片来源于[6]

多头美杜莎的单次迭代算法可以总结为5步：

1.得到原始模型（主预测头）和美杜莎头的logit输出（第一次迭代中为用户提示的序列长度；之后迭代中的序列长度为1），并将树状掩码加在模型注意力模块中。

2.从原始模型和美杜莎头的logit输出得到各自的top-k 候选tokens，并将其转换成树状组合格式。同时，按照树状组合结构来更新位置编码。

3.树状结构输入给模型并得到原始模型和美杜莎头的树状logit输出，并将其转换回原始组合的格式用于下一步的最大概率组合选择。

4.从候选tokens（由第1步中美杜莎头产生）索引找到与原始模型logits概率为1的索引能对上的数字（三个字的例子：[[1,0], [1,1], [0,1]]），并对其进行从左到右的累积乘积（cumulative product）操作（例如[[1,0], [1,1], [0,0]]），之后得到最佳候选tokens组合的索引以及其接受长度。

5.将最佳候选tokens组合与现有input_ids合并，并更新其索引位置对应的KV cache。最终，取出在接受长度索引位置的原始模型和美杜莎头的logit输出，用于下一次迭代。

这里，我们列出了多头美杜莎算法（温度超参设置为0）的大概思路，具体算法可参考官方源代码[7] 。

多头美杜莎的一大要点在于模型每次需要处理每个头的top-k树状分支数总和。例如，当我们在做贪婪预测时，意味着top-k = {1,1,1}，因此模型一次需要处理1+1+1*1=3个字；在之前的例子中top-k = {1,7,6}，这意味着模型一次需要处理1+7+7*6=50个字。从这里我们可以看到与投机采样类似的问题：当单词预测过多字时，预测成功的难度也随之提升。这里可以注意到，多头美杜莎还会面临的一个问题是随着美杜莎头数量增加，top-k的树状分支也将会以指数增长，造成庞大的计算开销。此外，许多基础和微调模型并没有开放其训练数据集，因此多头美杜莎面临的另一大问题是使用什么数据来训练美杜莎头。原文中使用了ShareGPT的数据来训练基于Vicuna模型的美杜莎头，根据参数规模的不同，训练时耗也仅仅为几个小时到一天。虽然ShareGPT只是微调训练Vicuna模型的一部分数据，然而训练好的多头美杜莎Vicuna还是能较好地保持原始模型的预测效果。

[1] Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv preprint arXiv:2305.13245.

[2] Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.

[3] Chen, C., Borgeaud, S., Irving, G., Lespiau, J. B., Sifre, L., & Jumper, J. (2023). Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318.

[4] https://github.com/FasterDecoding/Medusa

[5] Stern, M., Shazeer, N., & Uszkoreit, J. (2018). Blockwise parallel decoding for deep autoregressive models. Advances in Neural Information Processing Systems, 31.

[6] https://sites.google.com/view/medusa-llm

[7] https://github.com/FasterDecoding/Medusa/blob/main/medusa/model/utils.py

[8] https://github.com/dust-tt/llama-ssp

[9] Leviathan, Y., Kalman, M., & Matias, Y. (2023, July). Fast inference from transformers via speculative decoding. In International Conference on Machine Learning (pp. 19274-19286). PMLR.

[10] Spector, B., & Re, C. (2023). Accelerating LLM Inference with Staged Speculative Decoding. arXiv preprint arXiv:2308.04623.

本文中所引用的图片和视频来源均已标明，若有侵权请联系我们予以删除。