JetSpec：利用并行树草稿打破投机解码的扩展上限

提出因果并行树草稿框架，打破投机解码因果性-效率困境

前置知识

投机解码

投机解码是一种加速自回归大语言模型生成的方法，通过使用轻量级草稿模型先预测多个token，然后用目标模型并行验证这些预测。草稿模型提出N个草稿token后，目标模型并行验证它们并接受与自身预测一致的最长前缀，然后从第一个被拒绝的token恢复下一次迭代。这种方法的关键性能取决于接受率α和相对草稿成本c，预期加速倍数为 (1-α^{N+1})/((1-α)(Nc+1))。

理解投机解码是理解本文的基础，因为JetSpec正是为了解决现有投机解码方法在扩展草稿预算时面临的瓶颈问题而提出的。

自回归生成

自回归生成是指大语言模型按照从左到右的顺序逐个生成token的过程，每个token的生成都依赖于之前所有已生成的token。这种串行生成过程使得延迟成为部署在数学、编程等需要长生成任务中的主要挑战。自回归分布可以表示为连乘形式，其中x是前缀，y_{1:k}是长度为k的延续。

自回归生成是LLM的基础生成范式，也是投机解码要加速的对象。JetSpec的核心创新就是保持分支wise因果依赖，使草稿分布与目标模型的自回归因子化对齐。

树状投机解码

树状投机解码扩展了传统的线性投机解码，通过构建和验证多个候选分支来提高验证效率。在每一步，草稿器生成一个候选树，包含多个可能的延续路径，然后目标模型使用树注意力并行验证所有树节点。每个分支按照投机解码接受规则进行验证，接受的分支长度取决于草稿分布与目标模型分布的一致程度。

JetSpec的核心技术就是因果并行树草稿，理解树状投机解码的工作原理对于理解JetSpec如何突破现有方法的扩展瓶颈至关重要。

知识蒸馏

知识蒸馏是一种训练技术，通过让模型学习另一个教师模型的输出分布来传递知识。在JetSpec中，使用前向KL散度损失训练草稿头，使草稿分布保留目标模型在多个合理延续之间的相对偏好。对于每个活跃草稿位置m，最小化目标和草稿分布之间的前向KL散度，其中分布经过温度归一化处理。

JetSpec使用知识蒸馏来训练因果并行草稿头，理解蒸馏目标的选择对于理解JetSpec为什么能产生高质量的候选树分布很重要。

研究动机

现有的投机解码方法面临因果性-效率两难困境。自回归草稿器如EAGLE产生高质量的路径条件候选，但需要顺序草稿传递，随着树深度增加草稿成本增长。双向块扩散草稿器如DFlash一次通过生成所有位置，但它们的分支无关边缘可能形成个体合理但相互不一致的树，浪费预算并降低接受率。这种困境导致增加草稿预算只能在接受率保持高位且累积草稿开销保持较小时才能提高吞吐量，而随着草稿长度增长，即使接受率的轻微下降或不可忽略的草稿成本都可能抵消并行验证的好处。

本文的目标是本文的具体目标是设计一个投机解码框架，能够同时优化草稿成本和接受率，使投机解码能够更好地利用额外的解码计算。通过扩大草稿token预算，JetSpec旨在实现更高的端到端加速倍数，打破现有方法的扩展上限。

与已有工作不同的是，本文的独特切入角度是训练一个因果并行草稿头，在一次前向传递中预测多个树节点，同时通过对隐藏状态的块级因果注意力保留分支wise因果条件。这使每个分支具有与目标模型自回归因子化对齐的路径条件草稿分布，从而能够将更大的草稿预算转换为更长的接受前缀和更高的端到端加速。与现有方法不同，JetSpec不是单独优化接受质量或草稿成本，而是同时优化这两个关键因素。

核心方法

