Towards Optimal Multi-Draft Speculative Decoding

作者/机构: Zhengmian Hu, Tong Zheng, Vignesh Viswanathan, Ziyi Chen, Ryan A. Rossi, Yihan Wu, Dinesh Manocha, Heng Huang
(1Department of Computer Science, University of Maryland, College Park, MD, USA; 2Adobe Research, San Jose, CA, USA; 3Manning College of Information & Computer Sciences, University of Massachusetts Amherst, MA, USA; 4Department of Electrical and Computer Engineering, University of Maryland, College Park, MD, USA)

A1 主要贡献

大型语言模型（LLM）已成为自然语言处理任务中不可或缺的一部分，但其自回归采样过程已成为效率瓶颈。多草稿推测解码（Multi-Draft Speculative Decoding, MDSD）是一种新兴的加速方法，在生成每个token时，由一个小型草稿模型生成多个草稿，然后由目标LLM并行验证，以确保最终输出符合目标模型的分布。

核心问题与研究目标:
MDSD算法有两个主要的设计选择：草稿采样方法和验证算法。对于固定的草稿采样方法，其最优接受率是一个最优传输问题的解，但该问题的复杂性使得计算最优接受率和衡量现有验证算法与理论上限之间的差距变得困难。这一困难引出了两个悬而未决的问题：
1. 对于词汇量高达数千的现代LLM，其最优接受率据我们所知从未被计算过。简单的线性规划（LP）求解器只能处理小型的玩具模型，这使得在实际场景中衡量最优接受率充满挑战。
2. 尽管现有验证算法被认为是该最优传输问题的近似解，但它们的性能与理论上限之间的差距从未在真实文本分布上被量化过。不知道最优接受率，就难以评估这些算法的次优程度。

创新点与主要贡献:
本文旨在解决上述两个问题，其主要贡献如下：
* 新颖的理论视角：通过考虑最优传输问题的对偶问题并应用全幺模性（total unimodularity），我们将计算最优接受率的问题转化为一个子集选择问题。这为理解MDSD的效率提供了一个全新的视角。
* 高效的计算方法：对于某些特殊情况，包括有放回采样和无放回采样，我们通过发现集合函数中的类凸结构，提出了解决该子集选择问题的高效方法。这首次提供了一种实用的方法来计算MDSD在给定草稿分布下的理论接受率上限。
* 首次量化理论上限与差距：我们首次在真实文本上测量了MDSD效率的理论上限，并量化了现有验证算法与该上限的差距。通过比较不同草案采样方法的最优接受率，我们发现无放回采样优于有放回采样。我们还评估了现有的验证算法，包括针对有放回采样的K-SEQ和针对两种采样的RRS，发现它们与理论上限仍有显著差距。
* 新的贪心草稿方法：我们提出了一种新颖的草稿采样方法，该方法贪心地选择高概率的草稿，仅最后一个草稿是随机的。在某些情况下，它实现了比无放回采样更高的最优接受率。我们还提出了一个相应的验证算法，该算法能够完美达到其理论接受率上限。

A3 背景知识

2.1 用于加速LLM推理的推测解码

基础定义。令 Σ 表示词汇集。我们有一个目标模型 $P_{target}(·|x_1, x_2, ..., x_m)$，它是一个预测下一个词概率的概率模型。我们的目标是从这个模型中采样作为输出。

单步多草稿推测解码过程。该过程如下：
1. 对于一个草稿模型 $P_{draft}(·|x_1, x_2, ..., x_m)$，采样 n 个草稿 token $\hat{x}^{(1)}, . . . \hat{x}^{(n)}$。
2. 并行计算目标模型的概率：$P_{target}(·|x_1, x_2, ..., x_m)$， $P_{target}(·|x_1, x_2, ..., x_m, \hat{x}^{(1)})$，...，$P_{target}(·|x_1, x_2, ..., x_m, \hat{x}^{(n)})$。由于是并行计算，这一步并不比单独计算 $P_{target}(·|x_1, x_2, ..., x_m)$ 慢很多。
3. 运行验证算法 $x_{m+1} \sim P_{verify}(·|\hat{x}^{(1)}, . . . \hat{x}^{(n)})$。
4. 如果接受了某个草稿 $\hat{x}^{(i)}$，我们就有 $x_{m+1} = \hat{x}^{(i)}$。在这种情况下，我们可以执行另一次采样步骤 $x_{m+2} \sim P_{target}(·|x_1, x_2, ..., x_m, \hat{x}^{(i)})$，从而在一个步骤中生成两个 token，实现加速。

本文范围。尽管推测解码可以生成多个步骤，每步有多个草稿，所有草稿形成一棵树，但本文只考虑单步情况。为了便于后续分析，我们使用 $p(·) = P_{target}(·|x_1, x_2, ..., x_m)$ 来表示目标分布，用 $p_{draft}$ 表示草稿 token 的分布。

2.2 使用单个草稿Token的推测解码

问题形式化。非形式地说，验证算法依赖于两个分布 $p$ 和 $p_{draft}$，以及一个从 $p_{draft}$ 中采样的草稿 token $j$。目标是输出一个 token $i \sim p$，同时最大化目标 $max P(i = j)$，即最大化随机变量 $i$ 与随机变量 $j$ 相同的概率。更形式地说，给定两个在空间 Σ 上的概率分布 $p \in \Delta_{\Sigma}$ 和 $p_{draft} \in \Delta_{\Sigma}$，我们寻求一个联合分布 $\pi \in \Pi(p, p_{draft})$，使其边际分布分别为 $p$ 和 $p_{draft}$，并最大化目标 $max \sum_{i \in \Sigma} \pi(i, i)$。这构成了一个最优传输问题。最优传输表示为 $\pi^*_{p,p_{draft}} \in \Pi(p, p_{draft})$，最优目标函数值为 $\alpha^*(p, p_{draft}) = \sum_{i \in \Sigma} \pi^*_{p,p_{draft}}(i, i)$。

线性规划表示。该问题可以通过将联合分布表示为矩阵，形式化为一个线性规划（LP）问题：

$$\max_{C \in \mathbb{R}^{\Sigma \times \Sigma}} \sum_{i \in \Sigma} C_{i,i} \text{ s.t. } \sum_{j \in \Sigma} C_{i,j} = p(i), \sum_{i \in \Sigma} C_{i,j} = p_{\text{draft}}(j), C_{i,j} \geq 0 \, \forall i, j \in \Sigma.$$

最优解。最优传输具有以下封闭形式的表达式：

其最优目标函数值为：

$$\alpha^*(p, p_{\text{draft}}) = \sum_{i \in \Sigma} \min(p(i), p_{\text{draft}}(i)).$$

条件分布。当 $(i, j) \sim \pi$ 时，给定 $j$ 的 $i$ 的条件分布表示为 $\pi(·|j) \in G(\Sigma)$，其中：

$$\pi(i|j) = \frac{\pi_{i,j}}{\sum_{i \in \Sigma} \pi_{i,j}} = \frac{\pi_{i,j}}{p_{\text{draft}}(j)}$$

对于最优传输，这导致：

$$\begin{aligned} \pi_{p, p_{\text {draft }}}^*(i \mid j)= \begin{cases}\min \left(\frac{p(i)}{p_{\text {draft }}(j)}, 1\right) & i=j \\ \left(1-\frac{p(j)}{p_{\text {draft }}(j)}\right)_{+} \frac{(p(i)-p_{\text {draft }}(i))_{+}}{\sum_{z \in \Sigma}\left(p_{\text {draft }}(z)-p(z)\right)_{+}} & i \neq j\end{cases} . \end{aligned}$$

