TAPS：面向扩散模型草稿的感知目标前缀树选择策略

背景

推测解码（Speculative Decoding）旨在通过利用一个小而快的“草稿模型”并行生成多个 token，再由一个更大、更慢的“目标模型”进行验证，从而加速大语言模型的推理过程。近年来，使用扩散模型（Diffusion Models）进行并行草稿生成成为一种极具潜力的方向。与传统自回归草稿模型每次只预测下一个 token 不同，扩散草稿模型能够在单次前向传播中预测多个未来位置的 token，从而显著降低了草稿生成的延迟。

然而，这种并行化策略将性能瓶颈从“草稿生成”转移到了“验证”阶段。在推测解码中，验证过程必须严格遵循前缀条件（prefix-conditioned），即只有当前缀被目标模型接受时，后续的子节点才有资格被验证。现有的扩散树方法通常依据边缘概率（marginal probability）对节点进行排序，却忽略了验证的前缀依赖性。这导致了一个关键的不匹配问题：系统可能会花费大量计算资源去验证那些基于已拒绝前缀的“不可达后代节点”，从而增加了延迟，却未能带来相应的接受长度收益。

核心内容

为了解决上述瓶颈，研究团队提出了 TAPS（Target-Aware Prefix Tree Selection，感知目标前缀树选择）。TAPS 的核心思想是将扩散模型输出的边缘概率转化为路径条件化的接受估计（path-conditioned acceptance estimates），并在此基础上，在固定的验证预算下选择紧凑的前缀闭合子树（prefix-closed subtree）。

1. 现有方法的局限性分析

在传统的扩散草稿树构建中，节点重要性通常由其边缘概率决定。然而，在推测解码的验证阶段，验证器（Target Model）只能接受以特定前缀开头的序列。如果前缀被拒绝，其下的所有子节点在逻辑上都是无效的，无需验证。现有方法由于忽略了这一约束，往往构建出庞大且包含大量无效分支的草稿树，导致目标模型的验证延迟激增，抵消了扩散模型在草稿生成阶段的加速优势。

2. TAPS 的工作机制

TAPS 通过以下步骤优化草稿树的选择：

• 路径条件化估计：TAPS 不直接使用扩散模型的全局边缘概率，而是计算在给定前缀路径条件下的 token 接受概率。这意味着它评估的是“如果当前路径被接受，后续节点被接受的概率”。
• 预算约束下的子树选择：在目标模型有限的验证预算（即允许的验证 token 数量或时间）内，TAPS 算法会选择那些预期接受收益最高且构成前缀闭合结构的子树。前缀闭合意味着如果选择了一个节点，必须选择其所有祖先节点。
• 优化接受-成本权衡：TAPS 的目标不是简单地扩大草稿树的规模，而是优化“接受长度”与“验证成本”之间的权衡。通过剔除低概率或不可达的分支，它确保了每一单位的验证计算都用于最有希望的序列上。

3. 实验结果

研究者在多个数据集和模型家族上进行了广泛实验。结果表明，TAPS 相比传统的自回归解码实现了高达 7.9 倍 的无损端到端加速。与当前最先进的扩散推测解码方法 DFlash 和 DDTree 相比，TAPS 分别实现了 1.36 倍 和 1.74 倍 的性能提升。这证明了通过感知目标验证约束来优化草稿树结构，能更有效地释放扩散模型在并行草稿生成中的潜力。

关键要点

• 瓶颈转移：扩散模型用于推测解码时，瓶颈从草稿生成的延迟转移到了目标模型的验证延迟。
• 核心痛点：现有扩散树方法基于边缘概率排序，忽略了验证的前缀依赖性，导致大量计算浪费在基于已拒绝前缀的无效节点上。
• TAPS 创新：提出了一种感知目标的前缀选择方法，将扩散边缘概率转化为路径条件化的接受估计。
• 策略优化：在固定验证预算下，TAPS 选择紧凑的前缀闭合子树，优化接受长度与验证成本的权衡，而非盲目扩大树结构。
• 性能突破：相比 vanilla 自回归解码，TAPS 实现最高 7.9 倍无损加速；相比 SOTA 方法 DFlash 和 DDTree，分别提升 1.36 倍和 1.74 倍。

意义与影响

TAPS 的提出标志着扩散模型在语言模型加速领域的应用进入了一个更精细化的阶段。它不仅仅展示了扩散模型在并行生成上的速度优势，更深刻地指出了**“生成”与“验证”之间的协同优化**的重要性。

1. 理论贡献：TAPS 明确了扩散草稿树构建中的“前缀条件化”缺失问题，为后续研究提供了新的评估维度——即草稿树的选择必须考虑目标模型的验证逻辑，而不仅仅是草稿模型的概率分布。
2. 工程价值：通过显著减少无效的验证计算，TAPS 能够降低大模型推理的硬件资源需求（如 GPU 显存占用和计算时间），使得部署更大规模的扩散辅助推测解码成为可能。
3. 未来方向：这项工作启示我们，未来的推测解码优化不应仅局限于草稿模型的架构改进，还应深入探索草稿生成策略与目标验证机制之间的联合优化。TAPS 作为一种通用的前缀选择框架，有望被扩展到其他类型的并行草稿生成方法中，推动生成式 AI 推理效率的进一步提升。