Multi-Draft Speculative Sampling: Canonical Decomposition and Theoretical Limits

文章标题：多草稿推测采样：规范分解与理论极限
作者/机构：Ashish Khisti (Qualcomm AI Research & University of Toronto), M.Reza Ebrahimi (Qualcomm AI Research), Hassan Dbouk (Qualcomm AI Research), Arash Behboodi (Qualcomm AI Research), Roland Memisevic (Qualcomm AI Research), Christos Louizos (Qualcomm AI Research)

A1 主要贡献

本文研究了多草稿推测采样，其中提议序列由不同的草稿模型独立采样。在每一步中，一个令牌级别的草稿选择方案将一系列有效令牌作为输入，并产生一个其分布与目标模型相匹配的输出令牌。先前的工作已经证明，最大化接受输入令牌之一的概率的最优方案可以被看作是一个线性规划问题的解。本文的主要贡献如下：

• 提出规范分解：本文重新审视了【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】中引入的最优传输框架，并提出了一个规范分解。研究证明，最优接受概率可以通过一个两步方案实现：第一步使用一种重要性采样来从输入集合中选择一个令牌；第二步则使用所选令牌和目标分布进行推测采样。
• K=2时的理论分析：对于 K=2 个相同草稿模型的情况，本文为最优接受概率建立了一个解析表达式。在此设定下，还建立了接受概率等于1的充要条件。此外，还讨论了超过2个草稿情况下的数值证据。
• 提出新的选择方案：受理论分析的启发，本文提出了一类新的基于加权重要性采样的令牌选择方案。为了实现更快的实施，本文提出了启发式方法，这些方法在降低计算复杂度的同时，对接受概率的惩罚也较小。
• 实验验证：本文展示了使用 OPT 模型在多种任务上的实验结果。通过将提出的方案与基线方案进行性能比较，证明了在块效率和令牌生成速率方面取得了持续的改进。

A3 背景知识和相关工作

自回归采样的挑战与加速方法。大型语言模型（LLMs）的自回归采样本质上是顺序的，并且受内存带宽限制【27，Fast transformer decoding: One write-head is all you need，2019，arXiv】。文献中提出了多种加速LLM推理的方法，包括模型压缩技术（如量化【14，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv】、【3，Quantizable transformers: Removing outliers by helping attention heads do nothing，2024，NeurIPS】和稀疏化【19，Sparse is enough in scaling transformers，2021，NeurIPS】、【13，Sparsegpt: Massive language models can be accurately pruned in one-shot，2023，ICML】），这些技术以一定的解码质量下降为代价来降低LLM的整体复杂性。

无损加速方法：推测解码。作为一种有前景且正交的替代方案，推测解码【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】、【21，Fast inference from transformers via speculative decoding，2023，ICML】、【28，Blockwise parallel decoding for deep autoregressive models，2018，NeurIPS】实现了无损的LLM推理加速。早期的工作通过增强基础LLM【28，Blockwise parallel decoding for deep autoregressive models，2018，NeurIPS】或采用激进解码【16，Lossless acceleration for seq2seq generation with aggressive decoding，2022，arXiv】来草拟和预测多个令牌。然而，LLM文本生成通常需要从生成的logits中以非零温度进行采样。为此，推测解码【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】、【21，Fast inference from transformers via speculative decoding，2023，ICML】被提出。在该方法中，自回归采样被委托给一个较小的语言模型（草稿模型）来生成多个候选令牌。然后，LLM（目标模型）用于并行地对草稿中的所有令牌进行评分，并通过一系列令牌级拒绝采样来验证草稿令牌。推测解码保证最终序列遵循与目标模型相同的分布。推测方法的性能高度依赖于草稿模型的选择。【35，Distillspec: Improving speculative decoding via knowledge distillation，2023，arXiv】使用知识蒸馏【17，Distilling the knowledge in a neural network，2015，arXiv】来更好地对齐草稿和目标模型，从而获得更高的令牌接受率。

多草稿推测解码的最新进展。最近，【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】、【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】、【20，Recursive speculative decoding: Accelerating llm inference via sampling without replacement，2024，arXiv】等工作将推测解码扩展到多草稿设置，其中草稿模型在每个时间步生成多个令牌序列。具体来说，【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】将令牌级草稿选择问题表述为带隶属成本的离散最优传输问题，并提出了SpecTr，一种新的解码算法，允许草稿中每个令牌有多个候选。类似地，【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】和【20，Recursive speculative decoding: Accelerating llm inference via sampling without replacement，2024，arXiv】研究了基于令牌树的结构来改进草稿序列，以及不同于【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】的令牌级选择方法。此外，【5，Medusa: Simple llm inference acceleration framework with multiple decoding heads，2024，arXiv】提出不使用专门的草稿模型，而是通过增加额外的解码头来增强目标模型，这些解码头可以并行地草拟多个令牌。【29，Optimal block-level draft verification for accelerating speculative decoding，2024a，arXiv】则在单草稿设置中研究了块级验证，将其视为一个块级最优传输问题，并提出了一个计算高效的算法来最优地解决该问题，报告了相比于先前的令牌级验证【21，Fast inference from transformers via speculative decoding，2023，ICML】的速度提升。

A2 方法细节

3 令牌级最优草稿选择：理论分析

问题设定。本文专注于【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】中引入的令牌级最优草稿选择框架。假设词汇表为 $Ω = \{1, 2, . . . , n\}$，在给定步骤 $t$，考虑 $K$ 个有效令牌 $X_1, . . . , X_K$。每个令牌都以独立同分布（i.i.d.）的方式从分布 $p(\cdot)$ 中生成，该分布由底层的草稿模型和上下文序列 $u_t \in Ω^t$ 决定，即对于每个 $y \in Ω$，有 $p(y) = M_s(y|u_t)$，其中 $M_s$ 是由小（草稿）模型生成的分布。类似地，令 $q(\cdot)$ 为与大模型相关的分布，即 $q(y) = M_b(y|u_t)$，其中 $M_b$ 是由大模型生成的分布。在整个分析中，由于序列 $u_t$ 是固定的且对两个模型都通用，因此在讨论 $p(\cdot)$ 和 $q(\cdot)$ 时不显式指明。

令牌级选择规则（TLSR）的定义与优化目标。给定一个由 $K$ 个候选令牌组成的输入 $S \sim \prod_{i=1}^{K} p(X_i)$，一个令牌级选择规则（TLSR）是关于 $Ω$ 的条件分布 $P(\cdot|S)$。一个有效的TLSR必须满足约束：对于每个 $z \in Ω$，$\sum_S P(z|S)p(S) = q(z)$。TLSR的一个自然优化指标是接受其中一个令牌的概率，即如果 $Z \sim P(\cdot|S)$ 是TLSR的输出，我们希望最大化 $Pr(Z \in S)$。

问题1：最优令牌级选择规则。给定分布 $p(\cdot)$ 和 $q(\cdot)$，找到一个有效的TLSR来最大化接受概率：$P(acc) = Pr(Z \in S)$，并令 $P^⋆(acc)$ 为其最优值。问题1在【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】中被研究，并被证明是最优传输的一个实例，可以被构建为一个线性规划问题。作者利用此框架在单草稿（即$K=1$）情况下确立了推测采样【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】、【21，Fast inference from transformers via speculative decoding，2023，ICML】的最优性。对于$K>1$的情况，作者为$P^⋆(acc)$建立了一个信息论上的上界。本文重新审视了问题1，并对最优解的结构提出了新的见解，特别是发现多草稿情况下的最优解与重要性采样【31，Importance sampling: a review，2010，Wiley Interdisciplinary Reviews: Computational Statistics】有自然的联系。对于$K=2$个草稿的情况，本文精确地刻画了$P^⋆(acc)$，并给出了$p(\cdot)$和$q(\cdot)$使$P^⋆(acc)$等于1的充要条件。