改进方向。基础的单步单草稿推测解码可以在两个方向上改进。多步方法生成一个草稿序列，相关改进在附录B.3中讨论（Sun等人，2024d；Hu & Huang，2024；Sun等人，2024c）。本文专注于多草稿方向，即在每一步生成多个草稿 token。

2.3 多草稿推测解码

关联集定义。对于 $i \in \Sigma$，定义关联集 $A_i := \{\bar{i} \in \Sigma^n | \exists j \in [n], \bar{i}_j = i\}$。

草稿分布的构建。非形式地说，验证算法依赖于两个分布 $p$ 和 $p_{draft}$，其中 $p_{draft}$ 现在是 n 个 token 的联合分布。常见的 $p_{draft}$ 构建方法包括：
* 有放回采样：给定一个草稿模型，其输出分布为 $q(·)$，独立采样 n 次。对于 $\bar{i} = (\bar{i}_1, . . . , \bar{i}_n) \in \Sigma^n$，我们有 $p_{draft}(\bar{i}) = \prod_{j=1}^n q(\bar{i}_j)$。
* 无放回采样：$p_{draft}(\bar{i}) = \prod_{j=1}^n q_{\neg\bar{i}_1,...,\bar{i}_{j-1}}(\bar{i}_j)$，其中

$$q^{\neg \bar{i}_{1}, \ldots, \bar{i}_{j-1}}(x)=\begin{cases} \frac{q(x)}{1-\sum_{z \in\left\{\bar{i}_{1}, \ldots, \bar{i}_{j-1}\right\}} q(z)} & x \notin\left\{\bar{i}_{1}, \ldots, \bar{i}_{j-1}\right\} \\ 0 & x \in\left\{\bar{i}_{1}, \ldots, \bar{i}_{j-1}\right\} \end{cases}$$

• 不同草稿分布的乘积：$p_{draft}(\bar{i}) = \prod_{j=1}^n q_j(\bar{i}_j)$。

问题形式化。给定多个草稿 token $\bar{i} = (\bar{i}_1, . . . , \bar{i}_n) \sim p_{draft}$，目标是输出一个 $i \sim p$，使得目标 $max P(\exists j \in [n], i = \bar{i}_j)$ 或等价地 $max P(\bar{i} \in A_i)$ 被实现，即最大化随机变量 $i$ 与 $(\bar{i}_1, . . . ,\bar{i}_n)$ 中某个随机变量相同的概率。更形式地说，给定一个在空间 Σ 上的概率分布 $p \in \Delta_{\Sigma}$ 和一个在空间 $\Sigma^n$ 上的概率分布 $p_{draft} \in \Delta_{\Sigma^n}$，我们寻求一个联合分布 $\pi \in \Pi(p, p_{draft})$，使其边际分布分别为 $p$ 和 $p_{draft}$，并最大化目标 $max \sum_{i \in \Sigma} \sum_{\bar{i} \in A_i} \pi(i, \bar{i})$。最优传输表示为 $\pi^*_{p,p_{draft}} \in \Pi(p, p_{draft})$，最优目标函数值为 $\alpha^*(p, p_{draft}) = \sum_{i \in \Sigma} \sum_{\bar{i} \in A_i} \pi^*_{p,p_{draft}}(i, \bar{i})$。

线性规划表示与困难。该问题可以通过将联合分布表示为张量来形式化为一个LP问题：

$$\begin{aligned} \begin{aligned} \max_{C \in \mathbb{R}^{\Sigma \times \Sigma^n}} & \sum_{i \in \Sigma} \sum_{\bar{i} \in A_i} C_{i, \bar{i}} \\ \text{s.t. } & \sum_{\bar{i} \in \Sigma^n} C_{i, \bar{i}} = p(i) \quad \forall i \in \Sigma, \quad \sum_{i \in \Sigma} C_{i, \bar{i}} = p_{\text{draft}}(\bar{i}) \quad \forall \bar{i} \in \Sigma^n, \\ & C_{i, \bar{i}} \geq 0 \quad \forall i \in \Sigma, \bar{i} \in \Sigma^n. \end{aligned} \end{aligned}$$

其困难在于变量和约束的数量呈指数级增长。

现有近似方法。已有多草稿推测解码问题的一些近似解被提出，包括递归拒绝采样（Recursive Rejection Sampling, RRS）（Yang等人，2024b；Jeon等人，2024）和K-SEQ（Sun等人，2024e）。RRS 递归地验证草稿 token，而 K-SEQ 旨在提高有放回采样的接受率。由于篇幅限制，这些方法的细节（附录 A）、其他相关工作（附录 B）和所有证明（附录 C）都移至附录。

A2 方法细节

3 最优接受率作为子集选择问题

核心思想。我们证明了最优接受率可以表示为一个子集选择问题：

$$\alpha^*(p, p_{\text{draft}}) = 1 + \min_{H \subset \Sigma} \left( \sum_{i \in H} p(i) - \sum_{\bar{i} \in H^n} p_{\text{draft}}(\bar{i}) \right) .$$

3.1 对偶问题

等价公式。我们从线性规划公式(7)出发，推导出一个等价的公式：

$$\begin{aligned} \begin{aligned} \max _{S \in \mathbb{R}^{\Sigma \times \Sigma^{n}}} & \sum_{i \in \Sigma} \sum_{\bar{i} \in \Sigma^{n}} S_{i, \bar{i}} \\ \text { s.t. } & \sum_{\bar{i} \in \Sigma^{n}} S_{i, \bar{i}} \leq p(i) \quad \forall i \in \Sigma, \quad \sum_{i \in \Sigma} S_{i, \bar{i}} \leq p_{\text {draft }}(\bar{i}) \quad \forall \bar{i} \in \Sigma^{n}, \\ & S_{i, \bar{i}} \geq 0 \quad \forall i \in \Sigma, \bar{i} \in \Sigma^{n}, \quad S_{i, \bar{i}}=0 \quad \forall i \in \Sigma, \bar{i} \notin A_{i} . \end{aligned} \end{aligned}$$

引理1证明了公式(7)和(9)是等价的。

问题转换与对偶。这个等价公式将一个运输问题（Hitchcock, 1941）转换为了一个b-匹配问题，其对偶问题是一个w-顶点覆盖问题（Schrijver等人，2003）（详细推导见附录C.1）：

$$\min_{y \in \mathbb{R}^{\Sigma}, z \in \mathbb{R}^{\Sigma^n}} \sum_{i \in \Sigma} y_i p(i) + \sum_{\bar{i} \in \Sigma^n} z_{\bar{i}} p_{\text{draft}}(\bar{i})$$

3.2 全幺模性

矩阵属性。公式(10)中约束的系数矩阵是全幺模（totally unimodular, TUM）的。第一组约束形成了一个二分图的关联矩阵，其中图的一侧节点对应于Σ，另一侧对应于$\Sigma^n$。节点 $i$ 和 $\bar{i}$ 之间存在一条边当且仅当 $\bar{i} \in A_i$。因此，它是一个全幺模矩阵（Biggs, 1993）。第二和第三组约束的系数矩阵是单位矩阵，拥有 $|\Sigma| + |\Sigma|^n$ 个变量和 $|\Sigma| + |\Sigma|^n$ 个约束。一个TUM矩阵和一个单位矩阵的拼接仍然是TUM的（Commoner, 1973）。

整数解。由于约束的右侧是整数，对偶问题(10)总是有整数最优解（Hoffman & Kruskal, 2010）。

3.3 子集选择公式

整数变量约束。通过将公式(10)中的变量限制为整数，我们得到：

$$ \min_{y \in \mathbb{Z}^{\Sigma}, z \in \mathbb{Z}^{\Sigma^{n}}} \sum_{i \in \Sigma} y_{i} p(i) + \sum_{\bar{i} \in \Sigma^{n}} z_{\bar{i}} p_{\text{draft}}(\bar{i}) $$

