Featured image of post 推测解码:先猜几个 token,为什么答案仍然没变

推测解码:先猜几个 token,为什么答案仍然没变

推测解码让小模型先起草多个 token,再由大模型一次并行验证。本文从串行瓶颈、拒绝采样校正和收益模型出发,解释它为什么能提速而不改变目标模型的输出分布。

大模型生成一句话时,最难并行的不是理解整段 prompt,而是接下来的等待:第 2 个 token 必须等第 1 个 token 出来,第 3 个又必须等第 2 个。假如输出 100 个 token,目标模型至少要经历 100 个串行解码步。GPU 很快,却被一条依赖链拴住。

推测解码(speculative decoding)提出了一个像 CPU 分支预测的办法:让便宜的草稿模型先向前猜几个 token,再让昂贵的目标模型一次检查整段猜测。猜对就批量前进,猜错便在第一个错误处校正。真正关键的不是“小模型帮大模型写答案”,而是校正机制保证最终样本仍来自目标模型原本的分布。草稿只影响速度,不获得最终裁决权。

这使它成为一种少见的优化:不量化权重,不剪枝,不改目标模型,也不以输出质量换延迟。Google Research 的原始工作在 T5-XXL 上报告了 2–3 倍加速;DeepMind 的独立工作在 70B Chinchilla 的分布式设置中报告了 2–2.5 倍解码加速。倍率依赖模型、硬件和负载,但“把多个串行步折叠进一次并行验证”的思想已经进入主流推理引擎。

解码的瓶颈是一条串行链

