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

会议: NeurIPS 2026
arXiv: 2608.26112
代码: https://github.com/fjm9933/TreeGraft
领域: LLM 效率
关键词: 投机采样、树状解码、多草稿模型协同、在线调度、离线价值蒸馏

一句话总结¶

TreeGraft 打破了树状投机采样单一草稿模型的范式，通过扩展嫁接位置选择、非破坏性嫁接及离线价值蒸馏驱动的在线调度器，让低开销与高能力草稿模型协同构建共享候选树，在 10 组模型对与 6 项基准上相较最优单草稿端点平均提升 15.1% 的加速比。

研究背景与动机¶

自回归大语言模型（LLM）逐 token 生成的特性受限于 GPU 内存带宽瓶颈，使得大规模模型服务延迟极高。投机采样（Speculative Decoding, SD）通过“草稿-验证”范式加速推理，由较小的草稿模型先行生成候选序列，再由目标大模型并行校验。然而早期基于单序列的投机采样存在结构性瓶颈——只要某一个 token 被拒绝，后续所有投机结果全部失效，严重制约了平均接受长度。为此，以 SpecInfer 和 EAGLE-2 为代表的树状投机采样方法应运而生，通过将草稿候选组织成多路径的投机树，利用备选分支绕过局部拒绝，显著拓展了单轮接受长度。但高接受长度的前提是候选树必须紧密契合目标模型的生成偏好，而构建高质量树结构本身需要付出高昂的草稿计算代价。

这种架构面临的核心矛盾在于：现有树状投机方法在所有草稿步骤中均固定绑定单一草稿模型，迫使系统在树质量与单步计算延迟之间做极端的二选一。若采用参数极小的轻量草稿模型（如 n-gram 或超小模型），单步扩展极其迅速，但生成质量受限，候选树命中率低；若改用参数规模较大的高能力草稿模型，虽然树质量大幅提升，但其庞大的前向传播开销会贯穿整棵树构建的每一个步骤，抵消掉目标验证阶段带来的加速红利。

本文的切入角度是：构建一棵优质的候选树，并不意味着每一个扩展节点都需要重型模型参与；如果仅在决定树结构走向、纠偏分支的关键步骤唤醒高能力草稿模型，其余步骤沿用超轻量草稿模型，就能兼得高树质量与极低生成耗时。核心 idea：提出多草稿模型协同的共享树嫁接框架 TreeGraft，通过允许强草稿模型重选历史位置与打分实现纠偏，结合非破坏性分支嫁接保留有效分支，并借助离线价值系统蒸馏出的在线调度器动态决策调用时机，实现树质量与计算开销的自适应权衡。

方法详解¶

整体框架¶

TreeGraft 的核心在于让不同开销的草稿模型（以轻量免训练的 n-gram 小草稿模型和参数更大的神经网络中草稿模型为例）在同一棵共享草稿树 $\mathcal{T}_k$ 上协作扩展。整个流水线系统性地回答了三个核心问题：在树的什么位置嫁接（Where）、如何将新节点融入现有树结构（How）、以及在哪个时间步调用更昂贵的中草稿模型（When）。在单轮投机中，系统首先基于前缀初始化根节点；随后在每个草稿步 $k$（共 $D$ 步），在线调度器根据当前树状态特征和上下文决定采取跳过（由小草稿模型扩展）或调用（由中草稿模型扩展）；中草稿模型被唤醒时可重访并重估历史未访问节点，以非破坏性方式挂载新分支；所有扩展结束后，经过全局预算贪心剪枝，最终提交给目标模型进行并行树注意力验证与无偏拒绝采样。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["当前验证前缀<br/>初始化共享树根节点"] --> B["在线特征提取<br/>上下文 + 树结构信号"]
    B --> C["基于价值指导的边际蒸馏在线调度<br/>轻量 MLP 判定是否调用中草稿"]
    C -->|判定调用| D["扩展嫁接位置选择<br/>重估历史未访问节点得分并重选前沿"]
    C -->|判定跳过| E["轻量级小草稿局部推进<br/>基于上一步前沿进行极速扩展"]
    D --> F["非破坏性嫁接与全局预算剪枝<br/>挂载中草稿分支且保留已有子树"]
    E --> F
    F -->|未达到最大步数 D| B
    F -->|完成草稿步数 D| G["全局路径累积打分剪枝<br/>保留前 B_ver 节点送目标模型验证"]

