在线大模型服务中,自回归解码的每一步都要把整个模型参数从高带宽内存搬到计算核心,生成一个 token 就需要一次完整的前向传播。当请求并发不高、batch size 较小时,矩阵乘退化为矩阵乘向量,GPU 算力大量闲置,推理延迟主要由访存带宽决定。投机解码(Speculative Decoding)试图用一个小模型快速猜出多个 token,再让大模型并行验证,从而把多次串行前向合并成一次。这个思路听起来直接有效,但真正落地时,加速比并不只由草稿模型的接受率决定。本文从拒绝采样的机制出发,分析影响加速比的各项因素,并讨论在低接受率或长序列下投机解码为何会失效。
自回归解码的瓶颈与投机解码的直觉
自回归解码的串行特性是推理延迟的根本来源。生成第 t 个 token 时,模型需要读取全部参数和已有的 KV Cache,计算量随着序列长度增长而增加,但每一步只能产出一个 token。在 batch size 为 1 的在线场景中,前向传播是典型的访存密集型操作,计算单元大部分时间在等待数据从 HBM 搬入片上缓存。增大 batch size 可以提升算术强度,但单个请求的端到端延迟并不会因此下降。
投机解码的出发点很简单:既然每一步都要搬运全部参数,不如一次搬运后同时验证多个候选 token。草稿模型(draft model)以自回归方式快速生成 K 个候选 token,目标模型(target model)把这 K 个 token 连同前缀一起作为输入,并行计算每个位置的概率分布,然后通过拒绝采样决定接受哪些 token。如果平均每个解码步骤能接受超过一个 token,总的解码步数就会减少,端到端延迟随之下降。
这里的关键是,拒绝采样必须保证输出分布与直接使用目标模型采样完全一致。也就是说,投机解码不能改变生成结果的概率分布,否则用户会观察到与原始模型不同的输出。Leviathan 等人 2022 年提出的投机采样算法(Fast Inference from Transformers via Speculative Decoding)证明了这一点,并报告在 T5-XXL 上相比标准 T5X 实现获得 2~3 倍加速,且输出分布不变。
拒绝采样:如何保证分布一致
拒绝采样的核心是一个概率修正步骤。假设草稿模型给出的候选 token 序列为 x₁, x₂, …, x_K,目标模型在给定前缀时对每个位置的真实分布为 p(x),草稿模型的分布为 q(x)。对于第 i 个候选 token,目标模型计算出 p(xᵢ),然后以概率 min(1, p(xᵢ)/q(xᵢ)) 接受它。如果接受,继续检查下一个 token;如果拒绝,则从修正后的分布中重新采样一个 token 作为当前位置的输出,并停止后续验证。
这个修正分布的形式是 max(0, p(x) − q(x)) 的归一化版本。直观上,当草稿模型的概率高于目标模型时,接受概率被压缩;当草稿模型的概率低于目标模型时,接受概率为 1,但拒绝时从差额中采样,补足草稿模型低估的部分。这一机制保证了最终采样分布与目标模型完全一致,无论草稿模型的质量如何。
拒绝采样的一个关键性质是,即使草稿模型很差,输出分布也不会偏离目标模型。代价是接受率可能很低,导致加速比下降甚至变慢。因此,接受率是衡量草稿模型质量的重要指标,但它不是决定加速比的唯一因素。
加速比的构成:接受率、草稿成本与验证开销
加速比可以分解为三个部分:草稿模型生成 K 个 token 的耗时、目标模型一次并行验证 K 个 token 的耗时、以及拒绝采样带来的额外计算。理想情况下,目标模型验证 K 个 token 的耗时与验证一个 token 相差无几,因为并行计算增加了计算量但不会显著增加访存时间。草稿模型的成本则取决于其规模和生成 K 个 token 所需的串行步数。
假设目标模型单步前向耗时 T_target,草稿模型单步耗时 T_draft,验证 K 个 token 的并行前向耗时约为 T_target(忽略计算量增加带来的少量延迟)。一个解码周期的总耗时为 K × T_draft + T_target。如果平均接受长度为 A(包括被拒绝后重新采样的那个 token),则每个周期实际生成的 token 数为 A,有效加速比为 A × T_target / (K × T_draft + T_target)。
从这个公式可以看出,即使接受率很高,如果草稿模型生成 K 个 token 的成本过高,加速比也会被稀释。例如,K=4、接受率 0.8 时,平均接受长度 A=3.2,但如果草稿模型单步耗时是目标模型的 1/10,草稿阶段耗时 0.4 × T_target,总周期耗时 1.4 × T_target,加速比约为 2.3。如果草稿模型更慢,比如 1/5,加速比下降到约 1.8。
验证批大小 K 也是一个需要权衡的参数。K 越大,草稿模型需要生成的候选越多,草稿阶段耗时线性增长,但目标模型验证的并行度更高。当 K 超过某个阈值时,验证前向的计算量增加可能开始影响延迟,因为矩阵乘的规模变大,访存时间不再完全主导。工程上通常需要针对具体模型和硬件测量不同 K 下的加速比曲线,找到最优值。
从草稿模型到多头架构:Medusa 与 EAGLE 的取舍
投机解码的早期实现依赖一个独立的草稿模型,这带来了额外的部署和维护成本。Medusa 提出在目标模型上直接添加多个解码头,每个头预测后续不同位置的 token,从而免去单独的草稿模型。Medusa 使用树状注意力机制构造多个候选序列,并在每个解码步骤并行验证。Medusa-1 在冻结主干模型的情况下微调解码头,可达到 2.2 倍加速;Medusa-2 与主干模型联合微调,加速比提升到 2.3~3.6 倍,但需要特殊的训练流程以保持主干模型的能力。
EAGLE 则从特征不确定性的角度出发,认为草稿模型应该基于目标模型的特征层进行预测,而不是直接从 token 分布采样。EAGLE 的草稿模型输入是目标模型倒数第二层的隐藏状态,这比基于 token 的预测更准确,因为隐藏状态包含更丰富的上下文信息。EAGLE 的接受率通常高于基于独立小模型的方案,但代价是需要访问目标模型的内部特征,这限制了它在黑盒场景下的应用。
下表对比了投机解码的几种实现方式,从草稿来源、训练成本、加速比和适用场景四个维度给出决策参考。
| 方案 | 草稿来源 | 训练成本 | 加速比(资料支持) | 适用场景 |
|---|---|---|---|---|
| 独立草稿模型 | 单独训练的小模型 | 需要额外训练和部署 | 2~3 倍(T5-XXL) | 可接受双模型部署,黑盒目标模型 |
| Medusa-1 | 目标模型上的附加解码头 | 仅微调解码头,主干冻结 | 2.2 倍 | 目标模型可微调,希望免去独立草稿模型 |
| Medusa-2 | 解码头与主干联合微调 | 需要特殊训练配方 | 2.3~3.6 倍 | 允许联合训练,追求更高加速比 |
| EAGLE | 基于目标模型特征层的草稿头 | 需要访问特征层 | 接受率更高(定性) | 可访问内部特征,追求高接受率 |
选择哪种方案取决于部署约束。如果目标模型是黑盒 API,只能使用独立草稿模型;如果拥有模型权重,Medusa 和 EAGLE 可以省去额外模型的部署开销,但需要微调能力。EAGLE 的高接受率通常意味着更高的加速比,但实现复杂度也更高。
一个贯穿场景:在线客服机器人的长回答生成
考虑一个在线客服机器人,用户询问“如何申请退款”,模型需要生成一段 200 token 的说明。在标准自回归解码下,这需要 200 次串行前向,每次前向都要搬运全部参数。假设目标模型单步前向耗时 10ms,总延迟约 2 秒。
使用投机解码,草稿模型先生成 K=8 个候选 token,目标模型并行验证。如果接受率为 0.7,平均每个周期接受 5.6 个 token,加上拒绝后重新采样的一个,实际生成约 6.6 个 token。草稿模型单步耗时假设为 1ms,草稿阶段耗时 8ms,验证耗时 10ms,总周期 18ms,生成 200 token 需要约 30 个周期,总延迟约 540ms,加速比约 3.7。
但如果接受率下降到 0.3,平均接受长度只有 2.4,总周期数增加到 83 个,总延迟约 1.5 秒,加速比降到 1.3。更糟的是,如果草稿模型生成 8 个 token 的耗时接近目标模型验证时间,加速比可能低于 1。因此,在低接受率场景下,投机解码可能比直接解码更慢。
这个例子说明,接受率、草稿成本和验证批大小共同决定加速比,而接受率本身又受草稿模型质量、生成任务的难度和序列长度影响。
数据流与状态变化:一个解码周期的完整流程
下图展示了一个投机解码周期的数据流,从草稿生成到拒绝采样,再到输出 token 的拼接。
flowchart TD
A[输入前缀] --> B[草稿模型自回归生成 K 个候选 token]
B --> C[目标模型并行验证 K 个位置]
C --> D{逐个检查候选 token}
D -->|接受| E[保留 token,继续下一个]
D -->|拒绝| F[从修正分布采样一个 token]
E --> G{是否还有候选}
G -->|是| D
G -->|否| H[输出所有接受的 token 和重新采样的 token]
F --> H
H --> I[更新 KV Cache,作为下一周期前缀]
关键转折点在于拒绝分支:当某个候选 token 被拒绝时,后续所有候选 token 都失效,因为前缀正确性假设被打破。此时需要从修正分布中采样一个 token 作为当前位置的输出,然后整个周期结束。这个被拒绝的位置成为下一周期的起点。
在实现中,目标模型验证时需要使用合适的 attention mask,确保每个位置只能看到其前缀,而不是整个候选序列。对于树状候选,attention mask 需要更精细地控制,这是 Medusa 等方案的核心实现细节。
工程调优:观察指标与失效边界
生产环境中,需要监控的指标包括:接受率、平均接受长度、草稿模型耗时、验证耗时、端到端延迟和吞吐。接受率可以通过日志统计每个周期接受的 token 数除以候选数得到。平均接受长度直接反映加速效果,是比接受率更直接的指标。
当接受率低于某个阈值时,投机解码可能退化。这个阈值取决于草稿模型与目标模型的耗时比。如果草稿模型非常快(如 1/20),即使接受率只有 0.2,加速比仍可能大于 1;如果草稿模型较慢(如 1/5),接受率需要超过 0.5 才能获得收益。工程上通常会在部署前测量不同接受率下的加速比曲线,确定可接受的最低接受率。
长序列是另一个失效边界。随着序列长度增长,KV Cache 增大,目标模型验证前向的访存时间增加,但草稿模型生成候选 token 的成本也随前缀长度增长。更关键的是,长序列后段往往更难预测,接受率可能下降。例如,在代码生成任务中,函数体内部的 token 可能较容易预测,但跨行逻辑或 API 调用处的 token 很难猜中。
另一个常见问题是分布偏移。草稿模型在训练数据上表现良好,但线上请求的分布可能与训练分布不同,导致接受率下降。Medusa 的 typical acceptance scheme 通过调整接受标准来提升接受率,但需要权衡生成质量。EAGLE 的特征层预测对分布偏移的鲁棒性更好,因为隐藏状态比 token 分布更稳定。
与替代方案的比较:投机解码 vs 其他加速方法
投机解码并非唯一的推理加速手段。动态批处理(dynamic batching)通过合并多个请求提高吞吐,但不降低单个请求的延迟。量化通过减少参数位数降低访存压力,但可能影响生成质量。早期退出(early exit)在模型内部提前停止计算,但需要修改模型结构。
投机解码的独特优势在于它不改变输出分布,且可以与其他方法叠加。例如,量化后的目标模型仍然可以使用投机解码,进一步降低访存时间。但投机解码增加了系统复杂度,需要维护草稿模型或附加头,并处理拒绝采样的额外逻辑。
在延迟敏感的场景(如在线客服)中,投机解码是降低端到端延迟的有效手段;在吞吐敏感的场景(如离线批处理)中,动态批处理可能更简单有效。选择哪种方法取决于业务对延迟和吞吐的优先级,以及是否有能力微调模型。
尚未解决的问题
投机解码的接受率仍然难以预测。草稿模型的训练目标与目标模型不完全一致,导致接受率在不同输入上波动很大。如何设计更鲁棒的草稿模型,使其在长尾输入上保持高接受率,是当前研究的一个方向。
此外,拒绝采样的修正分布计算需要目标模型输出完整词表分布,这在大词表场景下可能成为瓶颈。一些实现采用 top-k 截断来降低计算量,但这会引入近似,破坏分布一致性。如何在保持分布一致的前提下降低验证开销,仍是一个开放问题。
最后,投机解码的加速比上限受限于草稿模型与目标模型的耗时比。如果草稿模型不够快,或者目标模型的前向已经非常高效,投机解码的收益就会有限。在实际部署中,需要结合具体硬件和模型特性,通过实验找到最优配置。