重要性加权采样的定义。本文首先定义了一族称为重要性加权采样的方案。一个重要性加权采样方案以候选令牌集 $S = \{X_1, . . . , X_K\}$ 为输入，并输出一个令牌 $Y_I \in S$，其条件分布定义为：

$$\begin{aligned} \operatorname{Pr}(Y_I = y|X_{1:K} = x_{1:K}) = \begin{cases} \beta_y(x_1, \dots, x_K), & y \in \{x_1, \dots, x_K\} \\ 0, & y \notin \{x_1, \dots, x_K\} \end{cases} \end{aligned}$$

其中，对于每个 $x_{1:K} \in Ω^K$，$\sum_{y \in \Omega} \beta_y(x_1, \ldots, x_K) = 1$，且 $0 \leq \beta_y(x_1, \ldots, x_K) \leq 1$。

重要性加权采样的说明。值得注意的是，可以不考虑公式(1)中选定令牌值的概率，而是考虑选择索引 $i$（在$1, \ldots, K$之间）的概率，即 $Pr(I = i|X_{1:K} = x_{1:K})$。这样的分布通过对所有 $x_i=y$ 的索引求和可以映射到公式(1)的形式。本文指出，公式(1)的形式在后续分析中更为方便。同时，经典的【31，Importance sampling: a review，2010，Wiley Interdisciplinary Reviews: Computational Statistics】重要性采样方案对应于 $Pr(I = i|X_{1:K} = x_{1:K}) \propto q(x_i)/p(x_i)$ 的情况。然而，定义1中的方案族不限于这种选择，而是将 $\beta_y(x_1, \ldots, x_K)$ 视为可以优化的自由参数。虽然重要性加权采样方案本身可能不是一个有效的TLSR，但它是令牌选择的规范分解中的一个关键构建块。

定理1：最优接受概率与规范分解。本文的第一个结果是关于最优令牌级选择规则的分解，它建立了与定义1中重要性加权采样的联系。令 $P^⋆(acc)$ 为问题1中最优令牌级选择规则的接受概率，则可以表示为：

$$P^{\star}(\text{acc}) = \max_{\{\beta_y(x_{1:K})\}} \left\{ \sum_{y \in \Omega} \min \left( q(y), \sum_{x_1, \dots, x_K \in \Omega} \beta_y(x_{1:K}) \cdot \prod_{i=1}^{K} p(x_i) \right) \right\},$$

其中，最大化是针对每个 $\{x_1, \ldots, x_K, y\} \in \Omega$ 的 $\beta_y(x_{1:K})$ 进行的，约束条件为 $0 \leq \beta_y(x_{1:K}) \leq 1$，以及

$$\sum_{y \in \Omega} \beta_{y}\left(x_{1: K}\right)=1, \quad \forall x_{1: K} \in \Omega^{K},$$

并且

$$\beta_{y}(x_{1:K})=0, \quad y \notin \{x_{1}, \ldots, x_{K}\}.$$

此外，如果 $\{\beta^⋆_y (x_{1:K})\}$ 是在公式(2)中实现最大值的参数，那么 $P^⋆(acc)$ 可以通过一个两步规范分解来达到：第一步，给定输入令牌列表 $\{x_1, \ldots, x_K\}$，应用定义1中的重要性加权采样，使用参数 $\beta^⋆_y (x_1, \ldots, x_K)$ 输出一个中间令牌 $y \in \{x_1, \ldots, x_K\}$；第二步，对选定的令牌 $y$ 应用单草稿推测采样方案【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】、【21，Fast inference from transformers via speculative decoding，2023，ICML】以生成最终输出令牌。图1展示了定理1中提出的系统，其中第一步涉及重要性加权采样以输出中间令牌，第二步涉及推测采样。该方法需要计算最优的 $\beta^⋆_y(x_{1:K})$。在实践中，可以使用计算速度更快但次优的选择，这将在后文讨论。

对定理1的扩展说明。尽管本文的理论发展主要集中在相同草稿模型的情况下，但定理1的结果可以自然地扩展到当 $K$ 个令牌是从一个联合分布 $p(x_1, \ldots, x_K)$ 中采样的情况，这在附录B.1中进行了讨论。特别是，定理1可以自然地扩展到当 $K$ 个令牌从不同的分布中独立采样，即 $S \sim \prod_{i=1}^{K} p_i(X_i)$ 的情况（如【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】），以及当令牌进行无放回采样的情况（如【20，Recursive speculative decoding: Accelerating llm inference via sampling without replacement，2024，arXiv】）。

定理2：K=2时接受概率为1的充要条件。本文基于定理1，为涉及 $K=2$ 个草稿的最优接受概率建立了新的分析结果。第一个结果是刻画了草稿和目标分布 $p(\cdot)$ 和 $q(\cdot)$ 使得 $P^⋆(accept) = 1$ 的充要条件。当 $K=2$ 时，定义1中的 $P^⋆(acc) = 1$ 的一个充要条件是：

$$\sum_{x \in \mathcal{S}} q(x) \geq\left(\sum_{x \in \mathcal{S}} p(x)\right)^2, \quad \forall \mathcal{S} \subseteq \Omega .$$

定理2的意义。值得注意的是，即使当 $p(\cdot)$ 和 $q(\cdot)$ 不完全相同时，接受概率也可以等于1。因此，当草稿模型的分布接近但并不等于目标模型时，接受概率仍可能为1。这与 $K=1$ 的情况形成对比，在 $K=1$ 的情况下，只有当 $p(\cdot)$ 和 $q(\cdot)$ 是完全相同的分布时，接受概率才能等于1【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】。此外，据我们所知，先前为多草稿设置提出的方案，如基于修正拒绝采样的SpecTr【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】和SpecInfer【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】，也要求 $p(\cdot) = q(\cdot)$ 才能使接受概率为1。定理1在图1的两步系统背景下很有趣。在这种情况下，重要性加权采样块的输出 $Y$ 匹配目标分布 $q(\cdot)$，因此不需要第二步的推测采样。

示例1。考虑词汇表 $Ω = \{1, 2\}$，草稿和目标分布分别为 $p = (p_1, p_2)$ 和 $q = (q_1, q_2)$。我们假设有 $K=2$ 个草稿。在这种情况下，公式(5)简化为 $q_1 \geq p_1^2$ 和 $q_2 \geq p_2^2$。如果 $p_1 = p_2 = 0.5$，那么当且仅当 $0.25 \leq q_1 \leq 0.75$ 时，$P^⋆(acc) = 1$。相比之下，对于 $K=1$ 草稿的最优方案，只有当 $q_1 = q_2 = 0.5$ 时，$P^⋆(acc)$ 才等于1。

定理2的证明思路。定理2的证明在附录C中，涉及分析定理1中重要性加权采样方案的输出分布 $p_I(\cdot)$，并证明当条件(5)满足时，存在一个可行的 $\beta_y(x_1, x_2)$ 选择，使得 $p_I(\cdot) = q(\cdot)$。证明基于傅里叶-莫茨金（FM）消元技术【36，Lectures on polytopes，2012，Springer Science & Business Media】。然而，直接应用该技术来满足每个 $i \in \Omega$ 的约束 $q(i) = p_I(i)$ 变得难以处理。我们的关键思想是证明，考虑一个形式为 $q(i) \geq p_I(i)$ 的松弛，可以得到与等式约束相同的解，并且适合使用傅里叶-莫茨金消元进行分析。我们在附录C中用一个涉及 $Ω = \{1, 2, 3\}$ 的例子进一步解释了这一点。

