EAGLE-3:大模型推理加速新突破,最高提速6.5倍且延续Scaling Law

EAGLE-3解锁大模型投机采样的Scaling Law能力,推理速度最高提升6.5倍,且保证无损。已被集成到SGLang。

原文标题:大模型推理无损加速6.5倍!EAGLE-3碾压一切、延续Scaling Law能力

原文作者:机器之心

冷月清谈:

EAGLE 团队发布 EAGLE-3,一种通过训练时测试优化投机采样的大模型推理加速方法。该方法通过增加草稿模型的训练数据量,并结合特征预测损失的优化,解决了传统投机采样方法在 scaling up 时效果不佳的问题。EAGLE-3 不仅重用目标模型的最后一层特征,还混合了低、中、高层的信息作为草稿模型的输入,显著提升了推理速度和接受率。实验表明,EAGLE-3 在多项任务上均优于其他投机采样方法,加速比可达 3.1x-6.5x。此外,EAGLE-3 已被集成到 SGLang 框架中,即使在大 batch size 下也能有效提高吞吐量。

怜星夜思:

1、EAGLE-3 在哪些场景下可能无法达到 6.5 倍的加速效果?例如,数据集特性、模型大小等等因素会如何影响?
2、EAGLE-3 中提到的“训练时测试”是如何模拟多步生成的?它解决了什么具体问题?
3、EAGLE-3 混合了目标模型不同层的信息作为草稿模型的输入,为什么这样做比只用最后一层特征更好?

原文内容

自回归解码已然成为大语言模型的推理标准。大语言模型每次前向计算需要访问它全部的参数,但只能得到一个 token,导致其生成昂贵且缓慢。


近日,EAGLE 团队的新作《EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test》通过一系列优化解锁了投机采样的 Scaling Law 能力,可以将大语言模型的推理速度提高 6.5 倍,同时不改变大语言模型的输出分布,确保无损。同时,随着训练数据的增加,加速比还能进一步提高。


  • 论文标题:EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test

  • arXiv 地址:https://arxiv.org/abs/2503.01840

  • 项目地址:https://github.com/SafeAILab/EAGLE

  • SGLang 版本:https://github.com/sgl-project/sglang/pull/4247


EAGLE-3 的加速效果(DeepSeek-R1-Distill-LLaMA 8B 在数学推理数据集 GSM8K 上测试,其他模型在多轮对话数据集 MT-bench 上测试):


不同方法的生成速度对比:


背景


投机采样使用一个小的模型快速生成草稿,一次生成多个 token。目标大语言模型通过一次前向计算并行验证草稿的正确性,输出正确的草稿,并确保无损。EAGLE 系列是投机采样的最快实现。


EAGLE-1 在更有规律的特征层面而不是 token 层面进行自回归,同时输入采样结果(超前一个时间步的 token)消除了不确定性,明显提升了草稿模型的准确率。EAGLE-2 利用草稿模型的置信度近似接受率,据此动态地调整草稿树的结构,进一步提升了投机采样的效率。 


机器之心之前已经报道了 EAGLE-1 和 EAGLE-2 的工作:


  • EAGLE-1:

  • EAGLE-2:


最新的大模型通过使用越来越多的训练数据以取得更好的性能。比如说,对于 LLaMA 系列 7B(8B)大小的模型,LLaMA 1、LLaMA 2 和 LLaMA 3 分别使用了 1T、2T、15T token 训练数据,模型结构和推理成本基本不变的前提下各项指标取得了明显提升。


EAGLE-3 的作者们希望通过增加草稿模型的训练数据量以提高接受率和加速比(蓝色曲线)。遗憾的是, EAGLE-1 从训练数据增加中得到的提升非常有限(红色曲线)。


EAGLE-1 和 EAGLE-3 的加速比随着训练数据的增加而变化的趋势:


动机


EAGLE 在特征层进行自回归,再使用目标模型的分类头得到草稿 token。这种方式利用了目标模型的中间结果和信息,相比 token 层自回归有更好的性能。因为草稿模型的最终任务是预测草稿 token,EAGLE 的损失函数包括两部分,一部分是特征预测损失,另一部分是 token 预测损失。特征预测损失也可以被认为是一种约束,限制了模型的表达能力。


