DistillSpec: Improving Speculative Decoding via Knowledge Distillation

📄 arXiv: 2310.08461v2 📥 PDF

作者: Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, Rishabh Agarwal

分类: cs.CL, cs.AI, cs.LG

发布日期: 2023-10-12 (更新: 2024-03-31)

💡 一句话要点

提出DistillSpec以解决草稿模型与目标模型对齐问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测解码 知识蒸馏 模型对齐 自然语言处理 推理效率

📋 核心要点

  1. 现有的推测解码方法在选择与目标模型对齐的草稿模型时面临困难，影响推理效率。
  2. DistillSpec通过知识蒸馏技术改善草稿模型与目标模型的对齐，采用在线数据生成和定制散度函数设计。
  3. 实验表明，DistillSpec在多个基准测试中实现了10%至45%的速度提升，并能显著降低解码延迟。

📝 摘要（中文）

论文提出了一种名为DistillSpec的方法，通过知识蒸馏技术改善草稿模型与目标模型之间的对齐，从而提升推理效率。现有的推测解码方法在选择紧凑的草稿模型时面临挑战，DistillSpec通过在草稿模型上进行在线数据生成和定制化的散度函数设计来解决这一问题。实验结果表明，DistillSpec在多个基准测试中相较于标准推测解码实现了10%至45%的速度提升，并且在不同模型规模的实际场景中，能够将解码延迟降低6至10倍，且性能损失最小。

🔬 方法详解

问题定义：论文要解决的问题是如何选择一个与目标模型良好对齐的紧凑草稿模型，以提高推测解码的效率。现有方法在这方面存在挑战，导致推理速度无法满足实际需求。

核心思路：论文的核心思路是通过知识蒸馏技术来增强草稿模型与目标模型之间的对齐性，进而提升推测解码的性能。具体而言，利用草稿模型生成在线数据，并根据任务和解码策略定制散度函数。

技术框架：整体架构包括两个主要阶段：首先，通过知识蒸馏提升目标模型的性能；其次，应用DistillSpec训练一个与目标模型对齐的草稿模型。该框架确保了草稿模型在生成过程中能够有效地与目标模型进行协同。

关键创新：最重要的技术创新点在于通过知识蒸馏实现草稿模型与目标模型的高效对齐，特别是在线数据生成和定制散度函数的设计，使得模型在推理时能够更快且更准确。

关键设计：在参数设置上，DistillSpec采用了特定的损失函数来优化草稿模型的输出与目标模型的分布一致性，同时在网络结构上进行了适当的调整，以适应不同任务的需求。整体设计确保了在速度和性能之间的良好平衡。

📊 实验亮点

实验结果显示，DistillSpec在多个标准基准测试中实现了10%至45%的速度提升，相较于标准推测解码方法表现出显著的性能优势。此外，在不同规模模型的实际应用中，解码延迟降低了6至10倍，且性能损失极小，验证了其有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和实时翻译等场景。通过提升推理速度，DistillSpec能够在需要快速响应的应用中提供更好的用户体验，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.