投机解码与 EAGLE

投机解码与 EAGLE

本节摘要:前沿 LLM 生成一个 token 需要对数十亿参数做一次完整前向传播,而这个前向传播是严重过度供给的——大多数时候,一个小得多的模型能猜对接下来 35 个 token,大模型只需「验证」这个猜测。猜对了,你就用一个前向传播的代价拿到了 5 个 token。投机解码(Speculative Decoding,Leviathan 等 2023)用一套精确的接受/拒绝规则把这件事变得严格等价于从目标模型采样;EAGLE-3(2025)则把每次验证的接受 token 数推到约 4.5 个,实现 45 倍加速且输出分布不变。本节讲透双模型 setup、精确性规则、期望加速比、用蒸馏训练草稿模型,以及 EAGLE 的树形草稿与特征复用两大改进。

对应原课程:Phase 10 · Lesson 23 · 25-speculative-decoding(原英文 phases/10-llms-from-scratch/25-speculative-decoding/docs/en.md)。

学习目标

阅读完本节,你应当能够:

1. 说清自回归生成的串行瓶颈与投机解码发现的并发机会。
2. 复述双模型 setup(目标模型 M_p + 草稿模型 M_q)与每步四动作。
3. 写出精确性接受/拒绝规则,并理解它为何能保持目标模型的采样分布。
4. 用期望加速公式 E[tokens]=(1-α^(K+1))/(1-α) 估算加速比。
5. 解释蒸馏训练草稿模型的标准配方(KL 散度对齐目标分布)。
6. 说清 EAGLE 的两大改进(特征复用 + 树形草稿)为何能把接受率推高。

一、问题与直觉

70B 级模型在 H100 上的解码吞吐通常只有 40~80 token/秒。每个 token 都要读全部权重做一次完整前向传播——你既不能缩小模型(会改变输出),也不能无限加大 batch(受显存所限)。看起来卡死了。

但自回归生成 x_{t+1} = sample(p(·|x_{1:t})) 虽然看似串行,却藏着一个并发机会:如果你有一个便宜的预测器说「接下来 4 个 token 大概是 [a,b,c,d]」,你就能在大模型的一次前向传播里并行验证全部 5 个位置,并接受最长匹配前缀。

Leviathan、Kalai、Matias(2023)用一套精巧的接受/拒绝规则让这件事严格等价于从目标模型 p 采样——同样的输出分布,2~4 倍速度。

二、从零实现

双模型 setup 与每步四动作

• 目标模型 M_p:你真正想要样本的大、慢、高质量模型,分布 p(x)。
• 草稿模型 M_q:小、快、低质量模型,分布 q(x),比目标小 5~30 倍。

每步:

1. 草稿模型自回归地提出 K 个 token:x_1,...,x_K ~ q。
2. 目标模型对全部 K+1 个位置跑一次前向传播,得到每个位置的 p(x_k)。
3. 用下面的修正拒绝采样规则从左到右接受/拒绝每个 token,接受最长匹配前缀。
4. 若某 token 被拒绝,从修正分布采样替换并停止;否则从 p(·|x_1...x_K) 采样一个奖励 token。

草稿与目标完全一致 → 一次目标前向拿 K+1 个 token;草稿第 1 位就错 → 只拿 1 个。

精确性规则(核心)

for each drafted token x_t:
    r ~ Uniform(0, 1)
    if r < p(x_t) / q(x_t):
        accept x_t          # 接受
    else:
        # 从残差分布采样替换,然后停止
        replacement = sample_from( positive_part(p - q) )
        stop
​

其中 positive_part(p-q) 是逐点差的正部。当草稿与目标一致(p≈q)时接受率接近 1;不一致时,残差分布的构造保证了整体样本仍精确等于 p。

💡 贪婪情形简化:temperature=0 时,只需检查 argmax(p)==x_t,是则接受,否则输出 argmax(p) 并停止。

