GRAFT：基于目标蒸馏边评分的自适应 DLM 草稿树构建

# 基于目标蒸馏边评分的自适应DLM草图树构建
来源：https://arxiv.org/html/2608.20375  
叶旭明1, 马泽明1, 余润杰1, 刘远1,3, 李天乐2  
白书翰1, 周坚1\\通讯作者, 吴飞1\\通讯作者  

###### 摘要  
基于树结构的投机解码通过验证多条草稿路径提升了标准投机解码的平均接受词元数。现有树构建器通常采用父节点条件扩展方式构建这些路径，即每个子词元基于其父路径条件生成。这种构建方式与DFlash等扩散语言模型草稿生成器不兼容，后者在单次前向传播中生成所有未来位置的分布。DDTree通过将每个未来位置分布中的高概率词元作为候选节点，并在固定节点预算下选择连续位置间的边来弥合这一差距。然而，其边选择仅依赖词元概率而未建模父子兼容性，导致目标兼容词元可能被连接到错误的父节点；此外，其固定预算忽略了吞吐量最优树规模随解码状态变化的事实。本文提出GRAFT，一种用于基于DLM投机解码的草稿树构建框架。GRAFT引入目标蒸馏边评分（TDES），从目标模型轨迹中蒸馏父子偏好以选择目标兼容边；以及状态感知预算分配（SABA），通过平衡预期草稿收益与验证成本来设置每轮树预算。在多模型和多任务实验中，GRAFT实现了自回归解码2.13×至6.36×的端到端加速，同时每轮仅增加低于0.5毫秒的开销，约占目标模型验证延迟的1.4%。

## 引言  
大型语言模型已成为从对话、代码生成到智能体系统等应用的核心基础设施（OpenAI 2026、DeepSeek-AI 2025、Hu et al. 2025）。然而，其推理效率受限于自回归解码——每次前向传播仅输出一个词元，导致GPU并行能力未被充分利用。投机解码通过“草稿-验证”范式缓解了这一瓶颈：轻量级草稿模型提议候选序列作为草稿，目标模型并行验证该草稿，并接受与目标模型一致的最长前缀（Chen et al. 2023, Leviathan et al. 2023）。投机解码的加速效果主要由平均接受词元数（MAT）驱动。基于树的投机解码通过用草稿树替代单一草稿序列进一步提升MAT（Miao et al. 2024, Cai et al. 2024, Li et al. 2026）。借助树注意力掩码，目标模型可在单次前向传播中验证整个草稿树。在大多数基于模型的系统中，树通过自回归的父条件节点扩展构建，草稿生成器通过在父路径下生成每个子词元来扩展或加深树（Li et al. 2026, Wang et al. 2025, Hong et al. 2025, Liu et al. 2026）。然而，扩散语言模型草稿生成器不适用于这种父条件构建方式。例如，最先进的DFlash草稿生成器在单次前向传播中预测未来词元分布（Chen et al. 2026）。该前向传播以当前目标模型状态为条件，因此扩展新父路径需要验证前不可用的目标模型状态。问题转变为如何从单次传播输出组装草稿树而无需为每条父路径重新运行草稿生成器。

为从DFlash构建草稿树，DDTree采用父节点无关的树组装策略：首先将每个未来位置的高概率词元视为候选节点，然后连接连续位置间的节点形成候选路径。每条路径仅通过乘以各位置词元概率来评分，DDTree在固定节点预算下选择高分路径。这种组装与父节点无关，因为子词元无论扩展哪条父路径都贡献相同分数。我们认为，父节点无关和固定预算的设计引入了两个挑战：(i) 父子不匹配，(ii) 预算误分配。

## 图1：基于DLM的树构建架构对比与实证分析  
(a) DDTree从位置维度的DFlash边际分布组装草稿树，因此目标兼容的子词元可能被选中并连接到错误的父节点。此处DDTree保留了of→is边并丢弃了目标兼容的fruits→is边，而GRAFT优先选择目标兼容边，比DDTree多接受三个词元。(b) 父子不匹配在多种任务中出现，GRAFT相比DDTree持续降低不匹配率（Qwen3-8B）。(c) 吞吐量最优树预算B*随解码状态变化。

**父子不匹配**  
路径评分仅反映草稿模型对每个词元的独立置信度，未考虑子词元与特定父节点的兼容性。结果导致目标兼容词元存在于树中但被连接到错误的父节点，即“父子不匹配”。如图1(a)所示，DDTree保留了高分的of→is路径，丢弃了目标兼容的fruits→is边——因为is在两个父节点下贡献相同分数，且通过of的路径总分更高。图1(b)进一步显示父子不匹配在多种任务中普遍存在，表明基于DLM的草稿树构建必须显式建模父子兼容性。

