Block Verification Accelerates Speculative Decoding

作者/机构：Ziteng Sun∗, Uri Mendlovic∗, Yaniv Leviathan∗, Asaf Aharoni∗, Jae Hun Ro∗, Ahmad Beirami∗, Ananda Theertha Suresh∗ (均为 Google Research)

A1 主要贡献

大型语言模型（LLMs）通常通过自回归采样进行解码，生成k个词元（token）需要k次昂贵的串行模型评估。为了改善生成延迟，推测解码（speculative decoding）被提出，它允许一个LLM并发地生成多个词元。该方法的核心思想是：在每次迭代中，使用一个快速的“草稿模型”（drafter）来猜测接下来的一块（block）$\gamma$个词元。然后，大型的“目标模型”（target model）会并行地评估这$\gamma+1$个生成的前缀。为了保证最终输出的分布与目标模型完全一致，一部分生成的词元会被接受，其余的则被拒绝。被接受的词元被追加到前缀中，然后重复此过程。标准推测解码算法（Token Verification）通过一系列独立的、逐词元的拒绝步骤来验证草稿。具体来说，对于每个草稿词元$X_i$，它以一定的概率被接受，这个概率基于目标模型和草稿模型对该词元的概率分配。这个过程会持续到某个词元被拒绝为止。

$$\min \left\{1, \frac{\mathcal{M}_{b}\left(X_{i} \mid \boldsymbol{c}, X^{i-1}\right)}{\mathcal{M}_{s}\left(X_{i} \mid \boldsymbol{c}, X^{i-1}\right)}\right\},$$

然而，本文做出了一个令人惊讶的发现：标准的逐词元验证算法并非最优。本文的核心观察是，通过联合验证整个草稿词元块，而不是独立验证每个词元，可以在保持分布一致性保证的同时，增加接受的词元数量。基于此，本文提出了一个名为块验证（Block Verification）的新算法，它带来了以下优势：

• 使用简单：该算法是标准推测解码验证算法的即插即用替代品，不增加额外的计算或代码复杂性成本。
• 分布一致性：该方法不是近似，它保持了推测解码的无损（lossless）保证，即最终输出与目标模型的分布完全相同（由定理1证明）。
• 最优改进：在使用相同草稿模型的情况下，块验证的加速效果不劣于标准的词元验证。更重要的是，本文证明了块验证是一种最优的验证程序（由定理2证明）。

实验证明，在一系列任务和数据集上，块验证相比标准的词元验证算法，在“块效率”（每次迭代生成的词元期望数量）上有一致的7%-10%的提升，在实际的“墙钟时间”（wall-clock time）上则有5%-8%的加速。鉴于块验证不增加代码复杂性，保持了无损保证，性能不会下降且总能带来提升，它可以作为推测解码实现中的一个良好默认选项。这些改进仅来自于验证阶段，因此可以与其他改进草稿阶段的工作相结合。

A3 动机示例

标准词元验证的原理与局限性。标准词元验证算法之所以需要随机拒绝某些词元，是因为草稿模型（Ms）对这些词元的生成概率高于目标模型（Mb）。这种拒绝机制是保证最终生成序列遵循Mb分布所必需的。本文的核心观察点在于，如果将整个词元块作为一个整体来联合决定是否拒绝，而不是逐个词元地做决定，最终可以接受更多的词元。

一个简单的双词元语言模型示例。为了说明这一点，我们考虑一个只包含两个词元（A和B）的极简语言模型。假设目标模型Mb和草稿模型Ms都是上下文无关的，其概率分布如下：

$$\mathcal{M}_{b}(\mathrm{A})=1/3, \quad \mathcal{M}_{b}(\mathrm{B})=2/3, \quad \mathcal{M}_{s}(\mathrm{A})=2/3, \quad \mathcal{M}_{s}(\mathrm{B})=1/3.$$

在这种设定下，标准的词元验证会独立地处理每个草稿词元：如果词元是B，接受概率为1；如果是A，接受概率为1/2。当块大小$\gamma=2$时，根据【索引18，Fast inference from transformers via speculative decoding，2022，ICML，作者Leviathan等人】的分析，使用词元验证时，从Ms中接受的词元期望数量是 $1 - 1/3 + (1 - 1/3)^2 = 10/9$。

一个拥有完全信息的理想算法。现在假设存在一个理想算法，它能够获取两个词元序列的完整联合分布信息，如下所示：

$$\begin{aligned} \begin{aligned} \mathcal{M}_{b}(\mathrm{AA})=1 / 9, & \mathcal{M}_{b}(\mathrm{AB})=2 / 9, & \mathcal{M}_{b}(\mathrm{BA})=2 / 9, & \mathcal{M}_{b}(\mathrm{BB})=4 / 9, \\ \mathcal{M}_{s}(\mathrm{AA})=4 / 9, & \mathcal{M}_{s}(\mathrm{AB})=2 / 9, & \mathcal{M}_{s}(\mathrm{BA})=2 / 9, & \mathcal{M}_{s}(\mathrm{BB})=1 / 9. \end{aligned} \end{aligned}$$

利用这些信息，该算法可以执行更优的接受逻辑：当草稿序列是AB、BA或BB时，由于$M_b(X_1X_2) \ge M_s(X_1X_2)$，总是接受这两个词元；当草稿序列是AA时，以$M_b(AA)/M_s(AA) = 1/4$的概率接受它（并将样本校正为BB）。这样一来，从Ms中接受的词元期望数量变为：$2 \times (M_s(AB) + M_s(BA) + M_s(BB) + 1/4 \times M_s(AA)) = 12/9$，这明显大于词元验证的10/9。这个理想化的例子揭示了联合考虑草稿块分布的潜在好处。

