Dynamic Delayed Tree Expansion for Improved Multi-Path Speculative Decoding

会议: NeurIPS 2026
arXiv: 2602.16994
代码: https://github.com/ritual-ai/dynamic-speculative-decoding
领域: LLM效率
关键词: 投机微调, 投机解码, 多路径验证, 最优传输, 动态树扩展

一句话总结

针对多路径投机解码中已有最优传输（OT）验证算法在根节点过早分叉导致算力浪费的缺陷，本文提出延迟树扩展与轻量级神经动态决策器，首次实现 OT 验证算法超越 Traversal Verification，平均端到端吞吐提升约 5%。

研究背景与动机

大语言模型（LLM）在推理阶段普遍受限于自回归逐 token 解码带来的高显存带宽瓶颈，投机解码（Speculative Decoding）通过轻量级草稿模型（Draft Model）预先生成候选 token，再由目标大模型并行前向验证，成为无损加速推理的关键范式。随着单路径向多路径投机解码的演进，系统借助草稿树结构和树状注意力掩码提升块效率（Block Efficiency）。现有验证策略主要分为自顶向下的最优传输（Optimal Transport, OT）线性规划求解器（如 SpecInfer、SpecTr、Khisti 等）以及自底向上的遍历验证（Traversal Verification），但在匹配实验设定下，各种验证算法的真实优劣格局与失效机理始终缺乏系统性的严格评测。

系统性评测揭示了一个出人意料的现象：在不同模型规模、任务与采样温度下，自底向上的 Traversal Verification 始终全面碾压所有主流 OT 验证算法，平均块效率与吞吐量领先达 15% 以上。深入分析表明，OT 算法并非数学机制不佳，而是被传统根节点独立同分布（i.i.d.）多路径展开策略所拖累。在树的前几层，草稿模型与目标模型的分布偏差（\(L_1\) 距离）极小，此时进行多路分支扩展带来的边际接受率收益微乎其微；而在树的深层，两模型分布剧烈发散、单路径极易被拒绝截断，急需多候选分支救场，然而传统固定树展开却无法将分叉预算后置，导致前向计算算力被严重浪费在浅层冗余分支上。

为解决这一分布发散与树拓扑不匹配的痛点，本文引入延迟树扩展（Delayed Tree Expansion），先沿单路径前向起步，在关键分叉点才展开为多路 i.i.d. 分支，并在数学上严格证明了其保留目标分布的无损性。核心 idea：通过构建预测块效率与硬件延迟的离线期望代理目标，训练上下文感知的轻量神经网络选择器（Neural Dynamic Selector），在推理时动态预测最佳主干长度、分支宽度与分支深度参数，使 OT 算法首次反超 Traversal Verification。

方法详解

整体框架

本文方法旨在打破静态分叉结构，将草稿树构建解耦为主干起步单路径与延迟多分支两阶段，并由根节点上下文状态驱动轻量决策网络。在每个解码步，系统首先根据当前提示与前一步隐藏状态提取不确定性、散度与硬件耗时统计量；随后由轻量分类 MLP 预测离散动作参数元组 \(a = (K, L_1, L_2)\)；草稿模型据此生成前 \(L_1\) 步单链并在叶节点发散生成 \(K\) 条长度为 \(L_2\) 的 i.i.d. 候选链；目标大模型执行一次树注意力批处理前向；最后调用相应 OT 求解器完成无损验证。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["上下文输入与根节点状态<br/>包含隐藏层状态/散度/耗时特征"] --> B["神经动态选择器<br/>轻量 MLP 预测 (K, L1, L2)"]
    B --> C["延迟树拓扑生成<br/>单路径主干 L1 步 + 延迟分支 K×L2"]
    C --> D["目标大模型批处理验证<br/>树状掩码并行前向计算分布"]
    D --> E["OT 验证与校正 Token 采样<br/>无损恢复目标模型输出"]

关键设计

1. 延迟树拓扑扩展机制：将分支多样性预算精准后置到分布高发散层

