caijizhuo

投机解码发展

发表于 更新于 分类于
本文字数： 3.1k 阅读时长 ≈ 11 分钟

投机解码与 EAGLE：步骤、发展历史与三代差异

一句话原理：用便宜的小模型先猜几个 token，再让大模型用一次前向把它们全部验完。 猜对了就白赚，猜错了就从出错的地方重新采样 —— 最终输出的分布与大模型自己逐字生成完全一致。

一、发展历史

时间 工作 一句话贡献
2023.01 Leviathan 等，Fast Inference from Transformers via Speculative Decoding 提出「草稿-验证」框架，并证明它无损
2023.02 Chen 等，Accelerating LLM Decoding with Speculative Sampling 独立提出同一思路，给出拒绝采样算法
2024.01 EAGLE 草稿从「token 级」改成「特征级 + 超前一步 token」，模型小 2~3 个数量级
2024.01 Medusa 不训独立草稿模型，直接在骨干上挂多个解码头
2024.02 Lookahead Jacobi 迭代 + n-gram 池，完全免训练
2024.06 EAGLE-2 草稿树从静态改成动态，按置信度分配节点预算
2025.03 EAGLE-3 training-time test + 多层特征融合，去掉特征预测约束

主线是 EAGLE 这一条，演进逻辑非常好记：

  1. EAGLE 解决「草稿用什么表示」—— token 换成特征；
  2. EAGLE-2 解决「节点预算怎么花」—— 静态树换成动态树；
  3. EAGLE-3 解决「训练和推理不一致」—— 单步训练换成 training-time test。

二、基础三步骤

投机解码的草稿、验证、接受三步流程

  1. 草稿（Draft）：小模型 Mq 自回归猜 4 个候选 token。一次只出一个，但模型小，4 次前向也不心疼。
  2. 验证（Verify）：把「上文 + 4 个候选」整段喂给大模型 Mp，一次前向并行算出每个位置的真实概率分布。原本要 4 次串行前向，现在 1 次。
  3. 接受（Accept）：从左往右逐个判定，遇到第一个不一致就截断；再从大模型的分布里补采 1 个 token（bonus token）。

以「中国的首都是」为例：草稿给出「北京 / 是 / 中国 / 的」，大模型验证后接受「北京」「是」，拒绝「中国」，并从残差分布补出「一座」—— 一轮产出 3 个 token，只花了 1 次大模型前向。

为什么这是无损的

判定只有两条规则：

  • 大模型比小模型更想要它（p ≥ q）→ 直接收下；
  • 小模型高估了它（p < q）→ 以 p/q 的概率收下，否则拒绝。

拒绝之后不是重来，而是从「残差分布」采样：把小模型多给出去的那部分概率扣掉，补给大模型真正想要的 token。一进一出刚好抵消，所以每个位置输出各 token 的概率，与大模型自己生成的分布完全相同。

注意：无损指的是分布一致，不是逐位一致，两次运行不会输出完全相同的结果。

什么时候能加速

一轮能多产出多少 token，只看两个量：

  • 接受率 α：草稿猜对的比例，越高越好；
  • 草稿相对开销 c：草稿一次前向比大模型一次前向便宜多少，越小越好。

草稿 token 数不是越多越好：多猜一个，能被全部接受的概率按 α 连乘迅速衰减，而草稿开销线性增长，所以存在一个最优长度。直觉上的临界条件是 α > c —— 接受率还不如草稿的便宜程度，就没得赚。

这也解释了为什么朴素的「7B 草稿 + 13B 目标」几乎零提升：草稿不够便宜，接受率也不够高。

三、EAGLE 的关键改动：从猜 token 到猜特征

朴素的草稿模型要单独训一个，还得让它和目标模型分布接近，很难同时做到「小」和「准」。EAGLE 的观察是：在特征层面做自回归，比在 token 层面更容易。 token 经过 softmax 之后信息已经被「压硬」了，而连续特征保留了更丰富的语义。

但特征级自回归有个麻烦：

位置 t+1 的特征长什么样，取决于位置 t+1 的那个 token 是什么 —— 而这个 token 还没定。只拿 f(t) 去猜 f(t+1)，是一对多的病态映射，学不准。

EAGLE 的解法很直接：把「超前一步」的 token 也一起喂进去。输入是「上一位特征的向量」加上「已经确定的真 token 的 embedding」，输出下一个特征和目标位置的 token 分布。条件一给，特征基本就被钉住了。

这个改动效果很大。论文消融实验（Vicuna-7B，MT-bench）：

草稿方式 加速比
token 级自回归 1.5×
特征级自回归 1.9×
特征级 + 超前一步 token 2.8×

