PRESTO: 前缀对齐的树状草稿生成用于扩散推测解码

# PRESTO：面向扩散推测解码的前缀对齐树草稿生成
来源: https://arxiv.org/html/2607.22634

Zheng Wang¹,∗ Zhifan Ye²,³,∗ Qi Cheng³ Yonggan Fu³ Ziyan Wang² Feng Zhu² Haozhe Zhao¹ Jan Kautz³ Pavlo Molchanov³ Humphrey Shi²,³ Minjia Zhang¹
¹伊利诺伊大学厄巴纳-香槟分校
²佐治亚理工学院
³NVIDIA
∗共同第一作者

###### 摘要

扩散大语言模型（dLLMs）近期已成为自回归（AR）大语言模型的一种有前景的替代方案，能够并行生成令牌。近期研究表明，dLLMs 特别适合作为推测解码（SD）中的草稿模型，因为它们可以在单次前向传播中高效地并行生成整个块状候选令牌。然而，现有的基于扩散的草稿生成方法主要依赖线性（单路径）草稿生成，尽管扩散模型同时为多个位置生成多个候选令牌，这自然会产生一个巨大的组合空间，包含多种可能的解码路径。因此，这些方法只能探索一小部分可用候选空间，从根本上限制了可实现的接受长度和解码效率。为了充分利用这种丰富的多候选结构，我们将基于树的草稿生成应用于扩散草稿模型，从而能够探索多样化的候选路径。然而，我们发现直接应用朴素的基于树的草稿生成效果不佳，原因是扩散草稿置信度与基于前缀的 AR 验证之间存在根本性不匹配：扩散边际分布本质上是前缀盲的，这可能导致不可靠的路径排序。我们提出 PRESTO，一个原则性的框架，它将基于树的草稿生成扩展到扩散草稿模型，同时通过*前缀对齐评分*和*基于优先级的树搜索*来解决扩散草稿置信度与基于前缀的 AR 验证之间的根本性不匹配，用于扩散推测解码。PRESTO 背后的关键原则是：(1) 候选排序应与 AR 验证的基于前缀特性保持一致；(2) 树构建应优先考虑具有高验证潜力的候选路径，以最大化接受长度。我们还证明 PRESTO 被设计成一个通用的树草稿生成框架，适用于专用扩散草稿模型 SD 和自推测 dLLM。大量实验表明，PRESTO 在多个基准测试上，对于最先进的专用扩散草稿模型 SD，平均获得 1.5 倍的端到端吞吐量加速；对于自推测扩散大语言模型，平均获得 1.12 倍的加速。

参见图注
图 1：PRESTO 在基于扩散的 SD 上的接受长度比较，包括专用扩散草稿模型 SD（Qwen3-8B-DFlash）和自推测 dLLM（Nemotron-Labs-Diffusion-8B；仅边际评分，无前缀信号）在多个基准测试上。PRESTO 在所有任务上一致地提高了接受长度，使得每轮验证能接受更多令牌。

## 1 引言

