PaperlayerView exact revision on arXiv

Source-linked AI summary

TALON: Confidence-Aware Speculative Decoding with Adaptive Token Trees

Tianyu Liu, Qitan Lv, Yuhao Shen, Xiao Sun, Xiaoyan Sun

arXiv:2601.07353v1cs.CL

TL;DR

Existing tree-based speculative decoding uses fixed-width and fixed-depth structures that do not adapt to varying token difficulty. TALON instead grows trees under a fixed token budget with confidence-aware expansion, and evaluations across 5 LLMs and 6 benchmarks report up to 5.16× speedup. Its current evaluation focuses on batch size 1, while large-batch serving remains an open engineering challenge.

  • Problem

    Existing tree-based drafting rigidly expands and prunes trees, failing to allocate computation to changing model confidence and token difficulty.

  • Method

    TALON is a training-free, budget-driven framework that iteratively grows adaptive trees using robust initialization and confidence-gated expansion.

  • Results

    Up to 5.16× speedup is achieved across 5 LLMs and 6 benchmarks, with TALON consistently outperforming state-of-the-art methods such as EAGLE-3.

  • Takeaways & Limitations

    TALON dynamically produces deep-and-narrow trees for deterministic contexts and shallow-and-wide trees for uncertain ones under a fixed budget.

  • Takeaways & Limitations

    The evaluation focuses on batch size 1; extending adaptive speculation to large-batch serving remains an open engineering challenge.

Abstract

from arXiv · show

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.

1 Introduction

Existing speculative-decoding trees use rigid expansion despite varying token confidence, motivating TALON’s budget-driven adaptive tree growth. TALON adjusts tree topology to context uncertainty and reports improved efficiency over EAGLE-3.

  • Autoregressive generation prevents LLMs from predicting multiple tokens in one step, making latency scale linearly with output length.
  • Speculative decoding separates inference into lightweight drafting and parallel target-model verification, allowing multiple proposed tokens to be accepted in one forward pass.
  • Tree-based methods mitigate early rejection by covering multiple plausible continuation paths, but existing approaches retain rigid, context-insensitive expansion.
  • TALON is a training-free, budget-driven framework that incrementally expands the draft tree until a fixed node budget is reached.It combines fixed-width first-layer initialization with confidence-gated expansion at subsequent layers.
  • TALON shapes trees deep-and-narrow for deterministic contexts and shallow-and-wide for uncertain contexts.The topology is intended to maximize draft length or enhance hit rate, respectively.
  • TALON significantly outperforms EAGLE-3 in draft efficiency and wall-clock speedup, particularly when generation difficulty fluctuates.

2 Related Work

Speculative decoding research progressed from chain-based proposals to tree-based verification and increasingly adaptive tree construction. TALON addresses the remaining gap by allocating token budgets according to real-time confidence.

  • Chain-based speculative decoding drafts one continuation but suffers when an early mismatch invalidates downstream tokens.
  • Tree-based speculative decoding verifies multiple candidate branches within one forward pass, mitigating early rejection.
  • Recent work improves tree-based speculation through calibrated drafters, context-aware trees, dynamic programming, and search-based construction.
  • Existing methods typically use rigid or heuristic expansion patterns rather than explicitly allocating token budgets from real-time confidence.
  • TALON contributes a training-free framework for adaptive token trees that become deep-and-narrow or shallow-and-wide according to uncertainty.

3 Background

Chain-based drafting wastes later tokens after early rejection, whereas tree-based drafting covers multiple paths but commonly relies on rigid expansion and pruning. These mechanisms motivate adaptive budget allocation.

  • Motivation: Static trees show an acceptance funnel: later-layer acceptance concentrates on high-confidence candidates despite wide expansion.
  • Chain-based Drafting: In chain-based drafting, rejecting token x_i causes all subsequent drafted tokens x>i to be discarded regardless of their correctness.
  • Tree-based Drafting: Tree-based drafting lets the verifier cover diverse candidate paths within a single token tree.
  • Tree-based Drafting: Static tree construction expands top-K children for each parent, selects parents using path scores, and ends at a predetermined depth D.
  • Motivation: Mean Accepted Tokens fluctuate across queries, including within the same task category, exposing the weakness of fixed-depth drafting.
  • Tree-based Drafting: The expand-then-shrink mechanism generates redundant nodes that are discarded during intermediate shrinking or final pruning.

4 Motivated Experiments

