从投机解码到 Rand5 模拟 Rand7：跨越十年的拒绝采样直觉

上一篇理清了 Prefill 和 Decode 的因果链之后，我开始翻 Decode 阶段的加速方案。结论已经很清楚了：在单请求、小 Batch 的典型场景下，Decode 阶段是 memory-bound 的——每生成一个 token 就要把整个模型权重从显存搬一遍，计算单元大部分时间在等数据（Batch 较大或模型结构不同时，瓶颈也可能转向 compute-bound，但单次请求的延迟优化场景下 memory-bound 是常态）。

一个直觉是：能不能让一次昂贵的 Target Model forward 验证多个 token，把算力利用率拉上去？

这就是 Speculative Decoding（投机解码）的出发点。我花了些时间理解它的核心机制，然后发现一个让我愣了几秒的事情——这套 accept/reject/resample 的数学内核，跟十年前打 ACM 和准备校招时反复练过的拒绝采样是同一个东西。

投机解码：猜了再验

#

Speculative Decoding 的思路很直接。既然大模型（Target Model）跑一次 forward 代价很高，那就找一个小模型（Draft Model）先快速串行猜 \(\gamma\) 个 token，再让大模型在一次 forward 中并行验证这些猜测。小模型参数量通常只有大模型的十分之一甚至更少，串行跑 \(\gamma\) 步的耗时可能还不如大模型跑一步。

设大模型的 next-token 概率分布为 \(p(x)\)，小模型为 \(q(x)\)。

如果只做 Greedy Decoding（永远选概率最高的 token），验证很简单：大模型算一遍，看小模型猜的 token 是不是 argmax，是就接受，不是就从该位置开始用大模型的结果。

但实际的 LLM 推理通常用 Temperature Sampling——按概率分布随机采样出 token，这样输出才有多样性。这时候问题变得微妙了：小模型的分布 \(q(x)\) 和大模型的 \(p(x)\) 不一样，直接接受小模型的采样结果，最终输出分布就不再是 \(p(x)\) 了。

怎么做到既用小模型加速，又保证最终输出分布严格等于 \(p(x)\)？ 这是 Speculative Decoding 要解决的核心问题。

验证规则是这样的：Draft Model 采样出一个 token \(x\) 后，再生成一个均匀随机数 \(u \sim \text{Uniform}(0, 1)\)，以概率

决定是否接受它。如果 \(u \le \min(1, p(x)/q(x))\)，接受；否则拒绝，从一个修正分布中重新采样（后面会推导这个修正分布）。

为什么生成必须串行，而验证可以并行？

#

写到这里有一个问题让我停下来想了想：Target Model 本身也是自回归模型，正常生成 token 时一步都不能跳。为什么 Draft 猜了一串之后，Target 就能在一次 forward 中把它们全部验证完？

关键在于：正常 Decode 的串行瓶颈不是计算本身的限制，而是输入的限制——下一步的输入依赖上一步的输出，在上一步完成之前根本不知道该算什么。

假设已确认的 Prefix 是 I love，Draft 给出三个候选 token：machine、learning、is。Target 需要验证的是这些条件概率：

• \(p(\text{machine} \mid \text{I love})\)
• \(p(\text{learning} \mid \text{I love machine})\)
• \(p(\text{is} \mid \text{I love machine learning})\)

每个概率所依赖的上下文 token 已经全部已知——Prefix 是确认的，候选序列是 Draft 提供的。Target 不需要等自己算出 machine 才去算 learning 的概率，因为 Draft 已经假设了这条路径。

这和 Prefill 处理整段 Prompt 的原理一样（之前的文章讨论过）：输入 token 全部已知时，Transformer 的 Causal Attention 可以在一次 forward 中并行计算所有位置的隐藏状态和 Logits。Causal Mask 限制的是每个位置能"看到"哪些前文，但不要求不同位置必须串行执行——同一个 Attention Layer 内，所有位置的 Query、Key、Value 可以同时计算，只是 Mask 确保位置 \(i\) 不会 attend 到位置 \(j \gt i\)。Transformer 的层与层之间有顺序依赖，但同一层内的多个位置是并行的。