扩散大语言模型（dLLMs）近期已成为自回归（AR）大语言模型的一种有前景的替代方案，引入了一种新的多令牌生成范式[1, 2, 3, 4, 5, 6]。然而，在实践中实现高质量生成通常需要迭代去噪过程和精心设计的令牌揭开策略，这引入了额外的计算量，并部分抵消了效率提升[3, 7]。有趣的是，近期研究表明扩散模型特别适合作为推测解码中的草稿模型[8, 9, 10, 11]。它们并行生成多个候选令牌的能力自然地与推测解码（SD）的提案阶段契合，从而实现了高效的草稿令牌生成。这一洞察催生了一系列越来越多地利用扩散模型作为有效草稿模型的工作，包括用于 AR 大语言模型的专用基于扩散的草稿模型（例如，dFlash[10]）和自推测 dLLM（例如，TiDAR[8]、I-DLM[11]、Nemotron-Labs-Diffusion[12]）。例如，dFlash 表明，训练良好的轻量级扩散草稿模型可以通过高效的并行草稿令牌生成显著加速推测解码，同时保持较高的草稿质量，突显了基于扩散的草稿生成的实际潜力。然而，现有的基于扩散的草稿生成方法主要依赖*线性（单路径）草稿生成*，即令牌沿单一轨迹生成并顺序验证。虽然这种设计在 AR 设置下有效，但对于扩散模型而言本质上是次优的。由于扩散生成的非自回归特性，多个可能的候选令牌可以同时存在，而具有最高接受长度的解码路径并不一定对应于 top-1 边际轨迹。根据经验，我们观察到即使在 dFlash 这样的强扩散草稿模型下，单路径草稿生成也只能获得有限的接受长度（见图 2(a)），这表明很大一部分提案空间未被探索。解决此局限性的一个自然方向是超越线性草稿生成，联合探索多个候选路径。在 AR 设置中，先前工作表明引入基于树的草稿生成可以通过并行扩展多个候选路径来显著增加接受长度，从而提高每一步验证的利用率[13, 14, 15, 16, 17, 18]。特别是，由于在小批量情况下目标模型通常不会完全受限于计算[8]，分配额外的草稿计算来探索更丰富的候选集可以使每步被接受的令牌更多。基于树的草稿生成在 AR 设置中的成功提出了一个关键研究问题：*我们能否在扩散草稿模型中利用基于树的草稿生成，通过增加接受长度来进一步提高基于扩散的推测解码的端到端吞吐量？* 虽然这一方向很有希望，但我们发现直接将基于树的草稿生成应用于扩散模型往往效果不佳。扩散草稿模型并行生成位置级的令牌分布，没有严格的前缀依赖关系，这使得高效灵活地构建树状结构候选集成为可能。然而，这一特性也意味着生成的分数是边际令牌概率，而非前缀条件似然。因此，这些分数无法准确反映早期令牌选择如何影响下游接受（例如，一个单独看来可能的令牌，一旦以所选前缀为条件，可能变得不可能），在草稿评分和基于前缀的 AR 验证之间造成了根本性不匹配。在这项工作中，我们提出 PRESTO，一个通过*前缀对齐评分*和*基于优先级的树搜索*实现的原则性基于树的扩散草稿生成框架。PRESTO 背后的关键原则是：(1) 候选排序应与 AR 验证的基于前缀特性保持一致；(2) 树构建应优先考虑具有高验证潜力的候选路径，以最大化接受长度。具体来说，我们通过引入前缀依赖校正来最小化调整扩散边际分布，从而使得得到的分数能更好地反映顺序验证下的接受行为。然后利用该分数指导基于优先级的扩展策略，在树构建过程中将计算分配到高质量的候选路径上，以最大化接受长度。因此，PRESTO 实现了有效的基于树的扩散草稿生成，并在用于 AR 大语言模型的专用基于扩散的草稿模型和自推测 dLLM 上显著提高了接受长度和吞吐量。具体来说，我们的贡献总结如下：

- • 我们识别出基于扩散的草稿评分与基于前缀的 AR 验证之间存在根本性不匹配。我们表明，虽然扩散概率提供了良好校准的边际接受信号，但它们并未完全捕捉 AR 验证的前缀条件特性，这可能导致树构建过程中的次优路径排序，并最终限制接受长度。
- • 我们设计了前缀对齐评分来解决上述不匹配。我们通过对扩散边际分布进行原则性的最小校正，纳入前缀依赖信号，得到一个既忠实于扩散草稿模型又与基于前缀的验证兼容的评分函数。
- • 我们形式化并实现了针对扩散草稿模型的基于树的推测草稿生成。在所提出的前缀对齐评分基础上，我们引入 PRESTO，一个原则性的基于树的扩散草稿生成框架，通过基于优先级的树扩展高效探索高质量草稿路径，大幅提高接受长度和解码吞吐量。
- • PRESTO 可推广到专用和自推测两种扩散 SD 框架。我们进一步将 PRESTO 扩展到自推测 dLLM，作为并行化验证与草稿管线中线性草稿生成的即插即用替代方案，使得 PRESTO 成为第一个既适用于专用扩散草稿模型也适用于自推测扩散大语言模型的基于树的草稿生成框架。在多样化的基准测试和模型规模上，与线性草稿生成基线相比，PRESTO 在 dFlash 上平均获得 1.5 倍的吞吐量加速，在 Nemotron-Labs-Diffusion 上平均获得 1.12 倍的吞吐量加速。