Pilot experiments with a fixed EAGLE tree reveal inefficient width allocation in deeper layers and substantial variation in the useful drafting depth across queries.

  • The pilot study evaluates EAGLE with width K = 10 and depth D = 8 on Qwen3-8B using MT-Bench.
  • Acceptance frequency forms a funnel-like pattern across the static tree.
  • In the initial layer, acceptance is relatively uniform across top-K candidates, indicating that wider search breadth is useful there.
  • For depths d ≥1, accepted tokens concentrate sharply in high-confidence regions such as top-1 and top-2 candidates.
  • Mean Accepted Tokens fluctuate drastically even within Math or Coding queries, exposing the limitation of fixed-depth drafting.

5 Method

TALON replaces fixed tree geometry with a budget-driven, confidence-aware expansion process. It combines robust root initialization with confidence-gated growth to form deep-and-narrow trees for deterministic contexts and shallow-and-wide trees for uncertain ones.

  • From Static Grids to Dynamic Budgets: TALON shifts draft-tree construction from fixed shape to a global token budget that is iteratively invested to maximize effective speculation length.The budget limits total draft-tree nodes while allowing topology to adapt to context difficulty.
  • Hybrid Expansion: The hybrid strategy uses robust Top-K initialization at the root, followed by confidence-gated expansion at deeper layers.Root initialization protects against draft-model over-confidence, while later layers allocate nodes according to relative confidence.
  • Confidence-Gated Expansion: At each deeper layer, TALON filters candidates by a confidence margin and retains the highest-scoring candidates when the remaining budget is insufficient.Candidates are ranked by cumulative path probability, and the threshold parameter µ controls the gating margin.
  • Adaptive Topology: High-confidence contexts produce deep-and-narrow chains, whereas uncertain contexts admit more candidates into shallow-and-wide layers.The mechanism links tree topology to entropy: concentrated confidence favors depth, while diffuse confidence favors coverage.
  • Efficiency Rationale: TALON’s draft-efficiency analysis relates speculation cost to accepted tokens and explains why adaptive expansion improves speed across changing context difficulty.Static methods pay a fixed cost despite context complexity, while TALON can stop early on hard cases and extend easy ones.
  • Method Overview: Figure 5 shows Top-K root expansion, confidence-gated deeper growth, budget-limited stopping, and parallel verification with a tree attention mask.The illustrated topology folds high-confidence paths into chains and branches uncertain paths more widely.

6 Experiments

Experiments evaluate TALON against EAGLE-3 across multiple models and six benchmarks using mean accepted tokens and wall-time speedup. TALON consistently improves speed, closely tracks the zero-waste efficiency baseline, and benefits from one layer of robust initialization.

  • Setup: TALON is evaluated on five listed LLM backbones and six benchmarks spanning instruction, mathematics, code, chat, question answering, and summarization.The benchmarks include MT-Bench, Alpaca, GSM8K, HumanEval, CNN/DM, and QA.
  • Setup: Experiments compare TALON with EAGLE-3 using a global token budget of 60, threshold µ = 0.03, mean accepted tokens, and wall-time speedup.The main experiments use a single NVIDIA H200 GPU with batch size 1.
  • Main Results: TALON consistently outperforms EAGLE-3 across all 8 models and 6 datasets, reaching 5.16× speedup on HumanEval with Vicuna-13B.On CNN/DM with Qwen3-8B, TALON reaches 2.30× versus the baseline’s 1.95×.
  • Main Results: TALON achieves substantial gains on reasoning tasks, including 2.67× versus 2.37× on Qwen3-8B GSM8K.The reported results attribute these gains to capturing correct paths in low-entropy reasoning steps where static trees under-explore.
  • Draft Efficiency: TALON closely approaches the Oracle relation τ = δ, while static methods incur computation waste from fixed cost allocation.Its adaptive trees shrink in uncertain regions and expand in deterministic ones.
  • Robust Initialization: The ablation identifies k = 1 initial Top-K layer as the strongest configuration across throughput and six-benchmark robustness comparisons.Removing initialization lowers speed, while extending it to k ≥2 adds overhead without proportional acceptance gains.
  • Sensitivity and Flexibility: The threshold µ trades exploration against exploitation, while changing N lets TALON adapt to different computation budgets.Larger µ favors deep-and-narrow trees for coding and math, whereas smaller µ can suit some creative tasks.

7 Conclusion

TALON replaces rigid speculative-decoding tree geometry with a training-free, budget-driven framework that combines robust initialization with confidence gating. Across 5 LLMs and 6 benchmarks, it consistently outperforms EAGLE-3, reaching up to 5.16× speedup.

  • TALON is a training-free, budget-driven framework that replaces rigid tree dimensions with flexible adaptive expansion.Its hybrid strategy combines robust initialization with confidence gating.
  • TALON shapes draft trees deep-and-narrow for deterministic contexts and shallow-and-wide for uncertain ones.
  • 5 LLMs and 6 benchmarks show consistent improvement over state-of-the-art EAGLE-3, with up to 5.16× speedup.