二进制变量简化。在最优解中，$y_i$ 和 $z_{\bar{i}}$ 不会超过1，所以它们只能取0或1。因此，问题可以进一步简化为：

$$ \min_{y \in \{0, 1\}^{\Sigma}} \min_{z \in \{0, 1\}^{\Sigma^n}} \sum_{i \in \Sigma} y_i p(i) + \sum_{\bar{i} \in \Sigma^n} z_{\bar{i}} p_{\text{draft}}(\bar{i}) $$

子集定义。定义 $H = \{i \in \Sigma|y_i = 1\}$，问题变为：

$$\min_{H \subset \Sigma} \min_{z \in \{0,1\}^{\Sigma^n}} \sum_{i \in H} p(i) + \sum_{\bar{i} \in \Sigma^n} z_{\bar{i}} p_{\text{draft}}(\bar{i})$$

最优解替换。对于z的最优解是：

$$\begin{aligned} z^{*}(H)_{\bar{i}}=\begin{cases}1 & \bar{i} \in \bigcup_{x \in \Sigma \backslash H} A_{x} \\ 0 & \bar{i} \notin \bigcup_{x \in \Sigma \backslash H} A_{x}\end{cases} \end{aligned}$$

代入此解，我们得到子集选择公式：

$$\min_{H \subset \Sigma} \sum_{i \in H} p(i) + \sum_{\bar{i} \in \bigcup_{x \in \Sigma \setminus H} A_x} p_{\text{draft}}(\bar{i}) .$$

最终推导。最后，注意到：

$$\sum_{\bar{i} \in \bigcup_{x \in \Sigma \backslash H} A_{x}} p_{\text {draft }}(\bar{i})+\sum_{\bar{i} \in H^{n}} p_{\text {draft }}(\bar{i})=1 .$$

这完成了子集选择公式(8)的推导。

4 在特殊情况下计算最优接受率

引言与定义。在本节中，我们讨论如何为草稿分布 $p_{draft}$ 的某些特殊情况高效地计算最优接受率。对于任何集合函数 $f$，我们定义元素 $x$ 相对于集合 $H$ 的边际值为 $f(x|H) = f(H \cup \{x\}) - f(H)$。我们还定义以下简写符号：$P(H) = \sum_{i \in H} p(i)$，$Q(H) = \sum_{\bar{i} \in H^n} p_{draft}(\bar{i})$，$f(H) = P(H) - Q(H)$。最优接受率可以表示为 $\alpha^*(p, p_{draft}) = 1 + \min_{H \subset \Sigma} f(H)$。

4.1 q-凸函数

q-凸函数定义。一个集合函数 $Q : 2^{\Sigma} \rightarrow \mathbb{R}$ 被称为 q-凸函数，如果存在一个函数 $q : \Sigma \rightarrow \mathbb{R}_{\ge 0}$，使得对于所有 $H \subset \Sigma$ 和 $x, y \in \Sigma \setminus H$ 且 $x \neq y$，我们有：

$$\frac{Q(x|H)}{q(x)} \leq \frac{Q(y|H \cup \{x\})}{q(y)}.$$

直观上，如果我们任意排序Σ的元素并逐个添加元素来构建一个集合序列$H_i$，那么$Q(H_i)$相对于q值之和的曲线总是凸的。

关键定理。定理3指出，对于有放回采样，函数 $Q$ 是一个q-凸函数。定理4指出，对于无放回采样，函数 $Q$ 也是一个q-凸函数。定理5表明，所有q-凸函数都是超模函数（supermodular functions）。

问题求解。对于有放回采样和无放回采样，计算 $\alpha^*(p, p_{draft})$ 都可以被形式化为一个无约束的子模最小化问题，该问题存在多项式时间算法（Iwata, 2008）。然而，通过充分利用q-凸函数的性质，我们可以更快地解决这个问题，如下一节所示。

4.2 高效计算

核心定理。定理6指出，假设 $Q$ 是一个q-凸函数，$Q$ 是单调递增的，并且对于所有 $x \in \Sigma$ 都有 $p(x) > 0$。对于所有 $H \subset \Sigma$ 和 $x, y \in \Sigma \setminus H$ 且 $x \neq y$，如果 $\frac{q(x)}{p(x)} \le \frac{q(y)}{p(y)}$ 并且 $f(x|H) \le 0$，那么 $f(y|H \cup \{x\}) \le f(x|H)$。上述定理要求 $p(x) > 0$。当 $p(x) = 0$ 时，存在一个公式(8)的最优集 $H^*$ 包含 $x$，因为 $Q$ 是单调递增的。

4.2.1 算法

算法步骤。受定理6的启发，我们可以高效地计算最优接受率，步骤如下：
1. 找到 Σ 的一个排序 σ，使得 $\frac{q(\sigma_1)}{p(\sigma_1)} \ge \cdot \cdot \cdot \ge \frac{q(\sigma_{|\Sigma|})}{p(\sigma_{|\Sigma|})}$。
2. 构建一个集合序列 $H_i = \{\sigma_1, . . . , \sigma_i\}$。
3. 计算 $\alpha^*(p, p_{draft}) = 1 + \min_i f(H_i)$。
直观地说，我们按照 $q$ 和 $p$ 的比率的非增序对元素进行排序，然后执行线性搜索。

4.2.2 计算Q和α*的复杂度

有放回采样。对于有放回采样，$Q$ 有一个简单的表达式 $Q(H) = (\sum_{x \in H} q(x))^n$。计算 $\alpha^*(p, p_{draft})$ 的时间复杂度为排序步骤的 $O(|\Sigma| \log |\Sigma|)$，加上线性扫描的 $O(|\Sigma|)$。

无放回采样。对于无放回采样，我们可以基于生成函数的系数来计算 $Q(H) = \frac{W_{n,H}}{W_{n,\Sigma}}$，其中 $W_{n,H} = \text{Coeff}_{t^n} G_H(t) = \text{Coeff}_{t^n} \prod_{i \in H}(1 + q(i)t)$，并应用动态规划，其递推关系为 $W_{n,H \cup \{x\}} = W_{n,H} + q(x)W_{n-1,H}$。时间复杂度为排序步骤的 $O(|\Sigma| \log |\Sigma|)$，加上使用动态规划计算生成函数系数的 $O(n|\Sigma|)$。

5 一种选择草稿Token的贪心方法

引言。在本节中，我们提出了一种构建草稿分布 $p_{draft}$ 的新方法，以及一个相应的验证算法，该算法能达到此分布的最优接受率。

5.1 草稿构建

构建过程。给定一个草稿模型输出分布 $q \in \Delta_{\Sigma}$，我们按如下方式构建草稿 tokens $\bar{i} = (\bar{i}_1, . . . , \bar{i}_n)$：
* 前 $n-1$ 个 token 被确定性地设置为 $q$ 中概率最高的 $n-1$ 个 token，即 $\bar{i}_1, . . . ,\bar{i}_{n-1} = \text{Top}_{n-1}(q)$，使得 $q(\bar{i}_1) \ge \cdot \cdot \cdot \ge q(\bar{i}_{n-1})$ 且 $\max_{i \in \Sigma \setminus \{\bar{i}_1,...,\bar{i}_{n-1}\}} q(i) \le q(\bar{i}_{n-1})$。
* 只有最后一个 token $\bar{i}_n$ 是从 $q$ 中无放回地随机采样的（即它与前 $n-1$ 个 token 不同）：$\bar{i}_n \sim q_{\neg \text{Top}_{n-1}(q)}(·) = \frac{q(·)}{1-\sum_{j=1}^{n-1} q(\bar{i}_j)}$。

草稿分布。最终的草稿分布为：