## 2 预备知识

### 2.1 基于树的推测解码

推测解码通过利用轻量级草稿模型加速目标大语言模型 \(p_T\) 的自回归生成。为了提高接受率，草稿模型可以构建一个编码多个候选序列的令牌树[13, 19, 18]。目标大语言模型以基于前缀的方式验证候选：从根开始，顺序接受令牌，直到出现不匹配为止。

**接受分解。** 对于候选路径 \(P = (x_1, \ldots, x_k)\)，令 \(a_i \in \{0,1\}\) 表示第 \(i\) 个令牌被接受的指示变量。接受从左到右进行，因此根据链式法则，
\[
\Pr(P \text{ accepted}) = \prod_{i=1}^k \Pr(a_i = 1 \mid a_1 = 1, \ldots, a_{i-1} = 1, x_1, \ldots, x_i).
\]
由于目标大语言模型以自回归方式计算下一个令牌分布 \(p_T(\cdot \mid x_{<i})\)，接受概率由草案覆盖 \(p_T\) 的方式决定。在树设置中，所有路径共享一个公共前缀，它们的接受概率可以通过一次前向传播计算，该前向传播使用树注意力掩码[18]。

**树注意力。** 标准推测解码中的一大关键加速是，一棵包含 \(B\) 个节点以及分布在树深度 \(k\) 上 \(k\) 个位置集合的树，可以使用 \(B\) 次树注意力计算来验证，而无需自回归式地处理每个位置。树注意力需要两个组件：原始序列中的位置索引 \(\{j_1, \ldots, j_B\}\)，以及指示每个节点在生成时可以看到哪些先前节点的因果掩码。这允许在单次前向传播中并行评估多个候选路径的接受概率。

### 2.2 基于扩散的推测解码

扩散大语言模型（dLLMs）通过从纯噪声 \(x_T\) 开始的迭代去噪过程生成文本。令 \(q_d(t)\) 表示扩散模型在去噪步骤 \(d\) 时在位置 \(t\) 上的边际令牌概率。在 SD 的上下文中，扩散草稿模型单次前向传播后直接提供边际分布 \(q_d(\cdot)\)，无需自回归生成。这对应于“一次全部分支”生成：每个位置 \(d\) 处的候选令牌 \(\{t_1, t_2, \ldots\}\) 来自 \(q_d(\cdot)\)，独立于其他位置，或者带有基于融合去噪状态的一些弱关联。线性扩散 SD 从每个 \(q_d(\cdot)\) 中采样 top-1 令牌，形成单一令牌序列 \(\hat{x}_1, \ldots, \hat{x}_N\) 供验证。由于 \(q_d\) 独立于先前令牌，基于树的扩展需要额外的结构来生成融合草案以进行条件采样。

## 3 扩散自回归验证中的路径选择问题

**设置。** 考虑一个扩散草稿模型，它在位置 \(d \in \{1, \ldots, L\}\) 处产生边际分布 \(q_d(\cdot)\)，以及一个 AR 目标模型 \(p_T\)。线性草稿生成选择 \(\hat{x}_d = \arg\max_t q_d(t)\)，形成单一路径。然后由 \(p_T\) 验证该路径，从 \(\hat{x}_1\) 开始，一旦 \(\hat{x}_i\) 被接受，下一个候选 \(\hat{x}_{i+1}\) 必须与 \(p_T(\cdot \mid \hat{x}_{\le i})\) 竞争。接受长度主要由沿该顺序路径的累积接受概率决定。