在部分信息下的验证。在实际应用中，计算下一个词元块的完整联合分布是不可行的。我们能获取的只是沿着草稿块采样路径的条件分布，即$M_b(\cdot | c, X_i)$和$M_s(\cdot | c, X_i)$。重要的是，理想的拒绝逻辑并不需要访问完整的分布，但需要仔细地分配残差分布。本文提出的块验证正是这样做的。对于上述的玩具示例，我们提出的改进算法（块验证的一般算法的简化版）执行如下操作：当草稿词元$X_1X_2$是AB或BB时，接受概率为1；当$X_1X_2$是AA时，接受概率为1/4，否则拒绝整个块并仅将第一个词元校正为B（因为算法无法访问$M_b(\cdot | B)$）；当$X_1X_2$是BA时，总是接受B，然后以1/2的概率接受A（否则将第二个词元校正为B）。该算法只使用基于草稿采样路径的条件分布，因此在部分信息设置下是可行的。最终，接受的词元期望数量是$2\times(M_s(AB)+ M_s(BB))+ (1+ 1/2)\times M_s(BA)+ 1/4\times 2\times M_s(AA) = 11/9$，优于词元验证的10/9。

标准词元验证并非最优的证明。这个例子证明了以下引理：

引理1：推测解码的标准词元验证算法并非最优。

值得注意的是，块验证在此例中的期望接受词元数（11/9）高于标准词元验证（10/9），但仍低于拥有完全信息的理想算法（12/9）。在第4节中，我们将证明，在具有自然假设的部分信息情况下，块验证确实是最优的。

A2 方法细节

3 块验证

将各词元接受决策进行耦合。本节将上述直觉扩展为一个通用的块验证算法，该算法适用于标准的部分信息推测解码场景。其高层思想是将每个草稿词元的接受决策与其他草稿词元耦合起来。为此，算法会考虑不同长度的草稿子块，并独立地决定是否接受每个子块。最终被接受的草稿块是上述过程中被接受的最长的子块。每个子块的接受概率和残差分布都经过精心选择，以维持最终输出的分布保证，并实现最优的加速效果。

算法1与算法2的对比。算法2展示了块验证的简要实现，算法1则为标准词元验证以供比较。请注意，两种实现遵循相同的整体结构（差异已高亮显示）。

算法1 词元验证
输入：草稿块 $X^\gamma$；小型模型分布 $\forall i < \gamma, M_s(\cdot | c, X^i)$；大型模型分布 $\forall i \le \gamma, M_b(\cdot | c, X^i)$。
1. 采样 $\eta_1, \dots, \eta_\gamma \sim U(0, 1)$。
2. 设置 $\tau = 0$。
3. 对于 $i = 1, \dots, \gamma$：
4. 设置 $h_{i}^{\text{token}} = \min\left\{\frac{M_s(X_i|c,X^{i-1})}{M_b(X_i|c,X^{i-1})}, 1\right\}$。
5. 如果 $\eta_i \le h_{i}^{\text{token}}$：
6. 设置 $\tau = i$。
7. 否则：
8. break。
9. 结束如果。
10. 结束循环。
11. 如果 $\tau = \gamma$：
12. 从 $M_b(\cdot | c, X^\gamma)$ 中采样 $Y$。
13. 否则：
14. 从方程(3)中的 $p_{\text{res}}^{\text{token}}(\cdot | c, X^\tau)$ 中采样 $Y$。
15. 结束如果。
16. 返回 $X^\tau, Y$。

算法2 块验证
输入：草稿块 $X^\gamma$；小型模型分布 $\forall i < \gamma, M_s(\cdot | c, X^i)$；大型模型分布 $\forall i \le \gamma, M_b(\cdot | c, X^i)$。
1. 采样 $\eta_1, \dots, \eta_\gamma \sim U(0, 1)$。
2. 设置 $\tau = 0, p_0 = 1$。
3. 对于 $i = 1, \dots, \gamma$：
4. 设置 $p_i = \min\left\{p_{i-1}\frac{M_b(X_i|c,X^{i-1})}{M_s(X_i|c,X^{i-1})}, 1\right\}$。
5. 设置 $h_{i}^{\text{block}}$ 如方程(5)。
6. 如果 $\eta_i \le h_{i}^{\text{block}}$：
7. 设置 $\tau = i$。
8. 否则：
9. continue。
10. 结束如果。
11. 结束循环。
12. 如果 $\tau = \gamma$：
13. 从 $M_b(\cdot | c, X^\gamma)$ 中采样 $Y$。
14. 否则：
15. 从方程(4)中的 $p_{\text{res}}^{\text{block}}(\cdot | c, X^\tau)$ 中采样 $Y$。
16. 结束如果。
17. 返回 $X^\tau, Y$。

Residual distribution in Algorithm 1 (Line 15): $\forall x \in \mathcal{X}$,

$$p_{\text {res }}^{\text {token }}\left(x \mid \boldsymbol{c}, X^{i}\right)=\frac{\max \left\{\mathcal{M}_{b}\left(x \mid \boldsymbol{c}, X^{i}\right)-\mathcal{M}_{s}\left(x \mid \boldsymbol{c}, X^{i}\right), 0\right\}}{\sum_{x^{\prime} \in \mathcal{X}} \mathcal{M}_{b}\left(x^{\prime} \mid \boldsymbol{c}, X^{i}\right)-\mathcal{M}_{s}\left(x^{\prime} \mid \boldsymbol{c}, X^{i}\right), 0} .$$

