[论文解读] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

Yuhui Li, Fangyun Wei+2 more|arXiv (Cornell University)|Jan 26, 2024

Natural Language Processing Techniques被引用 5

一句话总结

EAGLE 引入了一个基于特征级自回归的草拟阶段，利用一步前瞻的令牌信息来应对推测采样中的不确定性，在不进行微调的情况下实现显著的延迟加速（MT-bench 上的 2.1x–3.8x），同时保留目标 LLM 的输出分布。

ABSTRACT

Autoregressive decoding makes the inference of Large Language Models (LLMs) time-consuming. In this paper, we reconsider speculative sampling and derive two key observations. Firstly, autoregression at the feature (second-to-top-layer) level is more straightforward than at the token level. Secondly, the inherent uncertainty in feature (second-to-top-layer) level autoregression constrains its performance. Based on these insights, we introduce EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency), a simple yet highly efficient speculative sampling framework. By incorporating a token sequence advanced by one time step, EAGLE effectively resolves the uncertainty, enabling precise second-to-top-layer feature prediction with minimal overhead. We conducted comprehensive evaluations of EAGLE, including all models from the Vicuna and LLaMA2-Chat series, the MoE model Mixtral 8x7B Instruct, and tasks in dialogue, code generation, mathematical reasoning, and instruction following. For LLaMA2-Chat 70B, EAGLE achieved a latency speedup ratio of 2.7x-3.5x, doubled throughput, while maintaining the distribution of the generated text.

研究动机与目标

• 推动大型语言模型（LLMs）在不改变输出分布的前提下实现更快的自回归解码。
• 通过聚焦特征级自回归和不确定性处理，重新审视推测采样。
• 提出一个轻量级草拟框架（EAGLE），通过一步令牌移位来解决特征级不确定性。

提出的方法

• 提出 EAGLE，一种在草拟阶段对次顶层特征而非令牌进行操作的推测采样框架。
• 引入一个向前推进一个时间步的令牌序列作为草拟模型的输入，以解决特征中的采样不确定性。
• 使用包含嵌入、LM 头、自回归头的三模块草拟模型，以及具树状注意力的草拟结构。
• 使用固定数据集和轻量级数据增强来缓解特征误差累积，结合回归与分类损失来训练自回归头。
• 实现一个验证阶段，使用带树注意力的目标 LLM 以维持输出分布，遵循递归推测采样过程。
• 提供低成本的训练方案（70B 模型在 A100 上 1–2 天）并显示与诸如 gpt-fast 等其他加速方法的兼容性。

实验结果

研究问题

• RQ1 在 LLM 的推测采样中，特征级自回归是否可以胜过令牌级自回归？
• RQ2 使用一步前瞻令牌输入来解决特征级不确定性如何影响准确性和加速？
• RQ3 在不同模型和任务中，EAGLE 的时延增益和保持分布的保真性有哪些？
• RQ4 EAGLE 如何与其他加速技术（如量化、编译）和批量大小策略相互作用？
• RQ5 在草拟阶段，EAGLE 对训练数据选择和特征噪声的鲁棒性如何？

主要发现

• EAGLE 在 Vicuna 与 LLaMA2-Chat 家族中实现显著的加速，例如在 LLaMA2-Chat 70B 上获得 2.7x–3.5x 的时延加速，同时吞吐量翻倍并保持输出分布。
• 在特征级草拟并使用一步前瞻令牌输入，相较于仅令牌草拟及现有方法（Lookahead、Medusa）提升了准确性和速度。
• 草拟/验证阶段中的树状注意力相对于链式草拟带来温和的额外提升，但不会改变最大前向传播次数，增加每次传播的令牌数量。
• 在 MT-bench 上，EAGLE 的表现优于 Lookahead（快 1.7x–2.1x）和 Medusa（快 1.5x–1.6x）；在某些设置下还能实现 2x–3x 的吞吐提升，并在贪婪和非贪婪模式下仍保留分布。
• EAGLE 与 gpt-fast 集成，使 LLaMA2-Chat 7B 的解码在 RTX 3090 上达到 160.4 tokens/s；训练成本低（自回归头参数少于 1B；70B 模型在 A100 上 1–2 天）。
• 消融研究显示通过特征与移位令牌输入来处理随机性的重要性，并且固定数据集训练的表现几乎与目标-LM 生成数据相当，强调了对数据敏感性较低。