PaperlayerView exact revision on arXiv

Source-linked AI summary

TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding

Jiaming Fan, Daming Cao, Canchen Huang, Jiale Fu, Jin Zhang, Junjie Gao, Kai Yang, Xiangzhong Luo, Xu Yang

arXiv:2608.26112v1cs.CL

TL;DR

Tree-based speculative decoding must balance draft-tree quality against the latency of using a stronger drafter at every step. TreeGraft jointly builds and schedules a shared tree with multiple drafters, achieving a 1.60× average speedup across 10 model pairs and 6 datasets.

  • Problem

    Existing tree-based speculative decoding uses one drafter throughout, forcing a trade-off between lower-cost, lower-quality trees and higher-quality, higher-latency trees.

  • Method

    TreeGraft jointly constructs a shared draft tree with multiple drafters using expanded grafting selection, non-destructive integration, and value-guided scheduling.

  • Results

    1.60× average speedup across 10 model pairs and 6 datasets, exceeding both fixed single-drafter endpoints and yielding a 15.1% average gain over the better endpoint.

  • Takeaways & Limitations

    TreeGraft provides a controlled quality–cost trade-off for shared-tree speculative decoding and maintains speedup on unseen model pairs and tasks.

  • Takeaways & Limitations

    TreeGraft is training-free, so it does not replace training-based improvements to draft sources and remains complementary to them.

Abstract

from arXiv · show

Speculative decoding accelerates large language model inference through a draft-then-verify paradigm. Building on this, tree-structured methods improve inference by organizing proposals into multiple candidate paths, increasing the accepted length. However, existing tree-structured methods use a single drafter for all drafting steps, creating a dilemma: a smaller drafter is fast but yields lower-quality trees, whereas a larger drafter improves tree quality but suffers from high latency. To address this, we propose TreeGraft, a multi-drafter framework in which drafters of different costs jointly construct a shared draft tree. TreeGraft uses the stronger drafter to rescore candidates by updating scores assigned by the weaker drafter, reselect grafting positions, and recover promising paths left unexplored. It also integrates stronger drafter expansions non-destructively, preserving existing branches that may still be accepted by the target model. Together, these designs improve the quality of the shared draft tree. To control the drafting cost, TreeGraft introduces a lightweight scheduler distilled from an offline value system to decide when to call the stronger drafter. Across 10 model pairs and 6 benchmarks, TreeGraft outperforms the better of the two fixed single-drafter endpoint strategies by 15.1% on average, reaching a maximum gain of 26.6%. Our code is available at https://anonymous.4open.science/r/TreeGraft-E983.

1 Introduction

TreeGraft addresses the latency–quality trade-off in tree-based speculative decoding by using a lightweight and stronger drafter to jointly construct and schedule expansions of a shared draft tree. Its design combines expanded grafting-position selection, non-destructive integration, and online scheduling based on the current tree.

  • Motivation: Tree-based speculative decoding improves accepted length through multiple candidate paths, but existing methods typically use one drafter for every drafting step, forcing a speed–tree-quality trade-off.Early speculative decoding is also limited by single-sequence proposals, while tree construction expands highly ranked candidates from the current tree.
  • Framework: TreeGraft uses a lightweight small drafter and a stronger middle drafter to jointly build one shared draft tree, with each handling a subset of drafting steps.Nodes produced by either drafter can be added beneath nodes produced by the other through grafting.
  • Where to graft: Expanded grafting-position selection lets the middle drafter rescore previously under-ranked nodes instead of restricting expansion to the newest leaves.This addresses the problem that directly applying single-drafter candidate selection would force the stronger drafter to trust existing scores on the shared tree.
  • How to graft: Non-destructive grafting integrates middle-drafter children without overwriting existing subtrees, preserving branches that the target model may still accept.Overwriting is especially harmful when grafting occurs at earlier positions containing entire subtrees.
  • When to graft: An online scheduler decides at each decoding step whether the middle drafter’s quality gain justifies its extra latency, using signals from the current tree.Steps that do not warrant the stronger call remain with the small drafter.

2 Related Work

