如果你曾经盯着闪烁的光标,看着 ChatGPT 或 LLaMA 一个字一个字地写出回复,你就体验过大语言模型 (LLM) 固有的瓶颈。这种迟缓源于自回归生成 (autoregressive generation) : 模型必须先生成 Token A,才能生成 Token B,然后是 Token C。这是一个严格的串行过程,导致那些擅长并行计算的现代 GPU 在很大程度上未被充分利用。

为了解决这个问题,研究人员开发了投机采样 (Speculative Sampling) , 这是一种允许模型低成本地“起草”多个未来 Token 并行进行验证的技术。然而,目前大多数投机方法都依赖于僵化、静态的结构。它们假设无论上下文如何,预测下一个词的难度总是相同的。

这就引出了 EAGLE-2 , 这是一篇挑战这一假设的新研究论文。通过使起草过程具有动态性和上下文感知能力 , EAGLE-2 实现了比标准解码快 3 到 4 倍的推理速度,超越了以往的最先进方法,同时在数学上保持了无损。

在这篇文章中,我们将剖析 EAGLE-2 的工作原理,解释为何静态草稿树效率低下,以及动态树如何在无需重新训练基础模型的情况下,大幅提升 LLM 的性能。

瓶颈所在: 为什么 LLM 推理很慢?

在深入了解 EAGLE-2 之前,我们需要理解它所解决的问题。LLM 推理是受限于内存带宽的 (memory-bound) 。生成单个 Token 需要将数十亿个参数从 GPU 的高带宽内存 (HBM) 移动到计算单元。因为 LLM 一次只生成一个 Token,所以每生成一个词,我们都要付出这种内存访问的“税”。

解决方案: 投机采样

投机采样 (或投机解码) 利用“起草与验证”的方法减少了这种开销:

  1. 起草 (Draft) : 一个更小、更快的模型 (“草稿模型”) 快速猜测接下来的 \(K\) 个 Token (例如,接下来的 5 个词) 。
  2. 验证 (Verify) : 巨大的“目标模型” (原始 LLM) 在一次前向传播中处理所有起草的 Token。

因为目标模型并行验证 5 个 Token 的速度几乎与生成 1 个 Token 一样快,所以只要草稿是正确的,我们就能节省时间。

演变: 从链到树

经典的投机采样起草的是单一的 Token 链条 (Token A \(\rightarrow\) Token B \(\rightarrow\) Token C) 。如果 Token A 错了,整条链都会被丢弃。

较新的方法,如 EAGLE (EAGLE-2 的前身) , 使用的是树状结构 。 草稿模型不再只猜测一个序列,而是探索多个分支。如果草稿模型不确定下一个词是 “is” 还是 “was”,它可以为两者都起草分支。在验证过程中,目标模型使用特定的注意力掩码 (树状注意力) 一次性检查整棵树。

标准投机采样与 EAGLE 的对比。

图 3 所示,标准投机采样 (左) 是线性工作的。EAGLE (右) 在特征层级上运作——传递特征向量而不仅仅是 Token——并验证一个树状结构。这带来了更高的接受率,因为如果一个分支失败,另一个分支可能会成功。

静态树的问题

虽然 EAGLE 代表了一个巨大的飞跃,但它有一个隐藏的低效性: 它以相同的方式处理每一个句子。

EAGLE 以及其他像 Medusa 这样的方法,使用的是静态草稿树 。 这意味着预测树的形状是固定的。例如,它可能总是为第一个 Token 猜测 2 个选项,第二个猜测 2 个,第三个猜测 1 个。这种结构假设猜测下一个 Token 的难度仅取决于它的位置 (即紧邻的下一个 Token 比 3 步之外的 Token 更容易猜) 。

洞察: 上下文至关重要

EAGLE-2 背后的研究人员发现这一假设是有缺陷的。草稿 Token 的“可接受性”在很大程度上取决于上下文 , 而不仅仅是位置。

考虑两种情况:

  1. 语境 A: “The capital of France is…” (下一个 Token 几乎肯定是 “Paris”) 。
  2. 语境 B: “The best strategy for this game is…” (下一个 Token 可能是任何东西) 。

