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

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.

摘要： 投机解码通过“先草拟、后验证”的范式加速了大语言模型的推理。在此基础上，树状结构方法通过将建议组织成多条候选路径来提高推理效率，从而增加了被接受的序列长度。然而，现有的树状结构方法在所有草拟步骤中仅使用单一的草稿模型（drafter），这导致了一个困境：较小的草稿模型速度快但生成的树质量较低，而较大的草稿模型虽然提高了树的质量，却带来了高延迟。

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.

为了解决这一问题，我们提出了 TreeGraft，这是一个多草稿模型框架，其中不同成本的草稿模型共同构建一个共享的草稿树。TreeGraft 利用较强的草稿模型通过更新较弱草稿模型分配的分数来重新评估候选路径，重新选择嫁接位置，并恢复那些未被探索的有潜力的路径。它还以非破坏性的方式整合较强草稿模型的扩展，保留了目标模型可能仍会接受的现有分支。

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 this https URL.

总之，这些设计提高了共享草稿树的质量。为了控制草拟成本，TreeGraft 引入了一个从离线价值系统蒸馏而来的轻量级调度器，以决定何时调用较强的草稿模型。在 10 组模型对和 6 个基准测试中，TreeGraft 的表现平均比两种固定的单一草稿模型策略中较好者高出 15.1%，最高增益达到 26.6%。我们的代码可在该链接获取。