TreeFlash: Parallel AR-Approximation for Faster Speculative Decoding
2026-06-02 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study how to speed up generating text by making the whole draft at once instead of one word at a time. They point out a problem where this fast method ignores how previous words influence the next ones, which causes errors, especially in tree-based drafts. To fix this, the authors introduce TreeFlash, which adds a small step to consider the last word when predicting the next, making the generated text more accurate without slowing down the process. Their method improves the quality and speed of drafting compared to earlier approaches.
one-shot decodingspeculative decodingautoregressive modelsnon-autoregressive modelstree-based decodingMLP (multilayer perceptron)block efficiencydecoding speedlanguage modelingdrafting
Authors
Peer Rheinboldt, Frédéric Berdoz, Roger Wattenhofer
Abstract
One-shot block drafters for speculative decoding generate the full draft in a single forward pass, achieving strong throughput by eliminating sequential token generation. However, they predict each draft token conditioned only on the prefix context, with no dependence on previously drafted tokens. This non-autoregressive conditioning causes the drafter's distribution to diverge from the verifier's true autoregressive distribution as draft depth grows. This problem becomes more severe in tree-based drafting, where distinct branches are forced to share the same marginal distribution for subsequent tokens. We propose TreeFlash, which addresses this by incorporating an MLP layer conditioned on the drafter's hidden state and the previous token to approximate an autoregressive distribution. TreeFlash retains the $\mathcal{O}(1)$ decoding time complexity of one-shot drafters by employing a two-stage approximation mechanism. TreeFlash achieves state-of-the-art performance across a variety of tasks and models, improving over marginal tree drafting by $12\%$ higher block efficiency and $9\%$ higher speedup.