Prior tree-based speculative decoding improves candidate-tree quality through tree-side optimization or draft-source improvement [7] [11] [13] [14] [15] [16] [17]. TreeGraft extends tree construction to multi-drafter shared trees, contrasting with mainly sequence-based hierarchical multi-drafter methods such as Cascade Speculative Drafting [9] and TriForce [10].

  • Tree-based speculative decoding: Existing tree-based methods optimize construction [7], node allocation or pruning [11], and token or cache-tree verification [13] under a given draft source.
  • Tree-based speculative decoding: TreeGraft’s design addresses where, how, and when to graft within its shared-tree framework.
  • Tree-based speculative decoding: Training-based approaches improve draft quality with auxiliary draft modules or heads, including target-internal [14], sequentially dependent [15], feature-level self-drafting [16], and multi-layer feature-fusion prediction [17].
  • Tree-based speculative decoding: TreeGraft extends tree-side optimization from single-drafter construction to multi-drafter shared-tree construction while remaining largely orthogonal to draft-source improvement.
  • Hierarchical multi-drafter speculative decoding: Unlike mainly sequence-based hierarchical multi-drafter methods such as Cascade Speculative Drafting [9] and TriForce [10], TreeGraft studies hierarchical drafting on draft trees.

3 Shared-Tree Construction with Grafting

TreeGraft constructs a shared draft tree by letting drafters use asymmetric candidate-selection and scoring rules, while integrating stronger-drafter expansions without deleting existing branches. The resulting proposal tree is pruned before standard target verification, preserving the target-model output distribution.

  • Grafting framework: The shared-tree construction is framed as a grafting process in which multiple drafters with different capabilities iteratively build upon one shared tree.The design addresses where to graft and how to graft, which jointly determine improvements in tree quality.
  • Expanded grafting-position selection: TreeGraft lets the middle drafter reconsider any unvisited node in the shared tree, whereas the small drafter expands only newly generated nodes.This expanded selection allows the middle drafter to revisit historical grafting positions and correct earlier tree-growth decisions.
  • Asymmetric scoring: The drafters score candidates asymmetrically: the small drafter reuses stored scores, while the middle drafter recomputes edge and cumulative path scores from its own distribution.After scoring, frontier selection and child generation follow the standard expansion procedure.
  • Non-destructive grafting: TreeGraft grafts new children non-destructively, preserving existing descendants that might otherwise be removed before target verification.Although this can temporarily enlarge the tree, the constructed tree is pruned to verification budget B_ver by retaining high-scoring nodes and required ancestor chains.
  • Distribution preservation: TreeGraft changes only the proposal tree; final outputs still use standard target-side tree verification and rejection sampling, so the target-model output distribution remains unchanged.The method therefore improves proposal-construction efficiency without altering target-side generation semantics.

4 Online Scheduling via Value-Guided Distillation

TreeGraft learns an online scheduler that decides whether to call the middle drafter based on the current tree state and scheduling history. It distills an offline, future-aware value system into a lightweight runtime predictor of call-versus-skip preference.

  • Online Scheduling via Value-Guided Distillation: TreeGraft schedules each drafting step by choosing whether the middle drafter’s accepted-length gain justifies its added latency, using throughput as the objective.Throughput is defined as accepted length divided by total draft-construction and target-verification latency, but becomes observable only after the full scheduling trajectory.
  • Online Scheduling via Value-Guided Distillation: Offline trajectory evaluation records accepted length and latency separately, enabling future-aware comparison of call and skip actions across complete scheduling trajectories.The same current state and action can produce different outcomes under different future schedules, so the estimator incorporates the future suffix during offline reasoning.
  • Online Scheduling via Value-Guided Distillation: The scheduler state combines runtime context, drafting-step index, previous call/skip history, and selected signals from the current draft tree.Tree-signal subset selection uses offline fitting and out-of-fold identification scores to reduce noise from redundant or weak features.
  • Online Scheduling via Value-Guided Distillation: TreeGraft distills the value system into a lightweight predictor of a call/skip margin and uses its sign for runtime scheduling.This avoids directly comparing potentially error-prone absolute throughput predictions; future-suffix optimization is performed offline, leaving one lightweight forward pass per drafting step.

5 Experiments