针对传统根节点直接分支在浅层造成的算力空转，论文定义了 \((K, L_1, L_2)\) 延迟树（Delayed Tree）。草稿模型 \(M_q\) 在当前根上下文 \(\boldsymbol{c}\) 下，首先自回归采样一条长度为 \(L_1\) 的确定性单链 \(a_{1:L_1} \sim q(\cdot|\boldsymbol{c})\)；随后以此单链末梢为公共前缀分支点，并发采样 \(K\) 条独立的、长度为 \(L_2\) 的候选路径 \(a^{(k)}_{1:L_2} \sim q(\cdot|\boldsymbol{c}, a_{1:L_1})\)。当路径重合时，按前缀出现频次自然形成具备重数的子节点列表。

该机制在数学上严格保证了目标分布一致性：因为无论是否在前缀折叠，OTLP 求解器在任意树节点条件分布下生成的候选 token 始终服从目标模型分布 \(p(\cdot)\)。通过将分叉点后延 \(L_1\) 步，算法将计算算力集中在 \(L_1\) 深度之后——正是草稿分布 \(q\) 与目标分布 \(p\) 之间 \(L_1\) 距离显著增大、传统单路径容易遭遇拒绝的瓶颈区间，从而成倍放大深层多样化路径的命中几率。

2. 基于无偏条件期望的块效率解析建模：分离草稿树随机性与验证随机性

为了指导参数选择，必须在考虑验证随机性的同时衡量块效率。对于任意最优传输线性规划（OTLP）求解器 \(f_{p,q,k}\)，给定 \(k\) 个候选 token 列表 \(\boldsymbol{x}\) 和目标 token \(t\)，定义分支跃迁概率 \(B(f_{p,q,k}, \boldsymbol{x}, t) = \mathbb{P}(f_{p,q,k}(\boldsymbol{x}) = t)\)。根据条件期望全概率法则，草稿树 \(\mathcal{T}\) 上的预期块效率可以精确展开为所有可达节点的乘积边缘累加：

\[\mathbb{E}[\tau+1] = \mathbb{E}_{\mathcal{T}}\left[\sum_{(\boldsymbol{c}, t_1, \dots, t_d) \in \mathcal{T}} \prod_{j=1}^{d} B\left(f_{p,q,k}, \text{ch}(\boldsymbol{c}, t_1, \dots, t_{j-1}), t_j\right)\right]\]

该形式在理论上完全消除了验证过程中的采样方差。在离线构建监督信号时，仅需对 \(s=4\) 个独立采样的延迟树结构取均值，即可获得快速、无偏且数值极为稳定的真实期望块效率标签 \(\widehat{E}[\tau(\boldsymbol{c}, a)+1]\)。

3. 上下文与硬件感知的神经决策器：吞吐量导向与防严重回退优化

为了在推理时以微秒级开销做出最优拓扑决策，设计了一个紧凑的分类 MLP。网络输入包括：目标与草稿模型在上一步的隐层特征 \(h^p_{\text{prev}}, h^q_{\text{prev}}\) 及草稿模型当前根特征 \(h^q_{\text{cur}}\)（均投影至 \(d=128\) 维并做 LayerNorm）；标量不确定性特征（草稿与目标分布的熵、双向 KL 散度与 \(L_1\) 距离）；采样超参数（温度 \(T\)、Top-\(p\)、上下文长度）；以及微基准预热测得的单步草稿耗时 \(t_q\) 和目标前向耗时 \(t_p\)。总墙上耗时可解析近似为：

\[\widehat{T}(\boldsymbol{x}, a) = \sum_{j=0}^{L_1-1} t_q(l+j) + \sum_{j=0}^{L_2-1} t_q(l+L_1+j \cdot K) + t_p(l+L_1+K L_2)\]

