SpecBlock: Block-Iterative Speculative Decoding with Dynamic Tree Drafting¶

会议: NeurIPS 2026
arXiv: 2605.07243
OpenReview: https://openreview.net/forum?id=l6JgXCnnMq
领域: LLM 效率
关键词: 投机微调解码, 块迭代投机解码, 动态树推测, 层级状态传递, 成本感知在线自适应

一句话总结¶

提出兼顾单次多位置前向与严格路径依赖的块迭代投机解码框架，结合层级隐状态传递、秩预测头动态构树、有效前缀曲率训练与服务期成本感知自适应，在削减一半起草开销的同时超越 EAGLE-3 的推理加速比。

研究背景与动机¶

大语言模型在生成阶段严重受限于显存带宽瓶颈，投机解码（Speculative Decoding）通过轻量级起草器（Drafter）并行生成多个候选 token 并交由目标模型一次性批量验证，有效分摊了验证开销。然而现有的起草范式陷入了两难困境：以 EAGLE-3 为代表的自回归起草器能够维持单条候选链内严密的自回归语义依赖，使树搜索深度与接受长度保持较高水平，但每次树深度增加都必须串行调用一次起草器前向传播，使得起草本身在 8B 规模目标模型上占据了近 30% 的迭代延迟；而以 Medusa 为代表的并行起草器虽然将起草收敛至单次前向（起草耗时降至约 7%），却因各个未来位置彼此条件独立、完全丢失了局部转移约束，组合出来的候选树包含大量语义不连贯的无效路径，严重浪费目标模型的验证算力。

现有方法在“起草串行轮次过多”与“并行预测缺乏序列依赖”之间非此即彼。要使验证节省的计算真正转化为端到端的净吞吐收益，关键在于打破“一次起草调用仅生成一层树”与“多位置并行预测必须完全解耦”的固化设计，在单次起草前向内显式建立相邻 token 间的条件转移机制，并实现跨块迭代的动态分支管理。

核心 idea：将投机起草重构为块迭代（Block-Iterative）范式，单次前向通过层级隐状态移位联合生成 $K$ 个具有严格自回归依赖的位置块，再由协同训练的秩预测头动态决定分支宽度与跨块延展节点，配合有效前缀掩码训练与服务期成本感知在线微调，实现低延迟与高接受率的最佳平衡。

方法详解¶

整体框架¶

SpecBlock 将投机解码的起草过程设计为多块迭代展开过程。系统以目标模型多层特征、最近验证 token 与位置查询向量为输入；在单个块内部联合预测 $K$ 个连续位置，并在跨解码器层之间通过隐状态显式移位传递因果依赖；由轻量秩预测头动态预测各位置的目标词排序分桶，据此决定同级分支数并筛选跨块迭代的起始节点；最终汇总成验证树交由目标模型并行核验。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["目标模型多层特征融合与前缀广播<br/>融合低中顶层特征与已确认Token"] --> B["块内联合起草与层级隐状态移位<br/>跨解码层显式注入前一位置隐状态"]
    B --> C["秩预测头引导的动态树扩展<br/>预测分桶分配分支宽度并筛选跨块种子"]
    C -->|满足跨块条件且未达预算| D["跨块状态继承与迭代前向<br/>复用自身隐状态延展后续块"]
    D --> B
    C -->|构筑完整候选树| E["目标模型单次并行树验证<br/>接受最长匹配前缀并提交生成"]

关键设计¶