TreeGraft outperforms fixed single-drafter endpoints across model pairs and benchmarks, generalizes to held-out pairs and MT-Bench, and benefits from adaptive scheduling and its grafting components.

  • Overall performance: 15.1% average gain: TreeGraft reaches 1.60× speedup across 10 model pairs and 6 datasets, versus 1.32× for All Small and 1.39× for All Mid.Its maximum individual-result gain is 26.6% on Qwen3-32B/0.6B.
  • Generalization: On held-out data, TreeGraft achieves 1.48× speedup across 4 unseen model pairs and 1.60× on MT-Bench, exceeding both fixed endpoints in each setting.The corresponding endpoint values are 1.32× and 1.20× for held-out pairs, and 1.34× and 1.36× on MT-Bench.
  • Ablation: Ablations show non-destructive grafting is primary: removing it reduces speedup from 1.43× to 1.20× and accepted length from 2.17 to 1.47.Removing expanded grafting-position selection also reduces speedup to 1.35× and accepted length to 1.80.
  • Scheduling behavior analysis: TreeGraft improves over the better endpoint by up to 21.5% while falling behind it by at most 1.6% across model-pair averages.This indicates that scheduling captures available gains while limiting losses when stronger drafting is not beneficial.
  • Scheduling behavior analysis: When All Mid is superior, TreeGraft matches its 2.30× speedup through frequent middle-drafter calls; when All Small is superior, it reaches 1.27× with sparse calls.For Qwen3-8B/0.6B, this compares with 1.24× for All Small and 0.85× for All Mid.
  • Ablation: The online scheduler raises average speedup from 1.43× to 1.94×, while replacing TreeGraft with CSD lowers speedup from 1.43× to 0.94×.CSD also lowers average accepted length from 2.17 to 0.81; scheduler gains are sharpest on CNN/DM.

6 Conclusion

TreeGraft frames multi-drafter tree construction around where, how, and when to graft, using position selection, non-destructive grafting, and value-guided scheduling. Together, these components improve tree quality while controlling drafting cost, outperforming single-drafter endpoints on average and closely tracking the better endpoint.

  • 6 Conclusion: TreeGraft combines expanded grafting-position selection, non-destructive grafting, and value-guided online scheduling to improve tree quality while controlling drafting cost.These components address where, how, and when to graft, respectively, turning shared-tree grafting into a controlled quality–cost trade-off.
  • 6 Conclusion: TreeGraft outperforms both single-drafter endpoints on average and closely tracks the better endpoint when one endpoint dominates, with complementary gains confirmed by ablations.The paper notes limitations and future directions in Appendix A.

A Limitations · B Details of Online Scheduling

TreeGraft is training-free and operates at shared-tree construction and scheduling, complementing rather than replacing trained draft-source improvements. Its current implementation uses an n-gram low-cost drafter, while the online-scheduling appendix documents the method’s implementation and offline data construction.

  • A Limitations: TreeGraft does not train additional draft heads or auxiliary draft modules, enabling use with off-the-shelf draft sources.This design makes TreeGraft easy to apply but limits it as a replacement for training-based draft-source improvement methods.
  • A Limitations: TreeGraft complements training-based draft-source improvements, such as EAGLE-style draft heads [17], because it operates at shared-tree construction and scheduling.The two approaches are described as largely orthogonal.
  • A Limitations: The implementation uses an n-gram drafter as the low-cost small drafter because even the smallest neural LLMs introduce nontrivial forwarding latency.The n-gram choice is a practical implementation decision rather than an exhaustive selection of lightweight drafters.
  • A Limitations: Future TreeGraft implementations can explore stronger lightweight draft sources, including Engram-style conditional-memory lookup [26] and suffix-tree-based drafters.These alternatives are proposed within the same TreeGraft framework.
  • B Details of Online Scheduling: The online-scheduling appendix provides technical implementation details, hyperparameters, and network architectures for the scheduling method introduced in Section 4.It also documents the candidate signals used by the scheduler.
  • B Details of Online Scheduling: The appendix describes the candidate signals and offline data construction procedure used for TreeGraft’s online scheduling method.These details support the scheduler introduced in Section 4.

B.1 Fitting Data Collection … B.4 Margin Predictor Architecture