Residual distribution in Algorithm 2 (Line 15): $\forall x \in \mathcal{X}$,

Acceptance probability in Algorithm 2 (Line 5): $h_{\gamma}^{\text {block }}=p_{\gamma}$, and when $i<\gamma$,

核心差异在于验证逻辑。一个重要的区别是，词元验证一旦有词元被拒绝就会停止（算法1第9行的break），而块验证总是对整个块进行操作。这等价于，在词元验证中，接受长度 $\tau = \arg \min\{i | \eta_i > h_{i}^{\text{token}}\} - 1$，而在块验证中，$\tau = \arg \max\{i | \eta_i \le h_{i}^{\text{block}}\}$。这个差异使得块验证倾向于接受比词元验证更长的子块，从而带来更高的块效率。图3绘制了两种算法在方程(2)引入的玩具模型上的接受长度经验互补累积分布函数（complementary CDF），以展示这一现象。

推测解码外循环保持不变。算法3展示了推测解码算法的外循环，它对于两种验证方法都保持不变。图2提供了额外的定义。附录A提供了简要的Python实现。由于改动非常简单，该算法可以轻松地在实际系统中实现，而不会引入额外的代码复杂性。

算法3 推测解码 (SPECDEC) (Leviathan et al., 2022)
输入：前缀c，目标模型Mb，草稿模型Ms。草稿长度γ。验证算法VERIFY。
1. 当 E.O.S ∉ ($X^\tau$, Y) 时循环：
2. 使用自回归采样从 $M_s(\cdot | c)$ 采样 $X_1, \dots, X_\gamma$，并保留每一步的条件概率 $M_s(\cdot | c, X^i)$，其中 $i=0, \dots, \gamma-1$。{获取草稿块}
3. 调用大模型Mb并并行计算条件概率 $M_b(\cdot | c, X^i)$，其中 $i=0, 1, \dots, \gamma$。{并行评分}
4. 通过草稿验证获得接受的词元 {草稿验证与校正} $X^\tau, Y = \text{VERIFY}(X^\gamma, \{M_s(\cdot | c, X^i)\}_{i=0}^{\gamma-1}, \{M_b(\cdot | c, X^i)\}_{i=0}^\gamma)$。
5. c ← c, $X^\tau$, Y。{将解码的词元添加到前缀}
6. 结束循环。

理论保证。使用块验证的推测解码保留了其输出的分布（定理1）。此外，块验证在推测解码外循环（算法3）中，在所有有效的草稿验证算法中实现了最优的加速，从而比标准词元验证有了严格的改进（定理2）。我们将在第4节中详细阐述算法2中参数选择的形式化陈述和直觉。

4 理论保证

本节提供块验证的正式理论保证。本节我们将展示块验证的正式理论保证，即它能产生正确的分布，并且在期望生成词元数方面是最优的。我们用 $M^*(· | c)$ 表示在模型M和上下文c下，直到生成过程结束的序列分布。

有效草稿验证算法的定义。一个草稿验证算法VERIFY，其输入为草稿块 $X^\gamma$、沿采样路径的小模型分布（即对所有 $i < \gamma$, $M_s(\cdot | c, X^i)$）和大模型分布（即对所有 $i \le \gamma$, $M_b(\cdot | c, X^i)$），输出为 $X^\gamma$ 的一个前缀 $X^\tau$（其中 $\tau \le \gamma$）以及一个额外的词元 $Y$。如果对于任意的c、模型Ms、Mb和块长度$\gamma$，使用该VERIFY算法的SPECDEC（算法3）的输出满足以下条件，则称VERIFY为一个有效的草稿验证算法：

$$\text{SPECDEC}(\boldsymbol{c}, \mathcal{M}_{b}, \mathcal{M}_{s}, \gamma, \text{VERIFY}) \sim_{\mathrm{p}} \mathcal{M}_{b}^{*}(\cdot \mid \boldsymbol{c})^{3},$$

即，输出的分布得以保留。

标准词元验证算法的有效性。例如，标准的词元验证算法就是一个有效的草稿验证算法（见【索引18，Fast inference from transformers via speculative decoding，2022，ICML，作者Leviathan等人】的附录A.1）。

定理1：块验证的有效性。我们现在声明以下定理：
定理1. 块验证（算法2）是一个有效的草稿验证算法。
换句话说，使用块验证的推测解码保留了输出序列的分布。

定理2：块验证的最优性。我们进一步声明，块验证对于所有有效的草稿验证算法是最优的。
定理2. 对于 $i > 0$，令 $N^{(i)}$ 为算法3中经过 $i$ 次迭代后解码的词元数量。对于定义1中的任何有效草稿验证算法VERIFY，我们有，对于所有的 $c, M_s, M_b, \gamma$ 和 $i$：

$$\mathbb{E}_{\text{BLOCKVERIFY}}[N(i)] \geq \mathbb{E}_{\text{VERIFY}}[N(i)],$$

其中随机性来自于草稿块的随机性和算法的随机性。特别地，

$$\mathbb{E}_{\text{BLOCKVERIFY}}[N(i)] \geq \mathbb{E}_{\text{TOKENVERIFY}}[N(i)].$$