关键设计¶

1. 扩展嫁接位置选择：赋予强模型重访历史低估节点的纠偏能力 在传统的单草稿模型树构建（如 EAGLE-2）中，候选池 $\mathcal{C}_k$ 被严格限制在前一步刚生成的叶节点集合 $\mathcal{N}_k$。对于单一模型而言，历史节点早已完成打分，从最新前沿生长完全合理；但在多草稿协作场景下，若机械限制中草稿模型只能在 $\mathcal{N}_k$ 中扩展，就等于强迫中草稿模型盲目信任小草稿模型历史给出的低质量排序。一旦极富潜力的路径此前被小草稿模型低估淘汰，该路径将永久无法复活。TreeGraft 针对两类模型实施了不对称候选池机制：维护中草稿模型已访问节点集合 $\mathcal{P}_k$（初始 $\mathcal{P}_1 = \emptyset$）。当调度器调用中草稿模型时，候选池扩展为整棵共享树中所有该模型未曾评估过的节点全集，并由中草稿自身模型分布刷新转移概率 $p(v)$ 与累积路径得分 $s(v)$： $$\mathcal{C}_{k}=\begin{cases}\mathcal{N}_{k},&\text{小草稿模型步骤},\\ V(\mathcal{T}_{k})\setminus\mathcal{P}_{k},&\text{中草稿模型步骤}\end{cases}$$ 其中累积路径得分计算为路径上所有非根节点边权得分的连乘积 $s(v) = \prod_{u \in \pi(v)\setminus\{r\}} p(u)$。中草稿模型随后重新选取全局 Top-$W$ 节点作为扩展前沿 $\mathcal{F}_k$。统计显示，在中草稿步骤中，高达 61.47% 的前沿节点选中了历史更早深度的位置，候选池规模扩大至常规设计的 2.38 倍，成功重激活了被小草稿过早错失的高质量候选。

2. 非破坏性嫁接与全局预算剪枝：避免结构覆盖并控制验证开销 在序列级层次投机（如 CSD）中，高阶模型纠偏时普遍采用覆盖（overwrite）规则，即直接替换已有候选序列。然而在树结构中，中草稿模型重选的位置往往处于较浅层或历史分叉点，这些节点下方已附带了小草稿模型此前生长的整棵子树。若直接覆盖，将彻底抹除这些分支，而这些旧分支中很可能包含目标模型实际上认可的合法后续。TreeGraft 采用非破坏性嫁接原则：中草稿模型生成的新子节点直接以兄弟节点的形式追加挂载在父节点之下，严禁删除任何既有分支。为了防止候选树膨胀过大引发目标验证延迟激增，TreeGraft 将尺寸约束解耦至树构建完全结束后的全局剪枝阶段：根据最终所有非根节点的累积路径得分进行全局降序排列，贪心保留最高分的节点及其连接到根节点的祖先路径链，严格将非根节点总数截断至目标验证预算 $B_{\text{ver}}$（实验设为 63），兼顾了搜索空间包容性与验证延迟确定性。

