DistiLLM: Towards Streamlined Distillation for Large Language Models

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

Recent advancements in auto-regressive language models (LMs) such as large language models (LLMs) have significantly improved the quality of text generation in a variety of generative tasks such as task-agnostic instruction-following tasks. This success is often attributed to the increased scale of training data and model parameters (e.g., 175B parameters for GPT-3). However, expanding the parameter count brings associated costs, limiting the deployment of these models due to either high inference costs or large memory footprints. Therefore, a crucial objective for the practical use of these high-capacity models is to compress them by reducing the number of parameters while preserving their performance to the greatest extent possible.

As the necessity of reducing the demands on computational resources becomes important, KD emerges as a promising method. It involves the transfer of knowledge from a teacher model with large capacity (i.e., LLMs) to a smaller student LM. Most approaches of KD have employed the Kullback–Leibler divergence (KLD) loss, enforcing the student model to generate outputs that mirror the teacher model's outputs on a fixed dataset. Such methods of KD have significantly enhanced the performance of student models, making them competitive with teacher models while increasing efficiency, especially for various classification tasks.

These approaches using KLD on fixed datasets, however, have two primary shortcomings in applying to auto-regressive LMs. First, using the KLD can lead to sub-optimal results. Given the complexity of generative tasks compared to classification tasks, this can result in the student distribution becoming overly smooth and consequently failing to fully capture the teacher distribution, or becoming overly concentrated in high-probability distributions. This issue—referred to as mode averaging or mode collapse—arises due to the asymmetric nature of the KLD. Second, the use of a fixed dataset in the training phase can cause a distribution mismatch between the sequences observed during training and those generated by the student in auto-regressive inference, leading to exposure bias problems.

Recent studies have explored various divergence losses or the incorporation of student-generated output (SGO) to address the existing problems. However, these methods often lack standardized objective functions and are less efficient due to the continuous SGO generation. For instance, Agarwal et al. employed on-policy distillation with SGOs, but it faces low sample efficiency and high generation time as it constantly prompts the student to produce new training sequences. Also, their experiments indicated the optimal divergence seems to be task-dependent, which requires additional efforts to inconveniently select a proper loss function. Gu et al. introduced a policy optimization method to minimize reverse KLD between student and teacher distributions, yet this also compromises training efficiency by requiring generation from both models in every iteration.

Contributions. In this paper, we introduce DISTILLM, featuring a novel skew KLD loss and an adaptive off-policy approach, focusing on both distillation effectiveness and training efficiency. We provide both theoretical and empirical evidence that the components of DISTILLM work well individually and synergistically with each other. Our detailed contributions include: Skew KLD—we focus on two key issues of existing objective functions for auto-regressive LMs: instability from potential gradient explosions in optimizing the KLD loss, and lack of emphasis on generalizability and convergence; we introduce skew KLD, optimized for stable gradients and minimal approximation errors. Adaptive off-policy approach—while using SGOs in KD is generally effective, it significantly increases training time; we propose an adaptive off-policy module for adaptively and efficiently leveraging SGOs. Advanced performance and efficiency—DISTILLM accomplishes SOTA performance for student LMs on various generative tasks, while achieving 2.5×∼4.3× training speedup compared to recent KD techniques (MiniLLM, GKD).

We provide preliminary information on the KD for auto-regressive generative LMs. Given a source and target sequence pair, denoted as $(x,y)$, KD minimizes divergence $D$ between the distributions of a fixed teacher model $p(y|x)$ and a parameterized student model $q_\theta(y|x)$. The training data pairs $(x,y)$ are either sampled from a fixed ground-truth dataset or from teacher-generated outputs. Conventionally, KLD, denoted as $D_{\text{KL}}$, is the most widely used loss in KD due to its simplicity and tractability. The sequence-level distillation using the KLD is accurately decomposed into a sum of token-wise distillation:

We mathematically present our motivation that skewing such KLD is highly effective in improving the performance of student models with a more favorable optimization process. The definition of skew KLD (SKL) employs the parameter $\alpha$ that controls the mixing ratio of two distributions. The $\alpha$-SKL between $p$ and $q_\theta$ is defined as the KLD between $p$ and the mixture of distributions $\alpha p + (1-\alpha) q_\theta$:

We similarly define the $\alpha$-SRKL by $D^{(\alpha)}_{\text{SRKL}}(p, q_\theta) = D_{\text{KL}}\!\bigl(q_\theta,\,(1-\alpha) p + \alpha q_\theta\bigr)$. Here, following our thorough analysis, we present a comprehensive insight suggesting that S(R)KL is superior to other loss functions, owing to its more stable gradient and smaller approximation error.

Stable gradient. To provide stable optimization of SKL, we first analyze the gradients of KLD and SKL to parameter $\theta$. Given a context-target sequence pair $(x,y)$, we define the gradient of KLD w.r.t. $\theta$:

where $r_{p_1, p_2}$ is the ratio between arbitrary distributions $p_1$ and $p_2$. The result is the model probability's negative gradient, weighted inversely by its value. If $q_\theta(y|x) \approx 0$, the gradient norm grows large, causing a significant, potentially noisy parameter update step. These ingredients can adversely affect gradient updates, impacting the optimization process. We now compute the gradient of SKL w.r.t. $\theta$:

where $\tilde q_\theta(y|x) = \alpha p(y|x) + (1-\alpha) q_\theta(y|x)$. SKL offers a reduced gradient norm compared to KLD, due to $p$ and $q_\theta$ interpolation preventing the denominator of $r_{p,\tilde q_\theta}$ from reaching zero. This results in a more stable gradient for SKL. The gradient analysis for RKLD and SRKL reveals similar trends:

where $\tilde p(y|x) = (1-\alpha) p(y|x) + \alpha q_\theta(y|x)$. As $\alpha$ becomes large, the coefficient is effectively small for both SKL and SRKL.

Small approximation error. We show that the empirical estimator of SKL from mini-batch training has a bounded $L^2$ norm. This bounded norm ensures rapid convergence with minimal error between the estimator and true divergence, yielding high generalizability.

Theorem 1. Let $p^1_n$ and $p^2_n$ be empirical distributions of $n$ i.i.d. samples from $p^1$ and $p^2$. Under mild assumptions, we have an upper bound for the $L^2$ norm of the $\alpha$-SKL estimator:

for $c_1(\alpha) = \min\{1/\alpha^2,\;\chi^2(p^1,p^2)^2/(1-\alpha)^2\}$ and positive constants $c_2, c_3, c_4$ independent of $n, \alpha$. Thm. 1 states that a large $\alpha$ lowers the $L^2$ norm between empirical and true objectives. However, considering the gradient scale reduction in Eq. (6)–(7), the benefit of a reduced $L^2$ norm is negated by compensating the reduced gradient scale. Overall, selecting $\alpha$ involves a trade-off: the relationship between the upper bound of the normalized $L^2$ norm and $\alpha \in [0,1]$ appears to be convex. Our experiments indicate that $\alpha$-SKL and $\alpha$-SRKL are most effective with $\alpha = 0.1$, surpassing KLD, RKLD, and JSD.

In §2.2, we have discussed two main issues of naïvely using SGO: (1) the risk of noisy feedback due to the teacher model's unfamiliarity with the SGO and (2) the significant increase in training time. For instance, employing SGO at every training iteration (on-policy) can lead to a substantial increase in runtime, up to 5.5×, and may also result in performance degradation. To tackle these issues, we propose (1) an adaptive SGO scheduler to conservatively utilize SGO in KD, guided by the validation loss of student models, thus mitigating the risk of noisy feedback; and (2) an efficient off-policy strategy to improve the sample efficiency of SGO.

Adaptive SGO scheduler. We define the probability of using SGOs, denoted as $\phi$. We apply SGOs with a probability of $\phi$, i.e., using samples from a fixed dataset with a probability of $1-\phi$. Unlike previous methods that maintain a consistently high $\phi$, our scheduler starts with low $\phi$ value, gradually increasing during training. This strategy prevents student models from being overwhelmed by noisy feedback. To manage the increase of $\phi$, we primarily rely on validation loss as a metric. Our observations indicate that training on a diverse range of SGOs, rather than solely on a fixed dataset, mitigates training-inference mismatch and consequently lowers validation loss. We adjust $\phi$ by comparing the current and previous validation losses; an increase in validation loss leads to an increase in $\phi$. Concretely, if $L_{\tilde t} > L_{\tilde t-1} + \varepsilon$, then

This method effectively improves student model performance by striking a balance between managing noisy feedback and minimizing training-inference mismatch issues.