TreeGraft’s scheduler is fitted from exhaustive offline scheduling trajectories, uses selected pre-action tree signals and action-conditioned value heads, and distills their optimal throughput margins into a lightweight online predictor. The resulting policy makes each scheduling decision with a single forward pass while preserving exact suffix optimization offline.

  • B.1 Fitting Data Collection: The offline fitting set contains 960 complete scheduling trajectories from 6 target–middle model pairs and 5 datasets, with D = 5 binary drafting decisions per trajectory.The online evaluation additionally uses disjoint dataset subsets, held-out MT-Bench, and 4 held-out target–middle model pairs across LLaMA 3 and Qwen3 families.
  • B.1 Fitting Data Collection: Each fitting sample records the pre-action state, current action, future scheduling suffix, and complete-trajectory accepted length and latency labels rather than local one-step outcomes.The resulting value-system data uses action-conditioned return rows, with 2400 skip-return and 2400 call-return rows.
  • B.2 Candidate Tree Signals: The online state combines runtime context, scheduling history, and pre-action draft-tree signals, while candidate tree signals are computed before each action and filtered for finite values.The final scheduler uses six selected pre-action tree observations from the ten candidate signals.
  • B.3 Value System Architecture and Training: The value system predicts complete-trajectory accepted length and log-latency from state, action, and future suffix using four independent action-conditioned MLP heads.For each candidate signal subset, selection uses five-fold grouped out-of-fold validation by fitting report, preventing rows from one report crossing training and validation folds.
  • B.3 Value System Architecture and Training: The selected six-signal subset minimizes a validation score aggregating normalized errors for accepted length, latency, and throughput across skip and call branches.The final value-system input has 23 dimensions, comprising runtime context, scheduling history, the selected tree signals, and future-suffix encoding.
  • B.4 Margin Predictor Architecture: Margin targets compare the best predicted future throughput under call and skip, with exact exhaustive enumeration of all binary future suffixes because D = 5.This offline suffix planner provides the optimal action-value margin teachers for the canonical states.

B.5 Offline Evaluation Results · B.6 Ablation on the Value System · C Runtime Analysis

The distilled TreeGraft scheduler nearly matches the best enumerated scheduling trajectory offline, while direct online deployment of the fitted value system produces poor speedup despite a smaller offline gap. TreeGraft therefore uses the value system as an offline teacher for a lightweight state-based runtime scheduler.

  • B.5 Offline Evaluation Results: The selected 6-dimensional state minimizes grouped out-of-fold identification error relative to either the single best signal or all 10 candidate signals.The single signal lacks sufficient context, whereas all signals introduce redundancy and noise.
  • B.5 Offline Evaluation Results: 24.30 tok/s versus 24.78 tok/s shows that the TreeGraft online scheduler closely approaches the best enumerated call/skip trajectory on scheduling traces.The result supports using action-value margins from the offline value system to supervise the lightweight scheduler.
  • B.6 Ablation on the Value System: Direct planning is unreliable because small accepted-length and latency prediction errors can be amplified when optimizing over future suffixes.The planner compares predicted call and skip throughput after future-suffix optimization, requiring accurate absolute predictions online.
  • B.6 Ablation on the Value System: The offline mean actual gap averages each method’s throughput difference from the highest-throughput trajectory within the enumerated call/skip space.That enumerated oracle is not a theoretical global optimum, and the metric is distinct from real online speedup.
  • B.6 Ablation on the Value System: The direct value-system planner achieves a smaller offline mean actual gap than TreeGraft but only 0.95× online speedup.This shows that offline closeness to the enumerated oracle does not guarantee better real online performance.
  • B.6 Ablation on the Value System: TreeGraft uses the fitted value system only offline, converting trajectory outcomes into a relative call/skip margin for training a scheduler that uses only the current online state.This avoids deploying the future-conditioned value system directly during inference while improving end-to-end speedup.

C.1 Small Drafter Implementation and Cost Analysis · C.2 Online Scheduler Cost Analysis

