推理飙升2倍！普林斯顿北大校友祭出多头「美杜莎」，33B模型与13B一样快
大模型推理加速如何破局？普林斯顿、UIUC等机构华人团队提出一个简单框架Medusa，训练多个解码头，直接让LLM推理速度飙升2倍。
由于LLM架构固有的内存限制，使得生成又慢又贵。
对此，很多大佬都在寻找一种挽救的方法。Karpathy曾多次提出，大模型「投机采样」对于推理时间优化是一个出色的方案。
但是，尽管投机解码能够加速生成速度，但因其太过复杂，并没有被许多开发者采用。
今天，来自普林斯顿、UIUC等机构的华人团队提出了全新的简单框架：Medusa（美杜莎）。<但图片还是羊驼吧...>
没有额外的草稿模型，研究人员只是引入了几个额外的解码头，微调出「美杜莎头」，能够在单个A100-8G GPU，一天时间内完成训练。
结果发现，Medusa直接让模型推理加速约2倍。
Vicuna-7b与Medusa
为什么LLM生成效率低？
从系统角度来看，LLM生成遵循一种「以内存为限制」的计算模式，主要的延迟瓶颈来自内存读取/写入，而非算术计算。
这一问题的根源在于，自回归解码过程中固有的顺序性。
即每次前向传递都需要将整个模型的参数，从高带宽内存（HBM）传输到加速器的计算单元。
尽管每个样本只产生一个token，但这个操作未能充分利用现代加速器的算术计算能力，由此导致模型效率低下。
在LLM崛起之前，应对这种低效率的常见方法是，简单地「增加批大小」，从而实现更多token的并行生成。
但大模型的不断涌现，让情况变得更加复杂。
在这种情况下，增加批大小不仅会带来更高的延迟，还会大大增加Transformer模型的键-值缓存的内存需求。
此外，这种低效率也体现在「成本结构」上。
截止到2023年9月，与仅处理提示相比，GPT-4的生成成本约高出2倍，Claude 2的生成成本大约高出3倍。
研究人员主要关注的重点是，改善LLM生成的延迟，同时Medusa也可以适用于需要平衡延迟和吞吐量的LLM服务。
每个token的价格
「投机解码」是终极解决方案？
基于上述的挑战，加速文本生成的一种吸引人策略是：更高效地利用计算资源。
具体来说，通过并行处理更多的token。
这种方法，采用了一个简化的「草稿」模型，每一步都能快速生成一批token的候选项。
然后，这些候选token将通过原始的、全尺寸的语言模型进行验证，以确定最合理的文本延续。
这一基本逻辑基于一个有趣的假设：「草稿」模型虽然小，但应该足够娴熟，能够生成原始模型可以接受的序列。
如果这个假设成立，「草稿」模型可以迅速生成token序列，同时原始模型可以高效地并行验证多个token，从而最大化计算吞吐量。
最近的研究表明，通过精心调整的草稿模型，投机解码可以将延迟降低高达2.5倍，令人印象深刻。
然而，这种方法并非没有挑战：
1. 寻找理想的「草稿模型」（Draft Model）：找到一个「小而强大」的草稿模型，与原始模型很好地协调，说起来容易，做起来难。
2. 系统复杂性：在一个系统中托管两个不同的模型会引入多层的复杂性，不论是计算还是操作，尤其是在分布式环境中。
3. 采样效率低：使用投机解码进行采样时，需要使用一种重要性采样方案。这会带来额外的生成开销，尤其是在较高的采样温度下。
这些复杂性和权衡限制了投机解码的广泛采用。因此，虽然投机解码前景广阔，但并未被广泛采用。
Medusa：将简单与高效融合
为了满足对更加用户友好，且功能强大的解决方案的需求，普林斯顿研究团队推出了创新框架Medusa（美杜莎）。
它不仅加速了模型的生成，甚至让LLM能够让更多人去访问和使用。
据介绍，最新方法重新回顾了「Blockwise Parallel Decoding for Deep Autoregressive Models」论文中，一个被低估的宝藏：
回溯Transformer模型的发明，与其引入一个全新的「草稿」模型来预测后续token，为什么不简单地扩展原始模型本身呢？
论文地址：https://arxiv.org/abs/1811.03115
这就是「Medusa head」（美杜莎头）发挥作用的地方。
这些额外的解码头与原始模型无缝集成在一起，在每个生成的关键点产生token块。
与草稿模型不同的是，Medusa head可以与原始模型一起训练，而原始模型在训练期间保持冻结状态。
这种方法允许研究人员在单个GPU上微调大模型，充分利用强大的基础模型学习到的表征。
此外，由于新的头仅由一个与原始语言模型头类似的层构成，因此Medusa不会增加服务系统设计的复杂性，并且适用于分布式环境。
单靠Medusa head，并不能达到将处理速度提高一倍的目标。
但这里有一个小技巧：将其与基于树状注意力机制配对使用时，就可以并行验证由Medusa head生成的多个候选项。
这样一来，Medusa head的预测能力速度提升2倍-3倍。
另外，研究人员摒弃了传统的重要性采样方案，专门为Medusa head生成创建了一种高效且高质量的替代方案。
这种新方法完全绕过了采样开销，甚至进一步提升了Medusa的加速步骤。
简而言之，研究人员用一个简单的系统解决了投机解码的挑战：
1. 没有独立的模型：不是引入一个新的草稿模型，而是在同一个模型上训练多个解码头。
2. 轻松集成到现有系统中：训练参数效率高，即使是GPU性能较差的情况下也可以进行。而且由于没有额外的模型，无需调整分布式计算设置。
3. 将采样视为一种放松：放宽与原始模型分布相匹配的要求，使得「非贪心生成」甚至比「贪心解码」更快。
Medusa概述
具体来说，Medusa在LLM的最后隐藏状态之上引入了多个头，使其能够并行预测多个后续token。
在使用Medusa head扩充模型时，原始模型在训练期间被冻结，只有Medusa head经过微调。这种方法使得在单个GPU上对大型模型进行微调成为可能。
在推理过程中，每个头为其指定的位置生成多个顶级预测。这些预测被组合成候选项，并使用基于树状注意力机制并行处理。
最后一步是，使用典型接受方案选择合理的延续，被接受的最长候选项prefix将用于下一阶段的解码。
这样，Medusa通过同时接受更多token，从而减少所需的解码步骤，提高了解码过程的效率。
接下来，让我们深入了解Medusa的三个组成部分：Medusa head（美杜莎头）、tree attention（树状注意力机制）和典型接受方案。
Medusa head（美杜莎头）
那么，Medusa head究竟是什么呢？
它们类似于原始架构中的语言模型头（因果Transformer模型的最后一层），但有一个变化：它们预测多个即将出现的token，而不仅仅是下一个token。
受到块状并行解码方法的启发，研究人员将每个Medusa head作为单层前馈网络来实现，并增加了一个残差连接。
训练这些头非常简单。你可以使用训练原始模型的相同语料库，或者使用模型本身生成新的语料库。
在这个训练阶段，原始模型保持不变；只有Medusa head经过微调。
这种有针对性的训练会带来一个参数效率极高的过程，并能迅速达到收敛。
尤其是，与在推测解码方法中训练独立的草稿模型的计算复杂性相比，优势更加突出。
在研究人员测试的Vicuna模型上，Medusa head在预测next-next token方面准确率达到60％，位列第一。同时，它仍然有改进的空间。
tree attention（树状注意力机制）
测试中，团队发现一些引人注目的指标：尽管对于预测next-next token的第1名准确率大约在60%左右，但第5名的准确率却超过了80%。
这个显著的提高表明，如果可以有效利用Medusa head产生多个排名靠前的预测，就可以增加每个解码步骤生成的token数量。
为了实现这个目标，研究人员首先通过从每个Medusa head的顶级预测中，获取笛卡尔积来创建一组候选项。
然后，根据图神经网络的思路，将依赖关系图编码到注意力机制中，以便可以并行处理多个候选项。
比如，使用第一个Medusa head的前2个预测，以及第二个Medusa head的前3个预测，如下所示。
树状注意力
上图中的可视化效果，展示了使用树状注意力同时处理多个候选项的过程。
在这种情况下，第一个头的任何预测都可以与第二个头的任何预测配对，形成一个多级树结构。
这个树的每一层，都对应着一个Medusa head的预测。由此，能够产生2*3=6个候选项。
这些候选项中的每一个对应于树结构内的一个不同的分支。
为了确保每个token只访问其前置token，研究人员设计了一种注意力掩码，专门允许注意流从当前token返回到其前置token。
通过这样做，并相应地位位置编码设置位置索引，可以在不需要增加批大小的情况下，同时处理各种候选项。
研究人员还指出，一些研究也采用了非常相似的树状注意力思想。
与它们相比，最新的方法更倾向于一种更简单的树状注意力形式，在推理期间，树状模式是规则的且固定的，使得树状注意力掩码进行预处理，进一步提高了效率。
典型接受
在早期关于投机解码的研究中，「重要性采样」被用来产生与原始模型预测密切相关的各种输出结果。
然而，后来的研究表明，随着采样温度升高，这种方法的效率往往会降低。
简单来说，如果你的草稿模型和你原来的模型一样好，理想情况下，你应该接受它的所有输出，从而使这个过程变得超级高效。
然而，重要性采样可能会在中间环节，拒绝这个解决方案。
在现实世界中，我们经常调整采样温度，只是为了控制模型的创造力，不一定是为了匹配原始模型的分布。
那么为什么不把重点放在，可接受可信的候选项上呢？
对此，研究人员引入了「典型接受」方案。
从现有的截断采样（truncation sampling）中汲取灵感，普林斯顿研究人员目标是根据原始模型选择足够可能的候选项。
研究人员根据根据原始模型的预测概率设置一个阈值，如果候选项超过这个阈值，就会被接受。
用专业术语来说，研究人员采用硬阈值，以及熵相关阈值中的最小值，来决定是否像截断采样那样接受候选项。
这确保了在解码过程中，选择有意义的token和合理延续。
研究人员总是使用贪心解码接受第一个token，确保在每个步骤中至少生成一个token。最终输出是通过研究接受测试的最长序列。
这一方法的高明之处在于，适应性强。
如果你将采样温度设置为0，它就会恢复到最有效的形式——贪心解码。当温度升高时，最新研究的方法变得更加高效，可以接受更长的序列。
并且，研究人员还已经通过严格的测试证实了这一说法。
因此，从本质上讲，最新的典型接受方案提供了一种更有效的方式，让LLM进行更创造性地输出。
Llama「吐口水」可以有多快？
研究人员使用了专门针对聊天应用进行微调的Vicuna模型对Medusa进行了测试。
这些模型的大小不同，参数数量分别为7B、13B和33B。
研究的目标是，测量Medusa在实际聊天机器人环境中如何加速这些模型。
在训练Medusa head时，研究人员采用了一种简单的方法，使用了公开可用的ShareGPT数据集。这是最初用于Vicuna模型的训练数据的子集，仅进行了一个epoch的训练。
而且，关键是，根据模型参数的大小，整个训练过程可以在几小时到一天之内完成，仅需要一块A100-80G GPU。
值得注意的是，Medusa可以轻松与量化的基础模型结合使用，以减少内存需求。
研究人员利用这一优势，在训练33B模型时使用了8位量化。为了模拟实际环境，他们使用了MT bench进行评估。
凭借其简单的设计，Medusa在各种用例中，始终能够实现约2倍的绝对时间加速。
值得注意的是，通过Medusa的优化，33B参数的Vicuna模型可以像13B模型一样迅速运行。
消融实验
Medusa head的配置
在利用美杜莎头的预测能力时，研究人员灵活地选择了每个头应考虑多少个顶级候选项。
例如可以选择第一个头的前3个预测和第二个头的前2个预测。当对这些顶级候选项进行笛卡尔积运算时，就生成了6个延续选项供模型评估。
这种可配置性，是需要进行权衡的。
一方面，选择更多的顶级预测会增加模型接受生成token的可能性。另一方面，它还会增加每个解码步骤的计算开销。
为了找到最佳平衡，研究人员尝试了各种配置，并确定了最有效的设置，如附图所示。
典型接受的阈值
在典型接受方案中，一个关键的超参数，称为「阈值」，帮助研究人员根据模型自身的预测，来确定生成的token是否合理。
阈值越高，接受标准越严格，从而影响通过这种方法获得的整体加速。
研究人员通过在MT bench的2个以创造力为导向的任务上进行实验，来探讨质量和加速之间的这种权衡。
图中显示的结果表明，与贪心解码方法相比，典型接受能够加速10%。这种加速明显优于使用随机采样的投机解码方法，后者与贪心解码相比实际减慢了过程。
参考资料：
https://twitter.com/tianle_cai/status/1701272996328120408
https://sites.google.com/view/medusa-llm