**边际与条件概率的不匹配。** 令牌 \(\hat{x}_d\) 的边际概率 \(q_d(\hat{x}_d)\) 捕获的是它作为第 \(d\) 个令牌的单独可能性。然而，在 AR 验证中，\(p_T\) 计算的是条件概率 \(p_T(\hat{x}_d \mid t_1, \ldots, t_{d-1})\)，其中 \(t_{<d}\) 是先前接受令牌的序列。由于 \(q_d\) 并未对 \(t_{<d}\) 进行调节，因此两者之间的不一致可能导致排名错误。图 2(a) 展示了 dFlash 在单路径草稿生成下的微平均接受长度（所有可能深度的平均接受令牌数）。该图显示，即使对于强扩散草稿模型，在接受中值（中点）附近，接受曲线相对于位置也会加剧下降。然而，不同路径的微平均接受长度存在显著差异。图 2(b) 展示了黄金路径（手动标记为最高接受数的路径）与 top-1 路径之间的差距。这表明存在过欠接受问题，其中 top-1 路径并非总产生最长的接受序列；一条较低排名的路径可能产生更长的接受。当前的线性草稿生成无法利用这种多样化的路径，从而错过了潜在的接受长度提升。

**树草稿生成中的路径排名。** 假设我们构建一棵候选树 \(T\)。对于树中的每条路径 \(P = (x_1, \ldots, x_L)\)，我们需要一个分数 \(S(P)\) 以确定在树构建期间将计算扩展到哪些路径。理想情况下，\(S(P)\) 应与实际接受长度 \(\ell(P)\) 强相关。然而，简单的启发式方法，例如路径概率乘积 \(\prod_{d} q_d(x_d)\)，可能会产生误导，因为 \(q_d\) 未对前缀进行调节。例如，一个在前缀 \(c\) 下被 AR 模型高度接受的令牌 \(t\) 可能具有较低的 \(q_d(t)\)，而另一个具有高 \(q_d(\tilde{t})\) 但在前缀 \(c\) 下被拒绝的令牌 \(\tilde{t}\) 则可能被错误地优选。

**对树效率的影响。** 由不可靠排名导致的错误剪枝会降低树的效率。如果一棵树在低质量路径上分配了过多节点，其最大可实现接受长度会降低。由于树构建中固有的计算预算 \(B\)（节点数），我们必须决定如何将预算分配到路径和深度上。如图 2(b) 所示，不同路径的接受长度存在显著差异，优先考虑合适路径可以产生更高的接受率。

> **图 2**：...（原文此处有图，但未提供具体内容，保留原文标记）

## 4 PRESTO：前缀对齐树草稿生成

为了解决第 3 节中识别的不匹配，我们设计了 PRESTO，一个通过*前缀对齐评分*和*基于优先级的树搜索*将基于树的草稿生成适应于扩散草稿模型的框架。

### 4.1 前缀对齐评分函数

令 \(q_d(t)\) 为扩散草稿模型在深度 \(d\) 处关于令牌 \(t\) 的边际分布。对于给定部分前缀 \(c_d = (t_1, \ldots, t_{d-1})\)，我们通过结合来自目标 LLM 或来自扩散模型架构的内部前缀敏感分布的前缀条件信号来校正 \(q_d(t)\)。深度 \(d\) 处可用的最接近替代是来自 AR 验证模型的接受概率 \(\rho_d(t \mid c_d) = \min(1, \frac{p_T(t \mid c_d)}{q_d(t)})\)。然而，在草稿阶段访问 \(p_T\) 要么不可行（如果你有一个较小的草稿模型）要么成本过高（在自推测设置中，你尚未运行目标模型）。我们替代使用一个近似值：\(\rho_d(t \mid c_d) \approx \sigma(h_\phi(c_d, t))\)，其中 \(h_\phi\) 是一个轻量级可学习函数，用于预测验证接受率。在纯扩散草稿模型设置中，我们通过探索不同 \(c_d\) 值从目标模型（在离线验证期间）收集数据，并训练一个浅层预测器来估计 \(\rho_d(t \mid c_d)\)，其预测与新路径的接受概率相关。