在语境 A 中,我们应该采用深而窄的树,因为我们很自信。在语境 B 中,我们应该使用宽而浅的树来覆盖各种可能性。静态树强迫做出妥协,这对两者来说都不是最优的。

EAGLE 和 EAGLE-2 的区别。

图 4 完美地展示了这种直觉。

  • 左侧 (EAGLE): 树的形状是固定的。即使查询是 10 + 2 =,答案显然是 1,静态树仍然浪费资源去猜测一个极不可能的替代项 (3) 。
  • 右侧 (EAGLE-2): 模型识别出了上下文。对于 10 + 2,它扩展得很宽,因为操作数可能是任何数字。但对于 10 + 2 =,它识别出了确定性,并创建了一个包含 1, 2 的深度单分支。

证据: 依赖上下文的接受率

为了从数学上证明这一点,作者分析了 Token 的接受率。

不同位置 Token 的接受率。

图 5 (左) 展示了一个标准的静态树结构。 图 5 (右) 绘制了接受率。虽然平均而言较早的 Token (P1) 比靠后的 Token (P6) 更容易被接受,但点的垂直分布说明了另一个故事。每个位置上的巨大方差意味着,对于某些查询,位置 4 非常容易猜,而对于其他查询,位置 1 却很难。

这种差异证实了固定结构会导致性能损失。

方法: EAGLE-2 的动态架构

那么,我们如何构建一棵能够即时改变形状的树呢?我们需要一个信号来告诉我们何时变“宽”,何时变“深”。

关键在于,我们不能向目标 LLM 询问这个信号,因为运行目标 LLM 正是我们试图避免的事情。我们需要草稿模型进行自我评估。

第一步: 校准 (置信度信号)

研究人员发现,EAGLE 中的草稿模型具有非常好的校准性 。 这意味着模型给出的置信度分数 (它分配给某个 Token 的概率) 是该 Token 实际上被目标 LLM 接受的非常准确的预测指标。

草稿模型不同置信度区间的平均接受率。

图 6 所示,草稿模型的置信度 (x 轴) 与实际接受率 (y 轴) 之间存在近似线性关系。如果草稿模型说“我有 90% 的把握这个 Token 是对的”,那么它大概率有 90% 的时间会被接受。这使得 EAGLE-2 可以使用置信度分数作为接受率的代理。

第二步: 扩展阶段

EAGLE-2 逐层构建草稿树。在标准方法中,扩展是固定的 (例如,总是扩展前 2 个节点) 。在 EAGLE-2 中,我们根据一个称为价值 (Value, \(V_i\)) 的指标来选择扩展哪些节点。

一个 Token 的全局接受概率不仅仅关于该 Token 本身;它依赖于通向它的整条路径都是正确的。如果父节点错了,子节点就无关紧要了。因此,节点 \(V_i\) 的价值是其路径上所有置信度分数的乘积:

价值公式

这里,\(c_j\) 是草稿模型对节点 \(j\) 的置信度分数。

扩展阶段 , EAGLE-2 计算当前层所有节点的这个价值,并选择价值最高的 top-k 个节点进行进一步扩展。这确保了计算资源被用于延伸最有希望的分支,而不是死守僵化的几何形状。

第三步: 重排序阶段

一旦树完全扩展,我们就面临一个问题。我们需要发送固定数量的候选 Token 给目标 LLM 进行验证。

简单地从最终扩展层选取节点是有风险的。树深处 (深度 5) 的一个节点,其累积价值可能低于深度 2 处我们选择不扩展的一个节点。

  • 节点 A (深度 2) : 80% 置信度。
  • 节点 B (深度 5) : 40% 置信度 (累积) 。

为了最大化被接受 Token 的数量,EAGLE-2 执行重排序阶段 。 它会查看树中生成的每一个节点 (不仅仅是叶子节点) ,按其价值 (\(V_i\)) 对它们进行排序,并挑选前 \(m\) 个候选者。这创建了一个灵活的草稿,其中可能包含几条深度链和若干浅层替代项。

综合流程

让我们用图 7 来可视化整个流程。这是 EAGLE-2 算法的核心。

