TAPS：面向扩散草稿推测解码的目标感知前缀树选择

# TAPS：针对扩散草稿推测解码的目标感知前缀树选择

来源：https://arxiv.org/html/2606.00487

Zhuoyu Wang∗，Junnan Huang∗，Xinyu Chen†
香港科技大学（广州）
∗共同第一作者。†通讯作者。

###### 摘要

使用扩散模型进行并行草稿生成是推测解码中一种有前景的方法。通过在单次前向传播中预测多个未来位置的词元，扩散草稿生成器显著降低了草稿生成延迟。然而，这将瓶颈转移到了验证阶段：验证单一序列会限制接受长度，而验证大型草稿树则会引入过量的目标模型延迟。我们识别出现有草稿树方法中的一个关键不匹配：现有的扩散树方法根据边际概率对节点进行排序，忽略了验证是基于前缀条件的。因此，它们可能会验证被拒绝前缀下的不可达后代，从而增加延迟而只能获得有限的接受收益。为了解决这个问题，我们提出了 TAPS，一种目标感知的前缀选择方法，它将扩散边际概率转化为路径条件化的接受估计。然后，TAPS 在固定的验证预算下选择一个紧凑的前缀封闭子树，从而改善接受-成本权衡，而不是简单地扩展草稿树。跨不同数据集和模型族的实验表明，与标准的自回归解码相比，TAPS 实现了高达 7.9 倍的无损端到端加速，分别比最先进的 DFlash 和 DDTree 提升了 1.36 倍和 1.74 倍。我们的工作可在 https://anonymous.4open.science/r/TAPS-EMNLP2026-53DD 获取。

TAPS：针对扩散草稿推测解码的目标感知前缀树选择

Zhuoyu Wang∗，Junnan Huang∗，Xinyu Chen†
香港科技大学（广州）
∗共同第一作者。†通讯作者。

![参见说明](图1：整体吞吐量-接受权衡。我们在 Qwen3-4B 和 Qwen3-8B 设置下比较了 DFlash、DDTree 和 TAPS，平均结果来自 A40 GPU 上的所有基准测试。TAPS 实现了比先前方法更好的吞吐量-接受权衡，在提高吞吐量的同时保持了具有竞争力的接受长度。)

![参见说明](图2：跨模型和 GPU 平台的平均加速比。TAPS 在目标模型和硬件平台上始终优于 EAGLE-3、DFlash 和 DDTree。)

![参见说明](图3：洞察 1：(a) DFlash 边际分布中目标接受词元的排名分布。(b) 在按位置 Top-8 选择以及不同预算的选择树中，每个草稿块位置包含正确词元的概率。洞察 2：(c) 随着树节点预算增加，每轮时间分解和验证效率。所有测量均使用 Qwen3-4B 在多个数据集上收集。)

![参见说明](图4：DFlash、DDTree 和 TAPS 的比较。)

![参见说明](图5：TAPS 概述。扩散草稿生成器首先生成一个大型边际候选池；目标感知评分层将边际词元置信度转换为路径条件化的可达概率；自适应选择器随后保留一个紧凑的前缀封闭子树，用于高效的目标模型验证。)

## 1 引言

