用目标感知方法优化扩散采样树,显著提升生成速度与效率
TAPS: Target-Aware Prefix Tree Selection for Diffusion-Drafted Speculative Decoding

- 基于目标感知的前缀选择策略,将扩散概率转为路径条件估计
- 在固定验证预算下选出紧凑闭合子树,提升接受率与成本平衡
- 实测速度比自回归生成快7.9倍,优于当前最优方案1.74倍
使用扩散模型进行并行起草是推测解码的有前途方法。通过单次前向传播预测多个未来位置的词元,扩散起草器显著降低了起草延迟。然而,这将瓶颈转移到验证阶段:验证单个序列限制了接受长度,而验证大型起草树则带来过高的目标模型延迟。我们发现现有起草树方法存在关键不匹配:现有扩散树方法按边际概率排序节点,忽略了验证是前缀条件的。因此,它们可能验证被拒绝前缀的不可达后代,增加延迟但收益有限。为此,我们提出TAPS,一种目标感知的前缀选择方法,将扩散边际概率转化为路径条件接受估计。TAPS随后在固定验证预算下选择一个紧凑的闭合子树,改善接受-成本权衡,而非简单扩展起草树。在多种数据集和模型族上的实验表明,TAPS相比原始自回归解码实现高达7.9倍的无损端到端加速,优于最先进的DFlash和DDTree,分别提升1.36倍和1.74倍。
原文摘要 · Abstract (English)
Using a diffusion model for parallel drafting is a promising approach for speculative decoding. By predicting tokens at multiple future positions in a single forward pass, diffusion drafters substantially reduce drafting latency. However, this shifts the bottleneck to verification: verifying a single sequence limits acceptance length, while verifying large draft trees incurs excessive target-model latency. We identify a key mismatch in existing draft-tree methods: existing diffusion-tree methods rank nodes by the marginal probability, ignoring that verification is prefix-conditioned. As a result, they may verify unreachable descendants of rejected prefixes, increasing latency with limited acceptance gains. To address this, we propose TAPS, a target-aware prefix selection method that turns diffusion marginals into path-conditioned acceptance estimates. TAPS then selects a compact prefix-closed subtree under a fixed verification budget, improving the acceptance-cost tradeoff rather than simply expanding the draft tree. Experiments across diverse datasets and model families demonstrate that TAPS achieves up to 7.9x lossless end-to-end speedup over vanilla autoregressive decoding, outperforming state-of-the-art DFlash and DDTree by 1.36x and 1.74x respectively. Our work is available at https://anonymous.4open.science/r/TAPS-EMNLP2026-53DD
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。