EAGLE-2 流程示意图。

  1. 起草 (顶部) : 模型从 “It” 开始。
  • 它预测 “is” (高置信度) 和 “has” (低置信度) 。
  • 它计算下一层的价值。“is \(\rightarrow\) a” 的价值为 \(0.6 \times 0.8 = 0.48\)。
  • 它扩展价值最高的节点。注意树是如何不对称地生长的。“is” 分支延伸得更深,因为置信度更高。
  1. 重排序 (中间) : 我们生成了许多节点。现在我们按价值对所有节点 (蓝色块) 进行排序。
  2. 扁平化 (底部) : 选中的节点被扁平化为一维序列 (例如 [It, is, has, a, the, to...]) 以输入 GPU。
  3. 注意力掩码: 因为这是一棵被扁平化为列表的树,我们必须确保 “has” 不会关注到 “is”——它们属于不同的可能性 (分支) 。注意力掩码确保 Token 只能“看到”它们的直系祖先节点。

实验与结果

这种动态灵活性是否转化为现实世界的速度?作者在多个数据集 (MT-bench, HumanEval, GSM8K) 上使用 Vicuna 和 LLaMA 模型测试了 EAGLE-2。

加速比分析

结果显示,EAGLE-2 始终优于标准投机采样 (SpS) 和原始 EAGLE (EAGLE-1) 。

不同方法在温度 1 下的加速比。

图 1 展示了在温度=1 (创造性生成的常用设置) 下的加速比。

  • Vicuna 13B: EAGLE-2 实现了 3.80x 的加速,相比之下 EAGLE 为 2.32x,标准投机采样为 1.62x。
  • LLaMA2-Chat 13B: EAGLE-2 达到了接近 4x 的加速 (3.92x) 。

在温度=0 (贪婪解码) 时,差距进一步扩大,如图 2 所示。

不同方法在温度 0 下的加速比。

在这里,EAGLE-2 始终处于领先地位。在 Vicuna 13B 上,它达到了 4.26x 的加速。相对于基线自回归解码 (1.0x) ,这是一个巨大的提升。值得注意的是,它比 EAGLE-1 快了大约 20% 到 40%

为什么它更快?

加速是由平均接受长度 (Average Acceptance Length, \(\tau\)) 驱动的。该指标跟踪每次验证步骤平均有多少个起草的 Token 被接受。

如下面的表 1 所示,EAGLE-2 始终拥有最高的 \(\tau\)。对于 MT-bench 上的 Vicuna 13B,EAGLE-2 的接受长度为 4.83 , 而 EAGLE 为 3.98。这意味着对于庞大的目标 LLM 的每一次前向传播,EAGLE-2 有效生成了近 5 个可用 Token。

表 1: 不同方法的加速比和平均接受长度。

消融实验: 价值和重排序重要吗?

作者进行了消融实验 (移除功能以观察其影响) 来验证他们的设计选择。

表 3: 消融实验结果。

表 3 显示:

  • w/o value (无价值) : 如果仅根据局部置信度 (而非累积价值) 进行扩展,性能从 3.62x 降至 3.21x。
  • w/o reranking (无重排序) : 如果跳过全局重排序步骤,性能降至 3.48x。
  • w/o both (两者均无) : 如果两者都移除,实际上就回退到了接近基线的性能。

这证实了上下文感知扩展 (使用价值) 和全局选择 (重排序) 对于最大化效率都是至关重要的。

结论与启示

EAGLE-2 代表了 LLM 加速领域的一次重大改进。通过放弃静态草稿树的僵化约束,它使推理过程与并非所有预测都是生而平等这一语言现实相一致。

以下是主要收获:

  1. 无损加速: EAGLE-2 不进行近似。输出分布保证与原始 LLM 完全相同。
  2. 上下文感知: 它根据模型在当前上下文中的置信度,动态分配计算资源 (树的深度/宽度) 。
  3. 开箱即用: 如果你已经有为 EAGLE 训练的草稿模型,你可以无需任何训练直接切换到 EAGLE-2。这纯粹是草稿树构建和验证方式上的算法改变。

随着 LLM 体积的不断增长,推理延迟仍然是一个关键障碍。像 EAGLE-2 这样的技术,通过优化生成过程而不仅仅是压缩模型,对于使高级 AI 变得易用和响应迅速至关重要。通过将推理视为一个动态的决策过程,而不是一个静态循环,我们终于能够打破自回归生成的速度限制。