EAGLE-3 的作者们进行了实验验证,观察第一个草稿 token 的接受率 0-α,在数据集较小时,特征预测损失在训练数据较少时可以提高性能,但严重损害了草稿模型的 scaling up 能力。


不过,不使用特征预测损失会导致生成后续草稿 token 时的输入偏离训练分布,导致第二个草稿 token 的接受率 1-α 严重下降。为了解决这一问题,EAGLE-3 使用 “训练时测试” 模拟多步生成,从而兼顾了模型的 scaling up 能力和生成多个草稿 token 的能力


EAGLE、Medusa 等投机采样方法都重用目标模型的最后一层特征作为草稿模型的提示,但 EAGLE-3 的作者们发现这存在缺陷。大语言模型的最后一层特征经过线性变换就能得到下一个 token 的分布。对于满秩的分类头,最后一层特征和下一个 token 的分布一一对应。最后一层特征只有下一个 token 的信息,失去了目标模型的全局性质。


因此,EAGLE-3 不再使用目标模型的最后一层特征作为辅助信息,而是混合目标模型的低层、中层、高层信息来作为草稿模型的输入。


方法


与其他投机采样方法一致,EAGLE-3 交替执行草稿阶段和验证阶段。


草稿阶段以下图为例。在预填充或者上一个验证阶段,EAGLE-3 记录模型的低、中、高层特征序列,分别记为 l、m 和 h,拼接 k 维向量 l、m 和 h 得到 3k 维向量,然后通过一个全连接层将其降维到 k 维,就得到融合了不同层级信息的特征 g,其中 k 为目标模型的隐藏层维度。目标是生成以 “How can I” 为前缀的草稿 token 序列,只输入 g_how 和 g_can,草稿模型无法感知随机的采样过程。


因此 EAGLE-3 引入采样结果 I 的词向量嵌入 e_I。将 g 与对应的超前一个时间步的嵌入 e 拼接,即 g_how 与 e_can 拼接,g_can 和 e_I 拼接。将拼接后的向量通过全连接层降到 k 维,并将其输入进一个单层的 transformer 得到 a。然后,将 a 输入分类头,采样得到第一个草稿 token “do”。


在 step 1,以 “How can” 为前缀时,EAGLE-3 重用了来自目标模型的 g_how 和 g_can。


在 step 2,前缀是 “How can I”。理想的方式是重用来自目标模型的 g_how、g_can 和 g_I。但这是不可能的,因为 token “I” 还没有被目标模型检查,无法获取 g_I。EAGLE-3 用上一个 step 草稿模型的输出 a_I 替代 g_I,拼接 a_I 与采样结果 “do” 的词向量嵌入作为草稿模型 step2 的输入。


在 step 3,同样无法获得 g_do,所以使用 a_do 代替,拼接 a_do 与 e_it 作为草稿模型的输入。之后的 step 同理。


实验


EAGLE-3 在多轮对话、代码、数学推理、指令遵循、总结五项任务上分别使用 MT-bench、Humaneval、GSM8K、Alpaca、CNN/DM 数据集进行了实验,并与 7 种先进的投机采样方法(SpS、PLD、Medusa、Lookahead、Hydra、EAGLE、EAGLE-2)进行了比较。


该实验分别在 Vicuna (V), LLaMA-Instruct 3.1 (L31), LLaMA-Instruct 3.3 (L33), DeepSeek-R1-Distill-LLaMA (DSL) 上进行。


表格中的 Speedup 为加速比,τ 为平均接受长度,也就是目标模型每次前向计算能生成的 token 数。EAGLE-3 每次前向计算能生成大约 4-7 个 token,而自回归解码每次生成 1 个 token,因此 EAGLE-3 明显加速了大语言模型的生成,加速比为 3.1x-6.5x


在所有任务和模型上,EAGLE-3 的加速比和平均接受长度都是最高的,明显优于其他方法。


应用


EAGLE-3 发布第一天就被集成到 SGLang 中。在生产级框架中,EAGLE-3 也有数倍加速效果。以下实验由 SGLang 团队提供并以 LLaMA 3.1 8B(batch size=1, 1x H100)为例。


