TALON: Confidence-Aware Speculative Decoding with Adaptive Token Trees

TL;DR

This work tackles the latency challenges of auto-regressive LLM inference by rethinking speculative decoding. It introduces TALON, a training-free, budget-driven adaptive tree expansion that grows a token draft tree within a global budget $N$, using robust root initialization and confidence-gated expansion to morph the tree between deep-and-narrow and shallow-and-wide shapes based on context difficulty. TALON defines Draft Efficiency $\delta$, Mean Accepted Tokens $\tau$, and speedup $R$ with $R = \frac{\tau}{1 + c \cdot \delta}$, demonstrating substantial end-to-end speedups (up to $5.16\times$) over the state-of-the-art Eagle-3 across 5 models and 6 datasets, including reasoning tasks. The results show TALON achieves near-oracle efficiency by adapting the draft cost to the acceptance reward, reduces wasted computation, and maintains robustness under stochastic sampling, indicating broad practical impact for real-time and interactive LLM applications.

Abstract

Speculative decoding (SD) has become a standard technique for accelerating LLM inference without sacrificing output quality. Recent advances in speculative decoding have shifted from sequential chain-based drafting to tree-structured generation, where the draft model constructs a tree of candidate tokens to explore multiple possible drafts in parallel. However, existing tree-based SD methods typically build a fixed-width, fixed-depth draft tree, which fails to adapt to the varying difficulty of tokens and contexts. As a result, the draft model cannot dynamically adjust the tree structure to early stop on difficult tokens and extend generation for simple ones. To address these challenges, we introduce TALON, a training-free, budget-driven adaptive tree expansion framework that can be plugged into existing tree-based methods. Unlike static methods, TALON constructs the draft tree iteratively until a fixed token budget is met, using a hybrid expansion strategy that adaptively allocates the node budget to each layer of the draft tree. This framework naturally shapes the draft tree into a "deep-and-narrow" form for deterministic contexts and a "shallow-and-wide" form for uncertain branches, effectively optimizing the trade-off between exploration width and generation depth under a given budget. Extensive experiments across 5 models and 6 datasets demonstrate that TALON consistently outperforms state-of-the-art EAGLE-3, achieving up to 5.16x end-to-end speedup over auto-regressive decoding.

Figures (15)

• Figure 1: Illustration of chain-based drafting and tree-based drafting (Eagle-3 eagle3) with $K=2$. At each step, Eagle calls draft model forward on the selected $K$ parent nodes of last step, selects top-$K$ child nodes for each parent, and filters $K\times K$ child nodes. Then Eagle employs an additional top-$K$ operation to choose $K$ nodes as parents for next step.
• Figure 2: Limitations of tree-based drafting methods. (a) when the model is already confident to its prediction, the draft tree still grows $k$ child nodes. (b) when the model is very confused and highly uncertain, the top-$K$ draft tokens are still not sufficient.
• Figure 3: Visualization of token acceptance frequency within a static draft tree ($K=10, D=8$). The heatmap reveals an "Acceptance Funnel" effect: while the acceptance frequency of the first layer is relatively uniform, the acceptance in subsequent layers ($d \ge 1$) shows a funnel trend that the accepted tokens concentrate more on high confidence regions (e.g., top-1 and top-2), rendering the wide static expansion computationally wasteful. Note that the first layer only has $K$ nodes, while its subsequent layers have $K\times K$ nodes.
• Figure 4: Real-World Mean Accepted Tokens (MAT) distribution across different queries in Eagle. The results exhibit high volatility: even within the same task category (e.g., Math or Coding), the optimal generation length fluctuates significantly.
• Figure 5: (a) Talon’s Budget-Driven Tree. Talon dynamically allocates the token budget based on confidence. It uses Top-$K$ at the root for robustness and confidence gating at deeper layers, resulting in adaptive topologies—deep chains for high-confidence tokens (e.g., the folded “algo-rithm” sequence) and wide branches for uncertain ones. The expansion stops when the global budget is met (indicated by the icon). (b) Tree Attention Mask. The structural mask used by the target model to verify the adaptive tree in parallel. The verification process follows the standard token-tree verification protocol (see Appendix \ref{['app:verification']} for details).