大语言模型（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