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

投机采样（Speculative Sampling）用小模型猜、大模型验，能在不改变输出分布的前提下加速推理。

其中，“不改变输出分布”是整个方法的硬约束：外界无论用什么统计检验，都无法区分”大模型自己逐词生成”和”小模型猜测 + 大模型校验”两种模式的输出。

候选经过接受、拒绝和补采后，每个 token 的最终概率必须等于大模型给出的概率。

算法过程

问题与目标

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

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

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

执行流程

1. Draft 阶段：小模型（draft model）自回归生成 $K$ 个候选 token：$x_1, x_2, \dots, x_K$。
2. Verify 阶段：大模型（target model）对前缀 + 全部 $K$ 个候选做一次并行 forward，拿到每个位置上大模型自己的概率分布 $p_i$。
3. 逐位置接受/拒绝：从位置 1 开始依次判定。一旦在位置 $i$ 触发拒绝，丢弃 $x_i$ 及其后所有 draft token，在位置 $i$ 从调整后的分布中重新采样一个 token，本轮结束。
4. 全部通过的 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，不会被抽中，因此重新采样的结果必然不同于被拒绝的候选 token，导致大模型与小模型对于序列的预测发生分歧。而序列前缀一旦改变，后续的预测全都作废，此为 Causal Attention 的性质。

接受/拒绝规则

小模型在当前位置采样出一个候选 $x$ 后，以如下概率接受它：

两种情况：

• 情况一：若 $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)$ 时接受，否则拒绝。

概率推导

无损投机采样通过拒绝与补采，将分布 $q$ 调整为 $p$。由于两个分布的概率总和都为 1，且 $q$ 相对 $p$ 多出的概率总量，恰好等于 $q$ 相对 $p$ 缺少的概率总量，因此可以通过重新分配精确补齐。

联合概率与接受概率

固定当前前缀，draft 以概率 $q(x)$ 提出一个候选 $x$，再以概率 $A(x)$ 接受它。因此，$x$ 被提出并直接接受的联合概率为：

候选可能是词表中的任意 token。把各 token 被提出并接受的联合概率相加，就得到这次抽样与验证直接接受候选的总概率：

注意区别 $A(x)$ 是给定候选 $x$ 的接受概率，$\beta$ 则考虑了 draft 可能抽到的所有候选。综上，$1-\beta$ 等于触发补采的概率。

对比直接接受贡献的概率与目标 $p(x)$：

• 情况一：$p(x) \ge q(x)$ 时，联合概率 $= q(x) \cdot 1 = q(x) \le p(x)$。缺口为 $p(x) - q(x)$。
• 情况二：$p(x) < q(x)$ 时，联合概率 $= q(x) \cdot \frac{p(x)}{q(x)} = p(x)$，已经足额。

所以，补采只需填补 $p(x)>q(x)$ 的 token 的概率缺口。

拒绝后的补采

情况二的候选以 $p(x)/q(x)$ 的概率接受，以 $1-p(x)/q(x)$ 的概率拒绝。按这个比例随机拒绝，就把 draft 提出它的概率 $q(x)$ 降到直接输出它的概率 $p(x)$。

候选被拒绝后，当前位置仍需输出一个 token，因此系统从调整分布 $p'$ 中重新采样，补偿情况一留下的缺口。

对每个 token，目标概率 $p(x)$ 减去直接接受已贡献的 $\min(q(x),p(x))$，就是需要补偿的缺口。词表中目标概率的总和为 1，直接接受的总概率为 $\beta$，所以缺口总量为：

这也正是触发补采的概率。按缺口比例采样，将各 token 的缺口除以 $1-\beta$，便得到总和为 1 的补采分布：

• 情况一：$p(x)\ge q(x)$ 的 token 按缺口 $p(x)-q(x)$ 的比例参与补采，相等时权重为 0。
• 情况二：$p(x)<q(x)$ 的 token 已获得目标概率，补采权重为 0。

当 $p=q$ 时，候选全部接受，无需补采。

无损证明

对于词表中的任意 token $x$，最终输出它有两条路径：draft 提出它并通过验证，联合概率为 $\min(q(x),p(x))$；或者其他候选被拒绝后，补采抽到它，联合概率为 $(1-\beta)p'(x)$。两条路径互斥，概率相加得到：

当 $\beta<1$ 时，代入 $p'$ 并化简：

而当 $\beta=1$ 时，$q=p$，全部接受同样得到 $P_{\text{final}}(x)=p(x)$。因此，无论是否需要补采，每个 token 的最终输出概率都等于 $p(x)$，输出分布与大模型一致。

每一步在相同前缀下的条件分布都与大模型一致，因此整个序列的联合分布也一致。

案例

假设词表为 A、B、C、D，本轮 draft 依次提出 D、A、C、B 四个候选。

生成候选与计算分布

draft 依次提出候选，并保存各位置的分布 $q_i$。随后 target 对已有前缀和四个候选做一次 forward，得到 $p_1,\dots,p_4$ 和全部通过时使用的 $p_5$。

假设前两个位置的候选 D、A 均被接受，接下来验证第三个位置。该位置的完整词表分布如下：

接受、拒绝与补采

从第一个候选开始，使用对应位置的概率，依次判断是否接受，注意各位置下词表的采样概率不同：

第二个候选 A 的接受概率为 0.25，本轮随机数为 0.1，因此通过。第三个候选 C 被拒绝后，保留 D、A，在第三个位置补采。

按前面的概率定义，该位置各 token 被提出并直接接受的联合概率为：

四项相加得到 $\beta=0.2+0.1+0.1+0.1=0.5$。拒绝概率为 $1-\beta=0.5$，也等于各 token 的缺口总量 $0.3+0+0+0.2=0.5$。

补采时，调整分布 $p'_3$ 为：

随后从 $p'_3$ 采样。假设抽到 A，本轮输出 D、A、A 后结束。

核对输出分布

沿用前面计算的联合概率、$1-\beta$ 和 $p'_3$，将直接接受与补采的概率贡献相加：

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

因此，最终输出分布与目标分布一致。

为什么必须用随机性

为什么不用确定性阈值，比如 $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)$。

总结

无损来自概率补偿的精确性：拒绝的总概率恰好等于待补偿的缺口，按缺口比例补采，就能恢复目标分布。因此，小模型的准确性影响候选的接受概率，而输出分布的正确性由采样规则保证。