3. 基于价值指导的边际蒸馏在线调度：突破局部延迟不可见的决策瓶颈 中草稿模型能否带来净收益具有高度的状态依赖性：若在树结构平稳推进时滥用，将白白增加前向耗时；若在分支分歧严重时介入，则能挽救生成质量。然而，单步决策所引发的最终吞吐率 $\eta = A / (t_{\text{D}} + t_{\text{V}})$ 在当前步根本无法观测（必须等到整树构建和目标验证完成），存在显著的“可观测性鸿沟”。如果直接在推理时用离线价值网络对未来步骤进行穷举搜索，因微小绝对预测误差累积会导致加速比暴跌至 0.95x。TreeGraft 提出两阶段蒸馏框架：首先，在离线阶段执行完整的二进制调度轨迹 $\tau_{1:D} \in \{0, 1\}^D$，收集状态样本 $(z_k, a_k, \tau_{k+1:D}, A, t)$。从 10 项候选树信号中通过 5 折分组留一交叉验证选出 6 个最具预测力的紧凑几何与打分特征（前沿深度、叶子比例、塌陷债务、历史回退比例、重选重叠度及得分），拟合轻量多头 MLP 价值系统 $f_\theta$。其次，将价值系统作为离线教师，对每个标准状态计算未来最优动作边际标量： $$y_{k}=\max_{\tau_{k+1:D}}\hat{\eta}(z_{k}^{\star},1,\tau_{k+1:D})-\max_{\tau_{k+1:D}}\hat{\eta}(z_{k}^{\star},0,\tau_{k+1:D})$$ 最后，训练一个极轻量的三层 MLP 分类器 $g_\psi(z_k^\star)$（仅 3,265 个参数，单次前向耗时仅 0.318 ms）拟合该边际。在线推理时完全剥离未来轨迹，仅凭预测边际符号判决：$g_\psi(z_k^\star) > 0$ 时唤醒中草稿，否则跳过。这一设计将昂贵的未来规划彻底内化至离线边际中，使调度开销不足目标验证开销的 0.5%。

实验关键数据¶

主实验¶

实验在 2 张 NVIDIA A100 GPU 上评估，覆盖 LLaMA 3（70B/8B/3B/1B）与 Qwen3（32B/8B/4B/1.7B/0.6B）共 10 组目标-中草稿模型组合，以及 6 大基准数据集（含完全未见任务 MT-Bench）。基线包括全流程单一小草稿模型（All Small, 免训练 n-gram 检索，单步 0.92 ms）和全流程单一中草稿模型（All Mid，神经网络模型，单步 29-55 ms）。

目标 - 中草稿模型组合	All Small (S)	All Mid (M)	TreeGraft (本文)	相对 S 提升	相对 M 提升
Qwen3-32B / 1.7B	1.21x	1.26x	1.49x	+23.1%	+18.3%
Qwen3-32B / 0.6B	1.21x	1.19x	1.47x	+21.5%	+23.5%
Qwen3-32B / 8B	1.21x	1.06x	1.40x	+15.7%	+32.1%
LLaMA-70B / 1B	1.44x	2.30x	2.30x	+59.7%	0.0%
LLaMA-70B / 3B	1.44x	1.95x	1.95x	+35.4%	0.0%
LLaMA-8B / 1B	1.40x	1.30x	1.50x	+7.1%	+15.4%
Qwen3-32B / 4B (未见)	1.21x	1.15x	1.39x	+14.9%	+20.9%
Qwen3-8B / 0.6B (未见)	1.24x	0.85x	1.27x	+2.4%	+49.4%
LLaMA-70B / 8B (未见)	1.44x	1.86x	1.83x	+27.1%	-1.6%
LLaMA-8B / 3B (未见)	1.40x	0.94x	1.42x	+1.4%	+51.1%
整体全集平均	1.32x	1.39x	1.60x	+21.2%	+15.1%

消融实验¶

消融实验在固定调度策略 $[1,0,1,0,1]$（第 1, 3, 5 步调用中草稿）下，评测树状嫁接各核心机制的作用，并在 GSM8K 与 CNN/DM 数据集上记录端到端加速比与单轮平均接受长度 $A$。

