https://sites.google.com/view/eagle-llm
https://github.com/SafeAILab/EAGLE
- 比原始解码快3倍（13B）。
- 比Lookahead（13B）快2倍。
- 比Medusa（13B）快1.6倍。
- 在生成文本的分布上与原始解码保持一致，可以证明。
- 在8个RTX 3090 GPU上可训练（1-2天）和测试，即使GPU性能较差的用户也可以承受。
- 可与其他并行技术结合使用，如vLLM、DeepSpeed、Mamba、FlashAttention、量化和硬件优化。
大模型（LLM）如ChatGPT展示了令人瞩目的能力，并在各个领域中得到越来越广泛的应用。然而，它们的文本生成过程既昂贵又缓慢。这种低效性归因于自回归解码的性质：每个token的生成都需要进行一次前向传递，需要访问LLM的整个参数集，这可能达到数百亿甚至数千亿个参数。这导致自回归解码受到内存限制。
加速自回归解码的一种方法是投机采样。这种技术使用一个较小的草稿模型通过标准的自回归生成猜测下一个γ个token。随后，原始LLM验证这些猜测的token，只需要进行一次前向传递来进行验证。如果草稿模型准确预测了α个token，原始LLM只需进行一次前向传递即可生成α+1个token。
在投机采样中，草稿模型和原始LLM共享相同的任务：根据当前token序列预测下一个token。用参数明显较小的模型完成这个任务非常具有挑战性，并且通常会产生次优结果。此外，草稿模型的独立预测没有利用原始LLM提取的丰富语义信息，导致潜在的低效性。
这一限制启发了EAGLE的开发。我们的方法利用原始LLM提取的上层输出中的上下文特征（即在预测下一个token时收集的特征），而无需进行额外计算。EAGLE基于以下第一原则构建：
特征序列是可压缩的，从先前的特征向量预测后续特征向量很容易。
我们训练了一个轻量级插件，称为自回归头，在原始模型的冻结嵌入层中，根据原始模型第二层输出中的当前特征序列预测下一个特征。然后使用原始LLM的冻结分类头将特征映射到token，从而得到token。与token序列相比，特征更加抽象和清晰，因此回归特征比回归token的任务要简单得多。总之，EAGLE使用一个小型自回归头在特征级别上进行外推，然后使用冻结分类头生成预测的token序列。与类似作品如投机采样、Medusa和Lookahead一致，我们关注的是每个提示推理的延迟，而不是整体系统吞吐量。
看下图这个对比很清晰
下图展示了EAGLE的工作流程。在原始LLM的前向过程中，我们从第二层收集特征。从这些特征和原始LLM生成的token开始，特征外推器开始其“猜测”过程。特征外推器以自回归的方式整合嵌入并生成下一个特征。随后，使用冻结的LM头确定token的分布，使我们能够从该分布中进行采样。通过多次重复采样，我们进行了一种类似树的生成过程，如图右侧所示。在本例中，特征外推器的三次前向传递“猜测”了一个包含10个token的树。
我们使用轻量级的FeatExtrapolator来预测原始LLM的特征，但这些特征可能不总是准确的（如图中红色框中的特征所示）。为了确保生成的文本分布的一致性，我们随后验证预测的树结构。由于因果LM的属性，这个验证过程可以在单次前向传递中完成，同时还生成一个token。利用这个生成的token和收集到的特征，FeatExtrapolator可以进行进一步的“猜测”。通过这个预测和验证的循环，LLM能够快速生成token。
训练自回归头部非常简单。对于所有模型，我们使用了ShareGPT数据集进行训练，该数据集包含少于70,000个对话轮次。FeatExtrapolator还具有最小数量的可训练参数。如上图中的蓝色部分所示，大部分组件都是冻结的。唯一的要求是“一个自回归头部”，它是一个单层Transformer解码器，具有0.24B-0.99B个参数。即使在GPU性能较差的设置下，训练自回归头部也是可行的。例如，Vicuna 33B的自回归头部可以在RTX 3090 GPU上在24小时内训练完成。
为什么我们使用token嵌入？
与Medusa不同，Medusa仅使用第二层的特征来预测各个偏移处的token（例如，下一个token，其后的token等）。而FeatExtrapolator在预测时还将下一个token的嵌入信息结合到其输入中。这个额外的信息有助于FeatExtrapolator处理采样过程中固有的随机性。
考虑下面图中的例子，查询是"I"。LLM输出了"I"后面可能是"am"或"always"的概率。Medusa在预测时不考虑是采样到了"am"还是"always"，直接预测了下一个token的概率。因此，Medusa的目标是"I am"或"I always"后面的token的概率。由于采样过程的随机性，Medusa相同的输入可能具有不同的目标，导致输入和输出之间缺乏一致的映射。
相比之下，EAGLE的输入包括采样结果的嵌入，确保了输入和输出之间的一致映射。这种区别使得FeatExtrapolator能够更准确地预测后续的token，并考虑到采样过程所建立的特定上下文。
树生成说明
与其他猜测验证框架如推测抽样、前瞻和Medusa不同，我们的方法在“猜测”阶段采用了树状生成，从而实现了更高的效率。如图所示，推测抽样和前瞻的生成过程是线性或链式的。Medusa的方法由于在猜测阶段无法构建上下文，通过笛卡尔积生成一棵树，导致相邻层之间形成完全连接的树。这种方法经常导致无意义的组合，比如“我是开始”。另一方面，EAGLE创建了一种更稀疏的树结构。稀疏的树结构更加选择性和上下文感知，防止了无意义序列的形成，并将计算资源集中在更合理的token组合上。
Multi-Round Speculative Sampling 推测抽样通常用于链式猜测过程，保持分布的一致性。为了适应树状猜测场景，我们将这种方法扩展为多轮递归形式。下面是多轮推测抽样的伪代码。在树生成过程中，我们记录与每个抽样token对应的概率。通过多轮推测抽样，我们确保最终生成的每个token的分布与原始LLM的分布一致。