证明中的技术挑战与几何观点。定理2证明中的核心技术挑战在于确定一个线性方程组是否存在非负解。这类问题在文献中早有研究，【7，Algorithm for finding a general formula for the non-negative solutions of system of linear equations，1964，USSR Computational Mathematics and Mathematical Physics】、【12，On positive solutions of a system of linear equations，1926，Annals of Mathematics】提供了一种算法。这些考虑引出了一个涉及多面体锥的几何观点，我们在附录D中进行了讨论。我们解释了如何使用双重描述法【15，Double description method revisited，1995，Springer】来寻找多面体锥的对偶表示，从而数值上验证接受概率等于1的充要条件。实际上，这种方法被用来验证了类似于定理2的条件，适用于多达$K=6$个草稿和所有大小不超过$|\Omega| \leq 14$的字母表，尽管本文只提供了$K=2$个草稿情况的解析证明。原因在于我们的关键步骤，即附录中的引理2，它利用了傅里叶-莫茨金消元，但不容易推广到$K>2$个草稿的情况。

定理3：K=2时的最优接受概率表达式。我们的最后一个结果是给出了$K=2$个草稿情况下的最优接受概率的显式表达式。对于$K=2$个草稿，草稿分布$p(\cdot)$和目标分布$q(\cdot)$以及任意的令牌字母表$\Omega$，最优令牌级选择规则的接受概率$P^⋆(acc)$由下式给出：

$$P^\star(\text{acc}) = \min_{\mathcal{S} \subseteq \Omega} \left\{ \sum_{s \in \mathcal{S}} q(s) + \left( \sum_{s \in \mathcal{S}^c} p(s) \right)^2 + 2 \left( \sum_{s \in \mathcal{S}} p(s) \right) \left( \sum_{s \in \mathcal{S}^c} p(s) \right) \right\},$$

其中 $S^c = \Omega \setminus S$ 是 $S$ 的补集。

定理3的独特性与证明。据我们所知，定理3中的结果是首次提出的。【30，SpecTr: Fast speculative decoding via optimal transport，2024，NeurIPS】中给出了 $P^⋆(acc)$ 的上界，但这些上界不一定是紧的。相比之下，公式(6)为当 $X_1$ 和 $X_2$ 从 $p(\cdot)$ 中独立采样时，$K=2$ 个草稿情况下的接受概率提供了一个精确的表达式。定理3的证明在附录E中给出，它将傅里叶-莫茨金消元法应用于定理1中提出的线性规划，以在$K=2$个草稿的情况下刻画一个解析解。该证明建立在定理2证明的基础上，但需要消除额外的变量。

对定理3的推广猜想。值得注意的是，公式(6)可以表示为：

$$P^{\star}(\text{acc}) = \min_{\mathcal{S} \subseteq \Omega} \left\{ \sum_{s \in \mathcal{S}} q(s) - \left( \sum_{s \in \mathcal{S}} p(s) \right)^{2} + 1 \right\},$$

这使我们推测，在 $K>2$ 个草稿的一般情况下，最优接受概率可以通过将公式(7)中第二项的指数2替换为 $K$ 来获得。

最优接受概率的数值评估。我们在图2中提供了最优接受概率的数值评估。为了便于说明，我们假设 $Ω$ 的大小为3，并且假设 $p = [1/3, 1/3, 1/3]$。在左图中，我们考虑 $q = [1/3, q_2, 2/3 - q_2]$，在右图中考虑 $q = [1/6, q_2, 5/6 - q_2]$。$q_2$ 的值在x轴上变化。我们将定理3中的最优接受概率与两个基线方案 SpecTr【30，SpecTr: Fast speculative decoding via optimal transport，2024b，NeurIPS】和 SpecInfer【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】进行了比较。我们观察到，对于广泛的 $q_2$ 值，最优接受概率可以等于1。这与定理2一致。相比之下，基线方案似乎只有在 $q_2 = 1/3$ 使得 $q = [1/3, 1/3, 1/3]$ 的特殊情况下才能达到1的接受概率。

4 更快的加权重要性推测采样

加速计算的动机。我们讨论了在图1的规范分解中快速计算 $\beta_y(\cdot)$ 的方法。理想情况下，我们希望选择能够最大化公式(2)中接受概率的 $\beta_y(\cdot)$，但这涉及一个计算成本高昂的线性规划问题。幸运的是，我们可以开发出计算速度更快且仍能实现高接受概率的程序。在实践中，目标模型和草稿模型的分布通常集中在少数几个令牌上。研究也观察到，从高概率集合（如top-p集）中采样可以产生更连贯的输出【23，Locally typical sampling，2023，Transactions of the Association for Computational Linguistics】。经过这种采样后，有效字母表的大小（即具有非零概率的令牌数量）通常很小。我们在附录H中报告了OPT模型的有效字母表大小的一些测量结果。这促使我们开发一些方法，通过减少优化中所需的变量数量来加速我们提出的解决方案。本节我们考虑草稿模型具有相同分布的情况，并在附录I中讨论了非相同草稿分布的情况。我们首先关注 $K=2$ 个草稿的情况，然后讨论如何处理一般情况。

K=2草稿的线性规划。我们首先考虑 $K=2$ 个草稿的情况，并重新审视需要解决的线性规划问题，以计算公式(2)中的最优接受概率。然后我们将解释如何减少线性规划中的变量以加快计算。让 $X_1$ 和 $X_2$ 表示输入令牌，$Y$ 表示选定的令牌。注意，当 $X_1 = X_2 = i$ 时，我们有 $\beta_i(i, i) = 1$。此外，当草稿模型具有相同的分布时，根据对称性，我们有 $\beta_y(i, j) = \beta_y(j, i)$。引入一组新变量 $w_{i,j}$ 会更方便，定义如下：

$$w_{i,j} = \Pr(Y = i \mid \{X_1, X_2\} = \{i, j\}),$$

即，$w_{i,j}$ 表示在输入端看到无序对 $\{i, j\}$ 的情况下，输出令牌为 $i$ 的概率。接下来我们讨论如何构建一个涉及变量 $w_{i,j}$ 的线性规划来最大化公式(2)中的接受概率。为方便起见，我们令 $p = (p_1, \ldots, p_n)$ 作为草稿模型的概率向量，$q = (q_1, \ldots, q_n)$ 作为目标模型的概率向量。接受概率，即公式(2)中的目标函数，由 $\sum_{i=1}^{n} \min(p_I(i), q_i)$ 给出，其中

$$p_I(k)=p_k^2+\sum_{i=1,i\ne k}^n 2p_ip_kw_{k,i}$$

表示选定令牌为 $Y=k$ 的概率。此外，对于我们在(8)中对 $w_{i,j}$ 的定义，相关的约束是：$w_{i,j} + w_{j,i} = 1$ 和 $0 \leq w_{i,j} \leq 1$。这个优化问题可以被构建为一个具有 $O(n^2)$ 个变量的线性规划问题，在实践中可能很慢。接下来我们讨论在不显著损失接受概率的情况下减少优化中变量数量的技术。

线性规划的快速近似。为了提出一种更快的近似方法，我们注意到当草稿和目标分布集中在少数几个令牌上时，线性规划的解对大多数 $w_{i,j}$ 的选择不敏感。因此，可以启发式地设置大多数变量。我们将这种方法称为截断LP方案，并在附录G的算法2中给出了伪代码。假设词汇表 $Ω = {1, 2, \ldots, n}$ 的令牌已按 $q_i - p_i^2$ 的降序排列，即 $q_1 - p_1^2 \geq q_2 - p_2^2 \ldots \geq q_n - p_n^2$。我们将 $Ω$ 分为两个集合 $Ω_1 = \{1, 2, \ldots, s\}$ 和 $Ω_2 = \{s+1, \ldots, n\}$，其中 $s$ 是一个设计参数。我们按如下方式固定一部分权重：

$$\begin{aligned} w_{i,j}=\begin{cases}1, & i \in \Omega_1, j \in \Omega_2 \\ 1, & i \in \Omega_2, j \in \Omega_2, i < j\end{cases} \end{aligned}$$