**预算误分配**  
DDTree为每轮解码分配相同节点预算，而额外节点的价值随当前解码状态变化。简单轮次已获得长接受路径，额外节点主要增加验证成本；困难轮次起始接受路径较短，较大的树可暴露有用候选词元，但仅持续到接受长度增益饱和。图1(c)展示了这种状态相关的权衡：简单、中等和困难轮次的吞吐量最优预算各不相同，而验证成本随树规模持续增加。固定预算导致验证工作在各轮次间误分配，这促使我们采用每轮预算决策。

为此，本文提出GRAFT——一种用于基于DLM投机解码的树构建框架，包含目标蒸馏边评分（TDES）和状态感知预算分配（SABA）。GRAFT基于两项观察：目标模型轨迹揭示了父子转换的偏好，且吞吐量最优树预算随当前解码状态变化。基于第一项观察，TDES从目标模型轨迹蒸馏轻量级边评分器以选择目标兼容边，减少跨任务的父子不匹配。基于第二项观察，SABA建模草稿收益与验证成本的权衡，选择最大化解吞吐量的最优预算。总之，GRAFT确定哪些边进入树以及树应有多大，如图1(a)所示。

主要贡献如下：
- 识别基于DLM草稿树构建的两大挑战（父子不匹配和预算误分配），并证明其分别源于父节点无关的边际组装和固定预算树构建。
- 提出GRAFT框架，结合目标蒸馏边评分的父条件边选择与状态感知预算分配的每轮树规模调整。
- 在多模型和多数据集上进行实验，结果表明GRAFT持续将端到端推理吞吐量提升2.13–6.36倍，同时每轮仅增加低于0.5毫秒的开销。

## 预备知识  
### 投机解码  
设x⁽ᵗ⁾表示解码轮t开始时的前缀。标准投机解码使用轻量级草稿模型提议长度为K的候选序列ŷₜ,₁:ₖ，然后目标模型并行验证该序列。在贪婪解码表示中，候选序列以自回归方式生成：ŷₜ,ₐ ∼ P_draft(· | x⁽ᵗ⁾, ŷₜ,₀:ₐ₋₁)。目标模型从位置0到K计算每个词元的概率P_target(· | x⁽ᵗ⁾, ŷₜ,₁:ₐ)。设a* = max{a ≤ K | ŷₜ,₁:ₐ与P_target(· | x⁽ᵗ⁾, ŷₜ,₁:ₐ₋₁)一致}为最长接受前缀。该轮解码接受a* + 1个词元（包括当前轮首个词元x⁽ᵗ⁾₊₁）。

### 基于树的投机解码  
基于树的投机解码用草稿树T替代单一序列。T包含节点N（词元）和边E（转移）。目标模型使用树注意力掩码在单次前向传播中验证T。最长接受路径通过从根节点开始的贪心搜索确定。平均接受词元数（MAT）为所有轮次接受词元数的期望。

### 扩散语言模型草稿生成器  
扩散语言模型草稿生成器（如DFlash）通过单次前向传播同时预测所有未来位置的分布。设{Pₜ,ₐ⁽ᵏ⁾}ₐ₌₁ᴷ为DFlash在位置1至K的预测分布。在DDTree中，每条候选边v = (p, c, d)的分数为S(v) = log qₜ,ₐ(c)，其中qₜ,ₐ(c)为位置d的预测概率。

## 方法  
### 目标蒸馏边评分  
我们提出边评分函数f_θ(p, c, d)来预测边v = (p, c, d)的目标模型兼容性。该函数使用轻量级多层感知机（MLP）实现，将父子词元嵌入和深度索引映射到标量分数：  
f_θ(p, c, d) = wᵀ · ReLU(W · [g_p(e_p); g_c(e_c); e_d] + b) + b₀，  
其中e_p和e_c为词元嵌入，e_d为深度嵌入，g_p(·)和g_c(·)为可学习投影，W, w, b, b₀为可训练参数，s > 0为可学习缩放参数。由于投影词元嵌入g_p(e)和g_c(e)仅依赖词元ID，我们计算一次并在解码期间重用。

### 训练  
我们从使用DDTree进行投机解码时收集的目标模型轨迹训练f_θ。训练目标如下：  
*分布层面*：对于深度d-1的每个接受父节点p，设zₚ,ₐ₍ₜₙ₎为目标模型对子词元c ∈ Cₚ,ₐ的逻辑值。定义目标偏好分布和校正草稿分布为：  
πₚ,ₐ⁽ᵗᵍₜ⁾(c) = softmax_{c∈Cₚ,ₐ}(zₚ,ₐ₍ₜₙ₎)，  
πₚ,ₐ⁽ᶿ⁾(c) = softmax_{c∈Cₚ,ₐ}(log qₜ,ₐ(c) + f_θ(p, c, d))，  
并通过L_KL = 𝔼₍ₚ,ₐ₎[KL(πₚ,ₐ⁽ᵗᵍₜ⁾ ∥ πₚ,ₐ⁽ᶿ⁾)]对齐二者。  

