Accelerating Language Models with Diffusion Draft Trees

Summary & Methodology Analysis

Speculative decoding accelerates autoregressive language models by verifying multiple draft tokens in parallel. However, diffusion-based drafters predict position-wise distributions that are marginal rather than conditioned on tokens selected along each draft path. Existing methods for incorporating causal information along a single draft chain or constructing diffusion-based trees do not fully address the efficiency bottleneck caused by path-dependent scoring when combined with sequential heap operations. To solve this, DARTree constructs a fixed-width candidate tree by expanding and scoring all nodes at each depth in a single batch, using a pretrained causally corrected block-parallel drafter. All branches at the same depth are evaluated together, and their corrected scores are used to retain nodes for the next depth while carrying selected correction states forward. The system retains the global top-W child prefixes at each depth, where W is the layer width, and updates their states. After the supertree is complete, it performs deferred best-first pruning by selecting the global top-B nodes based on their scores, where B is the verification budget, to form the final verification tree. Finally, the resulting prefix tree is verified using the target model in a single forward pass with a tree-attention mask, which uses attention mechanisms to process tree-structured token hierarchies in parallel. Despite these efficiency gains, DARTree has notable limitations. It requires a pretrained diffusion drafter with a causal correction head, meaning it is not directly applicable to naive diffusion drafters without further training. As a speculative decoding method, DARTree does not reduce the total FLOPs, which are floating-point operations measuring total computational work, but uses additional computation to reduce inference latency. It incurs greater computational overhead than non-tree-based methods due to verifying a large tree of candidate tokens, making it potentially unsuitable for all deployment environments. Furthermore, DARTree is primarily designed for low-concurrency serving where decoding is memory-bandwidth-bound and compute resources are underutilized.