SpecBlock：具有动态树草拟的块迭代投机解码

# 基于动态树起草的块迭代投机解码

来源：https://arxiv.org/html/2605.07243
Weijie Shi1 Qiang Xu2 Fan Deng2 Yaguang Wu2 Jiarun Liu2 Yehong Xu1 Hao Chen1 Jia Zhu3 Jiajie Xu4 Xiangjun Huang2 Jian Yang2 Xiaofang Zhou1
1香港科技大学 2MetaX 3浙江师范大学 4苏州大学

###### 摘要

投机解码（Speculative Decoding）通过起草候选续写的树并在一次目标模型前向传播中进行验证，从而加速大语言模型（LLM）的推理。现有的起草模型分为两类，各自存在相反的弱点。自回归起草模型（如 EAGLE-3）保留了每条起草路径上的依赖性，但每增加一层树深度就需要调用一次起草模型，导致起草过程占用了每次迭代延迟中非微不足道的比例。并行起草模型通过在单次前向传播中预测多个未来位置来减少起草模型调用次数，但每个位置在预测时并未看到其他位置的信息，导致生成的路径常被验证器拒绝。

在本文中，我们提出了 SpecBlock，一种块迭代式起草模型，它将路径依赖性与低成本的起草相结合。每次起草模型的前向传播产生 $K$ 个相互依赖的位置，我们将此称为一个“块”。起草树通过反复的块扩展来增长。两种机制显式地传递路径依赖性，以保持后续起草位置的准确性：在块内部，逐层移位机制将前一个位置的隐藏状态传入每个解码器层；在块之间，每个新块可以从前一个块的任意位置开始，继承其隐藏状态以延长路径。为了将验证器预算分配在接受可能性较高的地方，一个联合训练的排名头（rank head）取代了固定的 top-$k$ 树，通过在起草过程中进行逐位置分支分配。为了避免在推理时永远不会产生的前缀上训练起草模型，有效前缀掩码（valid-prefix mask）会在较早位置出错时丢弃后续位置的损失。除了静态起草外，部署时的成本感知多臂老虎机（cost-aware bandit）利用免费的验证器反馈有选择地更新起草模型，仅当预期吞吐量增益超过更新成本时才执行更新。实验表明，SpecBlock 在 EAGLE-3 44%–52% 的起草成本下，将平均加速比提高了 8%–13%，而成本感知适应进一步将这一优势扩大到 11%–19%。

## 1 引言

由于大语言模型（LLM）的解码通常受限于内存带宽，投机解码（Leviathan et al., 2023; Chen et al., 2023; Miao et al., 2024; Sun et al., 2023; Zhou et al., 2023）通过使用小型草稿模型预测多个未来 token，并让目标模型并行验证这些候选 token 来解决这一瓶颈，使得一次目标前向传播可以接受多个 token，从而更充分地利用可用的计算能力。基于树的验证（Miao et al., 2024; Chen et al., 2024）进一步用起草树替换单一的起草序列，为验证器在未来位置提供多个备选方案，显著增加了接受长度。

基于树的验证的收益取决于一种平衡：树必须覆盖目标模型可能接受的续写，同时保持起草计算量足够小，使得节省的目标模型调用次数能转化为净加速。

**图 1：三种起草范式。** 自回归起草模型（左）每调用一次起草模型增加一层深度。并行起草模型（左下）在一次调用中独立预测所有深度。SpecBlock（右）每次调用产生 $K$ 个相互依赖的位置，并将来自多个起始位置的块批处理到随后的每次调用中，迭代地增长树。

自回归起草模型（如 EAGLE-3 (Li et al., 2024b, a, 2025)）逐层扩展起草树的深度。这保留了每条起草路径上的依赖性，平均接受长度可达近 6，但每增加一层树深度仍需额外的一轮顺序起草模型调用。尽管起草模型较小，但每轮调用本身受限于内存，因此串行调用积累了大量的权重加载开销，占 8B 级别目标模型每次迭代延迟的近 30%。并行起草模型（Cai et al., 2024）通过在一次调用中提出几个未来位置来减少这种开销，将起草开销缩减至约 7%。然而，一旦来自不同深度的备选方案组合成起草树，它们形成了一个巨大的组合空间，其中许多路径并非连贯的续写，验证器会在这些路径上浪费预算。