方法变体配置	Q32B/0.6B 平均 (加速比 / $A$)	L70B/1B 平均 (加速比 / $A$)	整体均值加速比	整体均值 $A$	核心结论与机制表现
CSD (序列级分层草稿)	0.77x / 0.77	1.11x / 0.84	0.94x	0.81	序列单路径无法绕过局部拒绝，加速比跌破 1.0
w/o EG (无扩展位置选择)	1.17x / 1.66	1.53x / 1.95	1.35x	1.80	局限在前步叶节点使接受长度下降 17.1%
w/o ND (无非破坏性嫁接)	1.13x / 1.44	1.27x / 1.51	1.20x	1.47	粗暴覆盖既有子树使接受长度大幅下挫 32.3%
TreeGraft (固定调度)	1.25x / 2.00	1.61x / 2.35	1.43x	2.17	扩展位置选择与非破坏性结构共同实现最优接受率
TreeGraft (自适应在线调度)	1.47x / 2.15	2.30x / 2.45	1.94x	2.30	自适应调度相较固定步调度加速比进一步暴增 35.7%

关键发现¶

非破坏性嫁接是保障接受长度的首要基石：去掉非破坏性规则（w/o ND）后，全集平均接受长度由 2.17 骤跌至 1.47，直接证实了“历史分支中仍蕴含大量目标模型青睐的 token”，盲目覆盖会导致树状优势大打折扣。
扩展位置选择带来纯粹的结构收益：引入扩展位置选择（EG）后，中草稿候选池均值扩大至 112.54 个节点，且 61.47% 的前沿扩展落在更早的历史节点上；由于每步最终保留的扩展前沿数 $W$ 严格一致，这一增益完全来源于“更优的历史嫁接锚点”，而非简单堆砌前沿节点数。
在线调度器具备强大的自适应调节能力：在 All Mid 占优的配置（LLaMA-70B/1B）中，调度器自动密集触发中草稿，加速比紧咬 2.30x；在 All Mid 发生严重减速的配置（Qwen3-8B/0.6B，All Mid 仅 0.85x）中，调度器自动稀疏触发，实现 1.27x 逆势加速；而在两端表现胶着的场景下（Qwen3-32B/0.6B），自适应混合调度取得了 1.47x，远超单一极值的 1.21x 与 1.19x。
泛化性优异：在 4 组完全未参与离线拟合的目标-草稿模型对上，TreeGraft 取得 1.48x 加速比（高于 All Small 1.32x 和 All Mid 1.20x）；在零样本迁移任务 MT-Bench 上达到 1.60x，证实所提取的 6 个几何树信号反映了普适的成本-收益规律。

亮点与洞察¶

将多模型分层投机从序列结构升维至树状拓扑：过去如 CSD 等多模型投机仅限于单链替换，TreeGraft 首次系统化建立了树状拓扑下的跨模型扩展与嫁接代数，在算法层面理清了 Where、How、When 三大支柱。
解耦生长探索与验证预算：在构建树时通过非破坏性嫁接宽容保留所有潜在分支，在验证前夕通过贪心路径累积打分做硬性预算裁剪，兼顾了搜索宽容度与硬件计算负载边界。
离线全局规划与在线轻量蒸馏的优雅闭环：巧妙避开“在线规划微小误差放大”的陷阱，利用全局遍历确定最优边际目标，仅训练一个 3K 参数的极简 MLP，以 0.318 ms 的极微开销完成了高维序列决策。

局限与展望¶

小草稿模型依赖上下文局部统计：当前实现中超轻量小模型采用了免训练的上下文 n-gram 检索，虽然单步开销只有 0.92 ms，但在少上下文或复杂长程逻辑推理场景下候选丰富度有所局限；未来可结合 Engram 或后缀树（Suffix Tree）增强轻量检索表达力。
正交性与训练型草稿头的结合：TreeGraft 聚焦于非侵入式的多模型树构建与调度，并不约束草稿源的内在生成机制；后续可与 EAGLE 等训练型专用轻量草稿头结合，探索草稿头与大模型的跨层共享嫁接。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出树状投机采样下的多草稿模型非破坏性嫁接与价值驱动在线调度机制。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 10 组主流大模型配对、6 类基准、严密的未见模型/任务泛化评测及完整的步进式消融。
写作质量: ⭐⭐⭐⭐⭐ 问题阐述精准透彻，符号体系清晰，实验对比维度完备。
价值: ⭐⭐⭐⭐⭐ 免训练、非侵入式、即插即用，为大规模 LLM 高吞吐低延迟部署提供了高实用价值的技术路线。