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