$$\begin{aligned} p_{\mathrm{draft}}(\bar{i})=\begin{cases}q^{\neg\operatorname{Top}_{n-1}(q)}(\bar{i}_n) & \bar{i}_1, \dots, \bar{i}_{n-1}=\operatorname{Top}_{n-1}(q) \\ 0 & \bar{i}_1, \dots, \bar{i}_{n-1} \neq \operatorname{Top}_{n-1}(q)\end{cases} \end{aligned}$$

5.2 验证算法

算法设计。对于这个草稿分布，相应的最优传输问题很简单，因为只有一个草稿 token 是随机的。我们可以设计一个验证算法，严格达到该草稿分布的最优接受率（展开的定义见附录D）：

$$\pi_{p, p_{\text {draft }}}^{\text {Greedy }}(i \mid \bar{i})=\pi_{p, q^{\neg \operatorname{Top}_{n-1}(q)}}^{*}\left(i \mid \bar{i}_{n}\right) .$$

最优接受率。定理7指出，贪心草稿分布的最优接受率为：

$$\alpha^*(p, p_{\text{draft}}) = \alpha^{\text{Greedy}}(p, p_{\text{draft}}) = \sum_{i \in \text{Top}_{n-1}(q)} p(i) + \sum_{i \in \Sigma} \min(p(i), q^{\neg \text{Top}_{n-1}(q)}(i))$$

我们的子集选择公式(8)为证明上述定理提供了一种便捷的方式。

5.3 与SpecHub的联系

SpecHub简介。SpecHub（Sun等人，2024b）是最近提出的一种MDSD方法，仅适用于 $n=2$ 的情况。SpecHub中的草稿构建如下：
* 首先，采样第一个草稿 token $\bar{i}_1$。
* 如果 $\bar{i}_1 = \text{Top}_1(q)$ 是 $q$ 中概率最高的 token，则无放回地采样第二个草稿 token $\bar{i}_2$，以确保它与 $\bar{i}_1$ 不同。
* 如果 $\bar{i}_1 \neq \text{Top}_1(q)$ 不是 $q$ 中概率最高的 token，则确定性地将第二个草稿 token 设置为概率最高的 token，即 $\bar{i}_2 = \text{Top}_1(q)$。

草稿分布。最终的草稿分布为：

$$\begin{aligned} p_{\mathrm{draft}}(\bar{i})=\begin{cases}q(\bar{i}_{1}) & \bar{i}_{2}=\mathrm{Top}_{1}(q) \\ \frac{q(\bar{i}_{1})}{1-q(\bar{i}_{1})}q(\bar{i}_{2}) & \bar{i}_{1}=\mathrm{Top}_{1}(q) \\ 0 & \mathrm{otherwise}\end{cases} \end{aligned}$$

联系与区别。我们注意到，当 $n=2$ 时，我们的贪心方法与SpecHub本质上是等价的，因为两种方法都确保了至少有一个草稿 token 是 $q$ 中概率最高的 token。然而，具体的草稿分布是不同的，这使得我们的贪心方法有更简单的验证算法。

A4 实验环境

我们的实验旨在测量各种MDSD方法在真实文本分布上的接受率，并将其与理论上限进行比较。

• 数据集:

  • Alpaca (Taori 等人, 2023): 用于指令遵循任务。
  • WMT’14 De-En (Bojar 等人, 2014): 用于翻译任务。
  • CNN-DailyMail (Hermann 等人, 2015): 用于摘要任务。
  • MT-Bench (Zheng 等人, 2023): 用于评估贪心采样方法的生成任务。
  • 每个任务使用LLM在1024个数据样本上生成响应，最大长度为128个token，以构建真实的目标分布 $p$ 和草稿分布 $p_{draft}$。

• 模型架构:

  • LLaMA系列: 目标模型为 LLaMA-7B，草稿模型为 LLaMA-68M。
  • OPT系列: 目标模型为 OPT-6.7B，草稿模型为 OPT-125M。
  • Vicuna系列: 目标模型为 Vicuna-7B-v1.3，草稿模型为 EAGLE 提供的 0.24B 参数模型。
  • Qwen系列: 目标模型为 Qwen2-7B-Instruct，草稿模型为 EAGLE 提供的 0.26B 参数模型。
  • 默认生成温度为0.7，草稿数量为3。

• 硬件配置:

  • 实验在 RTXA6000 GPU上进行，总计算成本低于50个GPU小时。

• 软件配置:

  • 贪心方法的实现基于 EAGLE 框架（Li 等人, 2024），该框架支持具有草稿树结构的多步MDSD。

A4 实验结果

6.1 主要实验

实验内容:
本实验比较了不同MDSD方法在多种LLM和任务上的接受率。比较的方法包括：有/无放回采样的RRS、有放回采样的K-SEQ、以及对应的理论最优接受率 $\alpha^*$，还有我们提出的贪心方法及其最优接受率。

实验结果与分析 (见表1):

• 现有方法与理论上限存在差距: 现有的验证方法，如RRS和K-SEQ，其接受率与理论上限之间存在明显差距。
• 无放回采样优于有放回采样: 无论是在理论上限还是现有验证算法上，无放回采样的接受率均高于有放回采样。这归因于有放回采样可能产生重复的草稿，对加速的帮助较小。
• 贪心方法的优势: 在大多数情况下，贪心方法获得了最高的接受率，但并非总是如此（例如在温度为1时性能会下降）。

Table 1: 不同MDSD方法在各种模型和任务上的接受率。∆α表示验证方法与理论上限之间的差距，统计上显著的差异用方向箭头表示。

6.2 消融研究 I：温度的影响

实验内容:
研究不同温度对接受率的影响。温度不仅影响目标模型和草稿模型的分布，还影响采样过程中的文本输出。本研究使用LLaMA-7B作为目标模型，LLaMA-68M作为草稿模型。

实验结果与分析 (见图1):
* 非单调影响: 温度的影响不是单调的，且不同方法对温度变化的响应也不同。
* 低温下的聚类: 在低温（如T=0）时，所有方法分为两类：允许重复token的方法（此时等效于只有一个有效草稿）和阻止重复token的方法（此时总是选择top-n草稿）。
* 差距随温度增大: 现有验证方法（RRS和K-SEQ）与最优接受率的差距随着温度升高而逐渐增大。
* 采样方法差距随温度减小: 随着温度升高，有放回采样和无放回采样方法之间的差距减小。这是因为在高温下，概率分布更平坦，有放回采样产生重复token的概率降低。

6.3 消融研究 II：草稿数量的影响

实验内容:
研究不同草稿数量对接受率的影响。

实验结果与分析 (见图2):
* 草稿越多，接受率越高: 随着草稿数量增加，对目标模型可能输出的覆盖范围提高，所有方法的接受率都随之提高。
* 采样策略影响收益: 无放回采样通常比有放回采样从增加草稿数量中获益更多，因为后者可能产生冗余草稿。
* 差距随草稿数增大: 现有验证方法（RRS和K-SEQ）与最优接受率的差距随着草稿数量的增加而逐渐增大。

6.4 在生成任务上评估贪心采样方法

实验内容:
在真实世界的生成任务中，评估我们提出的贪心草稿采样方法（第5节）的有效性和生成效率，并与其他MDSD方法进行比较。我们在EAGLE框架内实现了贪心方法，并测试了三种树结构：（1）2个草稿，4层深度；（2）4个草稿，3层深度；（3）EAGLE默认的稀疏树（最多4个草稿，5个步骤）。实验在MT-Bench数据集上进行，使用Vicuna-7B-v1.3作为目标模型。