TreeGraft uses a training-free n-gram small drafter to keep the low-cost endpoint inexpensive, while its online scheduler adds minimal overhead and selectively invokes the middle drafter. The implementation caches target-model top-m candidates and uses context n-gram lookup for proposals, but alternative small drafters require separate cost–quality analysis.

  • C.1 Small Drafter Implementation and Cost Analysis: The small drafter performs no neural forward pass, instead retrieving candidate tokens and scores through n-gram lookup over the available context.It serves as the small-drafter expansion step when the scheduler skips the middle drafter.
  • C.1 Small Drafter Implementation and Cost Analysis: During prefill, TreeGraft caches each position’s top-m target-model tokens and probabilities and builds an n-gram context index for later proposal lookup.At each drafting step, the small drafter matches the current-context suffix against this lookup-based representation.
  • C.1 Small Drafter Implementation and Cost Analysis: The All Small endpoint prioritizes genuinely low drafting cost rather than representing the strongest possible small drafter, clarifying why scheduling middle-drafter calls matters.The scheduler decides when a real middle-drafter call is worth paying for instead of comparing two expensive neural proposal sources.
  • C.1 Small Drafter Implementation and Cost Analysis: 0.92 ms per call: the n-gram drafter is far cheaper than neural alternatives costing tens of milliseconds, defining TreeGraft’s low-cost endpoint.Representative neural drafters require 29.1 ms (LLaMA 3.2-1B-Instruct) and 44.8 ms (Qwen3-0.6B) per call.
  • C.1 Small Drafter Implementation and Cost Analysis: TreeGraft can in principle pair its framework with neural or stronger lookup-based small drafters, but each replacement requires a separate cost–quality study and is left for future work.The reported implementation therefore does not establish the trade-offs of those alternatives.
  • C.2 Online Scheduler Cost Analysis: 0.318 ms per call: the online scheduler adds at most 1.59 ms per decoding turn with D = 5, so it is not a runtime bottleneck.The scheduler’s end-to-end benefit comes mainly from avoiding unnecessary middle-drafter calls rather than from negligible scheduling overhead.

D Supplementary Study on the Small Drafter Choice

Under a fixed schedule, replacing the neural small drafter with a training-free n-gram drafter consistently improves TreeGraft speedup. The results support using the cheaper n-gram endpoint while reserving neural forwarding for selective middle-drafter calls.

  • Experimental setup: The fixed schedule calls the middle drafter at steps 1, 3, and 5, while all other benchmark, budget, hyperparameter, and model-pair settings remain unchanged.The study compares the smallest pretrained neural model from the same family against the training-free n-gram drafter.
  • Results: 36.4% relative gain: replacing the neural small drafter with the n-gram drafter raises average speedup from 0.865× to 1.180× across all three configurations.Table 9 compares the choices under identical middle–target pairs and fixed experimental settings.
  • Design implication: The n-gram drafter is preferred because neural small-drafter forwarding can offset speculative-decoding benefits, whereas its lower cost preserves efficiency for selective middle-drafter invocation.In two Qwen settings, the neural small drafter produces speedups below 1×, making decoding slower than target-only autoregressive decoding.

E Online Evaluation Setting Details · F Algorithmic Details of Tree Construction

Online evaluation uses controlled hardware, sampling, budgets, and held-out settings, while TreeGraft and single-drafter baselines share target-side verification. The construction algorithms differ before verification: TreeGraft schedules small and middle drafters over a shared tree, followed by common score-based pruning with ancestor closure.

  • E Online Evaluation Setting Details: The formal evaluation covers 10 target–middle model pairs, with held-out model pairs excluded from offline selection, value fitting, and scheduler distillation, and MT-Bench held out from offline training.Table 1’s blue rows identify held-out model pairs, while its green MT-Bench column identifies held-out task evaluation.
  • E Online Evaluation Setting Details: Online evaluations use two NVIDIA A100 GPUs, batch size 1, identical budgets and decoding procedures, and the same 80 sampled examples per benchmark across methods.This controls hardware, batching, tree construction, target verification, and evaluation samples for TreeGraft, All Small, All Mid, and ablations.
  • E Online Evaluation Setting Details: Tree construction uses D = 5 drafting steps, frontier size W = 10, m = 10 child tokens per expanded node, verification budget Bver = 63, and max_new_tokens = 256.Bver is the maximum number of non-root draft nodes retained after pruning and submitted for target verification.
  • E Online Evaluation Setting Details: The scheduler’s runtime context contains log prompt length, target latency ttgt, middle-drafter latency tmid, and log(ttgt/tmid), with timing measured during online warm-up.Warm-up forces middle-drafter calls at all five drafting steps; the warm-up sample’s outcomes and scheduler decision are not used as supervision.
  • F Algorithmic Details of Tree Construction: Both methods apply global score-based pruning that ranks non-root nodes by cumulative path score, retains required ancestor chains, and produces a valid verification tree within Bver.The resulting tree also supplies the tree-attention mask, depth-based position ids, and padded root-to-leaf retrieval paths.
  • F Algorithmic Details of Tree Construction: Standard single-drafter construction expands the top-W newly generated nodes and selects each expanded node’s top-m next-token children for up to D drafting steps.The algorithm initializes the root, forms candidate pools from newly generated nodes, and runs the single drafter on each selected prefix.
  • F Algorithmic Details of Tree Construction: TreeGraft constructs a shared tree by scheduling either the small or middle drafter, reusing stored scores on small-drafter steps and tracking nodes visited by the middle drafter.Its inputs include small drafter qs, middle drafter qm, target model ptar, schedule actions, and the shared tree budget.

