用扩散模型构建树状草稿,提升自回归语言模型的推理速度。
Accelerating Speculative Decoding with Block Diffusion Draft Trees
- 基于扩散模型生成草稿块,并构建多路径草稿树。
- 在固定节点预算下,通过优先级算法选择最可能匹配目标模型的路径。
- 单次前向传播完成树结构验证,适合高吞吐场景使用。
推测解码通过轻量级草稿模型预测多个未来词元,由目标模型并行验证以加速自回归语言模型。DFlash表明,基于块扩散的草稿模型可在一次前向传播中生成整个草稿块,性能优于EAGLE-3等强自回归草稿模型。然而,原始DFlash每轮仅验证单一草稿轨迹,可能限制接受长度。本文提出DDTree(扩散草稿树),直接利用块扩散草稿模型各位置的概率分布构建草稿树。在固定节点预算下,采用简单优先队列算法选择最可能与目标模型匹配的延续路径(基于草稿模型输出定义的代理指标)。最终树结构可通过仅包含祖先的注意力掩码,在一次目标模型前向传播中高效验证。由于依托于领先草稿模型DFlash,DDTree在推测解码中达到顶尖水平。
原文摘要 · Abstract (English)
Speculative decoding accelerates autoregressive language models by using a lightweight drafter to propose multiple future tokens, which the target model then verifies in parallel. DFlash shows that a block diffusion drafter can generate an entire draft block in a single forward pass and achieve state-of-the-art speculative decoding performance, outperforming strong autoregressive drafters such as EAGLE-3. Vanilla DFlash, however, still verifies only a single drafted trajectory per round, potentially limiting its acceptance length. We introduce DDTree (Diffusion Draft Tree), a method that constructs a draft tree directly from the per-position distributions of a block diffusion drafter. Under a fixed node budget, DDTree uses a simple best-first heap algorithm to select the continuations that are most likely to match the target model according to a surrogate defined by the draft model's output. The resulting tree is verified efficiently in a single target model forward pass using an ancestor-only attention mask. Because DDTree builds on DFlash, a leading draft model for speculative decoding, these gains place DDTree among the leading approaches to speculative decoding.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。