实验结果与分析 (见表2):
* 理论验证: 实验证实了理论上的洞见，即当草稿数量为2时，贪心方法和SpecHub的接受率相等，两者在任何温度下都没有统计上的显著差异。
* 低温性能优势: 贪心方法在低温下表现出更好的性能。例如，在T=0.1时，它比无放回的RRS有更高的接受率，从而带来更快的生成速度。
* 高温性能减弱: 随着温度升高，贪心方法的性能增益减小，这与图1中的消融研究结果一致。

Table 2: 基于Eagle框架，在MT-Bench上不同MDSD方法的接受率。

A5 结论

本文研究了多草稿推测解码（MDSD）的接受率问题。

在理论方面，我们发现了最优接受率与一个子集选择问题之间的等价关系，并为常见的草稿分布提供了计算最优接受率的高效方法。

在实践方面，我们首次在真实文本分布下测量了最优接受率，并量化了现有算法与该最优值之间的差距。

此外，我们提出了一种实用的贪心草稿构建方法，在某些情况下，其接受率甚至高于无放回采样。

我们希望我们的工作能激发对提高大型语言模型推理效率的进一步研究，使这些强大的模型在现实世界场景中更易于访问和应用。

A6 附录

A 近似解

A.1 递归拒绝采样 (RRS)

残差分布定义。Yang等人（2024b）和Jeon等人（2024）使用了递归拒绝采样方法。对于$p, q \in \Delta_{\Sigma}$，定义残差分布$Res^{p-q} \in \Delta_{\Sigma}$为：

$$\operatorname{Res}^{p-q}(i)=\frac{(p(i)-q(i))_{+}}{\sum_{z \in \Sigma}(p(z)-q(z))_{+}}$$

当 $p=q$ 时，$Res^{p-q}$可以定义为任意分布。

有放回采样的RRS算法。对于来自有放回采样的$p_{draft}$，RRS算法被递归定义为：

$$\pi_{p, p_{\text {draft }}}^{\mathrm{RRS}, \mathrm{w}}(i | \bar{i})=\widetilde{\pi}_{p, q}^{\mathrm{RRS}, \mathrm{w}}(i | \bar{i})$$

其中

$$\begin{aligned} \widetilde{\pi}_{p, q}^{\mathrm{RRS}, \mathrm{w}}(i \mid \bar{i})= \begin{cases}\min \left(\frac{p\left(\bar{i}_{1}\right)}{q\left(\bar{i}_{1}\right)}, 1\right) & i=\bar{i}_{1} \\ \left(1-\frac{p\left(\bar{i}_{1}\right)}{q\left(\bar{i}_{1}\right)}\right)_{+} \widetilde{\pi}_{\mathrm{Res}^{p-q}, q}^{\mathrm{RRS}, \mathrm{w}}\left(i \mid \bar{i}_{2:}\right) & i \neq \bar{i}_{1}\end{cases} \end{aligned}$$

且

$$\widetilde{\pi}_{p,q}^{\mathrm{RRS,w}}(i|()) = p(i)$$

这里的 $\bar{i}_{2:}$ 表示移除了第一个元素的序列 $\bar{i}$。

无放回采样的RRS算法。对于来自无放回采样的$p_{draft}$，RRS算法被定义为：

$$\pi_{p, p_{\text {draft }}}^{\mathrm{RRS}, \mathrm{w}}(i \mid \bar{i})=\widetilde{\pi}_{p, q}^{\mathrm{RRS}, \mathrm{wo}}(i \mid \bar{i})$$

其中

$$\begin{aligned} \widetilde{\pi}_{p, q}^{\mathrm{RRS}, \mathrm{wo}}(i \mid \bar{i})= \begin{cases}\min (\frac{p\left(\bar{i}_{1}\right)}{q\left(\bar{i}_{1}\right)}, 1) & i=\bar{i}_{1} \\ (1-\frac{p\left(\bar{i}_{1}\right)}{q\left(\bar{i}_{1}\right)})_{+} \widetilde{\pi}_{\mathrm{Res}^{p-q}, q^{\neg \bar{i}_{1}}}^{\mathrm{RRS}, \mathrm{wo}}(i \mid \bar{i}_{2:}) & i \neq \bar{i}_{1}\end{cases} \end{aligned}$$

且

$$\widetilde{\pi}_{p,q}^{\text{RRS,wo}}(i|()) = p(i)$$

它们的接受率分别表示为 $\alpha_{RRS,w}(p, p_{draft})$ 和 $\alpha_{RRS,wo}(p, p_{draft})$。

A.2 K-SEQ

K-SEQ算法。Sun等人（2024e）提出了K-SEQ方法来验证有放回采样的草稿。定义

$$\beta_{p, q}(\rho)=\sum_{i \in \Sigma} \min (\frac{p(i)}{\rho}, q(i))$$

并让 $\rho$ 是以下方程的解：

$$1-(1-\beta_{p, q}(\rho))^{n}=\rho \beta_{p, q}(\rho)$$

K-SEQ算法定义为：

$$\pi_{p, p_{\text {draft }}}^{\text {K-SEQ }}(i \mid \bar{i})=\widetilde{\pi}_{p, q, \rho}^{\text {K-SEQ }}(i \mid \bar{i})$$

其中

$$\begin{aligned} \widetilde{\pi}_{p,q,\rho}^{\text{K-SEQ}}(i|\bar{i}) = (1 - \frac{p(\bar{i}_1)}{\rho q(\bar{i}_1)})_+ \widetilde{\pi}_{p,q,\rho}^{\text{K-SEQ}}(i|\bar{i}_{2:}) + \begin{cases} \min(\frac{p(\bar{i}_1)}{\rho q(\bar{i}_1)}, 1) & i = \bar{i}_1 \\ 0 & i \neq \bar{i}_1 \end{cases} \end{aligned}$$

且

$$\widetilde{\pi}_{p, q, \rho}^{\mathrm{K}-\mathrm{SEQ}}(i \mid())=\frac{p(i)-\min \left\{q(i), \frac{p(i)}{\rho}\right\}}{(1-\beta_{p, q}(\rho))^{n}} \frac{1-(1-\beta_{p, q}(\rho))^{n}}{\beta_{p, q}(\rho)}$$

接受率与保证。接受率表示为 $\alpha_{K-SEQ}(p, p_{draft}) = 1 - (1 - \beta_{p,q}(\rho))^n$，理论上保证能达到最优接受率的 $(1 - e^{-1})$-近似。

B 相关工作

B.1 草稿模型设计

设计原则。许多研究探讨了如何为推测解码设计更好的草稿模型。原则上，任何自回归概率模型都可以作为草稿模型。最简单的方法包括使用n-gram模型（Ou等人，2024）或以文档检索作为草稿模型（Yang等人，2023；He等人，2023）。小型的基于Transformer的语言模型也已被使用（Leviathan等人，2023；Chen等人，2023a），通常还会采用蒸馏技术来进一步增加草稿模型和目标模型之间的重叠（Zhou等人，2023）。

权衡与创新。设计一个好的草稿模型需要在其与目标模型的相似性和其计算复杂性之间进行权衡。更复杂的草稿模型因其与目标模型更接近而导致更高的接受率，但也会带来更高的计算开销。为了实现更好的权衡，一些工作提出重用目标模型的计算结果。例如，Monea等人（2023）使用带有“前瞻”token的原始模型，而Cai等人（2024）在原始模型的最后一个隐藏层添加新的头来预测更远的token。Li等人（2024）重用大模型的最后一层隐藏状态计算，并引入一个新的注意力层来预测下一个token。Sun等人（2024a）使用带有部分键值缓存的目标模型作为草稿模型。

B.2 多草稿推测解码

本文关注点。许多关于多草稿推测解码（MDSD）的相关工作已在其他部分介绍。本文专注于单步多草稿场景。当MDSD生成多个步骤，且每一步都涉及多个草稿时，它会形成一个树状结构。Sequoia（Chen等人，2024）提出了一种动态规划算法来搜索最优的树拓扑结构。

