Trees from Marginals: Autoregressive drafting with factorized priors

Autoregressive language models are inherently sequential, making them memory-bound and slow. While speculative decoding uses a fast draft model to generate multiple tokens in parallel, existing factorized drafters suffer from poor acceptance rates as the draft length increases due to their independence assumptions. The authors seek to improve speculative decoding efficiency by restoring conditional dependencies between draft tokens without incurring the high computational cost of a full-vocabulary projection.

The authors propose DFlash-TfM, a hybrid drafting architecture. It uses a factorized drafter (DFlash) to predict marginal distributions for future tokens in parallel. These predictions serve as a prior for Weaver, a lightweight autoregressive adapter. Weaver operates only on the top-K candidate tokens identified by the factorized drafter, which avoids the expensive full-vocabulary projection matrix multiplication. The resulting proposal tree is verified using a traversal-based algorithm. To support modern architectures like Gated Delta Nets (GDN), the authors derive a rollback-free tree-verification algorithm and implement optimized CUDA kernels in SGLang to handle the non-commutative state transitions inherent in these models.

DFlash-TfM significantly improves the mean acceptance length of speculative decoding compared to existing methods. By combining the parallel efficiency of factorized models with the conditional accuracy of autoregressive adapters, the system achieves a 4.37-fold speedup over standard autoregressive decoding. Furthermore, it outperforms the optimized DFlash baseline by 24.7% on the Qwen3.6-27B model. The authors demonstrate that their hybrid approach effectively bridges the gap between the speed of factorized drafting and the accuracy of autoregressive generation.