无损投机采样的概率论视角

投机采样（Speculative Sampling）用小模型猜、大模型验，能在不改变输出分布的前提下加速推理。“不改变输出分布”是整个方法的硬约束：外界无论用什么统计检验，都无法区分”大模型自己逐词生成”和”小模型猜测 + 大模型校验”两种模式的输出。

这篇文章只关注一个问题：这个”无损”是怎么在概率层面实现的？

问题的形式化

设词表中任意候选 token $x$ ，在当前已接受前缀条件下，大模型给出的概率为 $p(x)$ ，小模型给出的概率为 $q(x)$ 。

在某个位置上，小模型以概率 $q(x)$ 采样出候选 $x$ ，大模型决定是否接受。如果接受， $x$ 直接输出；如果拒绝，从调整后的分布中重新采样。

目标：让 $x$ 最终被输出的概率严格等于 $p(x)$ 。

执行流程

先看完整的一轮迭代是怎么跑的：

Draft 阶段：小模型（draft model）自回归生成 $K$ 个候选 token： $x_1, x_2, \dots, x_K$ 。
Verify 阶段：大模型（target model）对前缀 + 全部 $K$ 个候选做一次并行 forward，拿到每个位置上大模型自己的概率分布 $p_i$ 。
逐位置接受/拒绝：从位置 1 开始依次判定。一旦在位置 $i$ 触发拒绝，丢弃 $x_i$ 及其后所有 draft token，在位置 $i$ 从调整后的分布中重新采样一个 token，本轮结束。
全部通过的 bonus：如果 $K$ 个 token 全部被接受，大模型的 forward 已经顺带算出了位置 $K+1$ 的分布，直接从中采样一个额外 token。

因此每轮产出最多 $K+1$ 个 token（全部接受 + bonus），最少 1 个 token（第一个就被拒绝，重采样补一个）。

为什么拒绝后必须立刻停止？一个 token 能被拒绝，前提是 $p(x_i) < q(x_i)$ （ $p \ge q$ 时接受概率为 1，不可能触发拒绝）。而调整后的分布中 $p \le q$ 的 token 权重为 0，不会被抽中，因此重新采样的结果必然不同于被拒绝的候选。而前缀一旦改变，后续的预测全都作废。

接受/拒绝规则

给定小模型采样出的候选 $x$ ，接受概率为：

A(x) = \min\left(1,\ \frac{p(x)}{q(x)}\right)

两种情况：

若 $p(x) \ge q(x)$ ： $A(x) = 1$ ，无条件接受。大模型比小模型更看好这个 token，没有理由拒绝。
若 $p(x) < q(x)$ ： $A(x) = \frac{p(x)}{q(x)} < 1$ ，以此概率随机接受。小模型对这个 token 过度自信了，需要按比例压回来。

实际判定：生成均匀随机数 $r \sim U[0,1)$ ，当 $r < A(x)$ 时接受，否则拒绝。

联合概率与总接受率

$x$ 被输出需要两步都成功：先被小模型提出，再通过接受判定。联合概率：

q(x) \cdot A(x) = q(x) \cdot \min\left(1,\ \frac{p(x)}{q(x)}\right) = \min(q(x),\ p(x))

关键性质：每个 token 在第一阶段能拿到的概率上限是 $\min(q(x), p(x))$ 。

对所有 token 求和，得到总接受率：

\alpha = \sum_x \min(q(x),\ p(x))

$\alpha$ 是”小模型提出某个 token 且该 token 通过判定”的总概率。当 $q \ne p$ 时 $\alpha < 1$ ，互补的 $1 - \alpha$ 就是拒绝事件的总概率。

此时出现一个问题：

$p(x) < q(x)$ 的 token：联合概率 $= q(x) \cdot \frac{p(x)}{q(x)} = p(x)$ 。已经精确对齐，不多不少。
$p(x) > q(x)$ 的 token：联合概率 $= q(x) \cdot 1 = q(x) < p(x)$ 。存在缺口 $p(x) - q(x)$ 。

第二类 token 自身永远不会被拒绝（接受概率是 1），但小模型提议它的频率不够高，仅靠第一阶段无法凑满 $p(x)$ 。缺口必须从别处补回来。

调整分布：跨 token 的概率补偿

补偿来自别的 token 被拒绝时触发的重新采样。当小模型提出了一个 $p < q$ 的 token 并被拒绝，系统不会放弃这个位置，而是从调整后的分布中重新采样：

p'(x) = \frac{\max(0,\ p(x) - q(x))}{1 - \alpha}

这个分布的构造逻辑：

$p(x) > q(x)$ 的 token：权重为 $p(x) - q(x)$ ，归一化后成为重新采样的候选，正是那些被小模型低估、需要补偿的 token。
$p(x) \le q(x)$ 的 token： $p(x) - q(x) \le 0$ ，权重为 0，不参与重新采样。它们在第一阶段已经被足额录取。

因果链： $p < q$ 的 token 被过度提议 → 按比例拒绝，释放概率空间 → 触发重新采样 → 补偿 $p > q$ 的 token 的概率缺口。

闭环证明

$x$ 最终被输出的总概率 = 第一阶段直接通过 + 拒绝后从调整分布中被抽中：

P_{\text{final}}(x) = \min(q(x), p(x)) + (1 - \alpha) \cdot p'(x)

代入 $p'$ 的定义：

= \min(q(x), p(x)) + (1 - \alpha) \cdot \frac{\max(0,\ p(x) - q(x))}{1 - \alpha}