而我们将 $i < j$ 且 $i, j \in \Omega_1$ 的权重 $w_{i,j}$ 作为自由参数。公式(10)中权重选择的直觉是，在这些情况下，我们倾向于选择令牌 $i$ 而不是令牌 $j$ 来进一步增加 $p_I(i)$，这反过来可以减少 $q_i$ 和 $p_I(i)$ 之间的差异。注意，公式(9)简化为：

$$\begin{aligned} p_I(k) = \begin{cases} p_k^2 + \sum_{i=1,i \neq k}^s 2p_ip_kw_{k,i} + \sum_{i=s+1}^n 2p_ip_k, & k \in \Omega_1 \\ p_k^2 + \sum_{i=k+1}^n 2p_ip_k, & k \in \Omega_2 \end{cases} \end{aligned}$$

我们的目标是最大化 $\sum_{k=1}^{s} \min(p_I(k), q_k)$，这是关于变量 $w_{i,j}$ 的。因此，变量的数量减少到 $O(s^2)$。我们还在附录F中证明，如果 $P^⋆(acc)$ 是应用线性规划于所有 $O(n^2)$ 个权重变量得到的最优接受概率，而 $P˜(acc)$ 是截断程序的接受概率，则：

$$\tilde{P}(\text{acc}) \geq P^\star(\text{acc}) - \sum_{x \in \Omega_2} (q(x) - p^2(x))^+$$

因此，如果选择的 $\Omega_2$ 使得惩罚项很小，那么接受概率的下降可以保持在很小的范围内。

实验观察与第二种加速方案。在实验中，我们观察到对于训练良好的目标模型，即使 $s$ 的值很小，准确率的下降也可以忽略不计。因此，通过适当地截断线性规划中要优化的变量数量，我们期望能有更快的实现。我们在附录G中讨论了所提出方法的计算复杂性。我们的第二个提议是在执行重要性采样时截断字母表。具体来说，令 $Ω_0 \subseteq Ω$ 是目标分布下的一个高概率子集。我们将目标分布重新归一化为 $\tilde{q}$，使其完全支持在 $Ω_0$ 上，并在我们提出的方案中使用这个分布来生成一个输出 $Y \sim \tilde{q}(\cdot)$。由于我们希望输出遵循分布 $q(\cdot)$，我们执行以下后处理。我们以概率 $p_a = \sum_{x \in \Omega_0} q(x)$ 接受 $Y$，并以概率 $p_r = 1 - p_a$ 拒绝。如果被拒绝，我们从 $Ω \setminus Ω_0$ 中输出一个令牌，使得任何 $x \in \Omega \setminus \Omega_0$ 都以与 $q(x)$ 成正比的概率被选中。这确保了输出令牌遵循原始分布 $q(\cdot)$。我们将此方案称为截断字母表方案。

处理K>2草稿的情况。为了处理 $K>2$ 个草稿的情况，我们建议将输入令牌分成大小为2的组，然后以多阶段的方式应用双草稿重要性采样方案。例如，如果 $S = \{X_1, X_2, X_3\}$ 且 $K=3$，我们首先对组 $\{X_1, X_2\}$ 应用快速加权重要性采样，输出一个中间令牌 $Y_1$，其分布为 $p_1(\cdot)$。然后，我们对输入 $(Y_1, X_3)$ 应用加权重要性采样，此时令牌具有非相同的分布，并产生一个输出令牌 $Y$，然后对其应用推测采样。

A4 实验环境

• 硬件配置: 实验在一台配备80GB内存的A100 GPU上进行。
• 模型架构: 使用OPT模型【34，Opt: Open pre-trained transformer language models，2022，arXiv】。草稿模型拥有1.25亿参数，目标模型拥有130亿参数。
• 数据集: 评估使用了三个数据集及其相关任务：XSum【25，Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization，2018，arXiv】，Databricks-Dolly-15k【10，Free dolly: Introducing the world’s first truly open instruction-tuned llm，2023】和WMT18【2，Findings of the 2018 conference on machine translation (WMT18)，2018，WMT】。
• 软件配置: 所有实验都在4个任意选择的种子上进行，并对性能进行了平均。

A4 实验结果

5.1 使用top-p采样的实验

实验一：不同温度的双草稿模型。本节报告了在LLM任务中使用 top-p 采样（p=0.95）的实验。在图3的第一组实验中，我们考虑了 $K=2$ 个草稿模型的情况，它们共享相同的权重但在令牌生成时使用不同的温度。目标模型的温度设为1.0，一个草稿模型的温度设为1.2，而另一个草稿模型的温度在1.0到2.4的范围内变化。在所有实验中，每次调用草稿模型生成5个令牌。基线方案是单草稿推测采样【21，Fast inference from transformers via speculative decoding，2023，ICML】、【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】，其中仅使用一个采样温度为1.2的草稿。其他两个方案是SpecInfer【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】和我们提出的重要性采样（IS）方案。在IS方案中，我们同时采用了截断LP（截断参数s=5）和截断字母表（大小为40个令牌）的方法。图3报告了不同方案在三个任务上的性能。顶部图表报告了块效率（平均每个草稿模型调用接受的令牌数），底部图表显示了相对于基线单草稿方案的令牌速率百分比改进。结果显示，IS方案在所有三个任务上始终优于SpecInfer。随着第二个草稿温度的增加，SpecInfer的块效率提升微乎其微，而IS方案则取得了持续的改进。在令牌速率方面，SpecInfer由于计算开销未能转化为足够的块效率提升而出现负增益，而IS方案由于块效率的提高，即使增加了第二个草稿的计算量，令牌速率也表现出持续的改进。

实验二：截断参数的影响。表1研究了第4节中讨论的词汇表截断和LP截断不同参数的影响。我们使用与前一个实验相同的参数，但将第二个草稿的温度固定为2.0。结果显示，随着词汇表大小 $\Omega_0$ 从10增加到40，块效率和令牌速率都有所提高。超过40后，块效率的增益饱和，令牌速率下降。我们还发现，块效率对LP程序中的参数 $s$ 不太敏感，选择 $s=5$ 可以获得最佳的令牌速率。

表1：LP截断和字母表截断的影响

实验三：不同目标模型温度。在图4中，我们将两个草稿模型的温度固定为1.0，并在0.2到1.0的范围内改变目标模型的温度。其余参数保持不变。我们在Dolly任务上再次报告了块效率和相对于相同草稿温度的单草稿方案的令牌速率改进。由于两个草稿具有相同的温度，我们还能够与SpecTr方案进行比较。我们再次观察到，我们提出的方法通常能获得最佳性能。

5.2 使用top-k采样的实验

实验设定。接下来我们考虑在LLM任务中应用 top-k 采样的实验。我们再次使用与前一小节相同的OPT模型作为草稿和目标模型，并且每次调用草稿模型采样5个令牌。在Dolly任务中，目标和草稿模型都使用1.0的采样温度。在XSUM任务中，目标模型温度为0.9，草稿模型温度为0.3。

实验一：令牌级接受概率比较。在表2的第一个实验中，我们比较了不同方法的令牌级接受概率。我们首先使用目标模型通过自回归解码生成一个令牌序列。在每一步，我们计算草稿和目标模型生成的logits，然后用它们来计算不同方案的接受概率，并在100个实例上取平均值。在重要性采样（IS）方案中，我们使用截断线性规划，截断阈值 $s=5$。我们还考虑了 $K=2$ 草稿情况下的理论最优接受概率，并计算了在备注2中讨论的该表达式对 $K=4, 8$ 的自然扩展（用灰色字体报告）。在这个实验中，两个模型都使用 $k=5$ 的top-k采样。我们报告了 $K=2, 4, 8$ 个草稿时的接受概率。在将IS方案应用于处理 $K>2$ 个草稿时，我们使用了备注3中讨论的逐次选择方案。结果显示，随着草稿数量的增加，每种情况下的接受概率都会增加。

