TreeGraft：用于树基推测解码的自适应多草稿器嫁接技术

# TreeGraft：基于树结构投机解码的自适应多起草器嫁接
来源：https://arxiv.org/html/2608.26112  
贾明范<sup>1,2,3∗</sup>，曹达明<sup>4∗</sup>，黄灿尘<sup>1,2</sup>，付家乐<sup>1,2</sup>，张晋<sup>3</sup>，高俊杰<sup>3</sup>，杨凯<sup>3†</sup>，罗相众<sup>1,2</sup>，杨旭<sup>1,2‡</sup>  
<sup>1</sup>中国教育部，东南大学，新一代人工智能技术及其交叉应用重点实验室  
<sup>2</sup>东南大学  
<sup>3</sup>蚂蚁集团  
<sup>4</sup>南京信息工程大学  
∗同等贡献†项目负责人‡通讯作者  
jiaming\.fan@seu\.edu\.cn  
dmcao@nuist\.edu\.cn  
chengyue\.yk@antgroup\.com  
xuyang\_palm@seu\.edu\.cn  

###### 摘要
投机解码通过“先起草后验证”的范式加速大语言模型的推理。在此基础上，树结构方法通过将提议组织为多个候选路径来改进推理，增加了被接受的序列长度。然而，现有的树结构方法在所有起草步骤中都使用单个起草器，这造成了一种两难困境：较小的起草器速度快，但生成的树质量较低；而较大的起草器能提高树质量，却会带来高延迟。为解决这一问题，我们提出了 TreeGraft，一个由不同计算成本的起草器共同构建共享草稿树的多起草器框架。TreeGraft 利用更强的起草器通过更新由较弱起草器分配的分数来重新评估候选者，重新选择嫁接位置，并恢复先前未探索的有前景的路径。它还以非破坏性的方式整合更强起草器的扩展，保留可能仍会被目标模型接受的现有分支。这些设计共同提高了共享草稿树的质量。为了控制起草成本，TreeGraft 引入了一个轻量级的调度器，该调度器通过蒸馏自一个离线价值系统，来决定何时调用更强的起草器。在 10 个模型对和 6 个基准测试上，TreeGraft 平均比两种固定单起草器端点策略中较优的一个高出 15.1%，最高增益达到 26.6%。我们的代码在 https://anonymous.4open.science/r/TreeGraft-E983 开放。