*成对层面*：对于每条失配轨迹，我们形成目标兼容边v⁺和不兼容替代边v⁻。负面边要么将相同子词元连接到错误父节点，要么将竞争子词元连接到正确父节点。我们通过间隔损失强制f_θ(v⁺) > f_θ(v⁻)：  
L_pair = 𝔼₍ᵥ⁺,ᵥ⁻₎[softplus(γ - f_θ(v⁺) + f_θ(v⁻))]。  

*边层面*：对于每条验证边v = (p, c, d)，当c是父节点p下的目标下一个词元时设置yᵥ = 1，否则yᵥ = 0。我们使用类别平衡二元交叉熵，其中w₁抵消正边的稀疏性：  
L_BCE = -𝔼₍ᵥ,ᵧᵥ₎[w₁yᵥ log σ(f_θ(v)) + (1 - yᵥ) log(1 - σ(f_θ(v)))]。  

最终目标为L = λ_KL L_KL + λ_pair L_pair + λ_BCE L_BCE，其中λ_KL, λ_pair, λ_BCE控制三种监督信号的相对贡献。

### 推理  
推理时，TDES为每条候选边分配公式(4)的分数。对于包含边v₁, ..., v₍|ᵤ|₎的前缀u，路径分数为S_TDES(u) = Σₐ₌₁^{|ᵤ|} S_TDES(vₐ)。如图3所示，树构建器使用S_TDES(u)在最大堆中对候选前缀排序，并反复将最高分前缀添加到树中直至达到预算。测试中TDES每轮仅增加0.3毫秒，约占目标模型验证延迟的1.3%。

### 状态感知预算分配  
最优树预算随解码轮次变化，因为额外节点在不同解码状态下提供不同的吞吐量效用。较大的树可覆盖否则会错过的目标兼容序列，但进一步扩展会增加每轮解码成本。为捕捉这种权衡，SABA将预算分配建模为状态条件优化问题：  
Bₜ = argmax_B Uₜ(B), Uₜ(B) = Gₜ(B) / Cₜ(B)。  

**问题表述**  
给定当前解码状态，SABA将期望草稿收益Gₜ(B)和每轮解码成本Cₜ(B)建模为树预算B的函数。其比率定义吞吐量效用Uₜ(B)。SABA旨在找到最大化当前轮次Uₜ(B)的最优树预算Bₜ。

**解码状态**  
遵循先前的自适应树构建方法，草稿置信度和接受历史作为草稿质量的标准代理指标。我们还包括DFlash特有信号——表示退化。在DFlash中，相邻隐状态在后期位置趋于相似，使后期草稿词元更难区分。因此SABA使用三个互补信号表征每轮状态：草稿置信度、表示退化和接受历史。设hₜ,ₐ为深度d的DFlash隐状态，rₜ,ₐ = cos(hₜ,ₐ, hₜ,ₐ₊₁)表示d=1,...,D-1时的相邻位置相似度。定义：  
confₜ = (1/D) Σₐ₌₁ᴰ qₜ,ₐ(cₜ,ₐ⁽¹⁾),  
histₜ = (1/w) Σᵢ₌₁ʷ aₜ₋ᵢ,  
degenₜ = (1/(D-1)) Σₐ₌₁ᴰ⁻¹ rₜ,ₐ + (rₜ, D₋₁ - rₜ,₁)。  
其中degenₜ的两项捕捉平均相邻位置相似度及其从早期到晚期位置的增长，w为接受历史窗口大小。它们共同构成解码状态sₜ = [confₜ, degenₜ, histₜ]。

**收益模型**  
草稿收益随树预算增加但呈现边际收益递减。SABA因此将其建模为树预算B的函数：Gₜ(B) = aₜ · Bᵝ，其中aₜ > 0为状态相关缩放因子，β ∈ (0,1)为凹指数。aₜ从历史接受长度估计：aₜ = (1/w) Σᵢ₌₁ʷ (lₜ₋ᵢ + 1)，其中lₜ₋ᵢ为轮次t-i的接受长度。β固定为0.5以反映边际收益递减。

**成本模型**  
解码成本与树规模线性相关：Cₜ(B) = B · c_verify + c_fiexd，其中c_verify为每节点验证成本，c_fiexd为固定前向传播成本。c_verify从目标模型前向传播时间估计，c_fiexd为草稿生成时间。