有一个容易出错的细节：Next-Token Logits 的位置对齐。Transformer 输入位置 \(t\) 输出的 Logits 预测的是位置 \(t+1\) 的 token，不是位置 \(t\) 本身。也就是说：

• 输入位置 love（Prefix 末尾）→ 输出 Logits 预测 machine → 用于验证第一个 Draft Token
• 输入位置 machine → 输出 Logits 预测 learning → 用于验证第二个 Draft Token
• 输入位置 learning → 输出 Logits 预测 is → 用于验证第三个 Draft Token

验证第一个 Draft Token 所需的 Logits 其实来自 Prefix 最后一个位置，这个 Logits 可能在之前的 Decode 步骤中已经算好了。其余位置的 Logits 通过 Verification Forward 一次性得到。

需要强调的是：Target 只是在 Draft 给出的假设路径上计算概率分布，并不是确认所有候选都正确。验证仍然按顺序进行——从第一个候选开始逐个检查接受概率，一旦某个位置被拒绝，该位置之后的所有结果都必须作废，因为它们依赖的上下文已经不成立了。

一个两 token 的世界

#

为了把机制看清楚，我把问题压缩到最简的情况。假设词表只有两个 token：A 和 B。

小模型觉得 B 更可能（60%），大模型觉得 A 更可能（70%）。两者分歧明显。

代入接受公式：

• Token A：\(\frac{p(A)}{q(A)} = \frac{0.7}{0.4} = 1.75\)，取 \(\min(1, 1.75) = 1\)。100% 接受。 大模型比小模型更看好 A，小模型猜 A 的时候，大模型全盘认可。
• Token B：\(\frac{p(B)}{q(B)} = \frac{0.3}{0.6} = 0.5\)，取 \(\min(1, 0.5) = 0.5\)。50% 接受。 小模型过度高估了 B，大模型按比例砍掉一半。

这个操作在做什么？对于 \(q(x) \le p(x)\) 的 token（Draft 低估的），照单全收；对于 \(q(x) \gt p(x)\) 的 token（Draft 高估的），按 \(p/q\) 的比例削减。把小模型"多猜"的概率质量裁掉，“少猜"的直接放过。

到这里我停下来想了想——这个用概率阈值"裁剪"一个分布去逼近另一个分布的操作，好像在哪里见过很多次。

十年前的肌肉记忆

#

盯着 \(\min(1, p(x)/q(x))\) 越看越眼熟，突然反应过来：这不就是**拒绝采样（Rejection Sampling）**的接受准则吗？

经典拒绝采样的逻辑：从一个容易采样的提议分布 \(q(x)\) 里抽样，通过概率性地接受或拒绝，保证最终留下来的样本分布等于目标分布 \(p(x)\)。接受概率正比于 \(p(x)/q(x)\)。

这一下子把我拉回了很久以前的记忆。2017 年面京东实习的时候被蓄水池采样的问题卡过，回来补课写了一篇博客。蓄水池采样和拒绝采样的底层逻辑其实是一脉相承的：在随机流中，用一个概率阈值来决定元素的去留。蓄水池用 \(k/i\) 的概率决定是否替换已有样本，拒绝采样用 \(p(x)/(M \cdot q(x))\) 的概率决定是否接受当前抽样。形式不同，骨架相同。

打 ACM-ICPC 那几年和后来准备校招面试的时候，这类概率题反反复复练过。两道经典题最能说明这个共通性。

用不均匀硬币做公平决断

#

一枚硬币正面概率 \(p\)，反面 \(1-p\)，\(p\) 未知且不等于 0.5。怎么用它做出一个绝对公平的二元决断？

冯·诺依曼的做法：连抛两次。

HT 和 TH 的概率严格相等，都是 \(p(1-p)\)。丢弃对称的 HH 和 TT，保留的两种结果天然 1:1。