表2：不同任务中平均接受概率的比较。

实验二：块效率比较。在表3中，我们使用 $K=2$ 和 $K=3$ 个草稿比较了不同方法的块效率。我们应用了 $k=10$ 和 $k=5$ 的top-k采样，并对两个模型都使用了1.0的温度。在我们的IS方案中，我们将LP中的变量截断为 $s=7$。虽然我们提出的方案仍然比先前的工作实现了更高的块效率，但增益相对较小。这可以通过注意到在这种设置下，基线方案的接受概率已经接近理论极限来解释，如表2所示。

表3：在Dolly任务中使用top-k采样实现的块效率

A5 结论

本文重新审视了为令牌级多草稿选择最大化接受概率的问题【30，SpecTr: Fast speculative decoding via optimal transport，2024b，NeurIPS】。我们首先提出了一个分解结果，展示了其与重要性采样的联系。对于 $K=2$ 个（相同）草稿的情况，我们建立了一个接受概率的封闭形式解，并给出了接受概率等于1的充要条件。我们提出了一个实用的令牌选择方案，该方案模仿了理论分解原则，并提供了一种在计算速度和接受概率之间进行权衡的方法。我们使用OPT模型和三个任务（Dolly, XSUM, WMT）展示了实验结果，并证明了在块效率和令牌速率方面相比基线方案有所改进。

A6 附录

A 推测解码 – 背景

基本流程回顾。我们简要回顾一下【21，Fast inference from transformers via speculative decoding，2023，ICML】和【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】中提出的方案。设 $Ω$ 是语言模型允许的所有令牌的集合。给定一个上下文 $x_t = (x^{(1)}, \ldots, x^{(t)}) \in Ω^t$，我们的目标模型 $M_s$ 为每个 $y \in Ω$ 产生条件概率 $M_s(y|x_t)$。我们遵循先前工作并如【30，SpecTr: Fast speculative decoding via optimal transport，2024b，NeurIPS】中明确指出的那样，假设以下计算模型：1) 标准推理：给定上下文 $x_t$，自回归模型可以在 $O(1)$ 时间内输出每个 $y \in Ω$ 的 $M_b(y|x_t)$。2) 时间轴并行化：给定上下文 $x_t$，自回归模型可以在 $O(1)$ 时间内为所有 $i \in \{1, \ldots, t\}$ 和每个 $y \in Ω$ 输出 $M_b(y|x_i)$。由于标准推理是顺序过程，【21，Fast inference from transformers via speculative decoding，2023，ICML】和【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】的思想是利用时间轴上的并行化来加速推理。这是通过一个计算成本远低于 $M_b(\cdot)$ 的草稿模型 $M_s$ 来实现的，该模型能够生成多个令牌。主要步骤总结如下：
1. 草稿构建：草稿模型可以高效地以自回归方式采样 $L$ 个令牌，将上下文 $x_t$ 扩展为 $x^{(1)}, \ldots, x^{(t)}, \tilde{x}^{(t+1)}, \ldots, \tilde{x}^{(t+L)}$。此外，我们保留每个 $0 \leq i \leq L$ 和每个 $y \in \Omega$ 的条件概率 $M_s(y|x_t^1, \tilde{x}_{t+1}^{t+i})$。
2. 条件概率计算：给定 $M_s$ 生成的样本 $\tilde{x}_{t+1}^{t+i}$，我们利用时间轴并行化计算每个 $0 \leq i \leq L$ 和每个 $y \in \Omega$ 的条件概率 $M_b(y|x_t^1, \tilde{x}_{t+1}^{t+i})$。
3. 草稿选择：我们选择前 $L' \le L$ 个令牌，并设置 $x(t+i) = \tilde{x}(t+i)$ for $i \le L'$。选择基于下面的推测解码算法。

算法1：推测采样。

1: 输入: 分布 p(·) 和 q(·), 样本 X ∼ p(·)
2: 计算残差分布 p_res(x) = (q(x) - min(p(x), q(x))) / (1 - Σ min(p(x'), q(x')))
3: Accept = False
4: 以概率 β_p,q(X) = min(1, q(X)/p(X)) 设置 Accept = True
5: if Accept = True then
6:   Y = X
7: else
8:   Y ∼ p_res(·)
9: end if
10: 返回 Y

选择过程与接受概率。为方便起见，我们表示 $p(y) \equiv M_s(y|x_t)$ 和 $q(y) \equiv M_b(y|x_t)$。选择过程如下：给定上下文 $x_t$ 和一个候选草稿序列 $\tilde{x}_{t+1}^{t+L}$，算法在第一步中分别设置 $p(\tilde{x}_{t+1}) = M_s(\tilde{x}_{t+1}|x_t^1)$ 和 $q(\tilde{x}_{t+1}) = M_b(\tilde{x}_{t+1}|x_t^1)$。它以概率 $\beta_{p,q}(\tilde{x}_{t+1}) = \min(1, \frac{q(\tilde{x}_{t+1})}{p(\tilde{x}_{t+1})})$ 接受 $\tilde{x}_{t+1}$。如果令牌被接受，则 $x(t+1) = \tilde{x}_{t+1}$，我们继续处理 $\tilde{x}_{t+2}$。如果令牌被拒绝，我们从 $p_{res}(\cdot)$ 中采样 $Y$ 并设置 $x(t+1) = Y$。在这一步，我们以 $x_{t+1}^1$ 为输入，发起一次从草稿模型采样 $L$ 个新令牌的调用。这个过程持续进行，直到所有 $L$ 个令牌都被接受或有一个令牌被拒绝。所提算法的正确性可以通过直接计算 $Pr(Y=y)$ 来验证。我们还计算了接受概率：

$$\Pr(\text{accept}) = \sum_{x \in \mathcal{X}} \Pr(\text{acc} | X = x) p(x)$$ $$= \sum_{x \in \mathcal{X}} \beta_{p, q}(x) p(x)$$ $$\begin{aligned} \begin{aligned} &= \sum_{x \in \mathcal{X}} \rho_{p,q}(x) p(x) \\ &= \sum_{x \in \mathcal{X}} \min(p(x), q(x)) = 1 - d_{TV}(p, q), \end{aligned} \end{aligned}$$

其中 $d_{TV}(\cdot)$ 是总变分距离。注意 $Pr(\text{accept})$ 是一个关键指标，它决定了草稿模型中的一个令牌被接受的概率。$Pr(\text{accept})$ 的值越高，算法就越高效。在本文中，我们还研究了多草稿设置下的 $Pr(\text{accept})$。

B 定理1的证明

重要性加权采样方案分析。我们首先考虑重要性采样方案族，然后是推测采样，并推导出接受概率。假设 $Y$ 表示在重要性采样中选择的样本，令：

$$\begin{aligned} \operatorname{Pr}(Y=y|X_1^K=x_1^K) = \begin{cases} \beta_y(x_1^K), & y \in \{x_1, \dots, x_K\} \\ 0, & y \notin \{x_1, \dots, x_K\} \end{cases} \end{aligned}$$

其中，对于每个 $x_1^K \in \Omega^K$，$\sum_y \beta_y(x_1, \ldots, x_K) = 1$，且 $0 \leq \beta_y(x_1, \ldots, x_K) \leq 1$。由此可得：

$$\Pr(Y=y)=\sum_{x_{1}, \ldots, x_{K} \in \Omega} \beta_{y}(x_{1}^{K}) \cdot \prod_{i=1}^{K} p(x_{i})$$ $$= \sum_{x_1, \ldots, x_K \in \Omega} \beta_y(x_1^K) \cdot \mathbb{I}_y(x_1, \ldots, x_K) \cdot \prod_{i=1}^K p(x_i)$$