草稿模型有多小

EAGLE 草稿模型结构

  • Embedding 和 LM Head 直接复用目标模型的，冻结不动，只训练中间一个 FC + 一层 decoder。
  • 特征取自 LM Head 的输入（最后一层 Transformer 的输出，代码里是 hidden_states[-1]）。论文叫它 second-to-top-layer feature，是把 LM Head 当作「顶层」来数的，别误读成「倒数第二层」。
  • 可训练参数量约 0.24B（7B 目标）到 0.99B（70B 目标）；70B 的草稿头在 4×A100 上训 1~2 天。

为什么要树，以及树怎么一次验完

只猜一条链的话，中间错一个后面全废。EAGLE 让每个节点保留 top-k 个分支，长成一棵树。

草稿树与树形注意力掩码

难点是：这棵树怎么用一次前向验证？ 答案是树形注意力掩码 —— 把树展平成一条序列、给每个节点编号，然后规定每个节点只能看见自己 + 自己的所有祖先（看不见兄弟节点）。这样一次前向就能同时算出所有候选路径的概率。

EAGLE-1 默认 25 个节点、深度 5。因为同层节点互不依赖、可以并成一个 batch，这棵树只花 5 次草稿前向 + 1 次目标前向。

容易搞错的一点：网上流传的「EAGLE 用 KV Cache 重排让祖先相邻」并不在论文里。论文只有树形注意力掩码；真正的 KV 处理发生在接受之后（把被接受路径对应的 KV 行拷到一起、丢弃其余分支），属于代码实现细节。

四、一轮完整流程与伪代码

一轮 EAGLE 解码的完整流程

五个阶段循环执行：取特征 → 草稿建树 → 一次前向验证 → 沿树接受 → 更新 KV。

下面用伪代码把整条链路串起来（去掉工程细节，只保留主干）。

1
2
3
4
5
6
# ============ 对照：基线自回归解码 ============
def baseline_decode(prompt, Mp, n_tokens):
    for _ in range(n_tokens):
        p = Mp.forward(prompt)        # 大模型一次前向
        prompt.append(sample(p))      # 一次只能产出 1 个 token
    return prompt                     # n 个 token = n 次串行前向
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# ============ 朴素投机解码：单链草稿 ============
def spec_decode(prefix, Mp, Mq, gamma=4):
    # 1) 草稿：小模型自回归猜 gamma 个 token
    draft, q_prob = [], []
    for _ in range(gamma):
        q = Mq.forward(prefix + draft)
        tok = sample(q)
        draft.append(tok)
        q_prob.append(q[tok])

    # 2) 验证：大模型一次前向，算出 gamma+1 个位置的真实分布
    p_prob = Mp.forward_multi(prefix + draft)

    # 3) 接受：从左往右，第一个不通过就截断
    out = []
    for i, tok in enumerate(draft):
        if uniform() < min(1.0, p_prob[i][tok] / q_prob[i]):
            out.append(tok)                                     # 收下
        else:
            out.append(sample(residual(p_prob[i], q_prob[i])))  # 拒绝 -> 残差重采样
            return out
    out.append(sample(p_prob[gamma]))    # 全部接受 -> 白赚一个 bonus token
    return out
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# ============ EAGLE：特征级草稿 + 树 + 一次验证 ============
def eagle_round(prefix, Mp, Mq, depth=5, topk=10, budget=60):
    # 1) 取特征：大模型最后一次前向留下的隐状态（也就是 LM Head 的输入）
    f = Mp.feature_cache[prefix]

    # 2) 草稿建树：同一层的节点拼成一个 batch，所以 depth 次前向就能长 depth 层
    tree = Tree(root=prefix[-1])
    frontier = [Node(token=prefix[-1], feature=f, score=1.0)]
    for _ in range(depth):
        # 关键：输入 = (token embedding 并上 上一位的特征向量)
        f_hat, logits = Mq.forward(
            [(embed(n.token), n.feature) for n in frontier],
            mask=ancestor_mask(frontier),       # 只看自己 + 自己的祖先
        )
        for node, feat, dist in zip(frontier, f_hat, logits):
            for tok in topk_of(dist, topk):
                # 路径分数 约等于 整条路径被全部接受的概率
                tree.add_child(node, tok, feat, node.score * dist[tok])
        # 动态树（EAGLE-2）：全局挑分数最高的节点继续扩张，超预算就剪掉最低分
        frontier = tree.pick_top_nodes(count=topk, budget=budget)

    # 3) 验证：大模型一次前向，算完整棵树每个节点的真实分布
    p_all = Mp.forward(tree.flatten(), mask=tree.attention_mask())

    # 4) 沿树接受：找出最长可接受前缀
    accepted, node = [], tree.root
    while True:
        child = tree.best_child(node)
        if child is None:
            break
        p, q = p_all[child], child.draft_prob
        if uniform() < min(1.0, p[child.token] / q):
            accepted.append(child.token)              # 接受，继续往下走
            node = child
        else:
            accepted.append(sample(residual(p, q)))   # 拒绝 -> 残差重采样后结束
            break

    # 5) 收尾：补 1 个 bonus token，压缩双方的 KV Cache，回到第 1 步
    accepted.append(sample(p_all[node]))
    Mp.compact_kv(keep=node)
    Mq.compact_kv(keep=node)
    return accepted

