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

2026-05-28

AI Summary

为解决单个起草者在树结构推测解码中速度与质量的矛盾，TreeGraft采用多起草者框架联合构建共享草稿树，通过轻量级调度器控制成本，提高解码效率和质量。

Abstract

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.

Problem

• speculative decoding
• tree-structured methods
• draft-then-verify
• latency

Innovation

• multi-drafter framework
• shared draft tree
• non-destructive integration
• lightweight scheduler