Off-policy approach for sample efficiency. To enhance efficiency, we replace the recently adopted on-policy approach with an off-policy approach, employing a replay buffer. In this buffer, we store SGOs from student models at a probability of $\lambda_R$. Subsequently, we randomly draw samples from this pool. We also replace the oldest samples of $\mathcal{D}_R$ with new ones once it reaches its maximum capacity. This off-policy strategy significantly improves the sample efficiency of KD with SGO, saving more resources than on-policy which constantly requires new data.

Off-policy reinforcement learning is prone to high bias error, particularly when there is a significant divergence between past and current policies. To address this, we set $\lambda_R := \phi\,(1 - t/T)$, where $t$ represents the current training iteration and $T$ is the total number of iterations. Define the replay ratio $\zeta := (1 - t/T)$:

• Early training phase (small $\phi$): student parameters evolve rapidly, so we focus on current SGOs with a high replay ratio $\zeta$ to minimize bias error.
• Late training phase (larger $\phi$): student nears convergence, so we predominantly reuse stored SGOs from $\mathcal{D}_R$ with a small $\zeta$.

Hence we consistently maintain a small $\lambda_R$ across the entire training, effectively balancing bias error reduction with sample efficiency.

Synergy with SKL. Off-policy approach's success stems from the fast convergence speed of S(R)KL while other loss functions cannot be achieved. As shown in our validation curves, both SKL and SRKL have a significant early-stage improvement, effectively leveraging the off-policy approach without high bias issues. This advantage is also evident in our ablation: unlike other baselines (ImitKD, GKD) that suffer performance drops when switching from on-policy to off-policy, our method maintains its efficacy. Consequently, our off-policy approach significantly improves the training efficiency with a negligible performance drop.

The full pipeline is summarized in Algorithm 1: at each iteration $t$, sample $u \sim \mathrm{Unif}(0,1)$; if $u < \lambda_R$, generate fresh SGO from $q_{\theta_t}$ and push into $\mathcal{D}_R$; if $u < \phi$, use a mini-batch from $\mathcal{D}_R$, otherwise from the fixed dataset $\mathcal{D}$; update $\theta_t$ with $D^{(\alpha)}_{\text{SKL}}$ or $D^{(\alpha)}_{\text{SRKL}}$; on validation step, run the SGO scheduler to update $\phi$.

We evaluate DISTILLM on instruction-following, text summarization, and machine translation tasks. We apply DISTILLM with SRKL and the off-policy approach with initial probability as zero and replay buffer size of 1000 as SRKL with $\alpha$ of 0.1 as determined through our ablation studies. We compare our approach with previous KD: (1) SFT directly fine-tunes the student on fixed datasets; (2) KD (Hinton et al., 2015) uses KLD on fixed datasets; (3) SeqKD applies SFT to teacher-generated output; (4) ImitKD employs KLD on SGO; (5) MiniLLM utilizes a policy gradient approach on SGO; and (6) GKD uses JSD on a mixture of SGOs and a fixed dataset.

Implementation details. Our implementation follows the experiment setup of Gu et al. We first construct the training data from databricks-dolly-15k, wherein we randomly select 14K samples for training and equally leave 500 samples for validation and testing. We evaluate the trained models on five instruction-following benchmarks: Dolly evaluation, Self-Instruct, Vicuna evaluation, Super-Natural Instructions, and Unnatural Instruction. We add a language modeling loss to the OpenWebText corpus for all experiments. For evaluation, we adopt two metrics: ROUGE-L and GPT-4 feedback.

Main results. Our results demonstrate DISTILLM's superiority to state-of-the-art methods, across diverse teacher-student combinations and evaluation metrics. Notably, GKD and MiniLLM, which recently incorporated SGO in distillation, are less effective or even lead to a performance decline in smaller-sized student models, particularly in GPT-4 feedback and some ROUGE-L scores. For evaluating DISTILLM on larger-sized LLMs, we utilize LoRA for training, with OpenLLaMA2-7B as the teacher and OpenLLaMA2-3B as the student. The results reveal that DISTILLM significantly surpasses other baseline methods in performance. While other supervised KD techniques are less effective than MiniLLM in LLM applications, DISTILLM uniquely achieves superior performance compared to MiniLLM. Moreover, DISTILLM requires only 1.6× the training time of naïve KD, whereas other methods take 3∼7×.