8 Limitations

The paper identifies scalability and hyperparameter generalization as limitations of TALON. Its current evaluation targets batch size 1, while larger batches and specialized domains may require further engineering or tuning.

  • TALON is currently evaluated in latency-critical batch-size-1 scenarios, leaving large-batch serving as an open engineering challenge.Dynamic tree maintenance and memory management may become non-trivial in high-throughput batches.
  • Fixed threshold µ = 0.03 and budget N = 60 are robust across most tested settings, but specialized domains may require task-specific tuning.The paper proposes future auto-tuning based on acceptance history.

9 Ethic Statements

The paper states that TALON preserves the underlying models’ behavioral properties and safety risks, while its method uses adaptive tree construction followed by standard parallel verification. Supporting analyses characterize why static trees waste resources across models and tasks.

  • TALON does not modify model weights or introduce new harmful-content capabilities, but inherits biases and safety risks from the target and draft models.The paper recommends standard safety guardrails and alignment techniques.
  • Static EAGLE-style trees use fixed depth and width even though their top-K parent selection cannot adapt K to context difficulty.
  • TALON fills a global node budget iteratively, using top-K initialization at the root and confidence-gated expansion at deeper layers.
  • Tree Attention verifies the entire draft tree in one target-model forward pass, then selects an accepted prefix and resamples at divergence.
  • Across models and tasks, acceptance concentrates in deeper high-confidence branches while effective speculation length varies substantially with task difficulty.Low-entropy Math and Coding tasks support deeper speculation, whereas high-entropy tasks favor shallower trees.

C.2 Dynamic Accepted Length of Various Models

Accepted speculation length varies substantially across models and tasks, making fixed static depth a poor fit for changing generation difficulty. Deterministic Math and Coding contexts can support deeper speculation than static limits allow.

  • Math and Coding tasks often produce high accepted-token counts because the draft model is confident and accurate in deterministic contexts.
  • Static depth limits can cap speculation length in easy contexts, preventing the system from fully exploiting those contexts for speedup.
  • The observed variance supports dynamically adjusting tree depth using real-time confidence rather than relying on one static topology.

D Derivation of Draft Efficiency

This section defines wall-clock speedup and draft efficiency for speculative decoding, then explains why TALON preserves speedup as acceptance length varies. It also situates the evaluation against established baselines and reports the controlled experimental settings.

  • Derivation: Speculative decoding latency combines draft-model execution with parallel target-model verification.
  • Derivation: Wall-clock speedup R is defined as baseline latency divided by speculative-decoding latency.
  • Derivation: Draft efficiency δ measures computation invested in drafting per verification step, while mean accepted tokens τ measures accepted tokens per step.
  • Implication: Static trees fix draft efficiency at δ = D + 1, so speedup degrades when acceptance length τ falls in difficult contexts.
  • Implication: TALON reduces δ when τ is low, preserving robust speedup across varying generation difficulties.
  • Experimental setup: The evaluation compares TALON with chain-based, tree-based, and MLP-based methods, including EAGLE-3 and OPT-Tree.

E.2 More Visualization Results of Draft Efficiency

Across visualization and evaluation results, TALON aligns draft cost with accepted length more closely than static tree methods. Its confidence-gated expansion adapts tree width and depth across deterministic, uncertain, and stochastic settings.

  • Visualization: Figure 11 compares draft efficiency δ with mean accepted tokens τ across four LLMs, using the Oracle line τ = δ as the no-waste reference.
  • Visualization: Static EAGLE-3 follows a rigid cost trajectory, creating an efficiency gap when acceptance length is low in high-entropy contexts.
  • Visualization: TALON tracks the Oracle baseline by reducing draft investment in uncertain contexts and expanding tree size in deterministic contexts.
  • Comparative results: TALON consistently outperforms EAGLE-3 across evaluated models and benchmarks, including stochastic environments.
  • Comparative results: On Vicuna-13B HumanEval, TALON reaches 3.97× wall-time speedup versus 3.53× for EAGLE-3.
  • Budget behavior: Under constrained budgets, TALON prioritizes promising paths; increasing the budget lets it extend depth in deterministic regions.
  • Threshold behavior: Higher draft-target alignment favors stricter thresholds and deep-and-narrow trees, whereas weaker alignment favors lower thresholds and shallower, wider trees.
  • Construction latency: TALON reduces tree-construction latency by 1.18×–1.44× across vocabulary sizes from 32K to 152K, with gains especially pronounced in deterministic settings.
Loading 2601.07353v1…