投机采样往往被认为在大 batch size 下会降低吞吐量。但是在 SGLang 这一生产级框架下,EAGLE-3 在 batch size 为 64 时仍可以提高 38% 的吞吐量,而 EAGLE 在 batch size 为 24 时就导致吞吐量下降。这里 1.00x 以 SGLang (w/o speculative) 的吞吐量作为基准。以下实验由 SGLang 团队提供并以 LLaMA 3.1 8B(1x H100)为例。 


作者介绍


李堉晖:北京大学智能学院硕士,滑铁卢大学访问学者,受张弘扬老师和张超老师指导,研究方向为大模型加速和对齐。


魏芳芸:微软亚研院研究员,研究方向为具身智能、图像生成和 AI agents。


张超:北京大学智能学院研究员,研究方向为计算机视觉和大模型加速。


张弘扬:滑铁卢大学计算机学院、向量研究院助理教授,研究方向为大模型推理加速、AI 安全和世界模型。



© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

这个“训练时测试”有点像是在训练的时候就模拟真实推理过程,让草稿模型知道自己生成的内容会被怎么使用。它主要解决的是草稿模型在生成多个 token 的时候,因为输入的分布和训练时不一样,导致后面的 token 错误率上升的问题。通过在训练时就让模型接触到这种“不确定性”,可以提高其泛化能力,从而提升整体的接受率和加速效果。

最后一层特征虽然包含了预测下一个 token 的信息,但它也丢失了很多全局信息,就像只看到了树叶,没看到树干和树根。混合不同层的特征可以让草稿模型同时获得局部和全局的信息,更好地理解上下文,从而做出更准确的预测。毕竟,大语言模型的能力很大程度上来自于其对全局信息的理解和建模。

从信息论的角度来看,大语言模型的不同层级特征捕捉了不同抽象程度的语言信息。最后一层特征虽然包含了预测下一个token所需的信息,但可能损失了来自 lower layers 的上下文信息。通过混合不同层级的特征,EAGLE-3 为草稿模型提供了更丰富的上下文信息,使其能够更好地理解输入,从而提高预测的准确性。这种做法借鉴了人类在理解语言时,同时考虑字面意义和语境信息的思维方式。

从理论上讲,EAGLE-3的加速效果与草稿模型和目标模型之间的差距有关。如果目标模型本身已经非常高效,或者数据集的特性使得草稿模型很难做出有效的预测,那么加速效果就会受到限制。此外,硬件的并行计算能力也是一个重要因素,如果硬件无法充分利用投机采样带来的并行性,那么加速效果也会受到影响。

我觉得哈,这个加速比肯定是在特定数据集和模型上测出来的。实际用的时候,如果你的数据分布和训练数据差太远,效果肯定会打折扣。还有就是模型大小,像那种几百亿参数的大模型,可能加速效果更明显,小模型就不一定了。毕竟投机采样也是有成本的,太小的模型可能提升不大。

EAGLE-3 的加速比依赖于草稿模型的准确率和接受率。如果数据集的复杂性很高,导致草稿模型难以生成准确的草稿 token,接受率会下降,加速效果自然会打折扣。此外,对于非常小的模型,由于目标模型前向计算本身很快,投机采样的额外开销可能会抵消一部分加速效果。最后,如果硬件资源有限,无法有效支持草稿模型的并行计算,也会影响加速比。

我的理解是,“训练时测试”相当于给草稿模型做了一个“压力测试”。在训练过程中,不仅仅是预测下一个 token,还会模拟预测后续的 token,并根据预测结果调整模型的参数。这样做的目的是让模型更好地适应真实推理时的多步生成环境,减少误差积累,提高整体的准确性。这就像是运动员在比赛前进行模拟训练,以更好地适应比赛的节奏和压力。

可以把大语言模型想象成一个金字塔,底层是基础的语言特征,高层是抽象的语义信息。最后一层可能只包含了最直接的预测信息,但忽略了底层的支撑。混合不同层的特征就像是看到了金字塔的整体结构,更全面,更不容易出错。这有点像我们在做判断的时候,既要看眼前的事实,也要考虑历史的经验和背景信息。

“训练时测试”方法通过在训练阶段模拟多步生成过程,缓解了因草稿模型在逐步生成 token 时输入分布偏移的问题。具体来说,它允许草稿模型在训练时就接触到由自身预测产生的输入,从而更好地适应真实推理场景中连续生成 token 的过程。这种方法有助于提高草稿模型的泛化能力,进而提升投机采样的效率和准确性。