五、三代之间的变化

EAGLE-1 静态树与 EAGLE-2 动态树对比

EAGLE-3 的两个改动

EAGLE（2024.01） EAGLE-2（2024.06） EAGLE-3（2025.03）
草稿输入 特征 + 超前一步 token 同左 同左，但改成多层特征融合
预测目标 特征 + token（两个损失） 同左 只预测 token
草稿树 静态：每层都铺 top-k 动态：按置信度全局分配节点 沿用动态树，深度 6 → 8
训练方式 单步 teacher forcing 同左 training-time test：训练时就跑多步
解决的瓶颈 草稿模型不够「准」 静态树把预算浪费在低概率分支 特征预测约束封死了「堆数据」的收益
加速比 2.78× ~ 3.07× 3.05× ~ 4.26× 最高 6.5×

EAGLE-2：换成动态树。 静态树隐含假设「接受率只取决于位置」，但实测发现同一位置上不同分支的方差很大 —— 接受率是上下文相关的。论文举的例子：10+2= 后面几乎必然是 1，再铺一堆候选纯属浪费；而 10+2 很难猜，这时才该多分配节点。

EAGLE-2 用草稿模型的置信度来估计接受率（实测校准得很好：置信度低于 0.05 的候选，真实接受率约 0.04；高于 0.95 的约 0.98），给每个节点算一个「整条路径被全部接受」的分数，然后在整棵树里按分数挑最高的节点去扩张，超出预算就全局剪枝。同样 60 个节点的预算，树能扎得更深，平均接受长度更长。

注意：EAGLE-2 没有改变单 token 的接受率，涨的是每轮能捞回来的 token 数。

EAGLE-3：换训练方式。 作者发现把训练数据翻倍，收益却很有限，根因是特征预测约束：损失里要求回归大模型的特征，这是额外约束，限制了草稿模型的表达能力；而且顶层特征本身只包含「下一个 token」的信息，用它预测「下下个 token」天然吃亏。于是 EAGLE-3 直接砍掉特征回归，改成纯 token 预测，同时把单一层特征换成低/中/高三层（代码里取第 2、L/2、L−3 层）拼接后压回原维度。

但只砍损失会出新问题：训练时输入永远是真特征，推理从第二步起输入是草稿自己猜的特征，两边分布对不上。training-time test 就是让训练模拟推理 —— 把自己的输出喂回去跑多步（代码里 7 步），每步都算损失。大模型只在最开始算一次特征，反复跑的只有那一层草稿模型，所以并不贵。

六、工程上唯一要记住的事

三代加速比对比与 batch size 权衡

投机解码主要优化延迟，不是吞吐。batch 一大，GPU 本来就吃饱了，多出来的草稿计算只会挤占资源：

  • vLLM 实测（GPT-OSS-20B）：并发 1 时 1.55×，并发 64 时只剩 1.05×；
  • SGLang 默认在 batch size ≥ 32 时直接关掉多步投机。

其余几条经验：temperature 越高收益越低；代码补全、摘要这类「输入高度决定输出」的任务收益最大；MoE 模型（如 Mixtral 8x7B）收益明显更小。

框架支持：vLLM 用 method: "eagle" 或 "eagle3"；SGLang 用 --speculative-algorithm EAGLE3；TensorRT-LLM 用 decoding_type: Eagle3（EAGLE v1/v2 的权重不兼容）；HuggingFace Transformers 目前没有 EAGLE。

参考资料

  • Leviathan 等，Fast Inference from Transformers via Speculative Decoding
  • Chen 等，Accelerating Large Language Model Decoding with Speculative Sampling
  • Li 等，EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty（ICML 2024）
  • Li 等，EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees（EMNLP 2024）
  • Li 等，EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test
  • 官方代码：https://github.com/SafeAILab/EAGLE