网络通过预测每个离散动作 \(a \in \mathcal{A}\)（动作空间为 \(\{1,2,3,4\} \times \{0,\dots,8\}^2\)）的分类概率分布 \(\pi_\theta(a|\boldsymbol{c})\)，以最大化预测吞吐量比值。为防止偶发的大幅性能退化，训练损失在最大化相对基线吞吐量 \(-\log(\widehat{\text{TPS}}_\pi / \widehat{\text{TPS}}_{\text{base}})\) 的同时，引入了基于前 \(\alpha\) 分位数最差负样本的二次惩罚正则项：

\[\mathcal{L} = \frac{1}{B}\sum_{i=1}^B \left(-\log\frac{\widehat{\text{TPS}}_\pi(x_i)}{\widehat{\text{TPS}}_{\text{base}}(x_i)}\right) + \lambda \cdot \frac{1}{|\mathcal{I}_\alpha|}\sum_{i \in \mathcal{I}_\alpha}\left(\max\left\{1 - \frac{\widehat{\text{TPS}}_\pi(x_i)}{\widehat{\text{TPS}}_{\text{base}}(x_i)}, 0\right\}\right)^2\]

实验关键数据

主实验

主实验在 Llama-3 (70B/8B)、Gemma-3 (27B/270M) 和 Qwen-2.5 (32B/0.5B) 三组模型对上进行，涵盖不同参数量比例（9:1、100:1、64:1）。测试覆盖 MATH500、OlympiadBench、LiveCodeBench、LitBench、Opus 5 大数据集，包含 8 种采样配置（温度 0.2~1.2，Top-\(p\) 0.9/0.99）。下表展示了基线验证算法与引入神经延迟扩展（NDE）后的平均吞吐量（tokens/s）及相对基准表现对比：

验证方法	Qwen-2.5 (32B/0.5B)	Gemma-3 (27B/270M)	Llama-3 (70B/8B)	全局平均吞吐 (tokens/s)
NSS (Baseline)	16.88	5.94	12.91	11.91
Naive (Baseline)	16.55	9.20	11.58	12.44
Khisti (Baseline)	18.67	5.93	13.66	12.75
NaiveTree (Baseline)	19.16	6.00	14.71	13.29
SpecTr (Baseline)	19.01	6.09	14.86	13.32
Block Verification (BV)	17.30	10.57	12.20	13.36
SpecInfer (Baseline)	19.58	6.21	14.86	13.55
Traversal Verification	21.35	11.56	14.77	15.89
Khisti NDE (Ours)	18.41	12.84	13.24	14.83
NSS NDE (Ours)	20.29	12.40	12.45	15.05
NaiveTree NDE (Ours)	21.67	12.86	14.01	16.18
SpecTr NDE (Ours)	21.74	13.22	14.23	16.40
SpecInfer NDE (Ours)	22.54	13.26	14.27	16.69

消融实验

神经动态扩展（NDE）相较于各算法固定最优网格搜索基线的块效率与吞吐提升比例如下（数据来自 Table 4 & Table 5）：

配置 / 基础验证算法	块效率提升比 (Qwen / Gemma / Llama)	全局平均块效率比	吞吐提升比 (Qwen / Gemma / Llama)	全局平均吞吐提升比
Khisti + NDE	1.01 / 1.90 / 1.09	1.19×	0.99 / 2.17 / 0.97	1.16×
NaiveTree + NDE	1.07 / 0.99 / 1.17	1.09×	1.31 / 1.40 / 1.21	1.30×
NSS + NDE	0.97 / 1.68 / 1.01	1.11×	1.20 / 2.09 / 0.96	1.26×
SpecTr + NDE	0.97 / 1.68 / 1.01	1.10×	1.14 / 2.17 / 0.96	1.23×
SpecInfer + NDE	0.98 / 1.65 / 1.01	1.10×	1.15 / 2.14 / 0.96	1.23×

关键发现

