Trees from Marginals: Autoregressive drafting with factorized priors

Abstract: Speculative decoding greatly increases the interactivity of autoregressive LLMs by trading off computation for extra tokens generated in a single forward pass. Factorized draft models are especially efficient because they predict future-token marginals in parallel, but their independence assumption causes acceptance rates to degrade sharply as the speculative budget grows. We analyze this limitation and introduce Weaver, a lightweight autoregressive adapter that constructs proposal trees from the top-K marginals of a factorized drafter. Weaver restores conditional dependencies between proposed tokens while avoiding a full-vocabulary projection. To support fast verification for models with Gated Delta Net layers, we derive a rollback-free tree-verification algorithm and implement optimized CUDA kernels in SGLang. By combining these model and systems contributions we achieve a 4.37-fold speedup over autoregressive decoding, and outperform a highly optimized DFlash baseline by 24.7%.

• The paper introduces DFlash-TfM, a hybrid speculative decoding framework that uses factorized marginal predictions to guide an autoregressive adapter named Weaver, achieving up to 4.37× speedup in inference.
• The methodology employs top-K constrained decoding and a custom CUDA kernel for rollback-free tree verification, significantly reducing compute and memory overhead.
• Empirical results demonstrate a 7.1× speedup over naive implementations and notable improvements in interactivity across benchmarks like chat, math, and code generation.

Autoregressive Drafting with Factorized Priors: Trees from Marginals

Motivation and Problem Setting

Efficient inference in autoregressive LLMs is limited by the sequential dependence of token generation, resulting in suboptimal hardware utilization and low interactivity in single-query, low-batch regimes—critical for local deployments and interactive applications. Speculative decoding, particularly via auxiliary draft models, has emerged as a key approach to amortize generation over multiple proposed tokens per verifier pass, but current architectures are constrained either by scalability, acceptance rates, or memory bandwidth, especially as the speculative budget increases.

The dichotomy between factorized drafters (which parallelize generation via marginal predictions but suffer sharply degraded acceptance rates with increasing proposal length due to independence assumptions) and lightweight autoregressive drafters (which capture dependencies but are bottlenecked by full-vocabulary projections) drives the investigation of hybrid approaches.

Weaver and DFlash-TfM: Architectural Advancements

The paper introduces DFlash-TfM, a hybrid speculative decoding framework that leverages factorized marginal predictions as priors for an autoregressive drafter—Weaver—that restores dependency structure via an efficient, top-K constrained decoding procedure. Weaver operates only over the highest-scoring candidate tokens as determined by the DFlash factorized model, thereby obviating the primary bandwidth cost of traditional autoregressive drafters.

Proposal trees are constructed autoregressively, with Weaver utilizing both verifier and marginal drafter states to predict logit residuals over the candidate set. The result is a lightweight adapter with 56.7M parameters, implemented as a single transformer layer, sharing embeddings and output projection with the verifier but only querying the top-K subset (typically $K=512$ out of the full vocabulary), significantly reducing compute and memory pressure.

Tree Verification and GDN Compatibility

A central systems contribution is the derivation and implementation of a rollback-free tree verification kernel compatible with Gated Delta Net (GDN) layers, which pose challenges for attention-masked single-pass verification due to non-commuting state transitions. The authors propose a masked triangular solve algorithm operating in the dual chunk form of the GDN recurrence, avoiding speculative state updates and fusing all required operations in a single custom CUDA kernel within the SGLang framework. This approach yields constant-time verification with respect to the number of branches, scaling substantially better than per-branch recurrence.

Empirical profiling shows that the fused verification kernel provides up to $7.1\times$ speedup over naïve recurrent implementations at $T=128$ draft tokens, with the verification overhead becoming only a minor fraction of the overall decode step.

Theoretical Implications and Empirical Results

The paper provides a detailed analysis of the acceptance rate limitations inherent to factorized drafters, demonstrating empirically and theoretically that the independence assumption, rather than model capacity, fundamentally bounds their utility at high draft depths. It is established that DFlash-TfM’s acceptance rate for long proposal sequences surpasses the maximal achievable by any marginal-only drafter—even one trained to optimality—by conditioning proposal generation on realized token sequences within the marginal top-K support.

Strong numerical results are reported: DFlash-TfM achieves a macro-averaged $4.37\times$ speedup over standard greedy autoregressive decoding (Qwen3.6-27B, bfloat16) and $24.7\%$ improvement in interactivity over a highly optimized DFlash baseline. The advantage is attributed to both increased mean accepted length (a $77\%$ improvement over DFlash chain and $32\%$ over DDTree at the same tree size) and superior scaling with increased speculative budgets. The efficacy generalizes across chat, math, and code generation benchmarks (MTBench, GSM8K, MATH500, HumanEval, etc.).

Methodological and Practical Implications

By integrating a small, top-K constrained autoregressive adapter with a high-performance marginal predictor, DFlash-TfM closes the parallelism-quality gap limiting prior speculative decoding work. The acceptance-rate advantage relative to the theoretical bound of marginal drafters is a nontrivial claim, directly contradicting the prevailing assumption that only the marginal quality of the drafter determines scalable speculative utility.

Practically, the approach is especially impactful for local inference and interactive systems, where batching is infeasible and memory efficiency is paramount. The modularity of DFlash-TfM enables straightforward adaptation to future marginal drafters, dynamic tree construction heuristics, or scheduler-level throughput optimizations for multi-request scenarios.

Future Directions

The authors outline several research directions for extending the methodology:

• Tree construction guided by acceptance-probability estimates, rather than draft probability, to optimize resource utilization.
• Extending the marginal window or self-anchoring to surpass existing block-size limitations, exploiting longer-range dependency restoration.
• Joint training of the marginal drafter and autoregressive adapter to further align the marginal support and conditional corrections.
• Adapting the system for concurrent batch-serving paths to maximize throughput in multi-user server deployments.

Conclusion

DFlash-TfM advances speculative decoding by hybridizing factorized and autoregressive drafting within a memory- and compute-efficient architecture, empirically surpassing both practical and theoretical acceptance ceilings of prior factorized approaches. System-level innovations in verification for GDN-layered targets ensure this method is directly applicable to modern LLM backbones. These results have direct implications for accelerating LLM inference in latency-critical, single-query, or low-resource environments, and open avenues for further study in adaptive, conditional drafting, and scalable scheduling for LLM serving workloads.

Reference: "Trees from Marginals: Autoregressive drafting with factorized priors" (2607.06763)