Source-linked AI summary
Accelerating Speculative Decoding with Block Diffusion Draft Trees
Liran Ringel, Yaniv Romano
TL;DR
Speculative decoding seeks lower autoregressive inference latency without changing the target model’s outputs, but vanilla DFlash verifies only one trajectory despite producing per-position distributions. DDTree constructs and efficiently verifies a bounded draft tree from those distributions, and experiments show consistent improvements over vanilla DFlash across all tested settings.
Problem
Speculative decoding must reduce inference latency while retaining accepted target-model tokens, but vanilla DFlash verifies only one trajectory despite producing many plausible continuations.
Method
DDTree constructs a compact draft tree from one block-diffusion forward pass under a fixed node budget and verifies it in one target-model forward pass.
Results
DDTree improves every entry across 60 dataset-model-temperature settings, including all three target models, reasoning, code, general instruction tasks, and both temperatures.
Takeaways & Limitations
DDTree consistently improves over vanilla DFlash while preserving its one-pass drafting setup and enabling multiple continuations per round.
Takeaways & Limitations
DDTree optimizes a surrogate based on the drafter’s factorized distribution because one-pass block diffusion does not expose target-conditioned continuation probabilities.
Abstract
from arXiv · showhide
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.
1 Introduction
Speculative decoding reduces autoregressive inference latency by having a lightweight drafter propose tokens for parallel target-model verification. DDTree extends block-diffusion drafting by selecting and verifying multiple promising continuations under a fixed budget, with consistent gains over vanilla DFlash.
- Autoregressive decoding is latency-intensive because each token requires another forward pass through the large model.
- Speculative decoding uses a lightweight drafter and parallel target-model verification while preserving the target model’s output distribution.
- DFlash generates an entire draft block in one forward pass and outperforms strong autoregressive drafters such as EAGLE-3.
- DDTree addresses DFlash’s single-trajectory verification by constructing a draft tree from per-position distributions and selecting continuations under a specified node budget.
- DDTree verifies the selected tree in one target-model forward pass with tree attention and consistently improves over vanilla DFlash across models.
- DDTree’s tree provably maximizes expected accepted tokens under the draft model, using that quantity as a surrogate for target-model acceptance.
2 Related Work
Prior work combines speculative verification trees with parallel drafting, while DDTree specifically integrates tree construction with DFlash’s one-pass block-diffusion probabilities. Unlike DART, DDTree avoids external scoring and provides a provably maximized surrogate objective.
- Tree-based speculative decoding generalizes verification from one drafted continuation to a tree of candidate continuations.
- SpecInfer introduces tree attention for efficient target-model verification, and Medusa applies similar tree verification with multiple prediction heads.
- Block-diffusion research reduces drafting latency by predicting multiple future tokens in a single forward pass.
- DART constructs trees from one-pass parallel logits but uses continuity-aware pruning with an external N-gram score and runtime trie.
- DDTree uses the one-pass DFlash per-position probabilities directly, avoiding auxiliary external scoring and maximizing an explicit surrogate objective with best-first construction.
3 Background: Block Diffusion Drafting
Block diffusion drafting predicts a future token block in parallel from the current context and bonus token, producing one marginal distribution per position. DDTree uses these marginals to build a compact tree rather than collapsing them into one continuation, while recognizing they are not path-conditioned probabilities.
- Each decoding round begins with a target-model-produced bonus token that the drafter can use without the target model having yet processed it in the appended context.
- A block diffusion drafter predicts the next L masked positions in parallel from a context containing the bonus token and masked future positions.
- The drafter outputs per-position token distributions, with each q_i representing a marginal rather than a path-conditioned distribution.
- A one-pass drafter therefore provides per-position marginals q_i(y_i | c, b), not continuation-conditioned factors for every partial prefix.
- Unlike the target model’s autoregressive continuation distribution, the drafter does not condition position i on realized tokens at earlier positions within the same block.
- DDTree uses the factorized one-pass distribution to construct a compact draft tree for verification instead of collapsing the marginals into one continuation.
4 The Proposed Method: Diffusion Draft Tree
DDTree constructs a budgeted draft tree from one block diffusion drafter pass and verifies it with the target model. Its tree-selection objective is a surrogate based on the drafter’s factorized distribution, for which a best-first algorithm recovers an optimal valid tree.
- 4.2 Surrogate objective for draft-tree selection: DDTree seeks a valid draft tree with at most B nodes that maximizes expected speculative-token acceptance under a drafter-based surrogate.The true target-conditioned objective is unavailable from a single block diffusion pass, so the method optimizes expected acceptance under the factorized draft distribution Q.
- 4.1 Overview: One block diffusion drafter pass produces per-position distributions for the next L positions after the bonus token b.The bonus token is already available at the start of the round and serves as the tree root.
- 4.1 Overview: DDTree builds a B-node tree, verifies all drafted nodes in one target-model forward pass with tree attention, and walks the tree using the target model’s decoding rule.The matched path is accepted, while the first unmatched target token becomes the bonus token for the next round.
- 4.2 Surrogate objective for draft-tree selection: The surrogate expected acceptance length decomposes into an additive sum of prefix probabilities over the nodes in the tree.This reduces tree selection to choosing a prefix-closed set of high-probability nodes under the node budget.
- 4.2 Surrogate objective for draft-tree selection: The optimal surrogate tree consists of the B highest-probability prefixes and is valid under the prefix-closure constraint.Proposition 2 states that these prefixes maximize the surrogate expected acceptance length among valid trees with at most B nodes.
- 4.2 Surrogate objective for draft-tree selection: The optimality guarantee applies to the factorized draft distribution Q, not to the true target-model distribution.The construction is exact for the surrogate induced by information available from one block diffusion drafter forward pass.
- 4.3 Efficient and optimal tree construction: A simple best-first heap search returns an optimal valid draft tree for the surrogate objective under node budget B.The search explores each prefix’s next sibling and best child rather than enumerating exponentially many prefixes; logarithms preserve ordering while improving numerical stability.
5 Experiments
DDTree is evaluated for end-to-end decoding speed and acceptance behavior across models, domains, and temperatures, consistently outperforming vanilla DFlash. Its budget controls a tradeoff between longer accepted prefixes and verifier overhead.
- The benchmark evaluates Qwen3-4B, Qwen3-8B, and Qwen3-Coder-30B-A3B-Instruct across reasoning, code, instruction, and dialogue datasets at temperatures 0.0 and 1.0.
- DDTree improves every entry across 60 dataset-model-temperature settings, spanning three target models, reasoning, code, and general instruction tasks.
- Budget-quality tradeoff: On MATH-500 with Qwen3-8B at temperature 0.0, acceptance length rises with node budget while speedup peaks around budgets of 256 to 512.At budget 1024, additional verification overhead outweighs the gain from longer accepted prefixes.
- Acceptance length distribution: At budget B = 512, DDTree shifts the MATH-500 acceptance-length distribution toward longer prefixes and more full-block acceptances.
A Mathematical proofs
The proofs establish that selecting the highest-probability prefixes under the draft-model surrogate yields an optimal valid tree under node budget B, and that the heap algorithm returns it.
- Every ancestor has greater probability than its descendants, making the set of top-B prefixes prefix-closed and therefore a valid tree.
- The surrogate objective is additive over prefix probabilities, so the optimal valid tree consists of the highest-probability prefixes up to budget B.
- The heap algorithm generates sibling and child candidates so that each highest-probability remaining prefix is available when selected.
- After B pops, the algorithm returns exactly the top-B elements of the restricted prefix set and preserves parent inclusion.
- Consequently, Algorithm 1 returns an optimal valid draft tree for the surrogate objective under node budget B.
B Benchmark details
The benchmark fixes model, decoding, hardware, and timing conditions while comparing DDTree with baselines under attention implementations that favor the baselines.
- All runs use block size 16, DDTree budgets from 16 to 1024, temperatures 0.0 and 1.0, a 2048-token limit, and bfloat16 inference.Evaluation runs on 8 H200 GPUs with warmup prompts before timing.
- DDTree uses standard PyTorch scaled dot product attention because FlashAttention-2 does not support its required tree attention pattern.
- For the autoregressive baseline and vanilla DFlash, the faster result between standard attention and FlashAttention-2 is reported.This comparison can only improve the baselines relative to DDTree.
- Table 2 reports the number of evaluated examples for each dataset, following the original DFlash benchmark setup.