换句话说，在所有有效的验证算法中，使用块验证的推测解码在固定迭代次数内期望解码的词元数量最多。由于块验证增加的计算开销可以忽略不计，这确立了块验证算法的整体最优性。特别地，块验证提供了比标准词元验证更大的加速。证明见附录B。

参数选择和理论保证的直觉。实现加速和分布匹配保证的关键量是 $p_i$。在附录B.1的引理3中，我们证明了 $p_i$ 对应于子块 $X^i$ 将被保留在最终输出中的概率。这是通过恰当地选择每步的接受概率来保证的，因为块验证保留了最长的被接受子块。接下来我们讨论 $p_i$ 如何贡献于分布匹配和最优性保证。

• 分布匹配保证 (定理 1)：首先，我们忽略 $p_i$ 递归定义中的最小值操作。在这种情况下，每个 $p_i$ 只是 $M_b(X^i | c)/M_s(X^i | c)$，这是实际 $p_i$ 的一个上界。如引理3所示，对于任何 $X^i$，它在被接受块中的概率是 $M_s(X^i)p_i(X^i)$。由于草稿块 $X^i$ 是以概率 $M_s(X^i | c)$ 生成的，这保证了通过从草稿中接受它来得到 $X^i$ 的概率最多为 $M_b(X^i | c)$，因此算法不会过多地接受 $X^i$。剩下的部分是选择一个合适的残差分布 $p_{\text{res}}^{\text{block}}$，使得下一个词元的分布遵循 $M_b(\cdot | X^i)$。注意，对于任何可能的下一个词元 $x$，$(X^i, x)$ 也可能通过在 $X_{i+1} = x$ 时接受 $X^{i+1}$ 来获得，其概率为 $M_s(X^i, x)p_{i+1}(X^i, x)$，这个概率应该被减去以获得 $(X^i, x)$ 上的残差质量。这导致了在适当归一化后，方程(4)中 $p_{\text{res}}^{\text{block}}$ 的选择。
• 最优性保证 (定理 2)：对于最优性保证，主要证明是表明对于草稿块中的任何前缀 $X^i$，$p_i(X^i)$ 是一个有效验证算法可以接受 $X^i$ 的最大概率，这在引理4中陈述。这意味着在一次迭代中，块验证在期望上接受最多的词元，多轮迭代的情况可以通过归纳论证得到。为了理解为什么 $p_i(X^i)$ 是接受概率的上界，我们表明这是必要的，以保证对于任何可以从多个草稿采样路径获得的前缀 $X^i$，后续词元的分布总是相同的 $M_b(\cdot | c, X^i)$。这使得块验证可以作为词元验证在推测解码外循环（算法3）中的即插即用替代品。

关于贪心块验证的说明。最后，我们注意到最优性保证适用于所有可以按原样在算法3中使用的验证算法。具体来说，存在一些验证过程，它们强制解码逻辑依赖于先前的接受/拒绝决策，这些过程在一次迭代中平均能产生更多被接受的词元。然而，这会影响后续迭代的解码速度。在附录C中，我们介绍了这样一种算法，并称之为贪心块验证。我们通过实验观察到块验证始终优于它，因此我们主要将其作为一个理论结果包含进来。

A4 实验环境

实验中使用了PALM-2模型和Vicuna模型。

• 模型架构

  • PALM-2系列：目标模型（大模型）为PALM-2-S，草稿模型（小模型）为PALM-2-XXS或PALM-2-XXXS。模型尺寸排序为 PALM-2-XXXS < PALM-2-XXS < PALM-2-S。
  • Vicuna系列：目标模型为Vicuna-7B-v1.3，草稿模型为Vicuna-68M。

• 数据集与任务

  • PALM-2实验：评估了多种任务和数据集，包括：

    • 语言建模: one-billion language benchmark (LM1B)
    • 通用对话: ChatGPT prompts sourced from LearnGPT (GPT Prompt)
    • 推理问答: WebQA, PIQA (physical commonsense)
    • 对话: ShareGPT
    • 摘要: XSum
    • 数学: GSM8K
    • 翻译: WMT DeEn (德语到英语)

  • Vicuna实验：使用了Spec-Bench基准测试，涵盖多轮对话、检索增强生成、摘要、翻译、问答和数学推理等多种生成子任务。

• 硬件配置

  • Vicuna实验：使用单块NVIDIA H100 GPU进行。
  • PALM-2实验：硬件细节未在文中明确说明。

• 软件与实验配置

  • PALM-2实验：

    • 在每个数据集上解码前1000个提示。
    • 最大输入提示长度为512，最大解码输出长度为128。
    • 除附录D.1中的实验外，所有实验的批处理大小（batch size）均为1。
    • 温度（temperature）设置为1.0。

  • Vicuna实验：

    • 批处理大小为1。
    • 最大生成长度为1024。
    • 研究了不同温度的影响，测试了 {0.2, 0.6, 1.0} 三种设置。
    • 草稿长度 $\gamma$ 固定为8。

A4 实验结果

本文实验主要对比了块验证（block verification）和词元验证（token verification）。尽管近期有工作通过使用多个草稿块来提升块效率，但这些方法增加了大模型的验证计算量，在进行查询批处理（query batching）时并不理想。本文在附录中展示了，在拥有单个草稿块的大批量处理场景下，块验证方法优于这些多草稿方法。

6.1 在PALM-2模型上的实验结果

实验对比发现，无论是在理想化的块效率指标还是在真实的墙钟时间上，块验证在各种设置下都展现出虽小但一致的性能提升。