相关改进。随着树的深度增加，某些分支的接受概率会降低。Cascade Speculative Drafting（Chen等人，2023b）通过将最大的草稿模型分配给生成较浅层次（更可能被接受）的草稿token，并逐渐使用较小的模型为不太相关的分支生成草稿来解决此问题。

理论进展。Khisti等人（2024）研究了n=2个草稿的有放回采样的特殊情况下的最优接受率，并得到了以下结果：

$$\alpha^*(p, p_{\text{draft}}) = \min_{H \subset \Sigma} \left\{ \sum_{i \in H} p(i) + \left( \sum_{i \in \Sigma \setminus H} q(s) \right)^2 + 2 \left( \sum_{i \in H} q(s) \right) \left( \sum_{i \in \Sigma \setminus H} q(s) \right) \right\}$$

这与我们在这种特殊情况下的结果(8)本质上是相同的。然而，我们的理论更具普适性，对草稿采样方法或草稿token的数量没有任何假设。

B.3 多步推测解码

与多草稿解码的区别。在第2.1节中介绍的基本单步单草稿推测解码可以应用于多个步骤，每一步只有一个草稿和一个独立的验证过程（Leviathan等人，2023；Chen等人，2023a）。然而，这种重复应用单步验证的方法对于多步场景并非最优。一些工作，如Sun等人（2024d）、Hu & Huang（2024）、Sun等人（2024c），为多步设置设计了更好的验证算法。这些算法专为多步场景定制，同时与单步情况兼容，当应用于长度为1的草稿序列时，会退化为基本的推测采样算法。

不同的改进方向。多步推测解码和多草稿推测解码代表了不同的改进方向。如图3所示，Sun等人（2024e）、Khisti等人（2024）和我们的工作从多草稿的角度改进了推测解码。当只有一个草稿时，它就退化为Leviathan等人（2023）、Chen等人（2023a）的情况。另一方面，Sun等人（2024d）、Hu & Huang（2024）、Sun等人（2024c）从多步的角度增强了推测解码。当只有一个步骤时，它也退化为Leviathan等人（2023）、Chen等人（2023a）的情况。

未来展望。将多草稿和多步场景的改进结合起来将是理想的，并且可能成为未来研究的一个方向。

C 证明

引理1的证明。令 $f_1(C)$ 和 $f_2(S)$ 分别表示(7)和(9)的目标函数值。令 $v_1 = f_1(C^*)$ 和 $v_2 = f_2(S^*)$ 为最优目标函数值。首先，我们证明(7)的最优解对(9)是可行的。定义 $S_{i,\bar{i}} = C^*_{i,\bar{i}}$ 对于 $\bar{i} \in A_i$，以及 $S_{i,\bar{i}} = 0$ 对于 $\bar{i} \notin A_i$。该解保持了目标函数值，即 $f_2(S) = f_1(C^*)$。因此，$v_2 = f_2(S^*) \ge f_2(S) = f_1(C^*) = v_1$。接着，我们证明(9)的最优解对(7)是可行的。定义 $pres(i) = p(i) - \sum_{\bar{i} \in \Sigma^n} S^*_{i,\bar{i}} \ge 0$ 对于 $i \in \Sigma$ 以及 $pres_{draft}(\bar{i}) = p_{draft}(\bar{i}) - \sum_{i \in \Sigma} S^*_{i,\bar{i}} \ge 0$ 对于 $\bar{i} \in \Sigma^n$。我们有 $\sum_{\bar{i} \in \Sigma^n} pres_{draft}(\bar{i}) = \sum_{i \in \Sigma} pres(i)$。定义 $C_{i,\bar{i}} = S^*_{i,\bar{i}} + \frac{pres(i) pres_{draft}(\bar{i})}{\sum_{i \in \Sigma} pres(i)} \ge S^*_{i,\bar{i}}$。该解具有更大的目标函数值，即 $f_1(C) \ge f_2(S^*)$。因此，$v_1 = f_1(C^*) \ge f_1(C) \ge f_2(S^*) = v_2$。结合这两部分，我们有 $v_1 = v_2$，这证明了两个公式的等价性。

定理3的证明。对于 $\bar{i} = (\bar{i}_1, \dots, \bar{i}_n) \in \Sigma^n$，我们有 $p_{draft}(\bar{i}) = \prod_{j=1}^n q(\bar{i}_j)$。函数 $Q(H) = \sum_{\bar{i} \in H^n} p_{draft}(\bar{i})$ 表示所有n个有放回样本都在集合H中的概率。因此，$Q(H) = (\sum_{x \in H} q(x))^n$。考虑凸函数 $g(x) = x^n$。为了证明 $Q$ 的q-凸性，只需证明对于所有 $H \subset \Sigma$ 和 $x, y \in \Sigma \setminus H$ 且 $x \neq y$，我们有：

$$\frac{(Q(H)+q(x))^n-Q(x)^n}{q(x)}\leq\frac{(Q(H)+q(x)+q(y))^n-(Q(H)+q(x))^n}{q(y)}$$

这可以重写为：

$$\frac{g(Q(H)+q(x))-g(Q(x))}{q(x)} \leq \frac{g(Q(H)+q(x)+q(y))-g(Q(H)+q(x))}{q(y)}$$

注意两边都是凸函数g的有限差分。定义 $a = Q(H)$，$b = Q(H) + q(x)$，以及 $c = Q(H) + q(x) + q(y)$。只需证明：

$$\frac{g(b)-g(a)}{b-a} \leq \frac{g(c)-g(b)}{c-b}$$

这直接源于g的凸性。

定理4的证明。函数 $Q(H) = \sum_{\bar{i} \in H^n} p_{draft}(\bar{i})$ 表示所有n个无放回样本都在集合H中的概率。为了处理更复杂的无放回采样情况，我们使用生成函数。定义生成函数 $G_H(t) = \prod_{i \in H}(1 + q(i)t)$ 和系数 $W_{n,H} = \text{Coeff}_{t^n} G_H(t)$。注意 $Q(H) = \frac{W_{n,H}}{W_{n,\Sigma}}$。这些系数满足以下递推关系：

$$\begin{aligned} \begin{aligned} W_{n, H \cup\{x\}} & =\operatorname{Coeff}_{t^n} G_{H \cup\{x\}}(t) \\ & =\operatorname{Coeff}_{t^n} G_H(t)+q(x) t G_H(t) \\ & =W_{n, H}+q(x) W_{n-1, H} \end{aligned} \end{aligned}$$

为了证明Q的q-凸性，只需证明对于所有 $H \subset \Sigma$ 和 $x, y \in \Sigma \setminus H$ 且 $x \neq y$，我们有：

$$\frac{W_{n,H \cup \{x\}} - W_{n,H}}{q(x)W_{n,\Sigma}} \leq \frac{W_{n,H \cup \{x,y\}} - W_{n,H \cup \{x\}}}{q(y)W_{n,\Sigma}}$$

应用递推关系，只需证明：

$$W_{n-1, H} \leq W_{n-1, H \cup\{x\}}$$

再次应用递推关系，只需证明：

$$0 \le q(x)W_{n-2,H}$$

这是成立的，因为G的系数总为非负，即 $W_{n-2,H} \ge 0$。

定理5的证明。只需证明对于所有 $H \subset \Sigma$ 和 $x, y \in \Sigma \setminus H$ 且 $x \neq y$，我们有：

$$Q(x|H) \leq Q(x|H \cup \{y\})$$

根据 $Q$ 的q-凸性，我们有：

$$\frac{Q(x|H)}{q(x)} \leq \frac{Q(y|H \cup \{x\})}{q(y)}$$

因此，