Transformer 的训练可以同时处理一句话里的许多位置,因为每个位置的正确前缀都已知。生成时却没有这份答案:目标模型在时刻 \(t\) 先计算 \(p(x_t\mid x_{

KV cache 避免重复计算旧 token 的 key 和 value,却没有消除这种时间依赖。常见的小 batch 解码还经常受显存带宽限制:每次只为一个新位置做少量工作,却要读取大量模型权重和缓存。于是,一次目标模型前向虽然没有把计算单元填满,下一 token 仍然不能提前确定。

推测解码利用了两个事实。第一,目标模型一次处理一小段连续位置,通常不会比只处理一个位置慢同样的倍数;矩阵运算有机会获得更高并行度。第二,语言里有大量容易预测的局部片段,例如固定搭配、标点、代码缩进和从上下文复制的文字,小模型常能猜中。用便宜模型制造并行工作,再用大模型统一审核,便有机会用一次昂贵调用换来多个 token。

图 1:普通解码每个 token 都要等待一次目标模型;推测解码先串行起草,再把多个候选位置交给目标模型并行验证

一轮推测解码究竟发生什么

记目标模型的分布为 \(p\),草稿模型的分布为 \(q\),每轮让草稿模型提出 \(\gamma\) 个 token。

  1. 草稿模型按自回归方式生成 \(x_1,\ldots,x_\gamma\),并保留每一步的 \(q_i(x)\)。这部分仍是串行的,但草稿模型必须足够便宜。
  2. 目标模型把“原前缀 + 整段草稿”当作一个小批次,一次算出每个候选位置的 \(p_i(x)\),以及草稿全被接受后下一个位置的分布。
  3. 从左到右审核。每个草稿 token \(x_i\) 以 \(\min(1,p_i(x_i)/q_i(x_i))\) 的概率被接受;一旦拒绝,后面的草稿全部作废。
  4. 若在第 \(i\) 个位置拒绝,不是简单地从 \(p_i\) 重采样,而是从校正后的“剩余分布”采样;若全部接受,还可从目标模型已经算出的下一位置分布再取一个 token。

所以一轮至少前进 1 个 token,最多前进 \(\gamma+1\) 个。草稿并非并行地产生——经典小模型 drafter 仍逐 token 运行——真正被并行化的是昂贵的目标模型验证。

为什么“猜错再改”不会偏离目标分布

只保留目标模型也喜欢的草稿 token,听起来会偏向“两个模型都同意”的安全答案。避免这种偏差的关键,是拒绝后的补偿分布。

对一个候选 token \(x\),先从 \(q\) 采样。若 \(q(x)\le p(x)\),总是接受;否则以 \(p(x)/q(x)\) 的概率接受。这样经接受路径得到 \(x\) 的概率恰好是

\[ q(x)\min\left(1,\frac{p(x)}{q(x)}\right)=\min(p(x),q(x)). \]

接受路径只覆盖了两个分布重叠的部分。目标分布还缺少 \(p(x)>q(x)\) 的那部分质量,因此拒绝时从

\[ p'(x)=\operatorname{norm}(\max(0,p(x)-q(x))) \]

采样。重叠部分与剩余部分相加,正好还原 \(p\)。随后对每个位置条件化地重复同一论证,整段序列也保持目标模型的分布。

图 2:接受路径取得 p 与 q 的重叠概率质量;拒绝后的 residual sampling 补回 p 独有的部分

对 greedy decoding,逻辑更直观:只要草稿 token 等于目标模型的 argmax 就接受,否则在第一个不一致处改用目标 token。对随机采样,“输出不变”指的是概率分布不变,不是固定随机种子下每次字符串逐字一致。浮点精度、batch 形状和随机数消费顺序仍可能造成实际运行差异;vLLM 的官方文档也明确把理论无损、算法验证与数值稳定性分开讨论。

速度取决于两个数,而不是草稿模型单独的准确率

把平均接受率记为 \(\alpha\)。在论文采用的独立同分布近似下,一轮目标模型验证平均产出的 token 数为

\[ E[N]=\frac{1-\alpha^{\gamma+1}}{1-\alpha}=1+\alpha+\alpha^2+\cdots+\alpha^\gamma. \]

接受率越高,越可能连续吞下整段草稿。但这还不是速度。令 \(c\) 表示草稿模型单步时间与目标模型单步时间之比,理想化加速因子为

\[ S=\frac{1-\alpha^{\gamma+1}}{(1-\alpha)(1+\gamma c)}. \]

这个式子揭示了真正的工程权衡:更强的草稿模型可能提高 \(\alpha\),却也增大 \(c\);草稿长度 \(\gamma\) 太短,没吃到并行收益,太长则会在第一次拒绝后浪费后续起草与验证。原论文观察到,在其设置中,草稿模型比目标模型小约两个数量级常有较好平衡,但这不是通用配方。

图 3:收益由接受率、草稿成本和推测长度共同决定;只追求更长草稿或更强 drafter 都可能适得其反

更重要的是,公式假设目标模型验证 \(\gamma+1\) 个位置的耗时接近验证 1 个位置,并且机器仍有空余计算能力。低并发、memory-bound 的交互式生成较符合这个前提;高并发服务已经把 GPU 填满时,扩大验证 micro-batch 可能反而争抢算力。vLLM 因此把推测解码定位在中低 QPS、延迟敏感的 workload,并提供按 batch size 动态调整推测长度的机制。

草稿不一定是另一台小 Transformer

经典算法只要求 proposer 便宜且能给出候选,不要求它与目标模型同构。这个接口催生了几类实现:

  • 独立草稿模型:概念最直接,能给出完整 \(q\) 分布,但增加权重、KV cache 与调度成本,词表不一致也要额外映射。
  • n-gram / prompt lookup:从 prompt 或已生成文本中找重复片段,几乎没有模型成本。代码、结构化文本和改写任务常有可利用的复制模式,但开放式生成的命中率可能较低。
  • 轻量预测头与 EAGLE 类方法:复用目标模型的隐藏状态,用小型头预测多个后续 token,减少独立 drafter 的成本。它们需要匹配模型的额外权重或训练流程。
  • 原生 multi-token prediction:模型训练时就学习预测多个未来位置,部署时可把这些预测用于推测验证。它把“谁来起草”纳入模型本身,而不是外挂另一套完整模型。

这些方法改变的是候选生成器,验收思想仍相同。vLLM 的当前官方文档同时列出 draft model、EAGLE、MTP、n-gram、suffix decoding 等路线,说明推测解码已经从单一算法发展为一组共享“提出—验证”协议的系统设计。

什么情况下它不会更快

推测解码最容易被一个漂亮的“接受率”掩盖。部署前至少要测四件事:平均接受长度、草稿开销、验证阶段的 GPU 利用率,以及端到端的 time per output token。接受率高但 drafter 慢,仍可能亏损;接受率一般但 proposer 近乎免费,也可能值得。

它通常不利于以下场景:目标与草稿分布差异大;temperature 较高导致候选更难吻合;batch 已经很大、目标模型验证从带宽瓶颈转成计算瓶颈;输出很短,初始化和调度成本来不及摊薄;或多卡拓扑让草稿与验证引入额外通信。不同请求、不同生成阶段甚至会有不同最优 \(\gamma\),静态参数很难覆盖所有负载。

还要区分时延与吞吐。对一个等待回答的用户,少做几轮大模型串行调用很有价值;对已经充分 continuous batching 的平台,同一块 GPU 每秒完成的总 token 未必同步增加。推测解码可能用更多总 FLOPs 换更短墙钟时间——这不是矛盾,而是用闲置并行度购买低延迟。

昨天讨论 PagedAttention 时,核心是把更多活跃序列放进显存,扩大有效 batch;推测解码则是在一条序列内部,尝试一次跨过多个时间步。一个优化“横向并发多少请求”,另一个优化“纵向少走多少轮”。两者都没有改变模型回答什么,却分别重写了内存管理和时间调度。

推测解码最值得记住的也不是“小模型猜、大模型验”这句口诀,而是一条更普遍的系统原则:当昂贵计算被严格串行依赖限制时,可以先用便宜近似制造候选,再用一次可并行的精确计算收回正确性。猜测可以大胆,因为裁决仍属于目标模型;优化发生在执行路径,而不是答案标准上。

参考资料

  1. Leviathan, Kalman, Matias, Fast Inference from Transformers via Speculative Decoding, ICML 2023.
  2. Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling, 2023.
  3. vLLM Project, Speculative Decoding documentation.
  4. vLLM Project, Speculators: training and deployment library.
  5. Li et al., EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty, ICML 2024.
  6. Cai et al., Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads, ICML 2024.