期望加速比

设草稿的 token 级接受率为 α,则每次目标前向的期望产出 token 数:

E[tokens] = (1 - α^(K+1)) / (1 - α)      # K=草稿长度, α∈[0,1]
​

α=0.8、K=4 时:(1-0.8^5)/(1-0.8) = 3.36 个 token/前向。若 cost_p >> cost_q*K(目标远贵于 K 步草稿),吞吐加速约 3.36 倍。**唯一真正的参数是 α,它完全取决于草稿与目标的对齐度——好的草稿决定一切。**

训练草稿:蒸馏

随机小模型做草稿效果很差。标准配方是从目标蒸馏:

1. 选一个小架构(70B 目标用约 1B,7B 目标用约 500M)。
2. 在大语料上跑目标模型,存其 next-token 分布。
3. 用 KL 散度对齐目标的分布(不是对齐 ground-truth token)来训练草稿。

结果:代码场景 α 通常 0.60.8,自然语言聊天 0.70.85,生产加速 2~3 倍。

EAGLE:树形草稿 + 特征复用

Li 等(2024, EAGLE)指出标准投机解码的两个低效:

1. 特征复用:草稿做 K 步串行全栈前向,但其实可以复用目标最近一次验证算出的隐藏状态——目标已经算出了丰富表征,草稿却在从头重算。
2. 树形草稿:草稿输出的是线性链。若改成输出一棵候选树(每节点多个猜测),目标的一次前向就能用树形注意力掩码并行验证多条候选路径,选最长接受分支。

EAGLE-3(2025)把这两点做到极致,每次验证接受约 4.5 个 token,实现 4~5 倍加速,且分布严格不变。

三、框架对比

生产中,vLLM、TensorRT-LLM 等推理引擎均已内置 EAGLE-3 级别的投机解码。

四、可复用产物

本节产出一份 outputs/skill-speculative-decoding.md——何时用投机解码的决策清单:目标模型大(>7B)、解码是延迟瓶颈、能提供/蒸馏出对齐草稿时启用;反之小模型或 prefill 受限时收益有限。

五、练习

1. 实现贪婪投机解码:目标与草稿都用 numpy 小 MLP,temperature=0,验证接受规则与加速比公式。
2. 测 α 的敏感性:固定 K=4,扫 α∈{0.5,0.7,0.85,0.95},画出期望 token/前向曲线,解释为何 α 从 0.7 提到 0.85 收益巨大而提到 0.95 收益递减。
3. 残差分布的正确性:用 numpy 验证「接受+残差替换」产出的样本分布与直接从 p 采样在统计上不可区分(跑 10 万次,KL 散度趋近 0)。
4. EAGLE 树形掩码:实现一个 2 层、每节点 3 个候选的树,用树形注意力掩码让目标一次前向验证全部路径。

本节要点回顾

1. 并发机会:自回归看似串行,但小模型可猜多 token、大模型一次前向并行验证。
2. 双模型:目标 M_p(大慢准)+ 草稿 M_q(小快糙,小 5~30 倍),每步四动作。
3. 精确性规则:接受条件 r < p(x_t)/q(x_t),拒绝时从残差 positive_part(p-q) 采样,整体严格等于 p。
4. 期望加速:E=(1-α^(K+1))/(1-α),α(对齐度)是唯一关键参数。
5. 草稿靠蒸馏:小架构 + KL 对齐目标分布,代码 α 0.60.8、聊天 0.70.85。
6. EAGLE 两改进:复用目标隐藏状态(省重算)+ 树形草稿(一次验证多路径),接受数推到 4.5,加速 45×。
7. 分布不变:无论标准还是 EAGLE,输出分布严格等于目标 p,这是它区别于「草稿模型直接用」的根本。

至此「从零构建 LLM」章的推理加速部分完成。下一节(梯度检查点)转向训练侧的显存优化——用重算换显存。