= \min(q(x), p(x)) + \max(0,\ p(x) - q(x))

对任意实数 $a, b$ ，恒有 $\min(a, b) + \max(0, a - b) = a$ 。因此：

P_{\text{final}}(x) = p(x)

对词表中每一个 token 都成立。输出分布与大模型原始分布完全一致。证毕。

案例

假设词表有四个 token，概率如下：

Token	$q(x)$ （draft 给出）	$p(x)$ （target 给出）	关系
A	0.2	0.5	$p > q$ ，draft 低估了 A
B	0.4	0.1	$p < q$ ，draft 高估了 B
C	0.3	0.1	$p < q$ ，draft 高估了 C
D	0.1	0.3	$p > q$ ，draft 低估了 D

第一阶段：各 token 被 draft 提出且通过接受判定的联合概率 $\min(q, p)$ ：

A： $\min(q(A),\ p(A)) = \min(0.2,\ 0.5) = 0.2$ $min (q (A), p (A)) = min (0.2, 0.5) = 0.2$ 。
- $p > q$
- 无条件接受
- 但 draft 提议频率只有 0.2，缺口 $p(A) - q(A) = 0.3$ 。
B： $\min(q(B),\ p(B)) = \min(0.4,\ 0.1) = 0.1$ $min (q (B), p (B)) = min (0.4, 0.1) = 0.1$ 。
- $p < q$
- draft 以 0.4 的频率提议 B，接受概率 $p/q = 0.1/0.4 = 1/4$
- 剩下 $3/4$ 触发拒绝，即 B 贡献的拒绝概率 $= 0.4 \times 3/4 = 0.3$ 。
C： $\min(q(C),\ p(C)) = \min(0.3,\ 0.1) = 0.1$ $min (q (C), p (C)) = min (0.3, 0.1) = 0.1$ 。
- $p < q$
- draft 以 0.3 的频率提议 C，接受概率 $p/q = 0.1/0.3 = 1/3$
- 剩下 $2/3$ 触发拒绝。C 贡献的拒绝概率 $= 0.3 \times 2/3 = 0.2$ 。
D： $\min(q(D),\ p(D)) = \min(0.1,\ 0.3) = 0.1$ $min (q (D), p (D)) = min (0.1, 0.3) = 0.1$ 。
- $p > q$
- 无条件接受
- 缺口 $p(D) - q(D) = 0.2$ 。

总接受率 $\alpha = 0.2 + 0.1 + 0.1 + 0.1 = 0.5$ 。

总拒绝率 $1 - \alpha = 0.5$ ，由所有 $p < q$ 的 token 共同贡献：B 的 $0.3$ 和 C 的 $0.2$ 。

调整后的分布：拒绝发生时的重新采样池 $p'(x) \propto \max(0,\ p(x) - q(x))$ ：

A： $\max(0,\ 0.5 - 0.2) = 0.3$
B： $\max(0,\ 0.1 - 0.4) = 0$ （已足额，不参与）
C： $\max(0,\ 0.1 - 0.3) = 0$ （已足额，不参与）
D： $\max(0,\ 0.3 - 0.1) = 0.2$

归一化（除以 $0.3 + 0.2 = 0.5$ ）： $p'(A) = 0.6$ ， $p'(B) = 0$ ， $p'(C) = 0$ ， $p'(D) = 0.4$ 。

最终概率 = 第一阶段通过 + 总拒绝率 $\times$ 调整分布中的份额：

A： $0.2 + 0.5 \times 0.6 = 0.5 = p(A)$
B： $0.1 + 0.5 \times 0 = 0.1 = p(B)$
C： $0.1 + 0.5 \times 0 = 0.1 = p(C)$
D： $0.1 + 0.5 \times 0.4 = 0.3 = p(D)$

每个 token 精确对齐 target 分布。B 和 C 被拒绝释放出的总共 0.5 概率空间（ $0.3 + 0.2$ ），被 A 和 D 按各自缺口比例（ $0.3 : 0.2$ ）瓜分。

为什么必须用随机性

一个自然的疑问：为什么不用确定性阈值（比如 $p(x) < q(x)$ 就直接拒绝）？

因为确定性策略无法维持概率守恒。假设 $p(x) = 0.1$ ， $q(x) = 0.2$ 。一刀切拒绝的话， $x$ 在输出中的概率变成 0，大模型原本赋予的 10% 可能性被抹杀。这就是有损解码，输出分布不再等于 $p$ 。

而投机采样的做法是：小模型以 0.2 的频率提出 $x$ ，大模型以 $\frac{0.1}{0.2} = 0.5$ 的概率放行。最终 $x$ 出现的联合概率 $= 0.2 \times 0.5 = 0.1$ ，精确等于 $p(x)$ 。

小结

投机采样的”无损”本质上是一个概率论问题：对于 $q$ 超出 $p$ 的 token，通过拒绝采样按比例拒绝，释放出概率空间；再通过调整分布，将这些空间精确分配给被 $q$ 低估的 token。由于 $\sum p(x) = \sum q(x) = 1$ ，被高估 token 释放的概率总量必然等于被低估 token 的缺口总量（都等于 $1-\alpha$ ），因此补偿刚好闭合，输出分布恒等于 $p$ 。

整个机制不要求小模型准确。无论 $q$ 和 $p$ 差多远，走完接受-拒绝-重新采样流程，数学保证输出分布与 $p$ 完全一致。