[Paper Review] DistiLLM: Towards Streamlined Distillation for Large Language Models

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

Topic Modeling4 citations

TL;DR

DistiLLM proposes a streamlined knowledge distillation framework for autoregressive large language models, introducing a novel skew Kullback-Leibler divergence loss for stable, convergent training and an adaptive off-policy method to efficiently utilize student-generated outputs. It achieves state-of-the-art performance on instruction-following and summarization tasks with up to 4.3× faster training than recent KD methods.

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.

Motivation & Objective

• To address the lack of standardized, stable objective functions in knowledge distillation for autoregressive language models.
• To reduce training inefficiency caused by excessive student-generated output (SGO) generation in off-policy distillation.
• To improve distillation effectiveness by minimizing mode collapse and distribution mismatch in student model outputs.
• To achieve high-performance student models with significantly reduced training time and computational cost.
• To provide a theoretically grounded, efficient, and effective framework for distilling large language models.

Proposed method

• Introduces a skew Kullback-Leibler (KLD) divergence loss with strong theoretical foundations to stabilize gradients and reduce approximation errors during distillation.
• Employs an adaptive off-policy strategy that dynamically controls the proportion of student-generated outputs used in training, improving sample efficiency.
• Leverages both fixed-dataset supervision and on-policy student-generated sequences, balancing distributional alignment and training efficiency.
• Optimizes the distillation objective to minimize the asymmetric divergence between teacher and student output distributions, avoiding mode collapse.
• Uses a curriculum-like training schedule that gradually increases reliance on student-generated outputs to stabilize learning.
• Applies the skew KLD loss to sequence-level outputs, ensuring better alignment with the teacher’s generative behavior across diverse autoregressive tasks.

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 .

Experimental results

Research questions

• RQ1 Can a theoretically grounded, skew-symmetric KLD loss improve training stability and convergence in knowledge distillation for autoregressive LMs?
• RQ2 How can student-generated outputs be efficiently and adaptively incorporated into distillation without incurring high computational overhead?
• RQ3 Does combining fixed-dataset supervision with adaptive off-policy SGOs lead to better generalization and faster convergence than on-policy or fixed SGO strategies?
• RQ4 To what extent can the proposed method reduce training time while maintaining or improving student model performance on generative tasks?
• RQ5 Is the skew KLD loss robust across diverse instruction-following and text generation benchmarks compared to standard KLD and reverse KLD approaches?

Key findings

• DistiLLM achieves state-of-the-art performance on instruction-following and text summarization tasks, outperforming existing KD methods.
• The framework reduces training time by up to 4.3× compared to recent off-policy KD methods such as Gu et al. (2024) and Agarwal et al. (2024).
• The skew KLD loss enables faster convergence and more stable training, reducing gradient explosion risks observed in standard KLD.
• The adaptive off-policy module improves sample efficiency by dynamically adjusting the use of student-generated outputs, avoiding redundant or low-quality sequences.
• On the Self-Instruct and Vicuna benchmarks, DistiLLM correctly generates concise, accurate responses in 100% of tested cases where other methods failed or produced hallucinated outputs.
• DistiLLM achieves superior performance even with limited SGO usage, demonstrating that high-quality distillation does not require constant, expensive student generation.

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