Source-linked AI summary
GRAFT: Adaptive DLM-Based Draft Tree Construction with Target-Distilled Edge Scoring
Xuming Ye, Zeming Ma, Runjie Yu, Yuan Liu, Tianle Li, Shuhan Bai, Jian Zhou, Fei Wu
TL;DR
DLM-based draft-tree construction is challenging because DFlash produces all future-position distributions in one pass, while DDTree’s token-only, fixed-budget assembly can mismatch parents and waste verification capacity. GRAFT combines target-distilled edge scoring with state-aware budget allocation, achieving 2.13×–6.36× speedup over autoregressive decoding across models and tasks.
Problem
DFlash’s one-pass outputs are incompatible with parent-conditioned tree expansion, while DDTree’s parent-agnostic scoring and fixed budget leave parent–child mismatch and budget misallocation.
Method
GRAFT uses Target-Distilled Edge Scoring to select target-compatible edges and State-Aware Budget Allocation to choose each round’s tree size.
Results
2.13×–6.36× end-to-end speedup over autoregressive decoding is reported across models and tasks, with less than 0.5 ms overhead per round.
Takeaways & Limitations
GRAFT improves DLM draft-tree quality and budget efficiency by jointly choosing which edges enter the tree and how many nodes are verified each round.
Takeaways & Limitations
GRAFT still executes drafting, tree construction, and target verification sequentially, leaving pipeline bubbles for future concurrent-decoding work.
Abstract
from arXiv · showhide
Tree-based speculative decoding raises the mean accepted tokens of standard speculative decoding by verifying multiple draft paths, and existing tree builders typically construct these paths through parent-conditioned expansion, where each child token is generated conditioned on its parent path. This construction is incompatible with diffusion language model (DLM) drafters such as DFlash, which produces all future-position distributions in a single forward pass. DDTree bridges this gap by treating high-probability tokens from each future-position distribution as candidate nodes and selecting edges between consecutive positions under a fixed node budget. However, its edge selection relies on token probability alone without modeling parent--child compatibility, so target-compatible tokens can be attached to wrong parents; moreover, its fixed budget ignores that the throughput-optimal tree size varies with the decoding state. We propose GRAFT, a draft-tree construction framework for DLM-based speculative decoding. GRAFT introduces Target-Distilled Edge Scoring (TDES), which distills parent--child preferences from target-model traces to select target-compatible edges, and State-Aware Budget Allocation (SABA), which sets the per-round tree budget by balancing expected draft gain against verification cost. Across multiple models and tasks, GRAFT achieves $2.13\times$--$6.36\times$ end-to-end speedup over autoregressive decoding while adding less than $0.5$\,ms of overhead per round, approximately $1.4\%$ of the target-model verification latency.
Introduction
Tree-based speculative decoding improves accepted-token throughput by verifying multiple draft paths, but DLM drafters require tree assembly from one-pass outputs rather than parent-conditioned expansion. GRAFT addresses parent–child mismatch and fixed-budget inefficiency with target-compatible edge selection and state-aware tree sizing.
- DLM drafters such as DFlash predict future-token distributions in one forward pass, so new parent paths cannot be expanded using unavailable target-model states.
- DDTree assembles trees from position-wise high-probability tokens, but token-only scoring can attach target-compatible children to the wrong parents.
- Fixed node budgets misallocate verification work because the throughput-optimal tree size varies with the decoding state.
- GRAFT combines Target-Distilled Edge Scoring for parent-conditional edge selection with State-Aware Budget Allocation for per-round tree sizing.
- 2.13–6.36× end-to-end inference throughput improvement is reported across models and datasets, with less than 0.5 ms of overhead per round.
Preliminaries
Speculative decoding uses a draft model to propose tokens and a target model to verify them, while tree-based methods verify multiple paths in one forward pass. The preliminaries contrast parent-conditioned tree expansion with DFlash’s one-pass position-wise distributions and DDTree’s parent-agnostic assembly.
- Standard speculative decoding proposes a candidate sequence with a lightweight drafter and verifies it in parallel with the target model.
- Tree-based speculative decoding replaces one draft sequence with a prefix-closed tree whose nodes the target model verifies in one forward pass.
- Autoregressive tree builders expand children conditionally on their parent paths, widening or deepening the tree through recursive node expansion.
- DFlash produces a token distribution for every future depth in a single forward pass, with token probability serving as each position’s score.
- DDTree maintains candidate prefixes in a max-heap and selects high-scoring prefixes until the fixed node budget is reached.
Method
GRAFT constructs DLM draft trees by combining parent-conditional edge scoring with state-aware budget selection. TDES uses target-model traces and lightweight scoring, while SABA selects budgets by balancing expected draft gain against decoding cost.
- Tree Construction: The GRAFT tree builder ranks candidate prefixes with TDES scores and stops heap expansion at the SABA-selected budget.The resulting draft tree is then passed to the target model for verification.
- Target-Distilled Edge Scoring: TDES adds a residual parent–child compatibility term to token log-probability when scoring each candidate edge.The scorer operates on a parent token, child token, and depth, with centered compatibility removing parent-specific offsets.
- Target-Distilled Edge Scoring: TDES is trained from target-model traces using distribution-level, pairwise, and edge-level supervision signals.Pairwise supervision contrasts target-compatible and incompatible edges, while edge-level labels mark target next tokens under each parent.
- Implementation: 0.3 ms per decoding round is the reported TDES overhead, approximately 1.3% of target-model verification latency.The scorer uses a lightweight cacheable bi-tower architecture with frozen target-model token embeddings and reusable token projections.
- State-Aware Budget Allocation: SABA selects a per-round budget by maximizing throughput utility, defined from expected draft gain and decoding cost under the current state.It evaluates budgets from 32 to 512 in increments of 16 and adds approximately 0.1 ms per round.
- State-Aware Budget Allocation: SABA represents decoding state with draft confidence, representation degeneration, and acceptance history, then models gain saturation and near-linear budget-dependent cost.The gain model captures diminishing marginal returns, while the cost model includes tree construction and target-model verification.
Experiments
Experiments show that GRAFT improves throughput across models and task categories through complementary edge scoring and adaptive budget allocation, while its components generalize and track near-oracle budget choices.
- Main Results: GRAFT achieves the highest wall-time speedup in all 27 model–dataset groups, with best speedups of 6.36×, 6.17×, and 6.01× on Qwen3-4B, Qwen3-8B, and Qwen3-Coder-30B, respectively.The gains span mathematical reasoning, code generation, software engineering, and instruction following.
- Component Analysis: TDES raises MAT from 7.53 to 8.00 and TPS by 5.7% at DDTree’s fixed budget.This isolates improved edge scoring without changing tree size.
- Component Analysis: SABA reduces average tree size from 110.2 to 49.1 nodes while improving TPS by 12.8% with DDTree edge scoring.The result indicates that many fixed-budget nodes have limited throughput utility.
- Component Analysis: Full GRAFT reaches 263.5 TPS, restoring MAT from 7.18 to 7.62 compared with +SABA at nearly the same budget.The full system combines TDES’s tree-quality improvement with SABA’s budget efficiency.
- TDES Generalization: TDES trained only on GSM8K traces reduces parent–child mismatch and improves normalized TPS on MBPP, SWE-bench, and Alpaca.These evaluations cover code generation, software engineering, and instruction following.
- SABA Oracle-Tracking Analysis: 94.1% of oracle throughput is retained by SABA on average, with at least 95% reached in 64.1% of rounds.The oracle comparison replays budgets using the same prefix, KV cache, and DFlash outputs.
- Hyperparameter Sensitivity: At β = 0.2, GRAFT reaches 8.061 MAT and 276.17 TPS with a 54.71% mismatch reduction on GSM8K.Increasing β to 0.4 raises mismatch reduction to 79.32% but lowers MAT and TPS below β = 0.
Related Work
Related work distinguishes speculative decoding methods by their drafting strategy and tree-construction procedure, including parent-conditioned expansion for autoregressive drafters.
- Speculative decoding: Speculative decoding methods include model-free drafters based on reusable text patterns or retrieval sources and model-based drafters using auxiliary models.The related-work discussion presents these as broad categories of existing approaches.
- Tree-based speculative decoding: Tree-based methods for autoregressive drafters adapt tree shape or budget under parent-conditioned expansion using confidence and cost signals.The cited methods mainly differ in how they construct the draft tree.
Conclusion
GRAFT addresses DDTree’s two limitations by improving edge selection and adapting the number of verified nodes per decoding round. It achieves 2.13×–6.36× speedup over autoregressive decoding, while sequential pipeline stages remain a limitation.
- GRAFT uses target-distilled edge scoring and state-aware budget allocation to address DDTree’s parent-agnostic, fixed-budget tree assembly.The two components respectively determine which edges enter the tree and how many nodes are verified per round.
- 2.13×–6.36× speedup over autoregressive decoding is achieved by GRAFT in experiments.
- Sequential drafting, tree construction, and target verification create pipeline bubbles that concurrent speculative decoding could address in future work.