🤖 machine learning

Trees from Marginals: Autoregressive drafting with factorized priors

The paper introduces Weaver, a lightweight autoregressive adapter that reconstructs conditional dependencies from factorized draft marginals to enable efficient tree-based speculative decoding, achieving a 4.37-fold speedup over standard autoregressive decoding through a novel rollback-free verification algorithm and optimized CUDA kernels.

Original authors: Yuma Oda, Ryan Mathieu, Roman Knyazhitskiy, Artur Chakhvadze

Published 2026-07-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuma Oda, Ryan Mathieu, Roman Knyazhitskiy, Artur Chakhvadze

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a story with a very smart, but slow, librarian (the AI model). Every time you ask for the next word in the story, the librarian has to stop, think hard, check their entire massive library of books, and then whisper the next word to you. This is how current AI works: one word at a time, one step at a time. It's accurate, but it's slow.

The paper introduces a new way to make this librarian much faster without losing accuracy. They call their method "Trees from Marginals" (or DFlash-TfM). Here is how it works, broken down into simple analogies.

The Problem: The "Guessing Game" Limit

To speed things up, researchers invented a trick called Speculative Decoding.

The Old Way: A fast, junior assistant (the "drafter") guesses the next few words. Then, the slow librarian (the "verifier") checks if those guesses are right. If they are, the librarian accepts them all at once. If not, the librarian fixes the mistake and starts over.
The Problem with "Factorized" Drafters: Some assistants are super fast because they guess all the next words at once, ignoring how they connect to each other. It's like a chef who guesses the next three ingredients of a soup without tasting the previous ones.
- The Catch: As the list of guesses gets longer, the chef gets worse at guessing the sequence. The first guess might be right, but the third guess is usually wrong because it didn't account for the first two. This limits how many words can be accepted at once.

The Solution: The "Weaver" Assistant

The authors created a new system that combines the speed of the fast chef with the logic of a careful editor. They call the new editor Weaver.

The "Top-K" Shortlist: First, the fast assistant (DFlash) makes a quick, rough guess and provides a shortlist of the top 512 most likely words for the next spot. It's like a chef saying, "I think the next ingredient is probably one of these 512 spices."
The Weaver's Job: Instead of guessing blindly, the Weaver (a tiny, lightweight AI) looks at that shortlist. It acts like a smart editor who says, "Okay, if the first word was 'salt', then the next word is almost certainly 'pepper', not 'sugar'."
Building a Tree: The Weaver doesn't just make a straight line of guesses. It builds a tree.
- Imagine a family tree. The root is the current sentence.
- The Weaver branches out, creating different possible paths for the story (e.g., "The cat sat on the mat" vs. "The cat sat on the floor").
- Because the Weaver is small and only looks at the shortlist provided by the fast assistant, it is incredibly quick to build this tree of possibilities.

The Verification: Checking the Tree

Now the slow librarian has to check this tree of guesses.

The Old Problem: If the librarian uses a standard "recurrent" memory system (like the Gated Delta Net layers in modern AI), checking a tree is usually a nightmare. It's like trying to walk down every single branch of a tree one by one to see which path is real. This is slow.
The New Trick: The authors invented a special mathematical shortcut (a "rollback-free" algorithm).
- Instead of walking every branch, they use a masked triangular solve. Think of this as a magic map that lets the librarian look at the entire tree structure at once and instantly know which path is the correct one, without having to re-calculate the memory state for every single branch.
- This is like having a GPS that instantly highlights the correct route on a complex map without you having to drive down every dead-end street first.

The Result: Speed and Efficiency

By combining these ideas, the system achieves two major wins:

More Accepted Words: Because the Weaver fixes the logic errors of the fast assistant, the librarian accepts longer chains of words (up to 77% more than the previous best method).
Massive Speedup: The whole process is so efficient that the AI generates text 4.37 times faster than the standard slow method. It also beats the previous "fastest" method by about 25%.

Summary Analogy

Standard AI: A snail writing a story, one letter at a time, checking every letter against a dictionary.
Old Fast Method: A speed-reader guessing the whole next paragraph, but often getting the middle of the paragraph wrong because they didn't pay attention to the beginning.
This New Method (Weaver): A speed-reader who quickly picks the top 500 words that might fit, and a tiny, super-smart editor (Weaver) who instantly arranges those words into a branching tree of the most logical sentences. A special "magic map" (the new kernel) then checks the whole tree instantly to see which path is real.

The result is an AI that writes as fast as a speed-reader but with the accuracy of a careful editor, making interactions feel much more instant and responsive.

Technical Summary: Trees from Marginals (DFlash-TfM)

Problem Statement

Autoregressive Large Language Models (LLMs) are inherently sequential, generating one token per step. This sequential nature creates a memory-bound bottleneck where latency is dominated by DRAM bandwidth rather than compute, leading to poor energy efficiency and limited interactivity (tokens decoded per second per sequence). While speculative decoding offers a solution by using a fast "drafter" model to propose multiple tokens for parallel verification by the main "verifier" model, existing approaches face trade-offs:

Factorized Drafters: Models like DFlash predict future tokens as conditionally independent marginals in parallel. While efficient, their independence assumption causes acceptance rates to degrade sharply as the speculative budget (draft length) increases, because they fail to capture conditional dependencies between proposed tokens.
Autoregressive Drafters: These capture dependencies correctly but incur high overhead due to full-vocabulary projections at every step, negating speed gains.
Verification on Recurrent Models: Standard tree-verification algorithms rely on attention masks. However, modern models utilizing Gated Delta Net (GDN) layers (e.g., Qwen3.6) use non-commutative state transitions that do not support simple ancestor-restricted masking, making efficient single-pass tree verification an open problem.

Methodology: DFlash-TfM

The authors propose DFlash-TfM (Trees from Marginals), a hybrid architecture that combines the parallel efficiency of factorized drafters with the conditional accuracy of autoregressive models, alongside a novel system-level verification kernel.

1. The Weaver Adapter

The core innovation is Weaver, a lightweight autoregressive adapter (56.7M parameters) that operates on a constrained vocabulary.

Input: Weaver takes the hidden states from a pre-trained factorized drafter (DFlash) and the verifier.
Mechanism: Instead of projecting to the full vocabulary, Weaver extracts the Top- $K$ marginal tokens predicted by DFlash. It then autoregressively predicts the residual distribution over this small candidate set ( $K=512$ ).
Benefit: This restores the conditional dependencies between tokens (e.g., token $t+1$ depends on the realized token $t$ ) without the memory-bandwidth cost of a full-vocabulary projection. The model constructs a proposal tree where branches are sampled from these conditioned residuals.

2. Tree Construction

The authors modify the DySpec tree-construction algorithm to improve throughput. Instead of a strictly sequential best-first expansion, they extract the top- $w$ nodes from the candidate heap and perform $w$ expansions concurrently. This balances memory bandwidth and compute pressure, allowing the construction of the full draft tree in $\lceil B/w \rceil$ sequential operations.

3. Rollback-Free Tree Verification for GDN

To support models with Gated Delta Net layers, the authors derive a rollback-free tree-verification algorithm:

Challenge: GDN layers update state via $S_t = \alpha_t(I - \beta_t k_t k_t^\top)S_{t-1} + \beta_t k_t v_t^\top$ . The non-commutative nature of the matrix update prevents standard cumulative-product optimizations used in diagonal state-space models.
Solution: The authors utilize a chunked form of the linear recurrence. They formulate the verification as a masked triangular solve over the tree structure.
- The tree is sorted topologically.
- The algorithm computes interaction matrices ( $X$ and $Y$ ) representing ancestor relationships.
- It solves a linear system $(I+X)U = \beta V$ to compute outputs without speculatively updating the recurrent state.
- State commitment is delayed until the next decoding iteration, once the accepted branch is known. This eliminates the need for state rollback.
Implementation: This is implemented as a fused CUDA kernel in SGLang, which computes the Gram matrix, solves the system via tiled forward substitution, and applies the output stage in a single pass.

Key Contributions

Hybrid Architecture: Introduction of DFlash-TfM, which uses a small autoregressive adapter (Weaver) to condition factorized marginals, lifting the acceptance rate ceiling of pure factorized drafters at long depths.
Theoretical Insight: Analysis showing that the acceptance rate of factorized drafters is limited by the Total Variation (TV) distance between marginal and conditional distributions, and that conditioning on realized tokens (via Weaver) allows the hybrid model to surpass the theoretical limits of any marginal-only drafter.
System Optimization: Derivation of a novel, rollback-free tree-verification algorithm for Gated Delta Net layers, enabling efficient speculative decoding for models like Qwen3.6 that were previously difficult to optimize for tree-based methods.
Implementation: Optimized CUDA kernels for SGLang that fuse the verification steps, significantly reducing latency compared to per-branch recurrent baselines.

Experimental Results

Evaluated on Qwen3.6-27B (bfloat16) using an NVIDIA B200 GPU across chat, math, and code benchmarks:

Speedup: DFlash-TfM achieves a 4.37× speedup over standard autoregressive decoding.
Baseline Comparison: It outperforms a highly optimized DFlash baseline by 24.7% and a DDTree baseline by 10–30% depending on the dataset.
Acceptance Length: The method increases the mean acceptance length (MAL) by 77% relative to the chain DFlash baseline and 32% relative to DDTree at the same tree size.
Verification Efficiency: The fused kernel scales significantly better than per-branch recurrent verification. At a tree size of 128 tokens, the fused kernel is 7.1× faster than the recurrent baseline. The verification step accounts for only ~12% of the total decode time.

Significance and Claims

The paper claims that the primary limitation of factorized drafters is structural (independence assumption) rather than a lack of model capacity. By introducing a lightweight autoregressive adapter that operates on a constrained vocabulary, the authors demonstrate that it is possible to restore conditional dependencies and achieve high acceptance rates without incurring the overhead of full-vocabulary projections.

Furthermore, the work addresses a critical systems gap: enabling efficient speculative decoding for models with non-diagonal recurrent layers (GDN). The proposed rollback-free verification kernel removes the need for state rollback, making tree-based speculative decoding viable for a broader class of modern architectures. The authors conclude that these combined model and system contributions represent a state-of-the-art improvement in interactivity for local and server-side LLM inference.