## 1 引言
自回归大语言模型（LLMs）在多种任务上表现强劲\[1 (https://arxiv.org/html/2608.26112#bib.bib1),11 (https://arxiv.org/html/2608.26112#bib.bib15),22 (https://arxiv.org/html/2608.26112#bib.bib16)\]，但其逐词解码的本质使得推理速度缓慢。早期的投机解码（SD）\[14 (https://arxiv.org/html/2608.26112#bib.bib17),4 (https://arxiv.org/html/2608.26112#bib.bib18)\]通过使用一个较小的模型作为起草器来加速这一过程，该起草器构建草稿，然后由目标模型并行验证。然而，由于加速效果受限于被接受的序列长度，早期 SD 面临一个结构性限制：它将提议限制在单一序列上。因此，一旦目标模型拒绝一个词元，所有后续词元都会失效，从而限制了被接受的长度。基于树的 SD\[18 (https://arxiv.org/html/2608.26112#bib.bib2),15 (https://arxiv.org/html/2608.26112#bib.bib8),6 (https://arxiv.org/html/2608.26112#bib.bib10)\]通过将提议组织成包含多条路径的草稿树克服了这一问题，利用替代分支绕过局部拒绝，从而延长了被接受的长度。然而，长的被接受长度需要一棵高质量的树，其路径需要覆盖目标模型能够接受的内容，而生成这样的树需要巨大的起草成本。因此，核心目标是以低起草成本构建高质量的树。

为了实现这一目标，我们首先审视草稿树的构建过程。一个通用的范式\[15 (https://arxiv.org/html/2608.26112#bib.bib8)\]是在一系列起草步骤中增量式地构建它。在每个步骤，起草器从当前树中形成一个候选池（通常是上一步生成的新节点），并根据评分标准对这些节点进行排名。排名最高的节点随后被扩展以产生新的子节点。然而，现有的基于树的方法大多在所有起草步骤中依赖单一起草器\[18 (https://arxiv.org/html/2608.26112#bib.bib2),15 (https://arxiv.org/html/2608.26112#bib.bib8)\]，这迫使在树质量和每步成本之间进行二元权衡：参数较少的起草器生成质量较低的树，而简单地切换到一个更大的起草器虽然能提高树质量，但在每一步都会带来更高的延迟（图1 (https://arxiv.org/html/2608.26112#S1.F1)(a)）。这就引出了一个自然的问题：是否必须在每个起草步骤使用相同的起草器？如果仅在那些质量提升最关键的步骤调用更强的起草器，就可以在不支付所有步骤成本的情况下改进树。

受此观察启发，我们引入了两种计算成本不同的起草器——一个轻量级的小起草器和一个更强的中等起草器——它们共同构建一棵单一的草稿树，每个起草器负责一部分起草步骤。我们将这种方案称为**嫁接**：一个起草器产生的节点可以添加到另一个起草器产生的节点下方，形成一个共享的草稿树。

要使嫁接工作，需要回答三个问题：
1.  **嫁接到哪里？**：在一个步骤内，新节点应附着在树的哪个位置？它们应如何与现有树结合？
2.  **如何嫁接？**：在步骤间，应何时调用中等起草器？
3.  **何时嫁接？**：第一个问题关乎如何通过嫁接提高树的质量，第三个问题关乎如何控制因调用中等起草器而产生的起草成本。

**嫁接到哪里。** 在常见的单起草器树构建中\[15 (https://arxiv.org/html/2608.26112#bib.bib8)\]，候选池仅限于上一步生成的叶节点。对于单一起草器，这个规则是自然的：较早的候选节点已被同一个模型评估过，因此树构建从新生成的前沿推进。然而，直接将此应用于多起草器场景，会将更强的中等起草器限制在最新节点上，迫使它盲目信任共享树上的现有分数。如果一个有前景的节点之前被小起草器排名较低，该节点将被永久忽略。因此，在中等起草器步骤，我们将候选池扩展到超出上一步的范围，允许中等起草器重新评估较早的节点，重新选择嫁接位置，从而重新打开过早被忽略的路径（图1 (https://arxiv.org/html/2608.26112#S1.F1)(b)）。

**如何嫁接。** 现有的分层多起草器方法主要是基于序列的\[7 (https://arxiv.org/html/2608.26112#bib.bib13),19 (https://arxiv.org/html/2608.26112#bib.bib14)\]：当一个更强的起草器在一个父节点下生成新的子节点时，父节点现有的子节点会被覆盖。然而，在树结构中，结合我们上面“嫁接到哪里”的设计，覆盖操作会带来更大的危害：由于中等起草器节点可能附加在更早的、已经带有整个子树的位置，覆盖操作将丢弃整个子树，包括目标模型可能仍然接受的分支。因此，我们需要一个非破坏性的嫁接规则，即中等起草器节点在附加到现有父节点下方时，不移除任何现有分支（图1 (https://arxiv.org/html/2608.26112#S1.F1)(c)）。

**何时嫁接。** 只有在那些质量增益足以抵消额外延迟的步骤，调用中等起草器才划算。一个步骤是否满足这个条件主要取决于当前的树，而这只在解码过程中才能知道。因此，我们需要一个在线调度器：在每个步骤，它可以从当前树中读取几个信号，并决定调用中等起草器是否值得其成本；如果是，该步骤就交给中等起草器，否则继续由小起草器处理（图1 (https://arxiv.org/html/2608.26112#S1.F1)(d)）。

基于此观点，我们提出了 TreeGraft，一个用于树结构投机解码的多起草器框架。TreeGraft 专注于多个起草器如何共同构建共享的草稿树并调度其扩展。我们的贡献如下：
(1) 扩展的嫁接位置选择，允许中等起草器重新评估先前排名较低的节点并选择嫁接位置；
(2) 非破坏性嫁接，在不丢弃现有分支的情况下将中等起草器子节点集成到共享树中；
(3) 价值引导的在线调度，通过拟合一个离线价值系统来选择具有价值预测性的树信号，并将其调用/跳过阈值蒸馏为一个轻量级的运行时调度器；
(4) 在 10 个模型对和 6 个基准测试上的广泛实验，表明 TreeGraft 平均比两种固定单起草器端点策略中较优的一个高出 15.1%，最高增益达到 26.6%。

## 2 相关工作
**基于树的投机解码。** 现有关于基于树的投机解码及其草稿来源的研究主要从两个角度改进候选树的质量。第一是树侧优化，它在给定草稿来源的情况下改进树构建\[15 (https://arxiv.org/html/2608.26112#bib.bib8),6 (https://arxiv.org/html/2608.26112#bib.bib10)\]、节点分配或剪枝\[23 (https://arxiv.org/html/2608.26112#bib.bib11),25 (https://arxiv.org/html/2608.26112#bib.bib12)\]、以及词元或缓存树验证\[18 (https://arxiv.org/html/2608.26112#bib.bib2),20 (https://arxiv.org/html/2608.26112#bib.bib3)\]。第二是草稿来源改进，其中代表性的基于训练的方法通过训练辅助草稿模块或草稿头来提高草稿质量，包括目标内部草稿头\[3 (https://arxiv.org/html/2608.26112#bib.bib5)\]、顺序依赖草稿头\[2 (https://arxiv.org/html/2608.26112#bib.bib6)\]、特征级自草稿\[16 (https://arxiv.org/html/2608.26112#bib.bib7)\]以及使用多层特征融合直接预测词元\[17 (https://arxiv.org/html/2608.26112#bib.bib9)\]。TreeGraft 将第一个方向从单起草器树构建扩展到多起草器共享树构建。同时，TreeGraft 在很大程度上与第二个方向正交，包括基于训练的草稿来源改进，因为它在树构建层面操作，并不限制每个草稿来源的获取方式。

**分层多起草器投机解码。** 分层多起草器投机解码在起草-目标流水线中增加了中等起草器，但现有方法主要是基于序列的，如级联投机起草\[7 (https://arxiv.org/html/2608.26112#bib.bib13)\]和 TriForce\[19 (https://arxiv.org/html/2608.26112#bib.bib14)\]。相比之下，TreeGraft 在草稿树上研究分层多起草器投机解码。

## 3 基于嫁接的共享树构建
### 3.1 标准的单起草器树扩展
在引入嫁接之前，我们首先描述标准的单起草器扩展规则。在每个起草步骤，起草器从当前草稿树中选择一个候选池 $\mathcal{C}_{k}$，并选择一个扩展前沿 $\mathcal{F}_{k} \subseteq \mathcal{C}_{k}$ 进行扩展。在常见的单起草器树 SD 方法中\[15 (https://arxiv.org/html/2608.26112#bib.bib8)\]，这个候选池被限制为上一步生成的新节点。我们用 $\mathcal{N}_{k}$ 表示这些节点，其中第一步 $\mathcal{N}_{1} = \{r\}$，$r$ 是表示当前已验证前缀的根节点。因此，标准单起草器规则使用 $\mathcal{C}_{k} = \mathcal{N}_{k}$。起草器然后根据累积路径分数 $s(v)$ 对 $\mathcal{C}_{k}$ 中的节点进行排序，并选择得分最高的 $W$ 个节点作为扩展前沿，即 $\mathcal{F}_{k} = \mathrm{Top}_{W}(\mathcal{C}_{k}, s)$。只有 $\mathcal{F}_{k}$ 中的节点才会被扩展。对于每个前沿节点，起草器在以该节点结尾的前缀上执行前向传播，获取下一个词元的分布，并将 top-$m$ 个词元作为子节点附加。这些新附加的子节点形成下一步的 $\mathcal{N}_{k+1}$。通过在不同路径间共享前缀，生成的树可以通过树注意力机制被目标模型高效验证。

分数 $s(v)$ 定义为到达节点 $v$ 的路径上的累积概率。在起草步骤 $k$，令 $\mathcal{T}_{k}$ 表示当前树，$V(\mathcal{T}_{k})$ 为其节点集。每个非根节点 $u \in V(\mathcal{T}_{k}) \setminus \{r\}$ 对应于在其父节点 $\mathrm{pa}(u)$ 结尾的前缀下生成的一个草稿词元。当 $u$ 被生成时，起草器为该词元分配一个概率；我们将此概率存储为指向 $u$ 的父子边的边分数 $p(u)$。对于候选节点 $v$，令 $\pi(v)$ 表示从 $r$ 到 $v$ 的路径上的节点集。累积路径分数为 $s(v) = \prod_{u \in \pi(v) \setminus \{r\}} p(u)$。对于根节点，$s(r) = 1$。

### 3.2 在共享草稿树上的嫁接
第 3.1 节描述的扩展过程假设使用单个起草器。当多个能力不同的起草器协作时，它们在一个共享草稿树上迭代构建。我们将这种多起草器协作扩展构建为一个**嫁接**过程。为了形式化这一操作，我们必须先解决两个关键问题：**嫁接到哪里**和**如何嫁接**，这两个问题共同决定了 TreeGraft 如何提高树的质量。

图2：在中间-小-中间调度下的扩展嫁接位置选择示意图。面板 (1) 和 (3) 显示中等起草器步骤会重新考虑共享树中的早期节点，而面板 (2) 显示小起草器步骤仅扩展上一步生成的节点。

**扩展的嫁接位置选择。** 为了决定嫁接到哪里，TreeGraft 查看当前活跃的起草器考虑过哪些节点。我们说一个起草器**访问过**一个节点，如果该节点在其之前的步骤中已进入该起草器的候选池。这个定义恢复了单起草器情况下的标准规则：由于同一个起草器处理了所有较早的候选池，对于该起草器来说唯一**未访问**的节点是前一步生成的节点。在 TreeGraft 中，两个起草器被不对称地对待。中等起草器遵循相同的“已访问节点”规则。令 $\mathcal{P}_{k}$ 表示在步骤 $k$ 之前已进入中等起草器候选池的节点集，其中 $\mathcal{P}_{1} = \emptyset$。在中等起草器步骤，TreeGraft 使用 $\mathcal{C}_{k} = V(\mathcal{T}_{k}) \setminus \mathcal{P}_{k}$。因此，中等起草器可以重新考虑当前共享树中它尚未访问的任何节点，包括由小起草器生成或排序的历史节点。由于小起草器的能力低于中等起草器，TreeGraft 不使用它来重新访问旧节点进行重排序。相反，小起草器将所有先前的候选池（包括中等起草器的）视为已经评估过。因此，在小起草器步骤，我们将其候选池严格设置为上一步生成的新节点，即 $\mathcal{C}_{k} = \mathcal{N}_{k}$。在中等起草器步骤后，我们更新 $\mathcal{P}_{k+1} = \mathcal{P}_{k} \cup \mathcal{C}_{k}$；在小起草器步骤后，$\mathcal{P}_{k+1} = \mathcal{P}_{k}$。结合两种情况，得到 $\mathcal{C}_{k} = \begin{cases} \mathcal{N}_{k}, & \text{小起草器步骤} \\ V(\mathcal{T}_{k}) \setminus \mathcal{P}_{k}, & \text{中等起草器步骤} \end{cases}$。