We evaluate the effectiveness of task-specific LMs on summarization and translation tasks using two datasets, SAMSum and IWSLT 2017. For SAMSum we use T5-XL v1.1 as the teacher model and T5-Base/-Small v1.1 as the student models. For IWSLT we employ mT5-XL as the teacher and mT5-Base/-Small as students. ROUGE-L (summarization) and BLEU (translation) scores show that DISTILLM outperforms other baselines, but with a smaller performance margin in single-task scenarios than in general-purpose instruction-following tasks. In SAMSum, students trained with ImitKD outperform those trained with GKD; in IWSLT, GKD outperforms ImitKD—consistent with prior findings that the effectiveness of objective functions and the use of SGO are task-dependent. Despite these variations, DISTILLM consistently shows superiority across different tasks, attributed to its adaptive use of the SGO scheduler and skew divergence.

We compare the performance of various models trained with different objective functions: conventional KLD, RKLD, and JSD with a $\beta$ of 0.9, as well as our SKL and SRKL with $\alpha$ of 0.1. The results show that our proposed objective functions generally outperform the others. Notably, both SKL and SRKL achieve remarkably high validation ROUGE-L scores during the entire training phase, consistently demonstrating the rapid convergence and strong generalization capabilities of our proposed loss functions. These empirical results verify our theoretical analysis and indicate that even our simple modification leads to significant performance enhancements.

We compare adaptive off-policy with on-policy and a mixed strategy (on-policy with probability 0.5, fixed dataset otherwise). The results indicate that our adaptive SGO scheduler effectively balances the trade-off between the risk of noisy feedback and training-inference mismatch. Notably, while the baselines suffer from performance degradation when applying SKL, our proposed adaptive strategy consistently improves the performance for all datasets. Moreover, the off-policy approach leads to a minimal performance drop while significantly improving computational efficiency—achieving 2.2× to 3.4× faster training speed compared to the on-policy or mixed strategy. Replacing the on-policy in GKD and ImitKD with our off-policy method significantly hurts those baselines, validating the synergy of SKL and adaptive off-policy in DISTILLM.

Skew values $\alpha$. We empirically evaluate the performance of $\alpha$-SKL under a wide range of $\alpha$. Both SKL and SRKL achieve the best performance at $\alpha = 0.1$, highly consistent with the theoretical result where both normalized $L^2$ norms have the smallest values at $\alpha = 0.1$. SKL shows a mild performance reduction as $\alpha$ increases above 0.1, while SRKL's reduction is comparably severe. One-stage Distillation. A significant issue with previous KD methods using SGO is their reliance on beginning with fine-tuned student models. A key advantage of DISTILLM is its rapid convergence and the adaptive SGO scheduler, which eliminates the need for such fine-tuned student models. We evaluated the performance of students distilled from pre-trained parameters without any prior fine-tuning, and DISTILLM maintains efficiency with only a minor performance reduction—a feat not matched by other methods.

KD effectively compresses neural networks, allowing smaller student models to match the performance of larger teacher models. Recently, KD has been extended to compressing auto-regressive LMs to address the challenges posed by the large scale of current LLMs. One popular direction harnesses LLMs as supervising data generators where only the teacher predictions are accessible (e.g., ChatGPT APIs); these methods employ LLMs for guided annotations or for imparting reasoning capabilities. Another noteworthy approach, when the teacher model is accessible, entails matching the student's generation distribution with that of the teacher through divergence loss functions. ImitKD demonstrated the effectiveness of SGO in distillation, leading Agarwal et al. (GKD) to propose on-policy SGO with diverse objectives like RKLD and JSD. Wen et al. examined various $f$-divergences, while Gu et al. (MiniLLM) proposed a policy-gradient method addressing the high variance issues in RL-based methods. Building on this research, we introduce DISTILLM with comprehensive analysis in objective function and data utilization.

We have proposed DISTILLM to address the challenges of KD frameworks for auto-regressive LMs. Our approach incorporates two key components: (1) SKL which is based on mathematically in-depth analyses and empirical evidence; (2) an adaptive off-policy approach that enhances the utility of SGO by reducing the noisy feedback introduced by SGO and improving the sample efficiency with a replay buffer. Through extensive experiments on various generation tasks, we have demonstrated the superior performance of DISTILLM, achieving significant training efficiency and performance improvement.