$$\frac{Q(x|H)}{q(x)} \leq \frac{Q(x|H) + Q(y|H \cup \{x\})}{q(x) + q(y)} = \frac{Q(H \cup \{x, y\}) - Q(H)}{q(x) + q(y)} \leq \frac{Q(y|H \cup \{x\})}{q(y)}$$

类似地，通过对称性，我们可以交换x和y得到：

$$\frac{Q(y|H)}{q(y)} \leq \frac{Q(H \cup \{x, y\}) - Q(H)}{q(x) + q(y)} \leq \frac{Q(x|H \cup \{y\})}{q(x)}$$

因此，

$$\frac{Q(x|H)}{q(x)} \leq \frac{Q(H \cup \{x,y\}) - Q(H)}{q(x) + q(y)} \leq \frac{Q(x|H \cup \{y\})}{q(x)}$$

这意味着：

$$Q(x | H) \leq Q(x | H \cup\{y\})$$

定理6的证明。我们有：

$$\begin{aligned} \begin{aligned} f(x|H) &= p(x) - Q(x|H) \\ &= p(x)(1 - \frac{q(x)}{p(x)} \frac{Q(x|H)}{q(x)}) \leq 0 \end{aligned} \end{aligned}$$

因此，

$$\frac{f(x|H)}{p(x)}=1-\frac{q(x)}{p(x)}\frac{Q(x|H)}{q(x)}\leq0$$

根据假设，$\frac{q(x)}{p(x)} \ge \frac{q(y)}{p(y)}$。根据 $Q$ 的q-凸性，我们有 $\frac{Q(x|H)}{q(x)} \le \frac{Q(y|H \cup \{x\})}{q(y)}$。因此，

$$\frac{q(y)}{p(y)} \frac{Q(y \mid H \cup\{x\})}{q(y)} \geq \frac{q(x)}{p(x)} \frac{Q(x \mid H)}{q(x)}$$

由此得出：

$$\begin{aligned} \begin{aligned} \frac{f(y|H \cup \{x\})}{p(y)} &= 1 - \frac{q(y)}{p(y)} \frac{Q(y|H \cup \{x\})}{q(y)} \\ &\leq \frac{f(x|H)}{p(x)} \leq 0 \end{aligned} \end{aligned}$$

定理7的证明。我们首先证明贪心方法的接受率是：

$\alpha^{\text{Greedy}}(p, p_{\text{draft}})$

$$ = \sum_{i \in \Sigma} \sum_{\bar{i} \in A_i} \pi_{p, p_{\text{draft}}}^{\text{Greedy}}(i, \bar{i}) $$

$i \in \operatorname{Top}_{n-1}(q)$

注意 $\sum_{i \in \text{Top}_{n-1}(q)} \min(p(i), q_{\neg \text{Top}_{n-1}(q)}(i)) = \sum_{i \in \text{Top}_{n-1}(q)} \min(p(i), 0) = 0$。接下来，我们计算最优接受率。注意当 $\text{Top}_{n-1}(q) \not\subseteq H$ 时，我们必须有 $Q(H) = 0$。当 $\text{Top}_{n-1}(q) \subseteq H$ 时，我们有 $Q(H) = \sum_{i \in H} q_{\neg \text{Top}_{n-1}(q)}(i)$。因此，

$$\alpha^*(p, p_{\text{draft}}) = 1 + \min_{H \subset \Sigma} P(H) - Q(H)$$

最优集是 $H^* = \{i \in \Sigma | q_{\neg \text{Top}_{n-1}(q)}(i) \ge p(i)\} \cup \text{Top}_{n-1}(q)$。在这种情况下，

$\alpha^*(p, p_{\text{draft}})$

$$\begin{aligned} \begin{aligned} &=1-\sum_{i \in \Sigma}\left(q^{\neg \operatorname{Top}_{n-1}(q)}(i)-p(i)\right)_{+}+\sum_{i \in \operatorname{Top}_{n-1}(q)} p(i) \\ &=\sum_{i \in \Sigma} \min \left(p(i), q^{\neg \operatorname{Top}_{n-1}(q)}(i)\right)_{+}+\sum_{i \in \operatorname{Top}_{n-1}(q)} p(i) \end{aligned} \end{aligned}$$

C.1 对偶问题的推导

推导过程。我们从原始问题(9)开始：

$$\begin{aligned} \begin{aligned} \max_{S \in \mathbb{R}^{\Sigma \times \Sigma^n}} & \sum_{i \in \Sigma} \sum_{\bar{i} \in \Sigma^n} S_{i, \bar{i}} \\ \text{s.t.} & \sum_{\bar{i} \in \Sigma^n} S_{i, \bar{i}} \leq p(i) & \forall i \in \Sigma \\ & \sum_{i \in \Sigma} S_{i, \bar{i}} \leq p_{\text{draft}}(\bar{i}) & \forall \bar{i} \in \Sigma^n \\ & S_{i, \bar{i}} \geq 0 & \forall i \in \Sigma, \bar{i} \in \Sigma^n \\ & S_{i, \bar{i}} = 0 & \forall i \in \Sigma, \bar{i} \notin A_i \end{aligned} \end{aligned}$$

我们为每个约束 $\sum_{\bar{i} \in \Sigma^n} S_{i,\bar{i}} \le p(i)$ 引入对偶变量 $y_i$，为每个约束 $\sum_{i \in \Sigma} S_{i,\bar{i}} \le p_{draft}(\bar{i})$ 引入对偶变量 $z_{\bar{i}}$。拉格朗日函数是：

$$L(S, y, z)=\sum_{i \in \Sigma} \sum_{\bar{i} \in \Sigma^n} S_{i, \bar{i}}$$

对偶函数是：

$$\begin{aligned} \begin{aligned} g(y, z)= & \max _{S \in \mathbb{R}^{\Sigma \times \Sigma^n}} L(S, y, z) \\ & \text { s.t. } S_{i, \bar{i}} \geq 0 \quad \forall i \in \Sigma, \bar{i} \in \Sigma^n \\ & \quad S_{i, \bar{i}}=0 \quad \forall i \in \Sigma, \bar{i} \notin A_i \end{aligned} \end{aligned}$$

重新整理拉格朗日函数：

$$L(S, y, z) = \sum_{i \in \Sigma} \sum_{\bar{i} \in \Sigma^n} (1 - y_i - z_{\bar{i}}) S_{i,\bar{i}}$$

为了使对偶函数有界，我们必须有：

$$1-y_{i}-z_{\overline{i}} \leq 0, \forall i \in \Sigma, \overline{i} \in A_{i}$$

因此，对偶问题是：

$$\begin{aligned} \begin{aligned} \min_{y \in \mathbb{R}^{\Sigma}, z \in \mathbb{R}^{\Sigma^n}} & \sum_{i \in \Sigma} y_i p(i) + \sum_{\bar{i} \in \Sigma^n} z_{\bar{i}} p_{\text{draft}}(\bar{i}) \\ \text{s.t.} & y_i + z_{\bar{i}} \geq 1 & \forall i \in \Sigma, \bar{i} \in A_i \\ & y_i \geq 0 & \forall i \in \Sigma \\ & z_{\bar{i}} \geq 0 & \forall \bar{i} \in \Sigma^n \end{aligned} \end{aligned}$$

这完成了对偶问题的推导。

D 额外图示

单步草稿token生成和验证的图示：

使用任意树拓扑结构应用多草稿推测采样的多步伪代码。

