2023/12/14 14:26

大模型推理效率无损提升3倍，滑铁卢大学、北京大学等机构发布EAGLE

大语言模型（LLM）被越来越多应用于各种领域。然而，它们的文本生成过程既昂贵又缓慢。这种低效率归因于自回归解码的运算规则：每个词（token）的生成都需要进行一次前向传播，需要访问数十亿至数千亿参数的 LLM。这导致传统自回归解码的速度较慢。

近日，滑铁卢大学、加拿大向量研究院、北京大学等机构联合发布 EAGLE，旨在提升大语言模型的推理速度，同时保证模型输出文本的分布一致。这种方法外推 LLM 的第二顶层特征向量，能够显著提升生成效率。

技术报告：https://sites.google.com/view/eagle-llm
代码（支持商用 Apache 2.0）：https://github.com/SafeAILab/EAGLE

EAGLE 具有以下特点：

比普通自回归解码（13B）快 3 倍；
比 Lookahead 解码（13B）快 2 倍；
比 Medusa 解码（13B）快 1.6 倍；
可以证明在生成文本的分布上与普通解码保持一致；
可以在 RTX 3090 上进行训练（1-2 天内）和测试；
可以与 vLLM、DeepSpeed、Mamba、FlashAttention、量化和硬件优化等其他平行技术结合使用。

加速自回归解码的一种方法是投机采样（speculative sampling）。这种技术使用一个更小的草稿模型，通过标准自回归生成来猜测接下来的多个词。随后，原始 LLM 并行验证这些猜测的词（只需要进行一次前向传播进行验证）。如果草稿模型准确预测了 α 词，原始 LLM 的一次前向传播就可以生成 α+1 个词。

在投机采样中，草稿模型的任务是基于当前词序列预测下一个词。使用一个参数数量显著更少的模型完成这个任务极具挑战性，通常会产生次优结果。此外，标准投机采样方法中的草稿模型独立预测下一个词而不利用原始 LLM 提取的丰富语义信息，导致潜在的效率低下。

这个局限启发了 EAGLE 的开发。EAGLE 利用原始 LLM 提取的上下文特征（即模型第二顶层输出的特征向量）。EAGLE 建立在以下第一性原理之上：

特征向量序列是可压缩的，所以根据前面的特征向量预测后续特征向量比较容易。

EAGLE 训练了一个轻量级插件，称为自回归头（Auto-regression Head），与词嵌入层一起，基于当前特征序列从原始模型的第二顶层预测下一个特征。然后使用原始 LLM 的冻结分类头来预测下一个词。特征比词序列包含更多信息，使得回归特征的任务比预测词的任务简单得多。总之，EAGLE 在特征层面上进行外推，使用一个小型自回归头，然后利用冻结的分类头生成预测的词序列。与投机采样、Medusa 和 Lookahead 等类似的工作一致，EAGLE 关注的是每次提示推理的延迟，而不是整体系统吞吐量。

EAGLE——一种增强大语言模型生成效率的方法

上图显示了 EAGLE 与标准投机采样、Medusa 以及 Lookahead 关于输入输出的区别。下图展示了 EAGLE 的工作流程。在原始 LLM 的前向过程中，EAGLE 从第二顶层收集特征。自回归头以这些特征以及此前生成的词的词嵌入作为输入，开始猜下一个词。随后，使用冻结的分类头（LM Head）确定下一个词的分布，使 EAGLE 能够从这个分布中进行采样。通过多次重复采样，EAGLE 进行了类似树状的生成过程，如下图右侧所示。在这个例子中，EAGLE 的三次前向传播 “猜” 出了 10 个词组成的树。

EAGLE 使用轻量级的自回归头来预测原始 LLM 的特征。为了确保生成文本分布的一致性，EAGLE 随后验证预测的树状结构。这个验证过程可以使用一次前向传播完成。通过这个预测和验证的循环，EAGLE 能够快速生成文本词。

训练自回归头代价很小。EAGLE 使用 ShareGPT 数据集进行训练，该数据集包含不到 70,000 轮对话。自回归头的可训练参数数量也很少。如上图中的蓝色部分所示，大多数组件都是冻结的。唯一要额外训练的是自回归头，这是一个单层 Transformer 结构，具有 0.24B-0.99B 参数。即使是 GPU 资源不足的情况下，也可以训练自回归头。例如，Vicuna 33B 的自回归头可以在 8 卡 RTX 3090 服务器上在 24 小时内完成训练。

为什么使用词嵌入来预测特征？

