DistiLLM: Towards Streamlined Distillation for Large Language Models

📄 arXiv: 2402.03898v2 📥 PDF

作者: Jongwoo Ko, Sungnyun Kim, Tianyi Chen, Se-Young Yun

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

发布日期: 2024-02-06 (更新: 2024-07-03)

备注: ICML 2024; Code is available at https://github.com/jongwooko/distillm

💡 一句话要点

提出DistiLLM以解决大语言模型蒸馏效率低下问题

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

关键词: 知识蒸馏 大语言模型 自回归模型 偏斜KL散度 自适应策略 模型压缩 高效推理

📋 核心要点

  1. 现有的知识蒸馏方法在自回归序列模型中缺乏标准化的目标函数，导致效率低下。
  2. DistiLLM通过引入偏斜Kullback-Leibler散度损失和自适应离线策略，提升了蒸馏过程的效率和效果。
  3. 实验结果显示，DistiLLM在指令跟随任务中表现优异，速度提升可达4.3倍，显著优于现有方法。

📝 摘要（中文）

知识蒸馏（KD）广泛用于将教师模型压缩为更小的学生模型，以降低推理成本和内存占用，同时保留模型能力。然而，当前针对自回归序列模型（如大语言模型）的KD方法缺乏标准化的目标函数。此外，利用学生生成的输出以解决训练与推理不匹配的问题显著增加了计算成本。为了解决这些问题，我们提出了DistiLLM，这是一个更有效且高效的自回归语言模型KD框架。DistiLLM包括两个组件：（1）新颖的偏斜Kullback-Leibler散度损失，利用其理论特性；（2）自适应离线策略，旨在提高利用学生生成输出的效率。大量实验表明，DistiLLM在构建高性能学生模型方面有效，并且与近期KD方法相比，速度提升可达4.3倍。

🔬 方法详解

问题定义：当前的知识蒸馏方法在自回归序列模型中存在缺乏标准化目标函数的问题，导致训练效率低下。同时，利用学生生成的输出来解决训练与推理不匹配的问题，显著增加了计算成本。

核心思路：DistiLLM的核心思路是通过引入偏斜Kullback-Leibler散度损失和自适应离线策略，来提高蒸馏过程的效率和效果。偏斜KL散度损失利用其理论特性，能够更好地捕捉教师模型与学生模型之间的知识传递。

技术框架：DistiLLM的整体架构包括两个主要模块：偏斜KL散度损失模块和自适应离线策略模块。前者负责优化学生模型的学习过程，后者则提高了学生生成输出的利用效率。

关键创新：DistiLLM的关键创新在于引入了偏斜Kullback-Leibler散度损失，这一损失函数的设计使得知识蒸馏过程更加高效，与现有方法相比，能够更好地平衡计算成本与模型性能。

关键设计：在DistiLLM中，偏斜KL散度损失的具体实现方式和参数设置经过精心设计，以确保在训练过程中能够有效地传递知识。此外，自适应离线策略的设计使得学生生成的输出能够被更高效地利用，从而减少了计算开销。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DistiLLM在指令跟随任务中的实验结果显示，其性能显著优于现有的知识蒸馏方法，速度提升可达4.3倍。这一结果表明，DistiLLM不仅在模型性能上有所突破，同时在计算效率上也实现了显著的改进。

🎯 应用场景

DistiLLM的研究成果在大语言模型的知识蒸馏领域具有广泛的应用潜力，尤其是在需要高效推理的场景中，如智能助手、对话系统和自动文本生成等。通过提升蒸馏效率，该方法能够在资源受限的环境中部署更强大的模型，具有重要的实际价值和未来影响。

📄 摘要（原文）

Knowledge distillation (KD) is widely used for compressing a teacher model to a smaller student model, reducing its inference cost and memory footprint while preserving model capabilities. However, current KD methods for auto-regressive sequence models (e.g., large language models) suffer from missing a standardized objective function. Moreover, the recent use of student-generated outputs to address training-inference mismatches has significantly escalated computational costs. To tackle these issues, we introduce DistiLLM, a more effective and efficient KD framework for auto-regressive language models. DistiLLM comprises two components: (1) a novel skew Kullback-Leibler divergence loss, where we unveil and leverage its theoretical properties, and (2) an adaptive off-policy approach designed to enhance the efficiency in utilizing student-generated outputs. Extensive experiments, including instruction-following tasks, demonstrate the effectiveness of DistiLLM in building high-performing student models while achieving up to 4.3$\times$ speedup compared to recent KD methods.