• 块效率与墙钟时间提升：

  • 块效率 (Block efficiency)：衡量在理想情况下（忽略草稿模型评估时间且有足够并行计算能力）每次调用目标模型能解码的平均词元数。在使用PALM-2-XXS作为草稿模型，草稿长度$\gamma=8$时，块验证带来的块效率提升在7.00%到10.06%之间，平均为8.30%。
  • 墙钟时间 (Wall clock time)：衡量包含所有实际开销的真实加速效果。在相同设置下，墙钟时间加速在5.36%到8.14%之间，平均为6.49%。具体数据见表1。

  表1：词元验证（TOKENV）与块验证（BLOCKV）在γ=8时的加速比较。目标模型为PALM-2-S，草稿模型为PALM-2-XXS，在多个数据集和任务上进行。数据为3次不同种子运行的平均值和标准差。
  

• 草稿长度 $\gamma$ 的影响：

  • 实验对比了$\gamma = 4, 6, 8$的情况，均观察到一致的提升。
  • 如图5所示，随着$\gamma$的增加，块验证相对于词元验证的相对提升也随之增加。这符合直觉，因为更大的$\gamma$为块验证提供了更多协调接受规则的空间。
  • 如图4所示，虽然块效率随$\gamma$增加而持续提升，但墙钟时间加速在$\gamma=4$或$\gamma=6$时达到峰值，因为更大的$\gamma$会增加验证阶段的计算成本。

• 草稿模型质量的影响：

  • 实验对比了两个不同质量的草稿模型：PALM-2-XXS（更大，质量更好）和PALM-2-XXXS。
  • 两种草稿模型下，块验证都带来了提升。
  • 如图5所示，在使用质量更好的PALM-2-XXS时，块效率的相对提升更大。这表明块验证的改进可以与提升草稿模型质量的改进相结合，且在更好的草稿模型下效果可能更显著。

  图4：不同γ下，词元验证（TOKENV）和块验证（BLOCKV）在所有数据集上的平均块效率（BE）和墙钟加速（WS）。大模型为PALM-2-S，草稿模型为PALM-2-XXS（XXS）或PALM-2-XXXS（XXXS）。
  

  

6.2 在Spec-Bench上使用Vicuna模型的实验结果

在Spec-Bench基准上使用Vicuna模型的实验进一步验证了该方法的普适性。

• 不同温度下的性能：

  • 实验在温度 $T \in {0.2, 0.6, 1.0}$ 下进行，$\gamma$固定为8。
  • 如表2所示，块验证在所有非零温度设置下都获得了一致的提升（块效率最高提升8.7%，墙钟时间最高提升6.7%）。
  • 温度越高，块验证的额外改进越显著。这与算法的内在逻辑一致，因为它通过协调不同词元位置上接受决策的随机性来提升块效率。
  • 当温度为0（即贪心解码）时，块验证退化为词元验证，不提供额外加速。

  表2：在Spec-Bench上，词元验证（TOKENV）与块验证（BLOCKV）在不同温度T下的加速比较。目标模型为Vicuna-7B-v1.3，草稿模型为Vicuna-68M，$\gamma=8$。
  

A7 相关工作

并行解码背景。本文的工作改进了推测解码（speculative decoding）【索引18，Fast inference from transformers via speculative decoding，2022，ICML，作者Leviathan等人】，这是一个并发解码多个词元的框架。早期的工作包括【索引26，Blockwise parallel decoding for deep autoregressive models，2018，NeurIPS，作者Stern等人】提出的“草稿与验证”（draft and verify），该方法针对贪心解码（零温度）情况，并行地独立预测和解码多个词元。推测解码后来也由【索引6，Accelerating large language model decoding with speculative sampling，2023a，arXiv，作者Chen等人】提出。

单草稿改进工作。许多工作致力于在不使用多个解码草稿的情况下改进推测解码。表3列出了一系列在“草稿与验证”框架下的工作，并对其草稿和验证算法进行了分类。【索引41，Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding，2024，ACL，作者Xia等人】对此进行了全面的研究。在单草稿情况下，一些工作专注于改进推测解码的草稿阶段（例如，【索引15，Rest: Retrieval-based speculative decoding，2023，arXiv，作者He等人】；【索引7，Cascade speculative drafting for even faster llm inference，2023b，arXiv，作者Chen等人】；【索引27，Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding，2024，arXiv，作者Sun等人】；【索引36，Distillspec: Improving speculative decoding via knowledge distillation，2023，arXiv，作者Zhou等人】；【索引20，Online speculative decoding，2023，arXiv，作者Liu等人】；【索引14，Better & faster large language models via multi-token prediction，2024，arXiv，作者Gloeckle等人】；【索引35，Draft & verify: Lossless large language model acceleration via self-speculative decoding，2024，arXiv，作者Zhang等人】；【索引12，Layerskip: Enabling early exit inference and self-speculative decoding，2024，arXiv，作者Elhoushi等人】）。然而，这些算法都使用了相同的词元验证算法。本文提出的块验证算法可以与表3中的草稿技术结合使用，从而产生叠加的增益。对这些情况下块验证改进的系统性研究将留作未来工作。

表3：基于“草稿与验证”框架的近期工作。 温度0指贪心解码，非零温度指采样。