其中 $I_y(x_1, \ldots, x_K)$ 是指示函数，如果 $y \in \{x_1, \ldots, x_K\}$ 则等于1，否则等于0。公式(18)成立是因为如果 $y \notin \{x_1, \ldots, x_K\}$，则 $\beta_y(x_1, \ldots, x_K) = 0$。通过对选定的样本 $X_I$ 应用推测采样，接受概率由下式给出：

$$\begin{aligned} \begin{aligned} P^{\mathrm{M}-\mathrm{IS}}(\text {accept}=1) & =\sum_{y \in \Omega} \min \left(q(y), \operatorname{Pr}\left(X_{I}=y\right)\right) \\ & =\sum_{y \in \Omega} \min \left(q(y), \sum_{x_{1}, \ldots, x_{K} \in \Omega} \beta_{y}\left(x_{1}^{K}\right) \cdot \mathbb{I}_{y}\left(x_{1}, \ldots, x_{K}\right) \cdot \prod_{i=1}^{K} p\left(x_{i}\right)\right) \end{aligned} \end{aligned}$$

因此，在所提出的重要性采样方案类别中，我们可以将我们的目标表述为：

$$\max_{\{\beta_y(x_1^K)\}_{y,x_1,\dots,x_K}} \left\{ \sum_{y \in \Omega} \min \left( q(y), \sum_{x_1,\dots,x_K \in \Omega} \beta_y(x_1^K) \cdot \mathbb{I}_y(x_1, \dots, x_K) \cdot \prod_{i=1}^K p(x_i) \right) \right\}$$

约束条件为：对于每个 $y, x_1, \ldots, x_K \in \Omega$，有 $0 \leq \beta_y(x_1^K) \leq 1$，以及

$$\sum_{x_{1}^{K} \in \Omega^{K}} \beta_{y}(x_{1}^{K})=1, \quad \forall y \in \Omega,$$

并且

$$\beta_y(x_1^K) = 0, \quad y \notin \{x_1, \dots, x_K\}.$$

最优解分析。我们现在考虑在一般设置下为任何给定的 $p(\cdot)$ 和 $q(\cdot)$ 优化接受概率的问题。遵循【30，SpecTr: Fast speculative decoding via optimal transport，2024b，NeurIPS】中的框架，我们寻求找到每个 $y, x_1, \ldots, x_K \in \Omega$ 的 $p_{Y|X_1, \ldots, X_K}(y|x_1, \ldots, x_K)$，以最大化：

$$\Pr(\text{accept}=1)=\Pr(Y \in \{X_1, \dots, X_K\})$$

并受限于 $P_Y(\cdot)$ 的边际约束：

$$q(y)=\operatorname{Pr}(Y=y)=\sum_{x_{1}^{K}} \operatorname{Pr}(Y=y, X_{1}^{K}=x_{1}^{K})=\sum_{x_{1}^{k}} p_{Y|X_{1}^{K}}(y|x_{1}^{K}) \prod_{i=1}^{K} p(x_{i}).$$

接着我们考虑：

$$\begin{aligned} \begin{aligned} q(y) & = \sum_{x_1^k \in \Omega^k} p_{Y|X_1^K}(y|x_1^K) \prod_{i=1}^K p(x_i) \\ & = \sum_{x_1^k \in \Omega^k} p_{Y|X_1^K}(y|x_1^K) \mathbb{I}_y(x_1, \dots, x_K) \prod_{i=1}^K p(x_i) \\ & \quad + \sum_{x_1^k \in \Omega^k} p_{Y|X_1^K}(y|x_1^K) \bar{\mathbb{I}}_y(x_1, \dots, x_K) \prod_{i=1}^K p(x_i) \end{aligned} \end{aligned}$$ $$\ge \sum_{x_{1}^{k} \in \Omega^{k}} p_{Y | X_{1}^{K}}(y | x_{1}^{K}) \mathbb{I}_{y}(x_{1}, \ldots, x_{K}) \prod_{i=1}^{K} p(x_{i})$$

其中 $\bar{I}_y(x_1, \ldots, x_K) = 1 - I_y(x_1, \ldots, x_K)$ 表示 $I$ 的补集。现在注意到：

$$\begin{aligned} \begin{aligned} & \operatorname{Pr}\left(Y \in\left\{X_1, \ldots, X_K\right\}\right) \\ & =\sum_{x_1^K \in \Omega^K} \operatorname{Pr}\left(Y \in\left\{X_1, \ldots, X_K\right\} \mid X_1^K=x_1^K\right) p\left(X_1^K=x_1^K\right) \\ & =\sum_{x_1^K \in \Omega^K} \sum_{y \in \Omega} p_{Y \mid X_1^K}\left(y \mid x_1^K\right) \mathbb{I}_y\left(x_1, \ldots, x_K\right)\left(\prod_{i=1}^K p\left(x_i\right)\right) \\ & =\sum_{y \in \Omega} \sum_{x_1^K \in \Omega^K} p_{Y \mid X_1^K}\left(y \mid x_1^K\right) \mathbb{I}_y\left(x_1, \ldots, x_K\right)\left(\prod_{i=1}^K p\left(x_i\right)\right) \\ & =\sum_{y \in \Omega} \min \left(q(y), \sum_{x_1^K \in \Omega_K} p_{Y \mid X_1^K}\left(y \mid x_1^K\right) \mathbb{I}_y\left(x_1, \ldots, x_K\right)\left(\prod_{i=1}^K p\left(x_i\right)\right)\right) \end{aligned} \end{aligned}$$

我们使用了公式(27)，这意味着对于任何可行的 $p_{Y|X_1^K}(y|x_1^K)$：

$$\sum_{x_{1}^{k} \in \Omega^{k}} p_{Y \mid X_{1}^{K}}\left(y \mid x_{1}^{K}\right) \mathbb{I}_{y}\left(x_{1}, \ldots, x_{K}\right) \prod_{i=1}^{K} p\left(x_{i}\right) \leq q(y)$$

是满足的。

最优接受概率的上界。我们现在为公式(31)建立一个上界，并证明它与在重要性加权采样方案(21)中优化的接受概率相吻合。对于每个 $x_1^K \in \Omega^K$，我们定义：

$$D(x_{1}^{K})=\sum_{y \in \Omega} p_{Y | X_{1}^{K}}(y | x_{1}^{K}) \mathbb{I}_{y}(x_{1}, \ldots, x_{K})$$

并且，用 $N(x_1, \ldots, x_K)$ 表示 $x_1^K$ 中唯一元素的数量，

$$\begin{aligned} \tilde{p}_{Y|X_1^K}(y|x_1^K) = \begin{cases} \frac{p_{Y|X_1^K}(y|x_1^K)}{D(x_1^K)}, & y \in \{x_1, \dots, x_K\}, & D(x_1^K) > 0 \\ \frac{1}{N(x_1, \dots, x_K)} & y \in \{x_1, \dots, x_K\}, & D(x_1^K) = 0, \\ 0 & y \notin \{x_1, \dots, x_K\}. \end{cases} \end{aligned}$$

通过构造，对于每个 $x_1^K \in \Omega^K$，

$$\sum_{y \in \Omega} \tilde{p}_{Y \mid X_{1}^{K}}\left(y \mid x_{1}^{K}\right)=1$$

且

$$\tilde{p}_{Y|X_1^K}(y|x_1^K)=0, \quad y \notin \{x_1, \dots, x_K\}$$

并且：

$$\tilde{p}_{Y|X_1^K}(y|x_1^K) \cdot \mathbb{I}_y(x_1, \dots, x_K) \geq p_{Y|X_1^K}(y|x_1^K) \cdot \mathbb{I}_y(x_1, \dots, x_K), \forall y, x_1, \dots, x_K \in \Omega$$

将(37)代入(31)，我们得到，对于任何可行的 $p_{Y|X_1^K}(\cdot)$，都存在一个满足(35)和(36)的 $\tilde{p}_{Y|X_1^K}(\cdot)$，使得：

