Title: TreeFlash: Accelerating Speculative Decoding via Parallel AR-Approximation

Abstract:

Speculative decoding utilizing one-shot block drafters generates entire draft sequences in a single forward pass, thereby significantly boosting throughput by removing the need for sequential token generation. Nevertheless, these methods condition each draft token solely on the initial prefix context, ignoring dependencies on tokens previously generated within the draft. This non-autoregressive conditioning leads to a divergence between the drafter’s probability distribution and the verifier’s true autoregressive distribution as the depth of the draft increases. This issue is particularly pronounced in tree-based drafting architectures, where separate branches are compelled to rely on identical marginal distributions for subsequent tokens.

To resolve this, we introduce TreeFlash, a method that approximates an autoregressive distribution by integrating an MLP layer that takes both the drafter’s hidden state and the preceding token as inputs. By leveraging a two-stage approximation mechanism, TreeFlash preserves the $\mathcal{O}(1)$ time complexity characteristic of one-shot drafters. The proposed approach delivers state-of-the-art results across various models and tasks, demonstrating a $12\%$ improvement in block efficiency and a $9\%$ increase in speedup compared to marginal tree drafting.

Source: arXiv Generated at: 2026-06-03 00:00:00 UTC