DARTree: Speculative Diffusion Decoding with Autoregressive Draft Trees

DARTree: Speculative Diffusion Decoding with Autoregressive Draft Trees§

Core Idea§

DARTree combines the parallelism of diffusion-based drafters with the causal precision of autoregressive (AR) correction heads. Existing diffusion drafters predict entire token blocks in parallel but output marginal distributions that ignore the specific draft path. Conversely, recursive correction (e.g., DFlash) injects causal information along a single chain, limiting candidate diversity. DARTree generalizes a pretrained AR correction head from chains to trees, allowing the head to score multiple candidate draft branches simultaneously in a training-free manner.

Method§

Tree Construction

Let $D$ be a diffusion drafter and $H$ an AR correction head (from DFlash or similar). Starting from a prefix, DARTree builds a fixed-width candidate tree of depth $K$ and width $W$. At each depth $t$, it expands all nodes from the previous depth in a single batched call to $D$, generating $W$ candidate tokens per node. All resulting candidate sequences are then scored by $H$ in one batch:

$$ P_{\text{AR}}(v^{(t)}_i \mid \text{path}(v^{(t)}_i)) = H(\text{path}(v^{(t)}_i)) $$

Instead of sequentially maintaining a heap (which would serialize AR inference), DARTree applies best-first pruning after each depth: it selects the top-$W$ scored nodes across the entire frontier, discarding the rest. This decouples AR-head inference from heap operations, making the entire tree construction parallelizable.

Verification

The pruned tree—now containing multiple candidate branches with AR-compatible probabilities—is fed to the target LLM as a standard speculative decoding verification step. The model verifies all candidate tokens in parallel, accepting the longest prefix that matches its own distribution. Because the AR head has already injected causal conditioning along each branch, the acceptance rate improves dramatically over both chain-based (DFlash) and marginal tree-based (Domino) methods.

Implementation Sketch§

import torch

def build_dartree(prefix, draft_model, ar_head, width=4, depth=8):
    # prefix: prefix token ids (list)
    nodes = {0: [tuple(prefix)]}
    for t in range(1, depth + 1):
        candidates = []
        parents = []
        # Expand all nodes at depth t-1 in a batch
        for path in nodes[t - 1]:
            proposals = draft_model.batch_propose(list(path), num_drafts=width)
            for tok in proposals:
                candidates.append(tuple(list(path) + [tok]))
                parents.append(path)
        # Score all expanded candidates with AR head in one forward pass
        scores = ar_head.score_batch(candidates)  # shape: [num_candidates]
        # Best-first pruning: keep only top 'width' candidates across whole frontier
        top_idx = torch.argsort(scores, descending=True)[:width]
        nodes[t] = [candidates[i] for i in top_idx]
    return nodes

Results§

Across seven math, code, and chat benchmarks, DARTree achieves the highest average acceptance length and speedup in all four model–temperature configurations tested. Notably, it accepts up to 12.97 tokens per verification round—98.6% more than DFlash and 27.9% more than Domino—and reaches up to 9.73× lossless speedup over autoregressive decoding. The training-free nature and compatibility with existing drafters make it a drop-in acceleration layer for deployed LLMs.