推测解码已成为加速大型语言模型（LLM）推理而不牺牲输出质量的最有效方法之一 (Leviathan et al., 2023 (https://arxiv.org/html/2606.00487#bib.bib8); Chen et al., 2023 (https://arxiv.org/html/2606.00487#bib.bib9))。其关键思想是使用一个轻量级的草稿模型快速提出多个候选词元，然后目标模型通过单次前向传播并行验证这些词元。通过将目标模型推理的成本分摊到多个被接受的词元上，推测解码实现了显著的挂钟加速，同时可证明地保留了目标模型的输出分布。

基于扩散的草稿生成已成为推测解码中一种特别有前景的策略。与必须逐个生成候选词元的自回归草稿生成器不同，DFlash (Chen et al., 2026 (https://arxiv.org/html/2606.00487#bib.bib14)) 等扩散草稿生成器在单次前向传播中预测整个未来词元块，从而显著减少了草稿生成延迟。为了进一步提高接受长度，近期方法将候选词元组织成树结构，允许目标模型同时验证多个替代路径 (Miao et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib10); Li et al., 2024b (https://arxiv.org/html/2606.00487#bib.bib11))。结合树结构验证，DDTree (Ringel and Romano, 2026 (https://arxiv.org/html/2606.00487#bib.bib15)) 等方法从扩散草稿预测中构建大型候选树并并行验证所有分支，从而提高了平均接受长度。

然而，扩散草稿生成改变了瓶颈：一旦草稿生成变得廉价，目标模型验证就主导了端到端延迟。树验证可以通过在单次目标模型前向传播中检查多个连续路径来增加接受长度，但其成本随着所选节点数量的增加而增长。由于每轮解码最终只承诺一条前缀路径，被拒绝前缀下的节点增加了验证延迟却没有贡献接受的词元。为了解决这个问题，我们提出了 TAPS（目标感知前缀选择），该方法学习预测扩散草稿树中的哪些候选节点最有可能被目标模型接受，并动态选择一个紧凑的子树，在验证预算约束下提高吞吐量。

我们的主要贡献如下：

- **目标感知接受建模。** 我们引入了一个轻量级的学习评分器，用于估计每个候选词元沿其草稿路径的条件接受概率，用目标感知评分取代了基于草稿置信度的选择。
- **成本感知的动态树剪枝。** TAPS 将接受概率传播到候选树中，并贪婪地选择一个紧凑的前缀封闭子树，在验证预算约束下最大化期望接受的词元数。
- **显著的端到端加速。** 如图1 (https://arxiv.org/html/2606.00487#S0.F1) 所示，TAPS 在保持相对较高的平均接受长度 τ 的同时，实现了最高的端到端吞吐量。图2 (https://arxiv.org/html/2606.00487#S0.F2) 显示，TAPS 在三个目标模型和两个 GPU 平台上实现了平均 5.44 倍的加速，相比 DFlash 和 DDTree 分别提升了 1.36 倍和 1.74 倍。

## 2 相关工作

#### 推测解码与草稿树。
推测解码通过使用草稿候选词元加速 LLM 推理，这些候选词元由目标模型验证，同时保留目标分布 (Leviathan et al., 2023 (https://arxiv.org/html/2606.00487#bib.bib8); Chen et al., 2023 (https://arxiv.org/html/2606.00487#bib.bib9))。现有方法从不同角度改进候选词元生成，包括辅助草稿模型、自推测头、基于检索的草稿生成、循环草稿生成以及多头特征预测器 (Cai et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib21); Ankner et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib22); Cheng et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib23); He et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib24); Li et al., 2024b (https://arxiv.org/html/2606.00487#bib.bib11))。为了增加每次验证通过的接受长度，基于树的方法使用树注意力验证多个候选连续路径，代表性系统包括 SpecInfer (Miao et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib10))、Sequoia (Chen et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib12)) 和 EAGLE-2 (Li et al., 2024a (https://arxiv.org/html/2606.00487#bib.bib13))。这些工作表明，扩展候选结构可以提高接受长度，但更大的树也会增加目标模型验证成本。

#### 基于扩散的草稿生成。
扩散语言模型为非自回归的从左到右生成提供了一种替代方案，并在离散和连续文本生成设置中得到了研究 (Austin et al., 2021a (https://arxiv.org/html/2606.00487#bib.bib26); Li et al., 2022 (https://arxiv.org/html/2606.00487#bib.bib27); Nie et al., 2025 (https://arxiv.org/html/2606.00487#bib.bib28))。近期的推测解码方法利用这种并行性在一次前向传播中草稿生成多个未来词元。DFlash (Chen et al., 2026 (https://arxiv.org/html/2606.00487#bib.bib14)) 预测块级边际 logits，而 DDTree (Ringel and Romano, 2026 (https://arxiv.org/html/2606.00487#bib.bib15)) 则从这些 logits 构建大型验证树。其他扩散式解码方法探索了轻量级草稿生成、前瞻填充和自回归验证以实现更快的生成 (Liu et al., 2026 (https://arxiv.org/html/2606.00487#bib.bib16); Cheng et al., 2025 (https://arxiv.org/html/2606.00487#bib.bib17); Xu et al., 2025 (https://arxiv.org/html/2606.00487#bib.bib18))。这些方法显著降低了草稿生成延迟，使得目标模型验证成为端到端解码时间中越来越重要的组成部分。

#### 自适应验证效率。
另一条工作线通过根据运行时难度调整推测努力来改进推测解码。代表性方法通过早期拒绝、自适应停止、动态前瞻或基于熵的草稿控制来减少浪费的计算 (Pan et al., 2025 (https://arxiv.org/html/2606.00487#bib.bib19); Zhang et al., 2025 (https://arxiv.org/html/2606.00487#bib.bib20); Mamou et al., 2024 (https://arxiv.org/html/2606.00487#bib.bib25); Liu et al., 2025 (https://arxiv.org/html/2606.00487#bib.bib29))。这些方法表明，固定的推测策略通常不是最优的：简单的上下文可以从更激进的草稿生成中受益，而困难的上下文则需要更保守的验证。

## 3 动机

### 3.1 扩散草稿树的局限性

扩散草稿模型通过并行预测多个未来位置来减少草稿生成延迟。然而，一旦草稿生成变得廉价，瓶颈就转移到了目标模型验证上。如图4 (https://arxiv.org/html/2606.00487#S0.F4) 所示，DFlash (Chen et al., 2026 (https://arxiv.org/html/2606.00487#bib.bib14)) 验证单一草稿路径，保持验证成本较低但限制了接受长度。DDTree (Ringel and Romano, 2026 (https://arxiv.org/html/2606.00487#bib.bib15)) 将扩散草稿扩展为候选树，允许目标模型验证多个分支并接受更长的前缀。然而，现有的扩散草稿树仍然存在两个局限性。

**局限性 1：边际选择不考虑路径。** DDTree 主要根据扩散草稿生成器产生的边际概率来选择候选节点。这将候选质量视为按位置区分的属性，而树验证是基于前缀条件的：一个节点只有在它的所有祖先都被接受时才是有用的。因此，一个高边际概率的词元如果位于错误的前缀下，它在目标模型验证期间可能变得不可达，从而验证效用很低。图3 (https://arxiv.org/html/2606.00487#S0.F3)(a) 显示，在 99% 的情况下，目标接受的词元出现在草稿生成器的前 8 个边际候选词元中，表明局部候选池通常是足够的。然而，图3 (https://arxiv.org/html/2606.00487#S0.F3)(b) 显示，有预算限制的前缀封闭树仍然丢失了许多正确的词元，尤其是在后面的位置，下降了 10-15%。这表明正确的词元通常在局部可用，但被放置在与目标模型接受路径不同的前缀之下。因此，有效的树选择应该将每个词元与其前缀一起评估，而不是按位置的边际概率独立地对节点进行排序。

**局限性 2：更大的验证树带来的吞吐量收益递减。** 扩展树可以提高覆盖率，但每轮解码最终只能承诺一个被接受的前缀。因此，许多验证过的分支被丢弃，而树注意力和目标模型验证成本则随着所选节点数量的增加而增长。图3 (https://arxiv.org/html/2606.00487#S0.F3)(c) 显示，将树预算从 64 个节点增加到 512 个节点会使验证延迟从 33 毫秒增加到 81 毫秒，同时验证效率稳步下降。这表明，吞吐量最优的树应该在期望接受长度和验证成本之间取得平衡，而不是简单地使用最大的预算。

### 3.2 关键洞察

这些局限性激发了两个设计原则。首先，扩散草稿树的选择应该是路径感知的：由于正确的词元通常已经存在于局部候选池中，选择器应该保留那些前缀在目标模型验证中幸存可能性高的候选词元。其次，验证应该是成本感知的，而不是填满预算：选择器应该选择一个紧凑的前缀封闭子树，以最大化每单位验证成本的期望接受词元数。在这些洞察的指导下，TAPS 使用一个轻量级的目标感知评分器将扩散边际候选词元转换为路径条件化的可达估计，然后在验证预算下选择一个紧凑的前缀封闭子树。这使得候选选择与基于前缀条件的验证相一致，并在不简单地扩展草稿树的情况下改善了接受-成本权衡。

## 4 方法：TAPS

基于上述观察，我们引入了 TAPS，它估计每个词元在其前缀中被接受的可能性，并动态选择用于验证的紧凑子树。

### 4.1 推理流程

TAPS 在标准的推测解码循环中运行，在草稿生成和验证之间交替。图5 (https://arxiv.org/html/2606.00487#S0.F5) 说明了一轮解码过程，分为三个阶段。

在第一阶段，扩散草稿生成器在单次前向传播中为 d 个未来位置生成按位置的边际 logits；每个位置的前 K 个词元被组装成一个最多有 N_pool 个节点的候选树。

第二阶段使用一个轻量级的目标感知评分器对候选树进行评分，并将边级别的接受估计传播为路径级别的*可达概率*。（第4.2节 (https://arxiv.org/html/2606.00487#S4.SS2)）。

在第三阶段，TAPS 在验证预算下选择一个紧凑的子树。生成的子树适应于难度：当草稿生成器自信时子树较大，不自信时则较小，同时平衡草稿置信度和验证成本。然后，通过树注意力使用目标模型验证所选子树，保留精确的输出分布。（第4.3节 (https://arxiv.org/html/2606.00487#S4.SS3)）

### 4.2 路径条件接受算法

如第3节 (https://arxiv.org/html/2606.00487#S3) 所述，根据边际草稿概率对候选词元进行排序未能考虑树验证的顺序门控特性。TAPS 改而在*边*级别分解评分：对于每个父-子对 \((u \to v)\)，我们估计条件接受概率

\[
q_{\mathrm{cond}}(u \to v) = \Pr(\text{accept } v \mid \text{reach } u),
\]

这直接反映了目标模型如何评估每个草稿词元，条件是在同一路径上先前接受的词元前缀。

#### 评分器设计。
一个轻量级评分器将每个候选边 \((u \to v)\) 映射到一个标量 logit \(\ell_{u,v}\)。输入包括父词元和子词元、边的深度以及草稿模型统计量，例如