def multidraft_speculative_decoding(prompt,
                                    target_model,
                                    draft_model,
                                    tree_topology,
                                    max_steps):
    """
    用于加速语言模型推理的多草稿推测解码算法。
    
    示例树拓扑:
    tree_topology = [
        [0], [1], [2], [3],                                 # 第一层: 4个分支
        [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [2, 0], [2, 1], [3, 0], # 第二层
        [0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [0, 1, 1], [0, 2, 0], [0, 2, 1], [1, 0, 0], # 第三层
        [0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 0, 2],             # 第四层
        [0, 0, 0, 0, 0], [0, 0, 0, 0, 1]                      # 第五层
    ]
    这是默认的EAGLE树结构，其中每个数字代表在每个级别使用哪个草稿token。
    """

    # 初始化字典来存储草稿和分布
    drafts = {}
    draft_distributions = {}
    target_distributions = {}

    # 为树拓扑中的每个前缀生成草稿
    for prefix in [[]] + tree_topology:
        children = get_children(tree_topology, prefix)
        if not children:
            continue  # 如果没有可扩展的子路径，则跳过

        # 生成草稿和相应的分布
        # 例如，有/无放回采样或第5.1节的方法
        (drafts[tuple(prefix)], 
         draft_distributions[tuple(prefix)]) = generate_draft_tokens(draft_model, prefix)

    # 从目标模型计算所有草稿的概率分布
    target_distributions = compute_target_distributions(target_model, drafts)

    # 开始验证和生成过程
    prefix = []
    while True:
        if tuple(prefix) not in drafts:
            break  # 如果没有可用的草稿则结束

        # 例如，RRS或K-Seq或第5.2节的方法
        output_token = verification(
            drafts[tuple(prefix)],
            draft_distributions[tuple(prefix)],
            target_distributions[tuple(prefix)]
        )
        prefix.append(output_token)

    # 返回生成的序列
    return prefix

def get_children(tree_topology, prefix):
    """
    获取树拓扑中所有扩展给定前缀一个token的子路径。
    
    示例:
    tree_topology = [[0], [1], [0, 0], [0, 1], [1, 0]]
    get_children([], tree_topology) -> [[0], [1]]
    get_children([0], tree_topology) -> [[0, 0], [0, 1]]
    get_children([1], tree_topology) -> [[1, 0]]
    get_children([0, 0], tree_topology) -> []
    """
    return [
        path for path in tree_topology
        if len(path) == len(prefix) + 1 and path[:len(prefix)] == prefix
    ]

贪心草稿构建的验证算法的展开定义。

$$\begin{aligned} \begin{aligned} & \pi_{p, p_{\text {draft }}}^{\text {Greedy }}(i \mid \bar{i}) \\ = & \pi_{p, q^{\neg \operatorname{Top}_{n-1}(q)}}^{*}\left(i \mid \bar{i}_n\right) \\ = & \begin{cases}\min \left(\frac{p(i)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)}{q(i)}, 1\right) & i=\bar{i}_n \\ \left(1-\frac{p\left(\bar{i}_n\right)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)}{q\left(\bar{i}_n\right)}\right)_{+} \frac{p(i)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)}{\sum_{z \in \Sigma}\left(p(z)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)-\mathbb{I}\left(z \notin \operatorname{Top}_{n-1}(q)\right) q(z)\right)_{+}} & i \in \operatorname{Top}_{n-1}(q) \\ \left(1-\frac{p\left(\bar{i}_n\right)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)}{q\left(\bar{i}_n\right)}\right)_{+} \frac{\left(p(i)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)-q(i)\right)_{+}}{\sum_{z \in \Sigma}\left(p(z)\left(1-\sum_{j \in \operatorname{Top}_{n-1}(q)} q(j)\right)-\mathbb{I}\left(z \notin \operatorname{Top}_{n-1}(q)\right) q(z)\right)_{+}} & i \neq \bar{i}_n, i \notin \operatorname{Top}_{n-1}(q)\end{cases} \end{aligned} \end{aligned}$$

E 符号摘要

• Σ: 词汇集
• ∆Σ: 词汇 Σ 上的概率单纯形
• [n]: 集合 {1,...,n}
• $P_{target}(·|x_1, x_2, ..., x_m)$: 目标模型，一个预测给定上下文的下一个词概率的概率模型
• $P_{draft}(·|x_1, x_2, ..., x_m)$: 用于生成候选 token 的草稿模型
• $P_{verify}(·|x^{(1)}, . . . x^{(n)})$ : 从草稿 token 中选择最终输出 token 的验证算法
• p(·) = $P_{target}(·|x_1, x_2, ..., x_m)$: 目标分布的简写
• $p_{draft}$: 草稿 token 的分布
• π ∈ Π(p, $p_{draft}$): 一个边际分布为 p 和 $p_{draft}$ 的联合分布
• $\pi^*_{p,p_{draft}} \in \Pi(p, p_{draft})$: 最优传输联合分布
• $\alpha^*(p, p_{draft})$: 最优接受率
• $A_i := \{\bar{i} \in \Sigma^n|\exists j \in [n], i_j = i\}$: token i 的关联集
• q(·): 草稿模型输出分布的简写
• $q_{\neg\bar{i}_1,...,\bar{i}_{j-1}} (x)$: 无放回采样时 token x 的概率，排除了先前已选择的 token
• $Res^{p-q} \in \Delta_{\Sigma}$: 残差分布
• $\pi^{RRS,w}_{p,p_{draft}}, \pi^{RRS,wo}_{p,p_{draft}}$: 有/无放回采样的 RRS 验证算法
• $\alpha_{RRS,w}(p, p_{draft}), \alpha_{RRS,wo}(p, p_{draft})$: 有/无放回采样的 RRS 接受率
• $\beta_{p,q}(\rho)$: K-SEQ 算法中使用的一个函数
• $\pi^{K-SEQ}_{p,p_{draft}}$: K-SEQ 验证算法
• $\alpha_{K-SEQ}(p, p_{draft})$: K-SEQ 算法的接受率
• $P(H) = \sum_{i \in H} p(i)$: 集合 H 上的目标概率之和
• $Q(H) = \sum_{\bar{i} \in H^n} p_{draft}(\bar{i})$: 集合 $H^n$ 上的草稿概率之和
• $f(H) = P(H) - Q(H)$: 集合 H 上目标和草稿概率之差
• $\pi^{Greedy}_{p,p_{draft}}$: 贪心验证算法
• $\alpha^{Greedy}(p, p_{draft})$: 贪心草稿采样方法的接受率
• $C_{i,j}$: 联合分布的矩阵表示
• $(p(i) - p_{draft}(i))_+$: 差值的正部
• $G_H(t)$: 定义为 $\prod_{i \in H}(1 + q(i)t)$ 的生成函数
• $W_{n,H}$: $G_H(t)$ 中 $t^n$ 的系数
• $f(x|H)$: 元素 x 相对于集合 H 的边际值
• $\text{Top}_{n-1}(q)$: 根据 q 中的概率排名前 n-1 的 token
• $S \in \mathbb{R}^{\Sigma \times \Sigma^n}$: 等价 LP 公式中的变量
• $y \in \mathbb{R}^{\Sigma}, z \in \mathbb{R}^{\Sigma^n}$: 对偶变量
• $\pi(·|j)$: 给定 j 的条件分布
• $\pi_{i,j}$: 联合分布矩阵的单个元素
• $\bar{i}_{2:}$: 移除了第一个元素的序列 $\bar{i}$
• $H_i$: 逐个添加元素构建的集合
• σ: 在高效计算算法中使用的 Σ 的一个排序
• $\text{Coeff}_{t^n}$: 生成函数中 $t^n$ 的系数

F 额外实验

Table 3: 基于Eagle框架，在MT-Bench上不同MDSD方法的平均生成长度τ及其比率Δ。

备注8. 对于#草稿=2，#步骤=4，且T=0.6，三种方法——无放回RRS、SpecHub和Greedy——显示出相似的平均生成长度。截断到两位小数后，它们看起来是相同的。然而，它们实际上是不同的数字：3.51781, 3.51944, 3.51975。