大模型投机解码与拒绝采样为何无损

一、为什么需要投机解码

大模型自回归推理（autoregressive decoding）存在一个根本性的性能瓶颈：每次只生成一个 token，且每一步都要把整个模型跑一遍。对一个几十上百层的 Transformer 而言，生成 1 个 token 和一次 forward 把 batch 里所有位置都算一遍，计算量几乎一样——因为推理时 batch size 通常很小（甚至为 1），算力（GPU 的算术运算单元）远没有被喂饱，真正的瓶颈在显存带宽：每生成一个 token，都要把权重从 HBM 读到 SRAM，而权重读取这一步占据了绝大部分时间。

换句话说，现代大模型推理在 decode 阶段是 memory-bandwidth bound（访存受限），而不是 compute bound（算力受限）。GPU 的算术运算单元（FP16/BF16 tensor core）大量空闲，每次 forward 实际上只算出了概率分布里被采样到的那一个 token，剩下成千上万维的概率信息全部被浪费了。

投机解码（Speculative Decoding）正是为打破这个瓶颈而生，由 Leviathan 等人在 2023 年的论文 Fast Inference from Transformers via Speculative Decoding 和 Chen 等人的 Accelerating Large Language Model Decoding with Speculative Sampling 几乎同时提出。其核心思想是：

用一个又小又快的 draft model（草稿模型）先“猜”一串 token，再用大模型一次性并行验证这串 token，只接受其中被验证为正确的部分。

由于大模型验证一串 token 的代价和生成一个 token 几乎相同（一次 forward），如果草稿猜得准，就能在一次 forward 里产出多个 token，从而实现显著的加速——而且最终输出的分布与直接用大模型采样完全一致（无损）。

二、投机解码的整体流程

记大模型（target model）为 $q$，草稿模型（draft model）为 $p$。一次投机解码的步骤如下：

1. 草稿阶段：给定当前上下文 $x$，用小模型 $p$ 自回归地生成 $n$ 个草稿 token $x_1, x_2, \dots, x_n$。这一步串行，但因为 $p$ 很小很快，开销低。

2. 验证阶段：把这 $n$ 个草稿 token 拼到当前上下文后面，一次 forward 喂给大模型 $q$，并行得到每个位置的条件分布 $q(\cdot \mid x), q(\cdot \mid x, x_1), \dots, q(\cdot \mid x, x_{1:n-1})$。注意：大模型对这 $n$ 个位置的推理是并行的，只相当于一次正常 forward 的开销。

3. 接受/拒绝：从前往后逐个检查草稿 token $x_i$：

   • 若 $q(x_i \mid \cdots) \geq p(x_i \mid \cdots)$，接受 $x_i$。
   • 若 $q(x_i \mid \mid \cdots) < p(x_i \mid \cdots)$，以概率 $\frac{q(x_i \mid \cdots)}{p(x_i \mid \cdots)}$ 接受 $x_i$，否则拒绝。
   • 一旦某个 token 被拒绝，从该位置起丢弃后面所有草稿 token，并从修正后的分布中重新采一个 token 继续下一轮。

关键的“修正分布”是：当 $x_i$ 被拒绝时，不是直接从 $q$ 里采，而是从

$$q'(\cdot) = \text{norm}\left(\max(0,\ q(\cdot \mid \cdots) - p(\cdot \mid \cdots))\right)$$

中采样。也就是把 $q$ 中被 $p$ “超额使用”的那部分概率扣掉再归一化。下一节会证明，正是这一步保证了整体无损。

直觉：草稿 token 是小模型“免费”借给大模型的猜测，大模型决定要不要认账。如果某 token 在小模型那里出现的概率比大模型还高，说明小模型“抢跑”了，要把这部分超额概率扣除后再补采，才能让总采样分布回到与大模型独立采样一致。

三、为什么拒绝采样能保证无损

这是投机解码最精妙、也最容易被一带而过的部分。下面给出严格的推导。

3.1 目标

我们要证明：用上述接受-拒绝算法采出来的 token，其分布等于直接从大模型 $q$ 采样的分布。 即对任意 token $y$，

$$P(\text{最终输出} = y \mid x) = q(y \mid x).$$

证明采用全概率展开：最终输出 $y$ 要么是被接受的草稿 token，要么是被拒绝后从修正分布补采的 token。

3.2 单步证明

固定上下文，简记 $p(y) = p(y \mid x)$、$q(y) = q(y \mid x)$。$y$ 作为最终输出有两种情形：

情形 A：$y$ 作为草稿 token 被提出且被接受。

草稿提出 $y$ 的概率是 $p(y)$；在提出 $y$ 后被接受的概率是

$$\min\left(1,\ \frac{q(y)}{p(y)}\right).$$

所以这一路贡献

$$p(y) \cdot \min\left(1,\ \frac{q(y)}{p(y)}\right) = \min\big(p(y),\ q(y)\big).$$

