提升扩散模型生成效率,让多候选路径更高效验证
PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding

- 用树形结构探索扩散模型的多个候选生成路径
- 相比原有方法,平均提速1.5倍,最长生成长度显著增加
- 适合追求生成速度与推理效率的扩散模型研究者
扩散型大语言模型(dLLMs)作为自回归模型的替代方案,可并行生成多个词元,使其成为推测解码(SD)中理想的草稿模型,可在一次前向传播中生成一整块草稿词元。然而现有基于扩散的方法仍依赖线性草稿生成,尽管dLLMs在各位置输出多个候选词元,导致解码路径存在巨大组合空间,限制了接受长度与解码效率。为利用这一多候选结构,我们引入基于树的草稿机制,探索多样化候选路径。但发现原始树形草稿不理想:扩散模型的置信度分布对前缀无感知,与基于前缀的自回归验证不匹配,导致路径排序不可靠。为此提出PRESTO,一种系统性框架,通过前缀对齐评分与基于优先级的树搜索,解决扩散草稿置信度与前缀验证间的根本不匹配问题。核心原则是:(1) 候选排序需契合自回归验证的前缀特性;(2) 树构建应优先高验证潜力路径以最大化接受长度。大量实验表明,PRESTO在最先进专用扩散草稿器上实现平均1.5倍端到端吞吐加速,在自推测扩散模型上平均达1.12倍加速,覆盖多种基准测试。
原文摘要 · Abstract (English)
Diffusion Large Language Models (dLLMs) have emerged as a promising alternative to autoregressive (AR) LLMs, generating tokens in parallel. This makes them effective draft models for speculative decoding (SD), producing an entire block of draft tokens in a single forward pass. Yet existing diffusion-based drafting methods rely on linear drafting, even though dLLMs emit multiple candidate tokens across positions, inducing a large combinatorial space of decoding paths. Consequently, they limit acceptance length and decoding efficiency. To exploit this multi-candidate structure, we apply tree-based drafting to diffusion drafters, enabling exploration of diverse candidate paths. However, we find that naive tree drafting is suboptimal: diffusion marginals are prefix-blind, mismatching the prefix-based AR verification and yielding unreliable path ranking. We propose PRESTO, a principled framework that extends tree-based drafting to diffusion drafters while resolving the fundamental mismatch between diffusion draft confidence and prefix-based AR verification through PREfix-aligned Scoring and priority-based Tree search for diffusion speculative decOding. The key principles behind PRESTO are that (1) candidate ranking should align with the prefix-based nature of AR verification, and (2) tree construction should prioritize candidate paths with high verification potential to maximize acceptance length. Extensive experiments show that PRESTO achieves up to an average of $1.5\times$ end-to-end throughput speedup on the state-of-the-art dedicated diffusion drafter SD and an average of $1.12\times$ on self-speculative diffusion LLMs across diverse benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。