[论文解读] DistillSpec: Improving Speculative Decoding via Knowledge Distillation

Yongchao Zhou, Kaifeng Lyu+6 more|arXiv (Cornell University)|Oct 12, 2023

Topic Modeling被引用 4

一句话总结

DistillSpec 通过知识蒸馏提升推测解码性能，通过使用策略内生成的数据和任务特定的差异函数，使小型草稿模型与更大的目标模型对齐。相比标准推测解码，该方法实现了 10–45% 的加速，当先对目标模型进行蒸馏时，延迟降低高达 6–10 倍。

ABSTRACT

Speculative decoding (SD) accelerates large language model inference by employing a faster draft model for generating multiple tokens, which are then verified in parallel by the larger target model, resulting in the text generated according to the target model distribution. However, identifying a compact draft model that is well-aligned with the target model is challenging. To tackle this issue, we propose DistillSpec that uses knowledge distillation to better align the draft model with the target model, before applying SD. DistillSpec makes two key design choices, which we demonstrate via systematic study to be crucial to improving the draft and target alignment: utilizing on-policy data generation from the draft model, and tailoring the divergence function to the task and decoding strategy. Notably, DistillSpec yields impressive 10 - 45% speedups over standard SD on a range of standard benchmarks, using both greedy and non-greedy sampling. Furthermore, we combine DistillSpec with lossy SD to achieve fine-grained control over the latency vs. task performance trade-off. Finally, in practical scenarios with models of varying sizes, first using distillation to boost the performance of the target model and then applying DistillSpec to train a well-aligned draft model can reduce decoding latency by 6-10x with minimal performance drop, compared to standard decoding without distillation.

研究动机与目标

• 通过改进草稿模型与目标语言模型之间的对齐，解决推测解码中高接受率的挑战。
• 克服标准知识蒸馏的局限性，后者优化的是下游任务性能而非推测解码所需的分布对齐。
• 系统性地研究影响草稿模型质量及其与目标模型兼容性的设计选择——数据生成策略和差异函数选择。
• 通过将 DistillSpec 与有损推测解码结合，实现对延迟与性能权衡的细粒度控制。
• 开发一种实用的推理流水线：首先将大模型蒸馏为更小的高质量目标模型，然后应用 DistillSpec 训练高度兼容的草稿模型。

提出的方法

• 使用草稿模型生成的策略内数据训练蒸馏目标，确保学生模型从其自身生成的合理补全中学习。
• 根据特定的解码策略（贪婪 vs. 非贪婪）和任务，定制差异函数（如 KL 散度、交叉熵），以提升与目标模型输出分布的对齐度。
• 应用知识蒸馏训练一个紧凑的草稿模型，使其在一定程度上近似目标模型的标记和序列级分布，即使以牺牲独立任务性能为代价。
• 将 DistillSpec 与有损推测解码集成，允许对目标模型分布的可控偏离，从而实现更优的延迟-性能权衡。
• 提出两阶段流水线：首先将大模型蒸馏为更小的高质量目标模型，然后应用 DistillSpec 训练出高度对齐的草稿模型用于推测解码。

实验结果

研究问题

• RQ1 从草稿模型生成的策略内数据如何影响推测解码中草稿模型与目标模型之间的对齐？
• RQ2 在推测解码中，知识蒸馏的最优差异函数是什么？其选择如何随解码策略（贪婪 vs. 非贪婪）而变化？
• RQ3 DistillSpec 是否能在多种 NLP 任务和模型规模下提升推测解码的接受率和块效率？
• RQ4 将 DistillSpec 与有损推测解码结合后，与标准推测解码相比，其对延迟-性能权衡的影响如何？
• RQ5 当存在多个不同规模的模型时，最有效的推理流水线是什么？是否应在推测解码前或后应用蒸馏？

主要发现

• 在多个基准测试（包括 GSM8K、XSum 和 CNNDM）上，DistillSpec 在使用贪婪和非贪婪解码时，相比标准推测解码实现了 10–45% 的加速。
• 在蒸馏过程中使用策略内生成的数据显著提升了草稿模型与目标模型的对齐度，从而带来更高的接受率和块效率。
• 蒸馏的最优差异函数高度依赖于解码策略：在贪婪解码下，KL 散度表现最佳；在非贪婪采样下，交叉熵更有效。
• 将 DistillSpec 与有损推测解码结合后，可实现更平坦、更高质量的延迟-性能权衡曲线，实现更大的延迟收益且性能损失极小。
• 当按顺序应用——先蒸馏目标模型，再应用 DistillSpec 训练草稿模型——该方法相比标准解码将解码延迟降低了 6–10 倍，且性能退化可忽略。
• 在 GSM8K 上，使用非贪婪采样和最优超参数时，DistillSpec 相比标准推测解码，Rouge2 分数提升 15%，块效率提高 25%。