其他对验证步骤的改进工作。据我们所知，唯一另一项改进推测解码验证步骤的工作是【索引16，Accelerated speculative sampling based on tree monte carlo，2024，ICML，作者Hu和Huang】的独立工作，他们使用树蒙特卡洛方法在单草稿情况下改进推测解码，并证明了其算法优于词元验证。相比之下，我们证明了我们的算法在所有有效的验证算法中（包括他们的算法）实现了最优的加速。我们的算法还只需要对原始词元验证算法进行最小的改动，使其在实践中易于实现和适配。

多草稿扩展工作。最近，推测解码被扩展到多草稿情况（【索引28，Spectr: Fast speculative decoding via optimal transport，2023，arXiv，作者Sun等人】；【索引21，Specinfer: Accelerating generative llm serving with speculative inference and token tree verification，2023，arXiv，作者Miao等人】），并且为多草稿场景提出了新的验证算法（【索引19，Eagle: Speculative sampling requires rethinking feature uncertainty，2024，arXiv，作者Li等人】；【索引8，Sequoia: Scalable, robust, and hardware-aware speculative decoding，2024，arXiv，作者Chen等人】）。虽然增加草稿序列数量已被证明可以提高整体加速效果，但这是以更多计算为代价的。在附录D.1中，我们展示了在推理不那么受内存限制的大批量处理设置中，我们的方法在不增加草稿块数量的情况下优于这些方法。在所有这些工作中，验证算法都是词元验证过程的泛化。将块验证扩展到多样本情况是一个有趣的未来方向。

A5 结论

本文证明了推测解码所使用的标准词元验证算法并非最优。在此基础上，我们提出了一个新的验证算法——块验证，并证明了它是一个最优的验证算法。我们还通过实验证明，块验证在一系列任务中始终优于词元验证。尽管其理论证明有些复杂，但块验证的实际实现并不比标准算法更复杂（见附录A）。由于我们提出的算法性能只会更好，绝不会比标准词元验证算法差（见定理2），因此它可以作为推测解码实现中的一个良好默认选项。

A6 附录

A Python实现

本节我们提供了块验证（算法2）的一个Python简要实现。请注意，这仅用于说明目的，不适合实际使用。

设 V = |X| 为词汇表的大小。算法的输入是：

• ps: 一个(γ + 1) × V的numpy数组，包含大模型$M_b(\cdot | c, X^i)$的分布；
• qs: 一个γ × V的numpy数组，包含草稿模型$M_s(\cdot | c, X^i)$的分布；
• drafts: 一个长度为γ的numpy数组，包含草稿词元$X^\gamma$的id；

def block_verification(ps: np.ndarray, qs: np.ndarray, drafts: np.ndarray) -> list[int]:
    draft_length, vocab_size = qs.shape
    qs.resize((draft_length + 1, vocab_size))  # 追加一个零向量
    token_sequence = []  # 将包含我们返回的词元序列
    accept_probability = 1.0  # 每个子块的接受概率
    probability_ratios = ps / qs
    # 添加一个词元以表示拒绝该序列
    vocab_plus_one = np.arange(vocab_size + 1)
    for token_index, token_value in enumerate(drafts):  # xs应为drafts
        # 未归一化的残差概率
        sampling_weights = np.zeros(vocab_size + 1) # 初始化sampling_weights
        sampling_weights[:vocab_size] = np.maximum(0, ps[token_index] * accept_probability - qs[token_index])
        # 序列被拒绝的未归一化概率
        sampling_weights[vocab_size] = 1 - accept_probability
        sampling_weights /= np.sum(sampling_weights)
        chosen_token = np.random.choice(vocab_plus_one, p=sampling_weights)
        # 更新序列
        if chosen_token < vocab_size:
            token_sequence = list(drafts[:token_index]) + [chosen_token] # drafts[:token_index]
        # 更新接受概率
        accept_probability = min(1, probability_ratios[token_index, token_value] * accept_probability)
    return token_sequence

作为参考，这里是词元验证算法（算法1）的简要实现：

def token_verification(ps: np.ndarray, qs: np.ndarray, drafts: np.ndarray) -> list[int]:
    draft_length, vocab_size = qs.shape
    qs.resize((draft_length + 1, vocab_size))  # 追加一个零向量
    token_sequence = []  # 将包含我们返回的词元序列
    probability_ratios = ps / qs
    token_index = 0
    vocab_range = np.arange(vocab_size)
    for token_value in drafts: # xs应为drafts
        accept_probability = probability_ratios[token_index, token_value]
        if not np.isfinite(accept_probability) or np.random.random() > accept_probability:
            # 拒绝
            break
        token_index += 1
        token_sequence.append(token_value)
    # 计算残差分布
    sampling_weights = np.maximum(0, ps[token_index] - qs[token_index])
    sampling_weights /= np.sum(sampling_weights)

    token_sequence.append(np.random.choice(vocab_range, p=sampling_weights))

    return token_sequence

B 形式化证明

引理2。对于所有的c, Ms, Mb, $\gamma$，让 $X^\gamma$ 从 $M_s^\gamma(\cdot | c)$ 生成，并且

$$X^\tau, Y = \text{VERIFY}(X^\gamma, \{\mathcal{M}_s(\cdot \mid \boldsymbol{c}, X^i)\}_{i=0}^{\gamma-1}, \{\mathcal{M}_b(\cdot \mid \boldsymbol{c}, X^i)\}_{i=0}^\gamma).$$

让 $Z^{\gamma-\tau}$ 从 $M_b^{\gamma-\tau}(\cdot | c, X^\tau, Y)$ 生成。VERIFY是有效的草稿验证算法（定义1）当且仅当对于所有的c, Ms, Mb, $\gamma$，