G Qualitative Examples of Non-destructive Grafting

Two traced CNN/DailyMail examples show that non-destructive grafting preserves historical subtrees that the target later accepts while adding new middle-drafter children. Overwriting would discard these accepted continuations.

  • Grafting mechanism: Non-destructive grafting attaches new middle-drafter children under a selected historical node without removing its existing children or subtree.The traces use the Qwen3-32B/Qwen3-0.6B pair with fixed scheduling pattern [1, 0, 1, 0, 1].
  • Case 1: In Case 1, preserving the historical L → umber → Liquid → ators branch retains the target-accepted continuation . Lumber Liquidators, whereas overwriting removes it.New children such as The, They, and CBS are attached alongside the preserved branch.
  • Case 2: In Case 2, preserving the historical P → elle → gr → ini branch retains the target-accepted continuation mentions that Pellegrini, whereas overwriting discards it.New children such as Martinez, Hart, and the are added under the selected historical node.

H Stability of Online Evaluation … 2. Limitations

Repeated online evaluation yields nearly identical speedups, while the paper frames TreeGraft as a training-free shared-tree framework with explicit scope limitations. The study uses existing public assets, reports reproducibility and ethical safeguards, and introduces no new high-risk models or datasets.

  • H Stability of Online Evaluation: 1.884× to 1.881×: repeated online evaluation changes average speedup by only 0.003× across two representative target–middle pairs and six benchmarks.The repeated run uses the same scheduler, inputs, benchmarks, and decoding settings, differing only in independently collected wall-clock measurements.
  • I Existing Assets and Licenses: The evaluation uses publicly available pretrained checkpoints and benchmarks, cites their original sources, and follows the corresponding licenses and research-use terms.Examples include LLaMA, Qwen3, GSM8K, Alpaca, NQ, HumanEval, CNN/DM, and MT-Bench.
  • 2. Limitations: TreeGraft’s stated scope is a training-free multi-drafter framework, with the n-gram drafter serving as one practical low-cost small-drafter choice rather than an exhaustive option.The paper presents TreeGraft as orthogonal to training-based improvements of draft sources.
  • 2. Limitations: The paper reports no formal theoretical results, theorems, or proofs; its mathematical expressions define algorithmic components and the scheduling objective.Thus, the paper’s contribution is algorithmic rather than a formal theoretical development.
  • 2. Limitations: Reproduction is supported by disclosed model pairs, datasets, baselines, metrics, hyperparameters, scheduler procedures, evaluation protocols, hardware settings, and an anonymized code repository.The online evaluation uses two NVIDIA A100 GPUs with batch size 1, while appendices provide additional fitting, training, runtime, and latency details.
  • 2. Limitations: The paper discusses both societal benefits and risks: lower inference cost may improve accessibility and energy efficiency, while faster generation could facilitate harmful or misleading text.Responsible deployment is therefore tied to safeguards for the underlying language models.
  • NeurIPS Paper Checklist: The paper introduces no new pretrained model, dataset, benchmark asset, or high-risk release, instead contributing an inference-acceleration framework and evaluation protocol.It also reports no crowdsourcing or human-subject research and no important, original, or non-standard LLM component in the core method.
Loading 2608.26112v1…