**校正后的评分。** 我们定义一个代理分数 \(p_d^*(t \mid c_d)\)，它在忠实于扩散边际 \(q_d\) 的同时融入前缀依赖信号 \(\rho_d(t \mid c_d)\)。具体来说，我们求解以下 KL 正则化目标：
\[
p_d^* = \arg\min_{p} \left[ \text{KL}(p \| q_d) - \lambda_d \mathbb{E}_{p}[\log \rho_d(t \mid c_d)] \right], \quad \text{s.t.} \sum_t p(t) = 1,
\tag{9}
\]
其中 \(\lambda_d > 0\) 且 \(\rho_d\) 以非平凡方式依赖于 \(c_d\)。式 (9) 从而提供了一个候选代理，解决了 \(\tilde{p}_{\text{diff}}\) 的结构不匹配。该形式也可以解释为 KL 正则化目标函数（平衡与 \(q_d\) 的接近性和与 \(\rho_d\) 的对齐）的解（详见附录 B）。为简单起见，我们使用由式 (8) 导出的未归一化对数分数。具体来说，我们定义令牌级分数为
\[
s_{d,t}(c_d) = \log q_d(t) + \lambda_d \log \rho_d(t \mid c_d),
\]
以及路径分数
\[
S(P) = \sum_{d=1}^k s_{d, t_d}(c_d), \quad c_d = (t_1, \ldots, t_{d-1}).
\tag{10}
\]
式 (10) 在深度上可加分解，支持增量式的基于优先级的树扩展。

### 4.2 基于优先级的树构建

算法 1 带全局保留的束搜索
1: 过程 BeamSearch(\( s, B, b, W, D \))
2: \(\mathcal{P} \leftarrow \{\text{root}\}, \mathcal{T} \leftarrow \{\text{root}\}\)
3: for \(d = 1\) 到 \(D\) 且 \(|\mathcal{T}| < B\) do
4: 根据分数 \(s\) 选择 \(\mathcal{P}\) 中的候选
5: 扩展前 \(b\) 个候选
6: 将新节点添加到 \(\mathcal{T}\)
7: 更新 \(\mathcal{P}\)
8: end for
9: return \(\mathcal{T}\)
10: end procedure

**树构建算法。** 我们采用一种基于优先级的树构建算法，该算法将计算分配给最有希望的路径。给定最大节点预算 \(B\)，我们从根节点开始。在深度 \(d\)，我们维护当前路径的优先级队列，并根据其累积路径分数 \(S(P)\) 展开它们。对于优先级队列中的每条路径，我们根据代理分数 \(p_d^*(t \mid c_d)\) 展开 top-k 个令牌候选。该过程重复进行，直到分配完所有 \(B\) 个节点。由于 \(p_d^*\) 优先考虑那些在已知前缀下具有高接受概率的令牌，因此树更有可能生成具有较长接受序列的路径。

**计算效率。** 前缀对齐评分是一种轻量级计算：预测器 \(h_\phi\) 是一个具有隐藏维度的两层 MLP，可忽略不计的开销。我们仅需对树中的每个节点评估一次 \(\rho_d(t \mid c_d)\)，这与标准的树构建步骤一致。在自推测设置中，由于在验证期间相同的目标模型已经计算了条件分布，因此无需额外的前向传播。

### 4.3 自推测 dLLM 中的 PRESTO

