[论文解读] DistiLLM: Towards Streamlined Distillation for Large Language Models

Jongwoo Ko, Sungnyun Kim+2 more|arXiv (Cornell University)|Feb 6, 2024

Topic Modeling被引用 4

一句话总结

DistiLLM 提出了一种针对自回归大语言模型的简化知识蒸馏框架，引入了一种新颖的偏斜 Kullback-Leibler 散度损失，以实现稳定且收敛的训练，并采用一种自适应的离策略方法，高效利用学生模型生成的输出。该方法在指令遵循和摘要生成任务上实现了最先进性能，训练速度比近期的 KD 方法快达 4.3 倍。

ABSTRACT

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$ imes$ speedup compared to recent KD methods.

研究动机与目标

• 为自回归语言模型的知识蒸馏中缺乏标准化且稳定的损失函数提供解决方案。
• 减少因在离策略蒸馏中过度生成学生模型输出（SGO）而导致的训练低效问题。
• 通过最小化学生模型输出中的模式崩溃和分布差异，提升蒸馏效果。
• 在显著降低训练时间和计算成本的前提下，实现高性能的学生模型。
• 提供一个理论基础坚实、高效且有效的大型语言模型蒸馏框架。

提出的方法

• 引入一种具有坚实理论基础的偏斜 Kullback-Leibler（KLD）散度损失，以稳定梯度并减少蒸馏过程中的近似误差。
• 采用一种自适应的离策略策略，动态控制用于训练的学生生成输出的比例，提升样本效率。
• 同时利用固定数据集的监督信号和在线策略的学生生成序列，平衡分布对齐与训练效率。
• 优化蒸馏目标，以最小化教师与学生输出分布之间的非对称散度，避免模式崩溃。
• 采用类似课程学习的训练调度，逐步增加对学生生成输出的依赖，以稳定学习过程。
• 将偏斜 KLD 损失应用于序列级输出，确保在多种自回归生成任务中与教师模型的生成行为更好对齐。

Figure 1 : Examples of SGOs from GPT-2 (student) and their corresponding validation loss by GPT-2 XL (teacher). Since the teacher model may be not familiar with the SGO, using $p(\mathbf{y}|\mathbf{x})$ as a target distribution can misguide the student model, as shown in Tab. 7 .

实验结果

研究问题

• RQ1 一种理论基础坚实的偏斜对称 KLD 损失是否能提升自回归语言模型知识蒸馏中的训练稳定性和收敛性？
• RQ2 如何高效且自适应地将学生生成输出整合到蒸馏过程中，同时避免高昂的计算开销？
• RQ3 将固定数据集监督与自适应离策略 SGO 相结合，是否能比在线策略或固定 SGO 策略带来更好的泛化能力和更快的收敛速度？
• RQ4 所提出的方法在多大程度上能减少训练时间，同时保持或提升学生模型在生成任务中的性能？
• RQ5 与标准 KLD 和反向 KLD 方法相比，偏斜 KLD 损失在多种指令遵循和文本生成基准上是否具有鲁棒性？

主要发现

• DistiLLM 在指令遵循和文本摘要任务上实现了最先进性能，优于现有 KD 方法。
• 与近期的离策略 KD 方法（如 Gu et al. (2024) 和 Agarwal et al. (2024)）相比，该框架将训练时间减少了最多 4.3 倍。
• 偏斜 KLD 损失实现了更快的收敛速度和更稳定的训练过程，降低了标准 KLD 中观察到的梯度爆炸风险。
• 自适应离策略模块通过动态调整学生生成输出的使用方式，提升了样本效率，避免了冗余或低质量序列的生成。
• 在 Self-Instruct 和 Vicuna 基准测试中，DistiLLM 在所有测试案例中均能正确生成简洁且准确的回答，而其他方法则失败或产生幻觉输出。
• 即使在 SGO 使用量有限的情况下，DistiLLM 仍能实现卓越性能，表明高质量蒸馏无需持续且昂贵的学生模型生成。

Figure 2 : ( Left ): Normalized runtime according to the maximum response length of SGOs with GPT-2 XL teacher and GPT-2 student. ( Right ): Normalized runtime for various sizes of teacher and student models with a response length of 256. FWD and BWD denote forward and backward propagation, respecti