Command Palette
Search for a command to run...
Draft-OPD:面向投机草稿模型的在线策略蒸馏
Draft-OPD:面向投机草稿模型的在线策略蒸馏
摘要
推测解码通过将目标模型与轻量级草稿模型配对来加速大语言模型推理,该草稿模型提出的 tokens 会进行并行验证。构建草稿模型(如 EAGLE3 或 DFlash)的常见方法是在目标生成的轨迹上进行监督微调(SFT)。然而,我们观察到 SFT 会迅速达到瓶颈:草稿模型在测试数据上的接受长度停止提升。其原因在于离线训练与推理之间的不匹配:在 SFT 中,草稿模型从固定的目标生成轨迹中学习,而在推测解码过程中,它是在其自身策略下提出的块上进行评估。这催生了在线策略蒸馏(OPD),即目标模型在草稿诱导的状态下对草稿模型进行监督。然而,OPD 对草稿模型而言仍然难以实施,因为它们无法可靠地独立展开完整序列,而目标辅助生成会使收集的序列遵循目标分布,从而消除在线策略信号。因此,我们提出了 Draft-OPD,该方法使用目标辅助展开以实现稳定的序列延续,并从验证暴露的错误位置重放草稿生成过程。这使得草稿模型能够从目标模型对接受和拒绝提议的反馈中学习,将训练重点集中在限制推测接受率的草稿诱导错误上。实验表明,Draft-OPD 在多种任务上为思维模型实现了超过 5imes 的无损加速,相较于 EAGLE-3 和 DFlash 分别提升了 23% 和 13%。
一句话总结
作者提出了 Draft-OPD,这是一种面向投机草稿模型的在线策略蒸馏方法。该方法通过结合目标辅助 rollout 与在验证暴露的错误位置重放草稿,利用目标模型对接受和拒绝的提议提供的反馈来训练草稿模型,从而解决了监督微调中离线训练与推理阶段不匹配的问题。
核心贡献
- 所提出的 Draft-OPD 框架将草稿模型的优化从离线监督微调转向在线策略蒸馏,解决了限制投机解码效率的离线训练与推理不匹配问题。
- 该方法将目标辅助 rollout 策略与错误位置重放相结合,在生成稳定续写的同时保留了必要的在线策略训练信号。
- 通过将优化重点集中在验证暴露的错误位置,该方法使草稿模型能够从接受和拒绝的提议中学习,有效突破了以往方法中观察到的接受长度瓶颈。
引言
投机解码通过让轻量级草稿模型提出多个 token,由更大的目标模型进行并行验证,从而加速自回归语言模型的推理过程,在保持输出质量的同时显著降低延迟。近期的进展通过从冻结的目标模型中提取特征或设计更强的草稿架构来增强该流程,但这些方法均依赖于使用目标生成轨迹的离线监督微调。这种静态训练范式最终限制了草稿模型的接受率,并阻碍了进一步的效率提升。为突破这一瓶颈,作者引入了一种在线策略蒸馏框架,直接在推理策略上训练草稿模型,从而实现显著更长的接受序列,并将投机解码推向传统离线训练极限之外。
方法
作者提出了 Draft-OPD,这是一种专为投机解码中基于训练的草稿模型定制的在线策略蒸馏框架。所解决的核心挑战是离线训练与推理阶段的不匹配:标准的监督微调(SFT)在目标模型生成的固定轨迹上训练草稿模型,而投机解码则在草稿模型自身生成的序列上对其进行评估。这种不匹配导致 SFT 出现性能瓶颈,因为模型学习的是在实际推理中并未遇到的状态。在线策略蒸馏(OPD)旨在通过在草稿模型自身策略诱导的状态上进行训练来弥补这一差距。然而,将 OPD 直接应用于草稿模型存在问题。如框架图所示,强制草稿模型自我生成完整序列(a)会导致不可靠且重复的输出。另一方面,使用目标辅助 rollout(b)虽然能产生稳定且高质量的续写,但会丢弃草稿模型的被拒绝提议,而这些被拒绝的提议对于学习草稿模型的失败位置最具信息量。这会破坏在线策略信号。 
为克服这些限制,Draft-OPD 采用三阶段流程。首先,它利用投机解码从目标模型收集稳定且高质量的 rollout。在此过程中,它记录草稿模型提出的每个草稿块的起始位置作为锚点。这既保证了收集序列的可靠性,又保留了草稿模型提出提议时的确切状态。如下方插图所示,草稿模型提出一个草稿块,随后由目标模型进行验证。每个草稿块的起始索引均保存为锚点。该过程持续至序列结束。
其次,框架从每个锚点重放草稿生成过程。对于每个锚点,草稿模型提出的草稿块会被重新生成,并使用草稿模型和目标模型在同一前缀上计算各 token 的对数概率。这使得目标模型能够对草稿模型的提议提供反馈,包括在初始验证中被丢弃的拒绝 token。验证结果将生成的 token 划分为接受集和拒绝集,使模型能够从成功和失败的预测中学习。第三,Draft-OPD 采用感知接受状态的蒸馏目标。对于接受的 token,它使用前向 KL 散度促使草稿模型的分布覆盖目标模型的分布,从而强化可靠的共识。对于拒绝的 token,它使用后向 KL 散度对目标模型不赞同时草稿模型的高概率预测进行惩罚,将学习重点集中在限制接受长度的错误上。为优先处理早期错误,拒绝 token 会根据其在草稿块中的位置进行指数级加权。最终目标函数是接受损失与拒绝损失的加权平均,使草稿模型能够同时从稳定高质量的 rollout 和投机验证暴露的关键失败点中学习。
实验
评估设置在同一计算预算下,于数学、代码和通用推理基准测试中将 Draft-OPD 与现有基线方法进行对比,并在优化的推理引擎上进行部署效率测试。主要结果展示了解码速度与 token 接受率的持续改善,消融实验验证了这些提升源于针对验证阶段错误的在线策略蒸馏,而非简单的监督数据增强。分析实验进一步证实,在训练期间保留草稿引发的错误至关重要,并指出了高度不确定的思考模式场景下的性能差距。综合来看,这些发现确立了 Draft-OPD 作为一种稳健且专注于效率的后期训练框架,它在保持目标模型输出分布不变的前提下,增强了草稿模型与目标模型的对齐,并加速了投机解码过程。
作者在 Qwen3 模型上将 Draft-OPD 与 EAGLE-3 和 DFlash 进行对比,评估了不同任务和设置下的解码加速比与平均接受长度。结果表明,Draft-OPD 在加速比和接受长度方面均持续优于基线方法,无论是否启用思考模式,且在不同温度设置及部署推理环境中均能保持这些优势。Draft-OPD 在所有测试模型和任务上的解码加速比和接受长度均高于 EAGLE-3 和 DFlash。该方法在不同解码温度及思考模式开启或关闭的情况下均能维持其优势。性能提升转化为服务吞吐量的实际增长,且在并发级别变化中观察到一致的好处。
作者在 Qwen3-4B 模型上将 Draft-OPD 与基础 rollout 基线进行对比,重点关注解码效率指标。结果表明,与基础方法相比,Draft-OPD 在多个基准测试中实现了更高的加速比和接受长度,表明草稿模型与目标模型的对齐效果得到改善。性能差异在不同任务中保持一致,证明了所提方法的有效性。Draft-OPD 在所有评估基准上的加速比和接受长度均高于基础 rollout 基线。接受长度和加速比的提升在不同任务中表现一致,表明其性能稳健。结果证明,与基础 rollout 方法相比,所提方法有效增强了草稿模型与目标模型的对齐。
{"summary": "作者在多个基准测试上评估了 Qwen3 模型上的 Draft-OPD 与基线方法,重点关注解码速度提升和平均接受长度等解码效率指标。结果表明,与替代方法相比,Draft-OPD 持续改善了这两项指标。在思考模式与非思考模式下均观察到性能提升,且这些提升在不同模型规模与任务中表现稳健。", "highlights": ["与基线方法相比,Draft-OPD 在所有评估基准和模型规模上均实现了更高的解码速度提升和接受长度。", "该方法在思考模式与非思考模式下均能维持其优势,表明其针对不同生成风格的草稿模型与目标模型对齐效果显著。", "性能提升在数学、代码生成和域外基准测试中保持一致,证明了其广泛的适用性。"]}
作者在匹配的训练 FLOPs 预算下,于 Qwen3 模型上将 Draft-OPD 与 EAGLE-3 和 DFlash 进行对比,重点关注解码加速比与平均接受长度。结果表明,Draft-OPD 在不同模型、任务和解码模式下均持续改善接受长度与解码速度,且在非思考模式以及 SGLang 部署的高并发场景下提升更为显著。Draft-OPD 通过强调验证阶段错误的在线策略训练,有效对齐草稿模型与目标模型,从而维持了相对于基线方法的优势。与 EAGLE-3 和 DFlash 相比,Draft-OPD 在所有评估模型和任务上的解码速度与接受长度均得到持续改善。在非思考模式下,Draft-OPD 实现的加速比和接受长度高于思考模式,表明其与目标模型行为的对齐效果更佳。该方法在 SGLang 的更高并发级别下仍能保持性能增益,展现出实际的部署优势。
作者在不同任务和并发级别下,评估了 Qwen3 模型上 Draft-OPD 与基线方法的性能,重点关注解码加速比与平均接受长度。结果表明,Draft-OPD 在所有测试配置中均持续改善接受长度与加速比,且在思考模式与高并发条件下观察到更高的增益。该方法在不同解码温度与模型规模下均能维持其优势,证明了草稿模型与目标模型对齐的稳健性。在匹配的训练预算下,Draft-OPD 在所有测试模型和任务上持续改善接受长度与解码速度,优于基线方法。与非思考模式相比,该方法在思考模式下实现了更高的加速比和接受长度,且增益随并发级别提升而增加。Draft-OPD 在不同解码温度和模型规模下均保持优势,表明其在推理与非推理生成场景下的草稿-目标对齐具有高度稳健性。
作者在 Qwen3 模型上,针对多样化任务、模型规模与部署条件,将 Draft-OPD 与多种基线方法进行对比,以验证其实现稳健草稿-目标对齐的能力。实验一致表明,与替代方法相比,该方法能带来更优的解码效率与 token 接受率,且性能在不同温度、并发级别与推理配置下保持稳定。在高负载与非思考模式下,性能提升尤为显著,证实了该方法在实际推理服务中的实用价值。最终,结果确立了 Draft-OPD 作为一种广泛适用的解决方案,能够在保持与目标模型对齐的同时,可靠地加速生成过程。