$$X^\tau, Y, Z^{\gamma-\tau} \sim_{\mathrm{p}} \mathcal{M}_b^{\gamma+1}(\cdot \mid \boldsymbol{c}).$$

B.1 定理1的证明

根据引理2，只需证明块验证满足方程(7)。为简化，我们常将序列 $(X^\tau, Y, Z^{\gamma-\tau})$ 记为 $O^{\gamma+1}$。注意 $O_{\gamma+1} \sim M_b(\cdot | c, O_\gamma)$ 总是成立的。因此，只需证明以下即可：

$$\forall \ell \leq \gamma, \forall x^\ell \in \mathcal{X}^\ell, \quad \operatorname{Pr}\left(O^\ell=x^\ell\right)=\mathcal{M}_b^\ell\left(x^\ell \mid \boldsymbol{c}\right),$$

我们首先陈述以下关于块验证接受的词元数量分布的引理。

引理3。设 $X^\gamma \sim M_s^\gamma(\cdot|c)$，以及

$$X^{\tau}, Y = \text{BLOCKVERIFY}(X^{\gamma}, \{\mathcal{M}_{s}(\cdot \mid \boldsymbol{c}, X^{i})\}_{i=0}^{\gamma-1}, \{\mathcal{M}_{b}(\cdot \mid \boldsymbol{c}, X^{i})\}_{i=0}^{\gamma})$$

那么我们有，对于所有的 $i \le \gamma$ 和 $x^i \in \mathcal{X}^i$，

$$\operatorname{Pr}(\tau \geq i \mid X^{i}=x^{i})=p_{i}(x^{i}).$$

我们基于引理3证明定理1，并将该引理的证明推迟到附录B.3。我们通过对时间索引 $l$ 的归纳来证明方程(8)。当 $l = 1$ 时，通过计算可以证明 $Pr(O_1=x) = M_b(x|c)$。假设方程(8)对于所有小于$\gamma$的$l$都成立。对于$l = l+1$，通过将 $Pr(O^{l+1}=x^{l+1})$ 分解为 $\tau \ge l+1$, $\tau=l$ 和 $\tau < l$ 三种情况，并分别计算其概率，可以证明 $Pr(O^{l+1}=x^{l+1}) = M_b(x^{l+1}|c)$，从而完成归纳步骤，证明了方程(8)和定理1。

B.2 定理2的证明

我们首先陈述以下引理，它与引理3结合表明，在一次迭代中，在所有有效的草稿验证算法中，块验证以最高概率接受每个子序列。

引理4。对于满足引理2中约束的草稿验证算法，我们有，对于所有的 $i \le \gamma$ 和 $x^i \in \mathcal{X}^i$，

$$\operatorname{Pr}(\tau \geq i \mid X^{i}=x^{i}) \leq p_{i}(x^{i}).$$

我们将该引理的证明推迟到附录B.4，首先基于该引理证明定理2。
证明的核心是证明以下引理5：

引理5。对于所有满足引理2约束的草稿验证算法，我们有，对于所有的c和输出 $x^ \in \mathcal{X}^$

$$\Pr_{\text{VERIFY}}(O=x^*, N(i) \geq \ell \mid \boldsymbol{c}) \leq \Pr_{\text{BLOCKVERIFY}}(O=x^*, N(i) \geq \ell \mid \boldsymbol{c})$$

我们通过对迭代次数 $i$ 进行归纳来证明此引理。当 $i=1$ 时，利用引理3和4可以证明该不等式成立。假设该引理对所有直到 $i$ 的迭代都成立，对于第 $i+1$ 次迭代，我们通过分析在第 $i$ 次迭代后生成了 $\eta$ 个词元的条件期望，并利用一个关于随机占优和增函数的引理（【索引23，Admissibility and measurable utility functions，1962，The Review of Economic Studies，作者Quirk和Saposnik】），可以证明不等式对于 $i+1$ 也成立。这就完成了定理2的证明。

B.3 引理3的证明

证明通过反向归纳法进行。当 $i = \gamma$ 时，根据定义 $Pr(\tau \ge \gamma | X^\gamma=x^\gamma) = h_\gamma^{block} = p_\gamma(x^\gamma)$。假设该陈述对 $i \ge l$ 成立。当 $i = l-1$ 时，通过将 $Pr(\tau \ge l-1 | X^{l-1}=x^{l-1})$ 分解为在 $X^l$ 上的求和，并利用归纳假设和 $h_{l-1}^{block}$ 的定义，经过一系列代数运算，可以证明 $Pr(\tau \ge l-1 | X^{l-1}=x^{l-1}) = p_{l-1}(x^{l-1})$。因此，该引理通过归纳法成立。

B.4 引理4的证明

证明分两种情况。第一种情况是，对于所有的 $i<l$，都满足 $p_{i-1}(x^{i-1})M_b(x_i|c, x^{i-1}) \le M_s(x_i|c, x^{i-1})$。在这种情况下，可以推导出 $p_l(x^l) = M_b(x^l|c) / M_s(x^l|c)$，并利用 $Pr(O^l=x^l, \tau \ge l) \le M_b(x^l|c)$ 来证明 $Pr(\tau \ge l | X^l=x^l) \le p_l(x^l)$。第二种情况是，存在某个 $i$ 使得 $p_{i-1}(x^{i-1})M_b(x_i|c, x^{i-1}) > M_s(x_i|c, x^{i-1})$。在这种情况下，可以证明 $p_i(x^i)=1$ 且 $M_b(x^i|c) / M_s(x^i|c) > 1$。通过利用有效验证算法的约束条件，可以推导出如果 $Pr(\tau \ge l | X^l=x^l) > p_l(x^l)$ 将会导致矛盾。因此，引理成立。