JetSpec采用基于头的草稿架构，复用目标模型的丰富中间特征同时保持草稿生成便宜。JetSpec从冻结的目标模型中提取融合隐藏特征，并用这些特征条件化一个因果并行草稿头，在一次前向传递中生成高质量候选树。关键创新是使用树因果注意力掩码，确保所有树节点可以并行处理的同时，每个分支遵循类自回归的依赖结构。具体来说，对于两个树节点u和v，掩码允许每个节点可以关注原始前缀及其祖先，但不能关注后代或不相关的兄弟分支。

JetSpec的核心创新点是将因果性保持与并行草稿结合起来。现有的并行草稿技术如DFlash会破坏有效树构建所需的因果结构，独立预测的未来位置没有沿着每个分支选择的token进行条件化，导致不一致的延续。JetSpec通过训练因果并行草稿头解决这个问题，该头在保留分支wise依赖的同时保留并行预测的效率。树因果掩码诱导分支wise草稿因子化，这镜像了目标模型的自回归因子化，但将每个草稿节点条件化在其自身路径上的具体祖先token上。

方法步骤详情

JetSpec的完整流程包括训练和推理两个阶段。在训练阶段，首先准备目标对齐序列，从训练语料库或由目标模型重新生成的延续中获取。对于每个序列，采样锚点位置，在每个锚点处构造长度为N的训练块，包含N个连续的未来位置。然后在块因果掩码下使用前向KL散度损失训练草稿头，使草稿分布匹配目标模型在相同真实前缀上产生的每位置教师logits。在推理阶段，在每个解码步骤i，因果草稿头条件化在上一步中验证和接受的token的隐藏状态，构造具有最大草稿深度N、分支宽度W和总节点预算B的候选树。通过一次草稿头前向传递，获得所有草稿深度的logits，提取每个深度的前W个候选token。然后使用分支评分函数驱动树扩展，反复弹出最高分的可扩展节点，用最多W个子节点扩展它，并重新插入更新分数的子节点，直到预算填满或没有可扩展节点剩余。最后，目标模型使用树注意力并行验证所有树节点，并沿每个候选分支应用投机解码接受规则。

技术新颖性

JetSpec的技术新颖性在于它首次实现了因果并行树草稿，打破了投机解码的因果性-效率两难困境。与DFlash的分支无关块扩散草稿不同，JetSpec通过树因果注意力掩码保持了每个分支的路径条件依赖，使草稿分布与目标模型的自回归因子化对齐。与EAGLE-3的自回归草稿不同，JetSpec在一次前向传递中生成所有树节点，避免了顺序草稿传递的开销。实验结果显示，这种因果并行草稿架构在高预算制度下显著优于现有方法，在MATH-500上达到超过10的接受长度和接近9.64倍的加速，在MT-Bench上达到接近6的接受长度和超过4倍的加速。

实验结果

JetSpec在多个基准测试上显著优于现有方法。在低token预算制度下使用16个草稿token，JetSpec达到竞争性的端到端加速，在GSM8K上达到4.80倍加速和6.00的接受长度，在MATH-500上达到6.06倍加速和7.75的接受长度。在高token预算制度下使用256个草稿token，JetSpec展现强增益，在MATH-500上达到10.76的接受长度和9.64倍加速，在MT-Bench上达到5.94的接受长度和4.58倍加速。与DDTree相比，JetSpec在所有基准测试上都有更高加速，表明因果树草稿产生比分支无关树构建更高的接受率。在温度1的非贪婪设置下，JetSpec仍然有效，表明因果树草稿的好处在非贪婪解码下保持鲁棒性。系统性能评估显示，大树预算在中小批量时最有用，在批量大小1时，预算从16增加到128将吞吐量从443.3提高到968.2 TPS，加速从3.09倍提高到6.75倍。

**Table 1**: Low-budget regime comparison of JetSpec and baselines trained with the same Qwen3-8B model and data recipe

**Table 2**: High-budget comparison on Qwen3-8B with at least 64 draft tokens

**Table 3**: Learning-rate ablation with JETSPEC and without loss weighting training (gamma = 0)

