无损加速最高5x，EAGLE-2让RTX 3060的生成速度超过A100

AIxiv专栏是机器之心发布学术、技术内容的栏目。过去数年，机器之心AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

李堉晖：北京大学智能学院硕士，受张弘扬老师和张超老师指导，研究方向为大模型加速和对齐，正在寻找25届工作机会

魏芳芸：微软亚研院研究员，研究方向为具身智能、图像生成和AI agents

张超：北京大学智能学院研究员，研究方向为计算机视觉和机器学习

张弘扬：滑铁卢大学计算机学院、向量研究院助理教授，研究方向为LLM加速和AI安全

自回归解码已经成为了大语言模型（LLMs）的事实标准，大语言模型每次前向计算需要访问它全部的参数，但只能得到一个token，导致其生成昂贵且缓慢。

今日，一篇题为《EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees》的论文提出了动态草稿树投机采样，依据草稿模型的置信度动态调整草稿树的结构，最高可以将大语言模型的推理速度提高5倍，同时不改变大语言模型的输出分布，确保无损。

论文链接：https://arxiv.org/pdf/2406.16858

项目链接：https://github.com/SafeAILab/EAGLE

Demo链接：https://huggingface.co/spaces/yuhuili/EAGLE-2

EAGLE-2在多轮对话数据集MT-bench上的加速效果（上图为贪婪生成，下图为采样生成）：

使用EAGLE-2,2张RTX 3060（$300）的推理速度可以超过A100（$10000）。

背景

投机采样使用一个小的模型快速生成草稿，原始的大语言模型可以通过一次前向计算验证草稿的正确性，将正确的草稿作为输出，从而一次生成多个token，并确保无损。EAGLE是投机采样的一种改进。它在更有规律的特征层面而不是token层面进行自回归，同时输入采样结果（超前一个时间步的token）消除了不确定性，明显提升了草稿模型的准确率。

到目前为止，EAGLE在第三方测试Spec-Bench（https://github.com/hemingkx/Spec-Bench/blob/main/Leaderboard.md）中排名第一。

思路

EAGLE和Medusa等方法使用静态的草稿树，隐式地假设草稿token的接受率和上下文无关，下面是一个简单的例子

上文是“10+2”时，下一个token难以预测，EAGLE在这个位置添加两个候选token以增加草稿命中率，“10+2=”和“10+2+”有一个正确即可。当上文是“10+2=”时，下一个token明显是“1”，但是EAGLE使用静态的草稿结构，仍然添加两个候选“1”和“3”，“10+2=3”不可能通过大语言模型的检查，存在浪费。EAGLE-2旨在解决这一问题，如下图所示，当上文是“10+2=”时，EAGLE-2只增加一个候选token“1”，将节约出的token用于让草稿树更深，这样“10+2=12”通过大语言模型的检查，EAGLE-2可以一次生成更多的token。

EAGLE-2的作者们在Alpaca数据集上进行了简单的测试，下图显示了不同位置的草稿token的接受率，左图中的P1-P6代表位置，与右图的横轴坐标对应。实验结果显示，在相同的位置上的草稿token的接受率也有较大的差异，这说明了使用动态草稿树可能取得比静态草稿树更好的效果。

上述例子中，EAGLE-2根据预测草稿token的难易程度决定草稿树的结构，精确计算难易程度（接受率）需要原始大语言模型的计算结果，这违背了投机采样减少对原始大语言模型访问的初衷。幸运的是，EAGLE的草稿模型的置信度与接受率（难易程度）高度正相关。下图显示了草稿模型不同置信度区间的草稿token的平均接受率，红色虚线连接（0,0）和（1,1）。由此可见，草稿模型的置信度可以作为接受率的有效近似。

方法

EAGLE-2包括两个阶段，扩展和重排，扩展阶段加深加大草稿树，重排阶段修剪草稿树，丢弃部分节点（token）。

为了保证无损，一个草稿token被接受的前提是它的祖先节点都被接受，所以EAGLE-2将一个节点的价值定义为它和它祖先的接受率的乘积，用置信度的乘积来近似。

在扩展阶段，EAGLE-2选择草稿树最后一层价值最高的m个节点（token）进行扩展。这些token被送入草稿模型，然后将草稿模型的输出作为子节点连接到输入节点，加深加大草稿树。在重排阶段，EAGLE-2按照价值对整棵草稿树进行重排序，保留前n个节点（token）。草稿token的置信度在0-1之间，两个节点价值相同时优先保留浅层节点，因此重排后保留的草稿树一定是连通的，保证了语义上的连贯性。重排后草稿树变小，降低了原始大语言模型验证的计算量。为了保证计算结果的正确性，还需要调整attention mask，确保每一个token只能看到它的祖先节点，不受其他分支的影响。下面是一个简单的例子。

扩展（Expand）阶段的黄色框表示被选中进行扩展的节点，绿色框为以这些节点为输入时草稿模型的预测。重排（Rerank）阶段的蓝色框表示被保留的节点，之后它们被展平成一维作为原始大语言模型的输入。EAGLE-2根据树的结构调整attention mask，比如，”a”只能看到它的祖先“It”和“is”，看不到另一个分支的“has”。EAGLE-2也同时调整位置编码，确保和标准自回归解码的一致性。

实验

EAGLE-2在多轮对话、代码、数学推理、指令遵循、问答、总结六项任务上分别使用MT-bench、Humaneval、GSM8K、Alpaca、CNN/DM、Natural Questions数据集进行了实验，与6种先进的投机采样方法（SpS、PLD、Medusa、Lookahead、Hydra、EAGLE）进行了比较。

表格中的Speedup为加速比，τ 为平均接受长度，也就是原始大语言模型每次前向计算能生成的token数。EAGLE-2每次前向计算能生成大约4-5个token，而自回归解码每次生成1个token，因此EAGLE-2明显加速了大语言模型的生成，加速比为2.5x-5x。加速比和接受长度在代码生成任务（Humaneval数据集）上最高，这是因为代码中存在大量确定性的模板，草稿更容易命中。在所有任务和大语言模型上，EAGLE-2的加速比和平均接受长度都是最高的，明显优于其他方法。

应用

EAGLE-2也在工业界得到应用，集成至Intel/intel-extension-for-transformers等。