情形 B：草稿被拒绝后从修正分布补采得到 $y$。

令 $r = P(\text{拒绝})$。先把所有 token 的“接受贡献”求和，得到任意草稿被接受的总概率：

$$P(\text{accept}) = \sum_z \min\big(p(z),\ q(z)\big).$$

于是拒绝概率

$$r = 1 - P(\text{accept}) = 1 - \sum_z \min\big(p(z),\ q(z)\big).$$

利用恒等式 $\min(a, b) = a - \max(0, a - b)$（或等价地 $b - \max(0, b-a)$），有

$$\sum_z \min\big(p(z), q(z)\big) = \sum_z \big(q(z) - \max(0,\ q(z) - p(z))\big) = 1 - \sum_z \max(0,\ q(z) - p(z)),$$

其中用到 $\sum_z q(z) = 1$。因此

$$r = \sum_z \max\big(0,\ q(z) - p(z)\big).$$

这正是修正分布 $q'(\cdot) = \text{norm}(\max(0,\ q - p))$ 在归一化前的总量。

被拒绝后，从修正分布采到 $y$ 的概率是

$$q'(y) = \frac{\max(0,\ q(y) - p(y))}{r}.$$

所以情形 B 对输出 $y$ 的贡献为

$$r \cdot q'(y) = \max\big(0,\ q(y) - p(y)\big).$$

合并 A + B：

$$P(\text{输出}=y) = \underbrace{\min(p(y),\ q(y))}_{A} + \underbrace{\max(0,\ q(y)-p(y))}_{B}.$$

分两种情况验证：

• 若 $q(y) \geq p(y)$：$A = p(y)$，$B = q(y) - p(y)$，合计 $= q(y)$。✅
• 若 $q(y) < p(y)$：$A = q(y)$，$B = 0$，合计 $= q(y)$。✅

两种情况都得到 $q(y)$。故

$$\boxed{P(\text{输出}=y \mid x) = q(y \mid x).}$$

3.3 多步/序列无损性

上面的证明是“固定一个位置”的单步分布等价。由于投机解码每一步都在做这样的等价采样，且每一步的等价性不依赖于前面草稿是否被接受（拒绝后的补采依然落在大模型分布上），用数学归纳法即可得到：任意长度的生成序列，其联合分布与纯大模型自回归采样的联合分布逐 token、逐 token 地完全一致。

所以投机解码是**严格无损（lossless）**的——它的输出分布和大模型自己采样一字不差，只是把“一次 forward 产一个 token”改成了“一次 forward 产多个 token”，用空间换时间、用草稿模型的算力换大模型的访存等待。最终质量与大模型独立采样完全相同。

四、几个常被混淆的点

1. “无损”是对“采样”而言，不是对“贪心解码”而言。 如果目标本身是贪心（greedy / argmax），那投机解码里把采样换成“取最大”也成立——只要草稿 token 等于大模型 argmax 就接受，这也是投机解码能在贪心场景严格无损、且实现更简单的原因。但一般说的无损指的是采样分布等价。

2. 草稿模型越准，加速越多，但无损性不依赖草稿模型好坏。 哪怕草稿模型烂到完全乱猜，输出分布依然和大模型一致——只是接受率低、加速比差而已。无损是数学保证，加速是工程结果。

3. 修正分布 $q' = \text{norm}(\max(0, q-p))$ 不是可选项。 它是“把小模型多借的那部分概率还给大模型”的关键。若拒绝后直接从 $q$ 采，则接受路径和补采路径会双重计入 $q$ 中被 $p$ 覆盖的部分，分布就偏了。

4. 自投机解码（self-speculative / EAGLE 等） 不再依赖外部小模型，而是用大模型自身的浅层/一个 head 来做草稿，甚至把草稿阶段也并进大模型的一次 forward。接受-拒绝的数学不变，依然无损，只是把“小模型 $p$”换成了与大模型共享参数的某种近似分布。

五、总结

• 投机解码用 小模型起草稿 + 大模型一次并行验证，破解了 decode 阶段访存受限的瓶颈。
• 拒绝采样 + 修正分布 $q' = \text{norm}(\max(0, q-p))$ 保证每一步输出分布严格等于大模型采样分布。
• 无损性来自一个简洁的全概率恒等：$\min(p,q) + \max(0, q-p) \equiv q$。
• 加速来自草稿命中后“一次 forward 出多个 token”，命中越多越快；而无损性与命中与否无关。

一句话：投机解码把“猜测”留给了快的小模型，把“裁决”留给了慢的大模型，而拒绝采样的那一点概率修补，恰好让两者的合力在分布意义上等于大模型亲力亲为。

参考：

• Leviathan et al., Fast Inference from Transformers via Speculative Decoding, 2023.
• Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling, 2023.