**Table 4**: Loss-objective ablation with JETSPEC at LR 6e-4 and gamma = 0

**Table 5**: Model generalizability: JetSpec vs. DDTree on Qwen3-30B-A3B (MoE target)

**Table 6**: Training-data ablation: JETSPEC vs. JETSPEC-Corpus

**Table 7**: Architecture and gamma ablation on MATH-500 with JetSpec at LR 3e-4 with forward-KL distillation

**Figure 1**: End-to-end decoding speedup over standard autoregressive decoding on H100 GPUs across math, coding, and chat benchmarks

查看结构化数据

任务	指标	本文	基线	提升
MATH-500 (Temperature=0, Budget=256)	Speedup / τ	9.64× / 10.76	DDTree: 8.78× / 9.81	+0.86×加速, +0.95τ
GSM8K (Temperature=0, Budget=256)	Speedup / τ	7.82× / 8.62	DDTree: 7.04× / 7.77	+0.78×加速, +0.85τ
MT-Bench (Temperature=0, Budget=256)	Speedup / τ	4.58× / 5.94	DDTree: 4.26× / 5.41	+0.32×加速, +0.53τ
HumanEval (Temperature=0, Budget=256)	Speedup / τ	7.12× / 7.78	DDTree: 6.31× / 6.96	+0.81×加速, +0.82τ

局限与改进

JetSpec的主要局限性包括：大树预算在高负载场景下可能低效，因为添加的验证成本和计算压力减少了接受更多token每步在吞吐量改进中的好处。静态预算策略可能不是最优的，作者指出应该根据服务负载选择树预算，大预算在低到中等负载制度下更可取。实验主要在Qwen3模型上进行，虽然在MoE模型上也验证了有效性，但在其他架构上的泛化能力需要进一步研究。训练需要访问目标模型的内部隐藏状态，这可能限制了在某些部署场景中的应用。此外，大预算的制度需要更多的GPU内存来存储和验证候选树。

独立分析的弱点

JetSpec的独立分析弱点包括：训练成本较高，需要780K示例和8个H100 GPU，这对资源受限的研究团队可能构成障碍。在大批量服务场景下性能增益有限，限制了在高负载生产环境中的实用性。需要访问目标模型的内部隐藏状态，这可能不适用于某些黑盒部署或API服务场景。树预算的静态选择缺乏动态适应性，无法根据实时服务负载或输入复杂度自动调整。计算复杂度随预算增加而增长，特别是在树验证阶段，可能抵消部分加速收益。

未来方向

未来研究方向包括：开发动态服务时预算调度策略，根据服务负载、输入复杂度或其他启发式信息实时调整树预算。研究在小规模数据上训练有效草稿头的方法，降低训练成本和门槛。探索不依赖于访问目标模型内部隐藏状态的草稿策略，扩大适用范围。将JetSpec集成到更多不同的LLM架构中，验证泛化能力。研究与其他推理加速技术的组合，如量化、剪枝或缓存优化。开发更高效的树验证算法，降低大树预算的计算开销。探索在更长生成场景中的性能和优化策略。

复现评估

JetSpec的复现性评估显示：作者已开源代码和模型，提供了完整的实现细节。训练使用780K示例的混合数据集，包括Nemotron Post-Training Dataset V2的编码和数学分割、STEM和聊天分割的随机样本，以及来自CodeAlpaca的20K额外示例。在8个H100 GPU上使用特定学习率进行训练。实验在Qwen3-8B和Qwen3-30B-A3B模型上进行，涵盖了密集和MoE目标模型。评估基准包括多个数学、编程和对话任务。系统性能评估通过将JetSpec集成到vLLM中进行，使用了优化的Triton内核和自定义分页FlashAttention内核。总体而言，JetSpec具有良好的复现性，提供了足够的细节和资源。

JetSpec：利用并行树草稿打破投机解码的扩展上限 JetSpec: Breaking the Scaling Ceiling of Speculative Decoding with Parallel Tree Drafting