1. 块内联合起草与层级隐状态移位：单次前向捕捉局部序列依赖 并行起草器在单次前向中为所有预测位置分配独立查询，导致后续位置无法获知前续生成内容。虽然在 Transformer 内部施加因果注意力掩码可以在形式上限制位置 $k$ 仅能看到 $\le k$ 的位置，但随着序列增长，因果注意力的 softmax 权重被严重稀释，块内深层位置的接受率断崖式下跌。为此，SpecBlock 在每两个连续的解码器层 $\ell$ 与 $\ell+1$ 之间，显式将第 $k-1$ 个位置的输出隐状态与第 $k$ 个位置拼接并投影，以极低参数量在单次前向中还原出类似自回归串行调用的状态流动： $$ \tilde{h}^{(t)}{t,k} = W^{(\ell)}}} \left[ h^{(\ell){t,k}, h^{(\ell)} \right] $$ 这一机制在不增加串行核函数调用开销的前提下，强制深层解码器层感知前序 token 的特征演进，使得块内第 4 个位置的自回归一致性显著提升。

2. 秩预测头引导的动态树扩展：按需分配验证算力 固定分支因子的候选树要么在简单位置过度冗余，要么在复杂位置欠探索导致搜索链过早截断。SpecBlock 引入了一个协同训练的秩预测头 $g_\phi$，该头以起草器顶层隐状态和 15 维分布统计量（包含 Top-10 对数概率轮廓、前几名 Logit 差值、最高概率与熵）为输入，冻结主干梯度，输出四个粗粒度排名桶 $b_0$（目标在 Rank 1）、$b_1$（Rank 2–4）、$b_2$（Rank 5–10）及 $b_3$（Rank > 10）。系统依据分桶自适应分配兄弟分支数（如 $[2, 4, 10, 0]$ 或 $[2, 4, 6, 4]$），并仅挑选高质量位置作为下一块迭代的根节点，使整棵验证树的几何形态完全由模型对自身置信度的动态评估驱动。

3. 有效前缀掩码训练：消除误导性监督与曝光偏差 自回归起草器通常采用教师强制（Teacher-Forcing），每步均基于真实历史上下文；而在块推断时，若块内第 1 个 token 预测错误，目标模型在验证时会立即截断整条路径，其后 $k > 1$ 个位置的预测将毫无意义。若在训练时仍强行使用真实目标 token 去监督这些在错误上下文之上产生的位置，不仅会浪费网络容量，还会向起草器灌输推断时根本不可能出现的分布。SpecBlock 引入有效前缀二值掩码 $m_{t,k}$： $$ m_{t,k+1} = m_{t,k} \cdot \mathbf{1}\left[\arg\max p_{t,k} = y^\star_{t,k}\right] $$ 一旦当前路径前序出现错误，后续所有位置的起草损失与秩预测交叉熵损失立即被置零。此外，在训练跨块跳转时，随机均匀采样前序块的切分点作为新块输入，消除推断期特征来源差异带来的协变量偏移。

4. 成本感知服务期自适应：利用免费验证信号动态更新 起草器在离线预训练后面对线上领域迁移容易发生接受率退化，但盲目全量反向传播的延迟成本甚至会超过单次目标前向。SpecBlock 挖掘目标模型在验证被拒绝位置时天然产生的“免费反馈”分布，构建查询级误差信号 $s = \sum_{k \in \text{rejected}} (1 - r_{t,k})$。系统采用成本感知多臂老虎机（Bandit），基于历史 EWMA 吞吐增益预测评估当前信号：在“跳过更新”、“仅更新输出头（lm_head 与 rank_head）”与“全量主干更新”三者间决策。由于输出头反向传播成本极低且能快速修正对齐漂移，绝大多数微调被路由至输出头，并在后台 CUDA 流独立执行，配合锚定预训练模型的 KL 惩罚与三窗口滑动回滚机制，在零额外前向开销下恢复领域泛化能力。

损失函数 / 训练策略¶

起草器端到端总训练损失由各位置的有效前缀加权交叉熵与秩预测头分类损失构成： $$ \mathcal{L} = \sum_{k=1}^K \mathcal{L}_k^{\text{draft}} + \mathcal{L}^{\text{rank}} $$ 其中，$\mathcal{L}_k^{\text{draft}} = -\frac{1}{N_k} \sum_t m_{t,k} \sum_v p^\star_{t,k}(v) \log p_{t,k}(v)$。前 2,000 步冻结秩预测头以稳定主干特征分布，后续联合优化。服务期自适应则引入 KL 散度约束 $\mathcal{L}^{\text{adapt}} + \lambda \mathrm{KL}(p_\theta \parallel p_{\theta_0})$ 防止灾难性遗忘。

实验关键数据¶

主实验¶

在 NVIDIA A100-80GB 环境、Batch Size 1 下评估各方法在贪婪解码（$T=0$）时的加速比（Spd）、平均验证接受长度（$\tau$）以及起草耗时占比（$T_{\mathcal{D}}\%$）。

目标模型	方法	HumanEval (Spd / $\tau$)	MATH-500 (Spd / $\tau$)	Alpaca (Spd / $\tau$)	NQ (Spd / $\tau$)	MT-Bench (Spd / $\tau$)	Mean (Spd / $\tau$)	起草占比 $T_{\mathcal{D}}\%$
Llama-3.1-8B	SpS (标准投机采样)	1.55× / 2.91	1.17× / 2.29	1.48× / 2.70	1.06× / 2.09	1.38× / 2.45	1.27× / 2.39	38%
Llama-3.1-8B	Medusa (独立头)	2.16× / 2.70	1.53× / 2.30	2.01× / 2.44	1.44× / 2.00	1.88× / 2.59	1.73× / 2.35	6%
Llama-3.1-8B	ParallelSpec (掩码并行)	2.50× / 3.16	1.89× / 2.64	2.39× / 3.28	1.78× / 2.63	2.36× / 3.21	2.10× / 2.87	8%
Llama-3.1-8B	Falcon (半自回归块)	3.06× / 4.69	2.30× / 3.74	2.86× / 4.23	2.12× / 3.38	2.66× / 4.31	2.48× / 3.95	26%
Llama-3.1-8B	EAGLE-3 (自回归树)	3.59× / 6.98	2.66× / 5.93	3.35× / 6.27	2.42× / 5.39	3.22× / 6.16	2.92× / 5.89	31%
Llama-3.1-8B	SpecBlock (静态)	3.92× / 5.16	3.07× / 4.03	3.40× / 4.66	3.00× / 4.40	3.10× / 4.46	3.21× / 4.41	16%
Llama-3.1-8B	SpecBlock+adapt	—	3.14× / 4.23	3.47× / 4.72	3.51× / 5.41	—	3.24× / 4.54	15%
Qwen3-8B	EAGLE-3	2.45× / 4.53	1.94× / 4.46	2.35× / 4.29	2.50× / 4.20	2.40× / 4.23	2.24× / 4.14	29%
Qwen3-8B	SpecBlock (静态)	2.50× / 3.59	2.53× / 3.71	2.58× / 3.73	2.26× / 3.21	2.33× / 3.32	2.42× / 3.47	14%
Qwen3-8B	SpecBlock+adapt	—	2.61× / 3.90	2.64× / 3.78	2.69× / 3.75	—	2.56× / 3.70	17%
Qwen3-32B	EAGLE-3	2.54× / 4.36	1.82× / 4.27	2.28× / 4.09	2.38× / 3.96	2.18× / 4.09	2.15× / 3.96	24%
Qwen3-32B	SpecBlock (静态)	2.61× / 3.47	2.48× / 3.53	2.48× / 3.54	2.17× / 3.07	2.21× / 3.21	2.37× / 3.33	11%
Qwen3-32B	SpecBlock+adapt	—	2.55× / 3.73	2.51× / 3.61	2.57× / 3.61	—	2.47× / 3.54	13%

消融实验¶

基于 Llama-3.1-8B 在 $T=0$ 下逐步移除各关键模块的消融表现如下：

配置 / 变体	平均加速比 (Spd)	平均接受长度 ($\tau$)	相对基线变化	说明
SpecBlock (完整基线)	3.21×	4.41	—	完整模型（$K=4, M=2$）
− 层级隐状态移位 (layer-wise shift)	2.95×	4.02	Spd -0.26× / $\tau$ -0.39	贡献最大组件；破坏块内依赖，深层接受率崩溃
− 前缀广播 (prefix broadcast)	2.99×	4.13	Spd -0.22× / $\tau$ -0.28	丢失目标模型融合特征的强引导
− 秩预测头动态分支 (rank-guided)	3.13×	4.25	Spd -0.08× / $\tau$ -0.16	退化为固定均匀 Top-$k$ 树，算力分配失衡
− 有效前缀训练课程 (valid-prefix)	3.06×	4.23	Spd -0.15× / $\tau$ -0.18	错误前缀干扰损失函数导致曝光偏差
SpecBlock+adapt (自适应完整)	3.24×	4.54	—	在线自适应完整方案
− 成本感知老虎机 (cost-aware bandit)	3.14×	4.54	Spd -0.10× / $\tau$ ±0.00	退化为全量频密更新，反向传播吞噬吞吐
− 仅更新输出头策略 (head-only action)	3.00×	4.58	Spd -0.24× / $\tau$ +0.04	每次更新均强制全主干反向，延迟剧增

关键发现¶

起草延迟减半换取更高净吞吐：虽然 EAGLE-3 的平均接受长度较高（5.89 vs. 4.41），但其深度为 7 的串行起草导致起草耗时占迭代总时延的 31%（17ms）；SpecBlock 将其压缩为 2 次块前向（7ms），起草成本下降 52%，在目标模型验证耗时恒定（约 37ms）的约束下换取了整体 8%–13% 的净加速提升。
层级移位是保持块内因果依赖的核心支柱：消融显示去掉层级移位后平均加速比直接滑落 0.26×，在长文本与多轮对话基准（Alpaca 掉速 0.39×，MT-Bench 掉速 0.32×）上尤为剧烈，证明了单靠注意力掩码无法对抗长上下文下的注意力稀释。
服务期轻量更新的非对称收益：在线自适应中，仅更新输出头操作的微调成本较全主干低一个数量级以上，却能吸收绝大部分由任务分布漂移造成的校准偏差；配合成本感知 Bandit 剔除弱信号，在领域分布差异巨大的 Natural Questions 上实现了吞吐量自适应激增（加速比从 3.00× 跃升至 3.51×）。

亮点与洞察¶

层间隐状态横向传递规避串行调用：通过简单的层间线性变换将位置 $k-1$ 的隐状态显式注入位置 $k$，巧妙用微小的张量拼接算子在单次前向内模拟出了多轮自回归的特征演进轨迹。
输出端双头（LM Head + Rank Head）构成自适应防火墙：将语言建模头与排名预测头解耦，并在服务期自适应中优先微调输出双头，使系统能够在几毫秒的后台计算中修正输出概率空间，完全避开深层 Transformer 昂贵的全量反向。
有效前缀损失设计阻断负向梯度流：深刻抓住了投机验证“一旦截断后续全废”的推断特性，用简单的示性函数掩码切断无效推断轨迹上的梯度，精准杜绝了传统并行起草器因错误前缀导致的容量浪费。

局限与展望¶

块深与块数的超参数权衡固定：当前静态设定单块长度 $K=4$ 与块数 $M=2$，针对不同序列长度和特定领域生成任务（如高确定性的代码与高熵的发散创作），最佳块深几何可能存在动态调节空间。
依赖目标模型多层特征拼接：第一块输入仍需提取目标模型的低、中、高三层隐藏状态，在某些未开放中间层激活接口的黑盒推理引擎或特定硬件管线上集成成本较高。
显存与双流部署开销：服务期在线自适应依赖维护推理与训练两份起草器权重副本（$\theta_{\text{inf}}$ 与 $\theta_{\text{train}}$），在极端显存受限的单卡端侧设备上带来一定的显存压力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 块迭代架构与层级隐状态移位在保持路径依赖的同时大幅缩减起草轮次，思路精巧。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 种主流开源基座、6 大基准数据集、全套消融及单双卡在线自适应分析，数据详实自洽。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述清晰透彻，系统权衡分析严谨，实验图表极其丰富。
价值: ⭐⭐⭐⭐⭐ 为大模型投机解码推理系统提供了切实可行、兼顾吞吐与开销的高效解决方案。