核心操作是同一个：拒绝不满足条件的样本，保留满足条件的。

Rand5 模拟 Rand7

#

这道题在校招算法面试里出现的频率极高。手上有一个 rand5() 函数，等概率返回 \([1, 5]\)，用它构造等概率返回 \([1, 7]\) 的 rand7()。

三步走：

1. 进制扩容：5 * (rand5() - 1) + rand5()，等概率生成 \([1, 25]\) 中的任意一个整数。原理是把两次 rand5() 看成一个五进制两位数，25 个组合对应 25 个等概率值。
2. 截断拒绝：25 个等概率的数里，\([1, 21]\) 正好是 7 的最大整数倍。落在 \([22, 25]\) 的直接丢弃，重新来过。
3. 取模映射：(val - 1) % 7 + 1，把 \([1, 21]\) 均匀映射到 \([1, 7]\)。

 1import random
 2
 3def rand5():
 4    return random.randint(1, 5)
 5
 6def rand7_basic():
 7    while True:
 8        val = 5 * (rand5() - 1) + rand5()  # [1, 25] 等概率
 9        if val <= 21:
10            return (val - 1) % 7 + 1

落在 \([22, 25]\) 的概率是 \(4/25 = 16\%\)。每轮调用两次 rand5()，期望调用次数是 \(2 \times \frac{25}{21} \approx 2.38\) 次。

模式已经很清楚了：从一个"不完全匹配"的随机源中采样，通过拒绝不合规的结果来保证输出分布的精确性。 代价是被拒绝的那部分采样白做了。

变废为宝：被拒绝的熵不该扔掉

#

这是让我想得最久的部分——也是 Speculative Decoding 最精妙的地方。

Rand7 还能更快吗

#

基础版 Rand7 的问题在于：摇出 22~25 时直接丢弃，重新调两次 rand5() 从头来过。那 16% 的情况消耗的随机性（entropy）全部浪费了。

但仔细想想，摇出 22~25 本身就携带了信息。它是 \([22, 25]\) 中等概率的一个值，相当于手里已经有了一个 \([1, 4]\) 的均匀随机数（val - 21）。

与其扔掉重来，不如复用这个残余的随机数。只需要再调一次 rand5()，就能拼成 \([1, 20]\) 的均匀分布：

20 里面 7 的最大整数倍是 14。落在 \([1, 14]\) 就取模输出；落在 \([15, 20]\)（6 个数），又拿到了一个 \([1, 6]\) 的残余随机数，再调一次 rand5() 拼出 \([1, 30]\)，截断 \([1, 28]\)……每一级被拒绝的残余都被喂给下一级，递归压榨。

 1def rand7_optimized():
 2    val = 5 * (rand5() - 1) + rand5()  # [1, 25]
 3    if val <= 21:
 4        return (val - 1) % 7 + 1
 5
 6    # 第一级残余：val-21 ∈ [1, 4]
 7    val = 5 * (val - 22) + rand5()  # [1, 20]
 8    if val <= 14:
 9        return (val - 1) % 7 + 1
10
11    # 第二级残余：val-14 ∈ [1, 6]
12    val = 5 * (val - 15) + rand5()  # [1, 30]
13    if val <= 28:
14        return (val - 1) % 7 + 1
15
16    # 极小概率（~0.1%）走到这里，回退
17    return rand7_basic()

期望调用次数可以这样估算：

• 概率 84%（21/25）：2 次调用就返回
• 概率 11.2%（4/25 × 14/20）：3 次调用返回
• 概率 4.5%（4/25 × 6/20 × 28/30）：4 次调用返回
• 剩余不到 0.3%：回退到基础版

从 2.38 降到 2.21，看起来改善不大，但核心思想很重要：被拒绝不意味着信息归零。拒绝本身产生了残差，这个残差是有效的随机性，可以被下一轮利用。

回到投机解码：这 30% 怎么办

#

带着这个思路回到前面的两 token 算例。