$$Pr(Y \in \{X_1, \dots, X_K\}) \le \sum_{y \in \Omega} \min \left( q(y), \sum_{x_1^K \in \Omega_K} \tilde{p}_{Y|X_1^K}(y|x_1^K) \mathbb{I}_y(x_1, \dots, x_K) \left( \prod_{i=1}^K p(x_i) \right) \right)$$

因此，一般情况下的最优接受概率受限于通过对满足(35)和(36)的 $\tilde{p}_{Y|X_1^K}(y|x_1^K)$ 进行优化(38)得到的值。但这个问题恰好与所提出的IS方案类中的优化问题(21)-(23)相吻合，从而证明了后者的最优性。

C 定理2的证明

线性方程组的设定及其松弛。我们使用 $q = (q_1, \ldots, q_n)$ 作为目标模型分布，$p = (p_1, \ldots, p_n)$ 作为草稿模型分布。并定义 $\alpha_{i,j} = \text{Pr}(Y=i, \{X_1, X_2\} = \{i, j\})$。为了使输出分布 $\text{Pr}(Y=i)$ 与目标分布匹配，我们需要满足以下线性方程组：

$$q_{1}-p_{1}^{2}=\alpha_{1,2}+\ldots+\alpha_{1, n}$$ $$q_2 - p_2^2 = \alpha_{2,1} + \dots + \alpha_{2,n}$$ $$\vdots$$ $$q_{n}-p_{n}^{2}=\alpha_{n, 1}+\ldots+\alpha_{n, n-1}$$

其中 $\alpha_{i,j} \geq 0$ 且对于每个 $i \neq j \in \{1, \ldots, n\}$，有 $\alpha_{i,j} + \alpha_{j,i} = 2p_i p_j = 2p_{i,j}$。我们转而考虑一个松弛的不等式组：

$$q_{1}-p_{1}^{2} \geq \alpha_{1,2}+\ldots+\alpha_{1, n}$$ $$q_2 - p_2^2 \ge \alpha_{2,1} + \dots + \alpha_{2,n}$$ $$\vdots$$ $$q_n - p_n^2 \geq \alpha_{n,1} + \dots + \alpha_{n,n-1}$$

其中 $\alpha_{i,j} \geq 0$ 且对于每个 $i \neq j \in \{1, \ldots, n\}$，有 $\alpha_{i,j} + \alpha_{j,i} = 2p_i p_j = 2p_{i,j}$。我们注意到，不等式组(47)-(50)有解当且仅当不等式组(51)-(54)有解。事实上，假设(51)-(54)中的某个不等式为严格不等式，通过对左右两边求和并利用 $\alpha_{i,j} + \alpha_{j,i} = 2p_i p_j$，我们得到：

$$\sum_{i=1}^{n} q_{i}>\left(\sum_{i=1}^{n} p_{i}\right)^{2},$$

这是一个矛盾，因为两边都应该等于1。因此，我们只需考虑这个线性不等式组。

增强的不等式组。我们考虑一个增强的不等式组。引理1指出，原系统(51)-(54)有解当且仅当以下系统有解：

$$\sum_{s \in \mathcal{S}} q_{s}-\left(\sum_{s \in \mathcal{S}} p_{s}\right)^{2} \geq \sum_{s \in \mathcal{S}} \sum_{t \in \mathcal{S}^{c}} \alpha_{s, t} \quad \forall \mathcal{S} \subseteq\{1, \ldots, n\}$$

对于 $\alpha_{s,t} \geq 0$ 和 $\alpha_{s,t} + \alpha_{t,s} = 2p_{s,t}$，其中 $s, t \in \{1, 2, \ldots, n\}$ 且 $s \neq t$。证明过程通过对(51)-(54)中的 $s \in S$ 求和，并进行一系列代数变换，最终推导出增强系统中的不等式。

归纳论证。我们通过归纳法证明以下引理2：设 $V_r$ 表示经过 $r$ 轮傅里叶-莫茨金(FM)消元后被消除的变量的索引集合。那么剩下的约束条件由以下公式给出：

$$\sum_{s \in \mathcal{S}} q_{s}-\left(\sum_{s \in \mathcal{S}} p_{s}\right)^{2} \geq \sum_{s \in \mathcal{S}} \sum_{t \in \mathcal{S}^{c}} \alpha_{s, t} \cdot \mathbb{I}\left((s, t) \notin \mathcal{V}_{r}\right), \quad \forall \mathcal{S} \subseteq\{1, \ldots, n\}$$

当所有变量都被消除后，右侧将等于0，从而恢复定理2的结果。证明的核心在于展示，在每次FM消元步骤中，新产生的不等式都是冗余的，因为它们可以被已有的不等式所蕴含。这个过程依赖于对集合操作和代数表达式的仔细分析，最终证明了在消除所有 $\alpha_{s,t}$ 变量后，剩下的条件就是定理2中给出的充要条件。

E 定理3的证明

优化问题设定。我们令 $p = [p_1, \ldots, p_n]$ 为草稿模型分布，$q = [q_1, \ldots, q_n]$ 为目标模型分布。我们的优化问题可以表示为：

$$\text{maximize} \sum_{i=1}^{n} t_i,$$

服从于：

$$t_{i} \leq \min \left(q_{i}, p_{i}^{2}+\sum_{j \neq i} \alpha_{i, j}\right),$$ $$\alpha_{i, j} + \alpha_{j, i} = 2p_i p_j, 0 \le \alpha_{i, j} \le 1.$$

分析方法。为了解析地解决这个线性规划问题，我们引入一个额外的变量 $z$，满足单个不等式 $z \leq t_1 + \ldots + t_n$。我们提供 $z$ 的可行值范围并选择最大值。使用与定理2证明中类似的技术，我们首先通过傅里叶-莫茨金(FM)消元法消除变量 $\alpha_{i,j}$，得到一个只涉及变量 $t_i$ 和 $z$ 的不等式组（引理3）。

引理3。在应用FM消元技术消除变量 $\alpha_{i,j}$ 后，我们得到以下不等式系统：

$$t_{i} \leq q_{i}, i \in \Omega$$ $$\sum_{i \in \mathcal{S}} t_{i} \leq\left(\sum_{i \in \mathcal{S}} p_{i}\right)^{2}+2\left(\sum_{i \in \mathcal{S}} p_{i}\right)\left(\sum_{i \in \mathcal{S}^{c}} p_{i}\right), \quad \forall \mathcal{S} \subseteq \Omega, \mathcal{S}^{c}=\Omega \backslash \mathcal{S}$$ $$z \leq \sum_{i=1}^{n} t_{i}.$$

归纳证明。接下来，我们通过归纳法逐步消除变量 $t_1, \ldots, t_n$。我们证明（引理4），在消除了 $t_1, \ldots, t_{j-1}$ 之后，系统可以简化为一个关于剩余变量 $t_j, \ldots, t_n$ และ $z$ 的特定形式的不等式组。归纳步骤的核心是证明，每当消除一个变量 $t_j$ 时，通过FM方法新产生的不等式都是冗余的，即它们可以被现有的不等式所蕴含。这个过程一直持续到所有 $t_i$ 变量都被消除。

最终结果。当所有变量 $t_1, \ldots, t_n$ 都被消除后，我们得到了 $z$ 的上界，即最优接受概率 $P^*(acc)$ 的上界。由于我们是在寻找 $z$ 的最大值，这个上界就是我们所求的最优值。最终，我们得到：

$$z \leq \sum_{i \in \mathcal{S}} q_{i}+\left(\sum_{i \in \mathcal{S}^{c}} p_{i}\right)^{2}+2\left(\sum_{i \in \mathcal{S}} p_{i}\right)\left(\sum_{i \in \mathcal{S}^{c}} p_{i}\right), \quad \forall \mathcal{S} \in \Omega$$