Medusa 仅使用第二顶层的特征来预测下一个词，下下个词......与 Medusa 不同，EAGLE 还动态地将当前采样得到的词嵌入作为自回归头输入的一部分来进行预测。这额外的信息帮助 EAGLE 处理抽样过程中不可避免的随机性。考虑下图中的例子，假设提示词是 “I”。LLM 给出了 “I” 后面跟着 “am” 或 “always” 的概率。Medusa 不考虑是抽样了 “am” 还是 “always”，直接预测 “I” 下下个词的概率。因此，Medusa 的目标是，在只给定 “I” 的基础上，预测 “I am” 或 “I always” 的下一个词。由于抽样过程的随机性，Medusa 的相同输入 “I” 可能有不同的下下个词输出 “ready” 或 “begin”，导致输入和输出之间缺乏一致的映射。相比之下，EAGLE 的输入包括了抽样结果的词嵌入，确保了输入和输出之间的一致映射。这种区别使 EAGLE 能够考虑抽样过程建立的上下文，进而更准确地预测后续词。

树状生成结构

与投机采样、Lookahead 和 Medusa 等其他猜测 - 验证框架不同，EAGLE 在 “猜词” 阶段采用类似树状的生成结构，进而实现了更高的解码效率。如图所示，标准投机采样和 Lookahead 的生成过程是线性或链式的。Medusa 的方法由于在猜测阶段无法构建上下文，故通过笛卡尔积生成树，导致相邻层之间形成全连接图。这种方法经常导致无意义的组合，例如 “I am begin”。对比之下，EAGLE 创建了一个更稀疏的树结构。这种稀疏的树结构防止形成无意义的序列，将计算资源集中在更合理的词组合上。

多轮投机采样

标准投机采样方法在进行 “猜词” 的过程中保持了分布的一致性。为了适应树状猜词场景，EAGLE 将这种方法扩展成了多轮递归形式。下面呈现了多轮投机采样的伪代码。在树状生成过程中，EAGLE 记录了每个抽样词对应的概率。通过多轮投机采样，EAGLE 确保最终生成的每个词的分布与原始 LLM 的分布保持一致。

更多实验结果

下图展示了 EAGLE 在 Vicuna 33B 上关于不同任务中的加速效果。涉及大量固定模板的 “编程”（coding）任务显示出最佳的加速性能。

欢迎大家体验 EAGLE，并通过 GitHub issue 反馈建议：https://github.com/SafeAILab/EAGLE/issues

工程EAGLE

相关数据

词嵌入技术

词嵌入是自然语言处理（NLP）中语言模型与表征学习技术的统称。概念上而言，它是指把一个维数为所有词的数量的高维空间嵌入到一个维数低得多的连续向量空间中，每个单词或词组被映射为实数域上的向量。

来源：维基百科

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

伪代码技术

伪代码，又称为虚拟代码，是高层次描述算法的一种方法。它不是一种现实存在的编程语言；它可能综合使用多种编程语言的语法、保留字，甚至会用到自然语言。它以编程语言的书写形式指明算法的职能。相比于程序语言它更类似自然语言。它是半形式化、不标准的语言。

来源：维基百科

映射技术

映射指的是具有某种特殊结构的函数，或泛指类函数思想的范畴论中的态射。逻辑和图论中也有一些不太常规的用法。其数学定义为：两个非空集合A与B间存在着对应关系f，而且对于A中的每一个元素x，B中总有有唯一的一个元素y与它对应，就这种对应为从A到B的映射，记作f：A→B。其中，y称为元素x在映射f下的象，记作：y=f(x)。x称为y关于映射f的原象*。*集合A中所有元素的象的集合称为映射f的值域，记作f(A)。同样的，在机器学习中，映射就是输入与输出之间的对应关系。

来源：Wikipedia

生成树技术

在图论的数学领域中，如果连通图 G的一个子图是一棵包含G 的所有顶点的树，则该子图称为G的生成树(SpanningTree)。生成树是连通图的包含图中的所有顶点的极小连通子图。图的生成树不惟一。从不同的顶点出发进行遍历，可以得到不同的生成树。

来源：百度百科

文本生成技术

文本生成是生成文本的任务，其目的是使人类书写文本难以区分。

来源：paperswithcode

语言模型技术

统计式的语言模型是借由一个几率分布，而指派几率给字词所组成的字串。语言模型经常使用在许多自然语言处理方面的应用，如语音识别，机器翻译，词性标注，句法分析和资讯检索。

来源：维基百科

量化技术

深度学习中的量化是指，用低位宽数字的神经网络近似使用了浮点数的神经网络的过程。

来源：Medium

北京大学机构

北京大学创办于1898年，初名京师大学堂，是中国第一所国立综合性大学，也是当时中国最高教育行政机关。辛亥革命后，于1912年改为现名。2000年4月3日，北京大学与原北京医科大学合并，组建了新的北京大学。原北京医科大学的前身是国立北京医学专门学校，创建于1912年10月26日。20世纪三、四十年代，学校一度名为北平大学医学院，并于1946年7月并入北京大学。1952年在全国高校院系调整中，北京大学医学院脱离北京大学，独立为北京医学院。1985年更名为北京医科大学，1996年成为国家首批“211工程”重点支持的医科大学。两校合并进一步拓宽了北京大学的学科结构，为促进医学与人文社会科学及理科的结合，改革医学教育奠定了基础。

官网，http://www.pku.edu.cn/