先算一下各种情况的概率：

直接接受的总概率只有 70%。还有 30% 的概率——Draft 采样出 B 但被拒绝了——怎么处理？

如果像基础版 Rand7 一样直接"丢弃重来"，代价是灾难性的。Target Model 已经跑完了一次昂贵的 forward，算出了完整的 \(p(x)\)。如果丢弃，这次 forward 吃掉的显存带宽和算力全部白费，还得再跑一次。

Speculative Decoding 的做法和 Rand7 的进阶优化如出一辙——绝不浪费，从残差分布中直接采样。

定义修正分布（Residual Distribution）：

直觉是：\(p(x) - q(x)\) 就是 Target 还没被 Draft 覆盖掉的那部分概率质量。对于 Draft 已经高估的 token（\(q \gt p\)），残差为 0——Draft 的过剩概率质量已经在 accept/reject 阶段被裁掉了。对于 Draft 低估的 token（\(q \lt p\)），残差正好是 Target 期望中还差的那部分。

代入算例：

• \(\max(0,\; p(A) - q(A)) = \max(0,\; 0.7 - 0.4) = 0.3\)
• \(\max(0,\; p(B) - q(B)) = \max(0,\; 0.3 - 0.6) = 0\)

归一化后：\(p'(A) = 1.0\)，\(p'(B) = 0\)。

也就是说，一旦 Draft 的采样被拒绝（在这个算例里只有采样到 B 才会被拒绝），从修正分布采样的结果是——100% 输出 A。

闭环验算

#

把所有路径加起来：

Token A 的总概率：

• 路径 1：Draft 直接采样到 A 且被接受 = \(0.4 \times 1.0 = 0.40\)
• 路径 2：Draft 采样到 B 被拒绝，修正分布采样出 A = \(0.30 \times 1.0 = 0.30\)
• 合计：\(0.40 + 0.30 = \mathbf{0.70} = p(A)\) ✓

Token B 的总概率：

• 路径 1：Draft 直接采样到 B 且被接受 = \(0.6 \times 0.5 = 0.30\)
• 路径 2：Draft 采样到 B 被拒绝，修正分布采样出 B = \(0.30 \times 0 = 0\)
• 合计：\(0.30 + 0 = \mathbf{0.30} = p(B)\) ✓

最终输出分布严格等于 Target Model 的 \(p(x)\)。而且，每一轮 Verification 至少能推进一个有效 token——要么 Draft 的猜测被接受，要么从修正分布中重新采样。不过需要注意：如果 Draft 猜了 \(\gamma\) 个 token 但第一个就被拒绝了，后面 \(\gamma - 1\) 个位置对应的 Target 计算全部作废，只有修正分布产出的那一个 token 是有效收益。所以"至少推进一个"不等于"没有浪费计算"——被拒绝路径上的 FLOPs 和显存带宽确实白花了。

这和 Rand7 优化版的设计思想相通：基础版 Rand7 丢弃被拒绝的随机数从头再来，浪费了已有的 entropy；优化版把残余随机数喂给下一级继续用。Speculative Decoding 丢弃被拒绝的 token 意味着浪费已经算好的 forward 结果；用修正分布兜底意味着这一轮的 forward 至少会产出一个有效 token，不会完全白跑。

Bonus Token

#

还有一个细节值得提一下。假设 Draft 提出了 \(\gamma\) 个候选 token，而且全部被接受了——Verification Forward 已经计算出了最后一个候选位置对应的 Next-Token Logits，可以直接从 Target 分布再采样一个额外的 token。也就是说，一轮最多可以输出 \(\gamma + 1\) 个 token，这个额外的 token 就是 Bonus Token。

如果中途某个候选被拒绝了，那后面的候选全部作废，Bonus Token 自然也拿不到——只能用拒绝位置对应的修正 token。另外，Bonus Token 虽然已经采样出来了，但它自身的 KV Cache 还没有被计算，下一轮的 Draft 和 Verification 需要先把这个位置的 KV 补上。