这要求在这两类方法之间取得平衡：一种起草模型既能减少起草模型调用次数，又能保留每条起草路径上的路径一致性，如图 1 所示。为了实现这种平衡，我们提出了 **SpecBlock**，一种块迭代式起草模型，它将每次起草前向传播视为产生一个多 token 块，并通过重复的块扩展来增长起草树。生成的树节点可以作为后续块的起始点，因此一次批处理的起草前向传播可以并行扩展多个分支，而不是逐层扩展树深度。两种机制通过显式传递依赖性来保持此结构中后续位置的准确性。在块内部，逐层移位机制将前一个位置的隐藏状态传入每个解码器层。在块之间，每个新块可以从前一个块的任意位置继续，条件是该位置的隐藏状态。

**排名引导的树构建。** 不同的起草位置值得不同数量的分支，因为目标 token 在一个位置的起草分布中可能排在顶部，而在另一个位置则排在很后面。一个联合训练的排名头读取每个位置的隐藏状态，并预测目标 token 在该位置起草分布中的排名高低，表示为一个粗略的分桶（bucket）。该分桶设置了该位置的同级备选方案数量，并决定该位置是否开始一个后续块，因此树是在起草过程中动态构建的，而不是事后修剪。

**有效前缀课程学习。** 自回归起草教师强制（teacher-force）每一步都从真实的地面真值前缀开始，因此每一步的损失都在正确的上下文中进行监督。SpecBlock 无法做到这一点，因为它的 $K$ 个预测是在一次前向传播中共同产生的，后续位置读取的是实际的早期预测，而不是地面真值前缀。如果早期预测错误，验证器将拒绝整条路径。此时对后续位置基于地面真值目标进行监督，只会浪费起草模型在永远不会提交的那些 token 上的容量。因此，有效前缀掩码会在同一条路径上较早位置出错时，丢弃任何后续位置的损失。

**成本感知的服务时适应。** 离线训练的小型起草模型无法适应每个领域，当服务提示分布偏离训练混合分布时，接受长度会下降。验证器已经在每次查询时产生了一个免费的适应信号，即每个被拒绝位置的目标分布，这是在验证过程中计算的，无需额外成本。一个成本感知的多臂老虎机读取此信号，并决定是否跳过更新、仅更新输出头或更新整个起草模型，仅在预期吞吐量增益超过更新成本时才采取非跳过操作。

实验表明，SpecBlock 在 EAGLE-3 44%–52% 的起草成本下，将平均加速比提高了 8%–13%。成本感知的服务时适应在具有足够流式查询的基准测试上，将这一优势扩大到 11%–19%。代码可在 https://github.com/shiweijiezero/SpecBlock 获取。

## 2 相关工作

#### 自回归起草模型

标准投机解码（Leviathan et al., 2023; Chen et al., 2023）建立了一个无损失的先起草后验证框架，其中一个小起草模型提出一个未来 token 链，目标模型并行验证。SpecInfer（Miao et al., 2024）将链推广为 token 树，因此一次验证可以在多个候选者中选择最长的匹配路径，在起草模型不确定时提升接受长度。EAGLE 系列（Li et al., 2024b, a, 2025）通过在目标模型的特征空间中进行自回归，并从起草模型的置信度生成动态树，以起草模型容量换取保真度，而 HASS（Zhang et al., 2024b）则通过模拟起草模型自身的多步 rollout 进一步对齐训练和推理。其他变体通过蒸馏（Zhou et al., 2023）、复用目标模型自身的浅层（Zhang et al., 2024a; Liu et al., 2024a）或从数据存储中检索缓存的续写（He et al., 2024; Fu et al., 2024）来减少起草模型开销。所有这些方法都通过一次一个深度地传递依赖性来获得接受率，因此每增加一个深度仍需要另一次起草步骤。

#### 并行和块式起草模型