• OT 算法性能格局逆转：传统树展开下，自底向上的 Traversal Verification 以 15.89 tokens/s 领跑所有 OT 算法（领先第二名 SpecInfer 17.3%）；在结合 NDE 延迟展开后，SpecInfer NDE 提升至 16.69 tokens/s，超越 Traversal 达 5.0%，打破了 OT 算法长期弱于 Traversal 的固有认知。
• 极端异构模型收益最为惊人：在 Gemma-3（目标 27B 与草稿 270M 大小相差 100 倍）上，两模型分布偏差巨大，传统静态根分支极易折损，NDE 在该模型上实现平均吞吐翻倍以上（SpecInfer 提升达 2.14×，Khisti 提升达 2.17×）。
• 深层拒绝瓶颈被彻底化解：对 200,000+ 条解码轨迹的深度分析表明，树深度大于 3 之后 \(L_1\) 分布发散剧烈上扬；静态展开在浅层浪费分支，而 NDE 能够智能根据上下文发散度决定主干单路径长度 \(L_1\) 与分支参数，显著平抑了深层接受率崩溃。

亮点与洞察

• 机制诊断切中要害：没有盲目提出更繁复的 OTLP 求解器，而是从树深度节点接受率与 \(L_1\) 散度变化曲线入手，精准揪出 OT 方法受挫的根源在于“分叉时机过早”。
• 无偏方差消除的理论期望构建：提出基于分支概率的条件期望公式，将树展开随机性与 OTLP 验证随机性解耦，使得仅用 4 次树采样即可离线精准拟合连续块效率，为轻量网络提供了高质量监督信号。
• 鲁棒硬件自适应设计：在训练目标中显式纳入微基准耗时，并加入针对极值性能回退的 \(\mathcal{I}_\alpha\) 二次罚项，确保策略在不同 GPU 显存带宽与算力约束下均能稳健加速。

局限与展望

• Traversal 算法自身的动态拓展缺失：本研究的神经决策器主要针对 OT 类求解器定制，未能探索自底向上的 Traversal Verification 是否也能通过类似延迟展开获得叠加增益。
• 离线数据与硬件标定开销：策略训练依赖于目标模型在特定硬件上的轨迹收集与前向时间预热微基准，跨显卡平台部署时需要重新校准耗时函数。
• 动作空间离散性：当前搜索空间固定在网格 \(\mathcal{A}=\{1,2,3,4\}\times\{0,\dots,8\}^2\)，未能支持更非对称或自适应梯度的连续拓扑形态。

相关工作与启发

• vs Traversal Verification (Weng et al., 2025)：Traversal 采用自底向上的动态规划机制，天然利于叶节点长路径保留，但算法内部机制不透明且难与多样化 OTLP 结合；本文证明通过延迟扩展，标准 OT 方法（SpecInfer/SpecTr）完全可以补齐甚至超越 Traversal 的吞吐表现。
• vs Sequoia (Chen et al., 2024a)：Sequoia 同样关注硬件感知的树剪枝，但依赖离线静态动态规划，对在线温度和动态 context 缺乏感知；本文的 NDE 属于上下文驱动的在线逐 token 决策，泛化能力更强。
• vs EAGLE-2 (Li et al., 2024a)：EAGLE-2 依赖确定性草稿概率阈值启发式展开，且仅能适配最基础的 NSS 验证；本文的延迟扩展全面覆盖所有通用 OT 算法，具备坚实的无损分布理论保障。

评分

• 新颖性: ⭐⭐⭐⭐☆ 深入剖析了多路径投机解码中分叉层级与分布发散的冲突，提出延迟树展开与离线解析监督学习思路。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 大主流开源模型家族、5 项跨领域代表性评测集、8 种采样温度以及全套验证算法基准。
• 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰，数学定义严谨，实验图表信息密实且逻辑严密。
• 价值: ⭐⭐⭐⭐⭐ 为投机解码系统的工程落地与算法选型提供了决定性基准，给后续推理性大模型高效部署指明了方向。

相关论文

• [NeurIPS 2026] RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding
• [NeurIPS 2026] SpecBlock: Block-Iterative Speculative Decoding with Dynamic Tree Drafting
• [NeurIPS 2026] TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding
• [NeurIPS 2026] D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting
• [NeurIPS 2026] UniVer: A Unified Perspective for Multi-step and Multi-draft Speculative Decoding