验证失败时的 KV Cache

#

顺便提一下 KV Cache 的工程含义。假设 Draft 提出了 D1、D2、D3、D4 四个候选，Target 接受了 D1、D2 但拒绝了 D3。那么 D1、D2 对应的 Target KV 可以继续保留在 Cache 中；D3、D4 路径上的 KV 逻辑上作废，不能再用。Target 的修正 token（替代 D3 的那个）自身的 KV 通常需要下一轮 Forward 才能计算。如果使用了 PagedAttention，逻辑上的回滚只是修改 Block Table 的映射，不一定涉及物理显存的搬运或释放。

为什么增加了计算量，反而可能更快？

#

前面讲清楚了 Lossless 的算法原理——最终分布严格等于 Target。但还有一个朴素的疑问：Speculative Decoding 多跑了一个 Draft Model，又多给 Target 喂了一串 token 做 Verification，总计算量明显增加了，凭什么反而更快？

用一个简化的数字感受一下。假设：

• 普通 Target Decode：10 ms / token
• Draft 一轮串行生成 4 个候选 token：共 3 ms
• Target Verification（一次 forward 处理 4 个候选位置）：12 ms
• 平均每轮有效输出 3 个 token（含可能的 Bonus Token）

原本生成 3 个 token 需要 30 ms（三次串行 Decode），现在 15 ms 搞定。

这能成立需要三个条件同时满足：

1. Draft 足够便宜。小模型参数量通常只有 Target 的 1/10 甚至更少，单步 Decode 很快。
2. Verification 处理多个 token 的成本不能线性爆炸。这一点和之前讨论的 Arithmetic Intensity 直接相关——在小 Batch 下，普通 Decode 每步做的基本是 GEMV（矩阵-向量乘法），计算量很小但权重照搬不误，Arithmetic Intensity 极低。Verification 一次 forward 处理 \(\gamma\) 个 token 时，这些位置共享同一批权重读取，操作更接近 GEMM（矩阵-矩阵乘法），权重搬运成本被摊薄了。所以 Verification 处理 4 个 token 的耗时不是普通 Decode 的 4 倍，而可能只有 1.2 倍。
3. Draft 的接受率足够高。如果小模型和大模型分布差异太大，猜的全被拒绝，一轮只推进一个 token，Draft 和 Verification 的额外开销就白花了。

不过要注意：提高 Arithmetic Intensity 不等于必然降低延迟。如果 Target 在高 Batch 下已经接近 compute-bound，权重复用的边际收益很小，额外的 Verification FLOPs 反而可能拖累整体吞吐。Speculative Decoding 的甜区是延迟敏感、Batch 较小的在线推理场景。

写在最后

#

回头看一下这两个看似毫不相关的问题：

它们共享同一个设计思想：从一个不完全匹配的分布中采样，通过概率阈值裁剪来保证正确性，再通过残差复用来避免浪费。 当然，具体机制有差异——Rand5 → Rand7 的残余是均匀随机空间里的整数偏移量，通过进制扩张组合新的均匀分布；Speculative Sampling 的残差是非均匀的概率质量差 \(\max(0, p - q)\)，需要归一化后做一次加权采样。但"拒绝后不丢弃已有的信息"这个骨架是相通的。

2017 年面京东被蓄水池采样卡了，回来补课的时候把 Reservoir Sampling 和 Rejection Sampling 的关系捋了一遍。打 ACM 那几年做概率期望题的直觉，准备校招时白板上推 Rand5 → Rand7 的手感，在十年后看 LLM Inference 的论文时突然接上了。Speculative Decoding 的 accept/reject/resample-from-residual 和这些经典题共享同一层设计直觉——不确定性本来就是有价值的资源，对它的高效利用方式，十年前和十年后没什么不同。

参考资料

#

• Leviathan et al., Fast Inference from Transformers via Speculative Decoding, ICML 2023
• Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling, 2023
• 蓄水池抽样算法概述，111qqz’s blog，2017