并行和块式起草模型采取了相反的权衡，通过将起草模型压缩为一次前向传播，一次性预测几个未来位置。每个位置独立于其他位置进行预测，因此起草树未能捕获相邻 token 之间的依赖性，其路径在最初的几个深度之后就会偏离目标的续写。起草头方法（Stern et al., 2018; Cai et al., 2024）在固定偏移处附加独立的头。基于掩码的起草模型从可学习的掩码 token 预测每个未来偏移，BiTA（Lin et al., 2025）、ParallelSpec（Xiao et al., 2024）和 PARD（An et al., 2025）在掩码集成方式上有所不同，而 DART（Liu et al., 2026a）在其上层叠加了扩散风格的掩码预测目标。Hydra（Ankner et al., 2024）通过顺序连接头部重新引入依赖性，每个头部以前头部产生的候选续写为条件。块式和半自回归变体则通过层级联（Huang et al., 2026）、半自回归块起草（Gao et al., 2025; Liu et al., 2024b）或循环和块掩码架构（Kim et al., 2024; Gao et al., 2025; Cheng et al., 2024）来扩大起草单元，提高每次起草调用的平均 token 数。Falcon 是其中最近的，也起草半自回归块，但通过堆叠的 LSTM 层和松弛因果掩码注意力在块内携带依赖性，使块内所有位置都能看到彼此，并验证手工制作的静态解码树。SpecBlock 则通过逐层隐藏状态移位强制块内严格的从左到右依赖性，并通过联合训练的排名头动态塑造验证器树。

#### 树构建

在自回归起草模型之上，起草树是从起草模型信号外部塑造的，以决定目标验证哪些节点。Sequoia（Chen et al., 2024）解决了关于树大小和深度的离线动态规划问题，而 C2T（Huo et al., 2025）、OPT-Tree（Wang et al., 2025）、DySpec（Xiong et al., 2025）和 TALON（Liu et al., 2026b）则根据起草概率、置信度或预算信号调整树。SpecBlock 则通过排名头将树构建集成到起草模型中，设置逐位置分支以及哪些位置开始后续块。

#### 服务时适应

投机起草模型保持较小以使起草成本低廉，这也使它们对降低接受长度的服务时分布偏移敏感。一条路线保持起草模型权重冻结，仅适应投机超参数，如提议长度和树大小，要么通过对候选配置的多臂老虎机（Hou et al., 2025），要么通过对逐位置接受概率的学习阈值（Huang et al., 2024）。另一条路线通过验证器反馈蒸馏更新起草模型（Liu et al., 2023），扩展方案将起草模型重新表述为使用接受位置奖励的 KL-to-RL 时间表训练的自投机头（Bhansali and Heck, 2025），随时间调度更新（Park et al., 2026），或将训练更紧密地集成到服务堆栈中（Wang et al., 2026），但每种方法都提前确定了可训练参数子集。除了何时更新之外，SpecBlock 将刷新哪一部分起草模型参数作为一个每查询的决策，暴露了一个头与全起草模型的操作分裂，其中排名头和 `lm_head` 形成一个自包含的输出侧通道，可以吸收输出级错误而不触及解码器。

## 3 SpecBlock

令 $\mathcal{M}$ 为目标模型，$\mathcal{D}_\theta$ 为一个小型起草模型，它提出一个候选续写树供 $\mathcal{M}$ 在一次并行前向传播中验证。投机解码吞吐量是比率 $\Phi = \tau / (T_\mathcal{M} + T_\mathcal{D})$，其中 $\tau$ 是每次验证器调用平均接受的长度，$T_\mathcal{M}$ 是一次目标前向传播的时间，$T_\mathcal{D}$ 是组装该树所用所有起草模型调用的成本。自回归起草模型保持 $\tau$ 高，但每层树深度调用一次起草模型，支付与深度成比例的 $T_\mathcal{D}$。并行起草模型将 $T_\mathcal{D}$ 压缩为单次前向传播，但由于每个未来位置在预测时未看到其他位置，从而失去了 $\tau$。SpecBlock 通过介于这两种极端之间来提高 $\Phi$：每次起草前向传播产生 $K$ 个相互依赖的位置，深度超过 $K$ 的树通过在选择自早期块的起始点批次上重用起草模型来增长。我们进一步通过联合训练的排名头在起草过程中塑造树的分支，在推理时实际面临的前缀分布下训练起草模型，并在服务时使用验证器衍生的更新信号有选择地刷新它。

### 3.1 起草块

像之前的投机采样方法一样，SpecBlock 在起草和验证之间交替。与 EAGLE-3（Li et al., 2025）的不同之处在于起草阶段，每次起草前向传播预测 $K$ 个...