自推测 dLLM（例如，Nemotron-Labs-Diffusion、TiDAR）采用统一的模型架构，该模型既能充当草稿模型，也能充当验证模型。在此设置中，PRESTO 作为线性草稿生成的直接替代品，在单个前向传播中同时进行树草稿生成和验证。对于给定的决策步骤 \(d\)，扩散草稿模型现在不是为下一个位置生成单个令牌，而是为一组候选位置生成令牌分布，从中 PRESTO 构建一棵树。在验证阶段，目标大语言模型（在此情况下与草稿模型相同）使用树注意力在一次前向传播中同时验证所有路径。由于验证开销与树大小呈线性关系，但通常小于自回归基准的深度级开销，因此增加了接受长度。我们通过让目标模型为非候选位置提供草稿令牌来合并并行化，从而形成一种混合方法，该方法在保持高吞吐量的同时扩大接受长度（详见附录 C）。

### 4.4 从该 KL 正则化推导

我们从以下 KL 正则化目标开始：
\[
p_d^* = \arg\min_p \text{KL}(p \| q_d) - \lambda_d \mathbb{E}_{p}[\log \rho_d(t \mid c_d)], \quad \text{s.t.} \sum_t p(t) = 1,
\]
其中我们假设 \(q_d(t) > 0\) 且 \(\rho_d(t \mid c_d) > 0\) 在候选支持上成立。展开 KL 散度，
\[
\text{KL}(p \| q_d) = \sum_t p(t) \log \frac{p(t)}{q_d(t)},
\]
目标函数变为
\[
\mathcal{L}(p) = \sum_t p(t) \log \frac{p(t)}{q_d(t)} - \lambda_d \sum_t p(t) \log \rho_d(t \mid c_d).
\]
包含归一化约束 \(\sum_t p(t) = 1\) 及拉格朗日乘子 \(\mu\)，我们得到
\[
\mathcal{J}(p) = \sum_t p(t) \log \frac{p(t)}{q_d(t)} - \lambda_d \sum_t p(t) \log \rho_d(t \mid c_d) + \mu \left( \sum_t p(t) - 1 \right).
\]
对 \(p(t)\) 求导，
\[
\frac{\partial \mathcal{J}}{\partial p(t)} = \log p(t) - \log q_d(t) + 1 - \lambda_d \log \rho_d(t \mid c_d) + \mu.
\]
令导数为零，得到
\[
\log p_d^\star(t) = \log q_d(t) + \lambda_d \log \rho_d(t \mid c_d) + C,
\]
其中 \(C\) 吸收与 \(t\) 无关的常数。两边取指数，
\[
p_d^\star(t) \propto q_d(t) \, \rho_d(t \mid c_d)^{\lambda_d}.
\]
等价地，
\[
p_d^\star(t \mid c_d) = \frac{q_d(t) \rho_d(t \mid c_d)^{\lambda_d}}{Z_d(c_d)}, \quad Z_d(c_d) = \sum_{t' \in \mathcal{V}} q_d(t') \rho_d(t' \mid c_d)^{\lambda_d}.
\]
因此，最优代理呈现专家乘积形式，结合了校准后的边际信号 \(q_d\) 与一个前缀条件校正 \(\rho_d\)。

## 附录 C 混合 dLLM 中 PRESTO 的二次自推测解码模式

我们提供关于 PRESTO 如何在混合 dLLM 中实例化并采用二次自推测解码模式的详细说明，对应图 3。与标准推测解码设置（其中小草稿模型为大验证模型提出令牌）不同，自推测 dLLM 使用*同一个 dLLM* 同时作为草稿模型和验证模型：单次前向传播同时验证先前草稿的令牌并在尾部掩码位置推测新令牌。PRESTO 将这一单模型管线转变为基于树的草稿生成器，无需额外模型。在轮次 \(r\) 开始时，输入序列由三个连续区域组成：

- • 一个**已接受前缀**，包含先前轮次中接受的令牌（图 3 中带有 ✓ 的绿色令牌，例如 C、D）；
- • 一个**验证区域**，包含在轮次 \(r-1\) 中提出的草稿令牌（例如 E、F、G），其中一部分将在本轮次中被接受；