这正是定理3中给出的表达式。整个证明过程依赖于对FM消元法在特定线性规划问题上的系统应用，以及证明在消元过程中产生的许多新约束是冗余的。

I 针对非相同草稿分布的LP和快速LP版本，K=2

问题设定。对于 $K=2$ 个草稿的情况，我们解释了当两个令牌从不同的分布独立采样时，即 $X_1 \sim p_1(\cdot)$ 和 $X_2 \sim p_2(\cdot)$，如何扩展重要性加权采样方案及其快速变体。我们用 $p_1 = (p_{1,1}, \ldots, p_{1,n})$ 和 $p_2 = (p_{2,1}, \ldots, p_{2,n})$ 表示草稿模型的分布，用 $q = (q_1, \ldots, q_n)$ 表示目标分布。

变量定义与目标函数。在这种情况下，令牌的顺序很重要。对于 $i < j$，我们定义：

$$\begin{aligned} \begin{aligned} w_{i,j} &= \Pr(Y=i | X_1=i, X_2=j), \bar{w}_{i,j} = 1 - w_{i,j} = \Pr(Y=j | X_1=i, X_2=j) \\ w_{j,i} &= \Pr(Y=i | X_1=j, X_2=i), \bar{w}_{j,i} = 1 - w_{j,i} = \Pr(Y=j | X_1=j, X_2=i) \end{aligned} \end{aligned}$$

如果 $Y$ 表示所选令牌，那么考虑到所有令牌 $i$ 作为输入令牌之一出现的情况，我们有：

$$p_{I}(i)=p_{1, i} p_{2, i}+\sum_{j=1}^{i-1} p_{1, i} p_{2, j} \bar{w}_{i, j}+\sum_{j=i+1}^{n} p_{1, i} p_{2, j} w_{i, j}+\sum_{j=1}^{i-1} p_{1, j} p_{2, i} \bar{w}_{j, i}+\sum_{j=i+1}^{n} p_{1, j} p_{2, i} w_{j, i}$$

我们需要找到最大化 $\sum_{i=1}^n \min(q_i, p_I(i))$ 的 $w_{i,j}$ 和 $w_{j,i}$。这是一个关于变量 $w$ 的线性规划问题，其中 $0 \le w \le 1$。

截断LP版本。LP的截断版本通过根据 $q_i - p_{1,i}p_{2,i}$ 对 $\Omega$ 中的令牌进行排序获得，同样考虑集合 $\Omega_1 = \{1, 2, \ldots, s\}$ 和 $\Omega_2 = \{s+1, \ldots, n\}$。如果 $i, j \in \Omega_1$，我们将 $w_{i,j}$ 视为需要优化的变量。如果 $i \in \Omega_1$ 且 $j \in \Omega_2$，我们设置 $w_{i,j} = 1$。如果 $i, j$ 都在 $\Omega_2$ 中，若 $i < j$ 则设置 $w_{i,j} = 1$，若 $i > j$ 则为0。最终的分布如下。对于 $i \in \{1, \ldots, s\}$：

$$\begin{aligned} \begin{aligned} \tilde{p}_I(i) = & p_{1,i}p_{2,i} + \sum_{j=1}^{i-1} p_{1,i}p_{2,j} \bar{w}_{i,j} + \sum_{j=i+1}^{n} p_{1,i}p_{2,j} w_{i,j} + \sum_{j=1}^{i-1} p_{1,j}p_{2,i} \bar{w}_{j,i} \\ & + \sum_{j=i+1}^{n} p_{1,j}p_{2,i} w_{j,i} + \sum_{j=s+1}^{n} (p_{1,j}p_{2,i} + p_{1,i}p_{2,j}) \end{aligned} \end{aligned}$$

对于 $i = s+1, \ldots, n$，我们有：

$$\tilde{p}_{I}(i)=p_{1, i} p_{2, i}+\sum_{j=i+1}^{n}\left(p_{1, j} p_{2, i}+p_{1, i} p_{2, j}\right)$$

理论保证。通过遵循导致(194)的步骤序列，我们可以证明：

$$\tilde{P}(\text{acc}) \geq P^\star(\text{acc}) - \sum_{i \in \Omega_2} (q_i - p_{1,i} p_{2,i})^+. $$

截断字母表方案可以通过类似的方式应用，即考虑一个高概率子集 $\Omega_0 \subseteq \Omega$，并只保留属于 $\Omega_0$ 的输入令牌。我们生成截断分布 $\tilde{p}_1(\cdot)$ 和 $\tilde{p}_2(\cdot)$，并在这些分布上应用线性规划，然后使用目标分布 $q(\cdot)$ 进行推测采样。

J 额外的实验结果

J.1 具有相同温度的两个草稿模型

实验设置。本节中，我们考虑使用相同的草稿模型来生成候选令牌序列的情况。我们将我们方法的性能与SpecTr【30，SpecTr: Fast speculative decoding via optimal transport，2024b，NeurIPS】、SpecInfer【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】以及单草稿推测采样【21，Fast inference from transformers via speculative decoding，2023，ICML】、【6，Accelerating large language model decoding with speculative sampling，2023，arXiv】进行了比较，其中我们为草稿模型使用了相同的温度。

性能比较。在图8中，我们将目标模型的温度设为1.0，同时将 $K=2$ 个草稿模型的温度在1.2到2.4的范围内变化，并报告了不同方案在之前讨论的三个任务上的性能。顶部图表报告了块效率，底部图表显示了相对于基线单草稿方案的令牌速率百分比改进。

结果分析。我们观察到，IS方案在所有三个任务上都持续优于SpecTr和SpecInfer。实际上，当第二个草稿的温度增加时，SpecTr和SpecInfer的块效率相对于单草稿基线的改进可以忽略不计，而令牌速率实际上低于单草稿基线。另一方面，我们提出的IS方案能够在所有三个任务中实现持续的改进。

块效率随草稿数量变化。表4比较了使用 $K=2$ 到 $K=6$ 个草稿时不同多草稿推测采样方法的块效率，所有草稿都相同且使用1.2的采样温度。我们再次看到所提出的重要性采样方案的持续改进。

表4：在Dolly任务中不同草稿模型数量下实现的块效率。

J.2 三个草稿模型

实验设置。在本节中，我们考虑了 $K=3$ 个使用不同温度进行令牌生成的草稿模型的情况。如前所述，在非相同草稿模型的情况下，唯一可行的多草稿采样方案是SpecInfer方案【24，Specinfer: Accelerating large language model serving with tree-based speculative inference and verification，2024，ASPLOS】。我们将目标模型的温度设为1.0，两个草稿模型的温度分别设为1.2和1.4，同时将第一个草稿模型的温度在1.0到1.6的范围内变化。

结果分析。如图9所示，所提出的IS方案在块效率和实现的令牌速率方面均优于SpecInfer和单草稿推测采样方法。

J.3 ROUGE-L和BLEU分数

相同草稿模型评估。在涉及XSum和WMT任务的实验中，我们分别测量了ROUGE-L【22，ROUGE: A package for automatic evaluation of summaries，2004，ACL】和BLEU【26，Bleu: a method for automatic evaluation of machine translation，2002，ACL】分数。表5和表6显示了相同草稿模型情况下的ROUGE-L和BLEU分数，对应于J.1节中图8的实验。

表5：在XSum任务上，不同解码器和采样温度下的ROUGE-L分数。

表6：在WMT数据集上，不同解码器和采样温度下的BLEU分数。

非相同草稿模型评估。类似地，表7和表8显示了非相同草稿模型情况下的ROUGE-L和BLEU分数，对应于5.1节中图3的实验。

表7：在XSum任务上，不同解码器和采样温度下的ROUGE-L分数。

表8：在WMT数据集上，不同解码器和采样温度下的BLEU分数。