C 贪心块验证

本节中，我们展示了通过修改算法3中的推测解码框架，允许解码逻辑依赖于先前的接受/拒绝决策，可以在一次迭代中接受比块验证（算法2）更多的词元。然而，如表4所示，由此产生的算法——贪心块验证，并没有比块验证更好。我们将其描述和分析作为一个理论结果。

我们首先介绍算法4，然后讨论为保持分布一致性保证所做的必要修改。贪心块验证算法的过程与块验证类似，但在接受概率和残差分布的设置上有所不同。与块验证中的 $p_i$ 相比，贪心块验证中的 $\tilde{p}_i$ 的递归定义没有最小值项，因此总是 $p_i$ 的上界，这导致了更高的接受概率（定理3）。然而，算法4不能直接用于推测解码的迭代实现，因为它会违反分布一致性保证。为了解决这个问题，我们引入了一个分布修改算法（算法5），它可以与算法4一起使用以维持保证。由此产生的完整算法是算法6。

算法4 贪心块验证
(伪代码省略，其核心是使用没有min(1, ...)项的$\tilde{p}_i$来计算接受概率，并使用相应的残差分布。)

$$p_{\mathrm{res}}^{\mathrm{greedy}}(x \mid \boldsymbol{c}, X^{i}) = \frac{\max \{ \tilde{p}_{i} \cdot \mathcal{M}_{b}(x \mid \boldsymbol{c}, X^{i}) - \mathcal{M}_{s}(x \mid \boldsymbol{c}, X^{i}), 0 \}}{\sum_{x^{\prime} \in \mathcal{X}} \max \{ \tilde{p}_{i} \cdot \mathcal{M}_{b}(x^{\prime} \mid \boldsymbol{c}, X^{i}) - \mathcal{M}_{s}(x^{\prime} \mid \boldsymbol{c}, X^{i}), 0 \}}.$$

算法5 分布修改
(伪代码省略，其核心是在贪心块验证拒绝了某些草稿后，修改后续词元的目标分布$M_b$，以补偿之前可能产生的分布偏差。)

算法6 使用贪心块验证的推测解码
(伪代码省略，其外循环与算法3类似，但在每次迭代后增加了调用算法5来修改目标分布的步骤。)

C.1 与块验证的比较

定理3 (非正式)：在具有相同模型和草稿长度的一次草稿迭代中，贪心块验证解码的词元数至少与块验证一样多。
然而，由于分布修改步骤，目标分布可能会在第一次迭代后改变，这可能会影响后续迭代中接受的词元期望数。理论上不清楚哪种方法占优。

经验比较：我们对贪心块验证进行了与第6节相同的实验。表4列出了当使用PALM-2-XXS作为草稿模型且$\gamma=8$时的块效率比较。可以看出，虽然贪心块验证仍然比词元验证有所改进，但改进幅度不如块验证显著。因此，我们推荐使用块验证而非贪心版本。

表4：词元验证、块验证和贪心块验证在γ=8时的块效率比较。

D 额外的实验结果

D.1 与使用多个草稿的推测解码的比较

最近的工作通过使用多个草稿块来提高块效率，但这增加了大模型的验证计算量。在LLM服务系统中，查询批处理是常用技术，这种情况下推理受内存限制较小，没有足够的额外并行计算能力来评估增加的草稿而不降低延迟。

我们通过实验比较了块验证与SpecTr和SpecInfer在查询批处理（批大小 B=8）下的性能。我们使用PALM-2-XXS作为草稿模型，$\gamma=8$，SpecTr和SpecInfer的草稿块数设为2。结果如表5所示。我们观察到，尽管SpecTr和SpecInfer能实现更高的块效率，但由于评估更多候选者所需的计算量增加，我们的方法实现了更好的墙钟时间加速，证明了我们的方法在常见的实际设置中的优势。

表5：γ=8, B=8。使用PALM-2-XXS作为草稿模型时，词元验证（TOKENV）、块验证（BLOCKV）、SpecTr和SpecInfer在不同数据集和任务上的加速比较。

D.2 其他参数设置下的详细结果

本节提供了与第6节描述的实验集相同的实验结果，但使用了不同的块长度（$\gamma = 4, 6, 8$）和不同的草稿模型（PALM-2-XXS和PALM-2-XXXS）。这些结果进一步证实了块验证在各种参数配置下都能提供一致的性能提升。

表6：当γ=4且草稿模型为PALM-2-XXS时，词元验证（TOKENV）与块验证（BLOCKV）的加速比较。

表7：当γ=6且草稿模型为PALM-2-XXS时，词元验证（TOKENV）与块验证（BLOCKV）的加速比较。

表8：当γ=4且草稿模型为PALM-2-XXXS时，词元验证（TOKENV）与块验证（BLOCKV）的加速比较。

表9：当γ=6且草稿模型为PALM-2-XXXS时，词元验证（TOKENV）与块验证（BLOCKV）的加速比较。

表10：当γ=8且草稿模型为PALM-2-XXXS时，词元验证（TOKENV）与块验证（BLOCKV）的加速比较。