MiniLLM: On-Policy Distillation of Large Language Models

Knowledge Distillation (KD) is a promising technique for reducing the high computational demand of large language models (LLMs). However, previous KD methods are primarily applied to white-box classification models or training small models to imitate black-box model APIs like ChatGPT. How to effectively distill the knowledge of white-box LLMs into small models is still under-explored, which becomes more important with the prosperity of open-source LLMs. In this work, we propose a KD approach that distills LLMs into smaller language models. We first replace the forward Kullback-Leibler divergence (KLD) objective in the standard KD approaches with reverse KLD, which is more suitable for KD on generative language models, to prevent the student model from overestimating the low-probability regions of the teacher distribution. Then, we derive an effective on-policy optimization approach to learn this objective. The student models are named MINILLM. Extensive experiments in the instruction-following setting show that MINILLM generates more precise responses with higher overall quality, lower exposure bias, better calibration, and higher long-text generation performance than the baselines. Our method is scalable for different model families with 120M to 13B parameters.

Figure 1. The comparison of MINILLM with sequence-level KD (SeqKD) in terms of the average GPT-4 feedback score on our evaluation sets. Left: GPT-2-1.5B as the teacher model and GPT-2 125M / 340M / 760M as the student models. Middle: GPT-J 6B as the teacher and GPT-2 760M, 1.5B, GPT-Neo 2.7B as the students. Right: OPT 13B as the teacher and OPT 1.3B / 2.7B / 6.7B as the students.

With the rapid development of large language models (LLMs), a common technique to reduce their high computational resource demand is knowledge distillation (KD), where we train a small student model with supervision from a large teacher model. Two categories of KD are commonly applied: black-box KD, where only the teacher-generated texts are accessible, and white-box KD, where the teacher model's output distribution or intermediate hidden states are also available. Recently, black-box KD has shown promising results in fine-tuning small models on the prompt-response pairs generated by LLM APIs. With the emergence of more open-source LLMs, white-box KD becomes more valuable for both research communities and industry sectors because student models receive better signals from the output distribution and hidden states of teacher models, thereby potentially resulting in higher performance. However, white-box KD approaches are mostly studied for small ($<$ 1B parameters) language understanding models, while white-box KD for LLMs is yet to be explored.

In this work, we investigate white-box KD of LLMs where the output distribution of the teacher model is available. We argue that the standard KD objectives are sub-optimal for LLMs that perform tasks in a generative manner. Given the teacher distribution $p(y|x)$ and the student distribution $q_\theta(y|x)$ parameterized by $\theta$, standard KD objectives (including several variants for sequence-level models) essentially minimize the approximated forward Kullback-Leibler divergence (KLD) between the teacher and the student distribution, termed as $\mathrm{KL}[p\,\|\,q_\theta]$, which forces $q_\theta$ to cover all modes of $p$. For text classification tasks, $\mathrm{KL}[p\,\|\,q_\theta]$ works well because the output space usually consists of a finite number of classes such that both $p(y|x)$ and $q_\theta(y|x)$ have few modes. However, for open-ended text generation tasks, which is usually the case of LLM applications, the output spaces are much more complex and $p(y|x)$ can contain many more modes than what $q_\theta(y|x)$ can express due to the limited model capacity. Minimizing forward KLD causes $q_\theta$ to assign unreasonably high probabilities to the void regions of $p$ and produces very unlikely samples under $p$ during free-run generation.

Figure 2 (toy): we fit a Gaussian mixture distribution with a single Gaussian using forward KLD and reverse KLD. To alleviate this problem, we propose to minimize reverse KLD, $\mathrm{KL}[q_\theta\,\|\,p]$, widely used in computer vision and reinforcement learning. Compared to $\mathrm{KL}[p\,\|\,q_\theta]$, minimizing $\mathrm{KL}[q_\theta\,\|\,p]$ causes $q_\theta$ to seek the major modes of $p$, and assign low probabilities to $p$'s void regions. In text generation, this means that the student avoids learning too many long-tail variants in the teacher's distribution to focus on the generation correctness, which is critical in practical scenarios that require truthfulness and reliability.

To optimize $\min_\theta \mathrm{KL}[q_\theta\,\|\,p]$, as shown in Section 2.2, we derive its gradient with Policy Gradient and adopt an on-policy training approach. To further stabilize and accelerate training, we propose (1) single-step decomposition to reduce variance, (2) teacher-mixed sampling to alleviate reward hacking, and (3) length normalization to eliminate the length bias. Finally, we introduce the overall KD algorithm in Section 2.3. Our student models are named MINILLM, indicating our method is suitable for compressing large (generative) language models. We apply our method to various generative language models with sizes ranging from 120M to 13B in the instruction-following setting that covers a large range of NLP tasks. We use 5 datasets with Rouge-L, the GPT-4 feedback, and human judgment for evaluation. Experiments show that MINILLM consistently outperforms standard KD baselines on all the datasets and scales up well from 120M to 13B models (see Figure 1). More analysis shows that MINILLM yields lower exposure bias, better calibration, and higher long response generation performance, with negligible loss of diversity.

We consider conditional text generation where the model produces a response $y = \{y_t\}_{t=1}^T$ conditioning on a prompt $x$ sampled from the distribution $p_x$, which is typically how LLMs perform tasks. We formulate KD as an optimization problem to minimize the difference between a fixed teacher model distribution $p(y|x)$ and a student model distribution $q_\theta(y|x)$ parameterized by $\theta$. The standard KD methods approximately minimize the forward KLD:

where $p'$ can be the real data distribution (word-level KD) or the teacher distribution $p$ (sequence-level KD). Though widely used, $\mathrm{KL}[p\,\|\,q_\theta]$ tends to overestimate the void regions of $p$ in text generation tasks when $q_\theta$ is insufficiently expressive. KD for LLMs fits this case because LLMs perform tasks in a generative manner, such that the low-capacity student models cannot perfectly imitate the complex text generation distribution of the teacher models or humans.

Figure 3: Comparison between sequence-level KD (left) and MINILLM (right). Sequence-level KD samples $y \sim p$ (the teacher) and forces the student to memorize all samples generated by the teacher model — it minimizes forward KLD with $y$-expectation under $p$. MINILLM instead samples $y \sim q_\theta$ (the student) and uses the teacher's distribution as feedback to improve the student's own generations — it minimizes reverse KLD with $y$-expectation under $q_\theta$. The key change is which distribution we draw $y$ from.

We consider minimizing the reverse KLD between the student and teacher model distributions as the learning objective for MINILLM:

Minimizing reverse KLD has been shown to cause the mode-seeking behavior in generative modeling, where $q_\theta$ assigns high probabilities to $p$'s large modes and ignores the small ones (illustrated in a toy experiment in Figure 2). In this work, we first study this property for KD of LLMs in text generation. Minimizing forward KLD causes $q_\theta$ to place large probability masses on the zero-probability regions of $p$, corresponding to the generation of low-quality text in practice, while reverse KLD focuses on $p$'s major modes, which is crucial to ensure the correctness and faithfulness of text generation. As illustrated in Figure 3, unlike sequence-level KD that minimizes forward KLD, MINILLM that minimizes reverse KLD does not force $q_\theta$ to fit all $y$ sampled from the teacher distribution $p$. Instead, it encourages the student to generate samples preferred by the teacher within its own capacities, which is more possible to achieve. Interestingly, we also find another perspective of understanding MINILLM motivated by Inverse Reinforcement Learning. We present the related derivation in Appendix A.1.

Gradient Derivation. We notice that the gradient of the objective function $\mathcal{L}(\theta)$ in Equation (1) can be derived using the Policy Gradient Theorem for on-policy optimization:

We take instruction-following as the conditional text generation task, where models are trained to generate responses according to the instructions. We fine-tune a large model on the dataset $\mathcal{D}$ consisting of instruction-response pairs as the teacher model. Then, we compare different KD methods on $\mathcal{D}$ by evaluating the student model's instruction-following performance.

Base Models. Our student models come from three model families with various sizes: GPT-2 (120M, 340M, 760M), OPT (1.3B, 2.7B, 6.7B), and LLaMA (7B). For teacher models of each family we use GPT-2-1.5B, OPT-13B, and LLaMA-13B respectively. These teachers are fine-tuned on $\mathcal{D}$ in advance. Results with GPT-J as the teacher are presented in Appendix C.1.

Training. We construct the training data from databricks-dolly-15K, consisting of 15K human-written instruction-response pairs. We filter out samples that exceed the context length, then randomly split 1K / 0.5K samples for validation / testing, leaving about 12.5K examples for training. For $\mathcal{D}_\text{PT}$ we use OpenWebText for the GPT-2 family and the RoBERTa training corpus for the others. We set the teacher-mix-in strength $\alpha = 0.2$ throughout the experiments in Eq. (4). We search hyper-parameters using validation Rouge-L because it aligns with human preference better than validation loss.

Evaluation Datasets. We evaluate on five instruction-following sets: ① DollyEval — the 500-sample test split from databricks-dolly-15k; ② SelfInst — a user-oriented set with 252 samples; ③ VicunaEval — the 80 challenging Vicuna questions; ④ S-NI — the SUPER-NATURALINSTRUCTIONS test set (9K samples across 119 tasks), split by ground-truth length into $[0,5]$, $[6,10]$, $[11,+\infty]$ subsets (we use $[11,+\infty]$ in §3.2 and analyze all subsets in §3.3); ⑤ UnNI — 10K samples from UNNATURALINSTRUCTIONS, evaluated with the same length-split protocol.

Metrics. Three metrics: ① R-L — Rouge-L score, suitable for large-scale instruction-following evaluation; ② GPT4 — GPT-4 feedback: ask GPT-4 to score model responses vs. ground-truth answers on a 1–10 scale and report the ratio of model-total to gold-total; ③ Human Evaluation — volunteers compare two responses on SelfInst and annotate Win / Tie / Loss. For all test sets we sample with $\text{temperature}=1$ and average over 5 generations per prompt (different random seeds).

Baselines. Three baselines: SFT w/o KD directly fine-tunes the student on $\mathcal{D}$ with golden responses; KD (word-level) uses the teacher distribution as token-step supervision; SeqKD fine-tunes the student on teacher-generated data.

We present the R-L and GPT-4 evaluation results in Table 1, from which we have three observations. First, by comparing the overall performance of MINILLM with the baselines, we observe that the model distilled by our KD method outperforms the baselines in almost all cases — across different base models, evaluation sets, and both Rouge-L and GPT-4 feedback. This verifies the good generalization and high overall performance of our KD method. We also find that MINILLM generally works much better than the baselines on datasets other than Dolly, indicating its good out-of-distribution generalization.

Second, the Rouge-L scores show that MINILLM produces the most precise responses that have high overlaps with the ground-truth responses. In some cases — especially on Vicuna, S-NI, and UnNI — student models reach even higher Rouge-L scores than the teacher. We conjecture that the standard teacher-forcing fine-tuning on $\mathcal{D}$ brings a training-inference discrepancy (i.e., exposure bias) to the teacher, while MINILLM is optimized with policy optimization that samples responses from the student itself during training, thereby alleviating exposure bias. Further analysis on exposure bias is in §3.3.

Third, comparing the results across model sizes and families, MINILLM's improvement is consistent when base model sizes range from 120M to 13B across three model families (also illustrated in Figure 1), showing excellent scalability and generalization in the era of LLMs. The human evaluation on SelfInst (LLaMA family) is shown in Figure 4: MINILLM obtains better human preference than all baselines, performing comparably to the teacher.

Table 1 highlights. Best per-size scores boldfaced; an asterisk marks rows where the student beats its teacher. Selected numbers: GPT-2-1.5B → 120M MINILLM 44.7 (GPT4 on Dolly) vs. SeqKD 41.2 / KD 40.3; OPT-13B → 1.3B MINILLM 60.7 vs. SeqKD 51.0 / KD 52.7; LLaMA-13B → 7B MINILLM 76.4 vs. SeqKD 73.6 / KD 73.7.

Scaling Law of Teacher. Although it is intuitive that we can distill better students from larger teachers, prior work has shown that increasing teacher size does not guarantee improvement, sometimes even harming distillation. We therefore compare MINILLM and SeqKD with a fixed student and varying teacher sizes. Figure 5 (GPT-2 family) and Appendix C.2 (OPT family) show that MINILLM consistently outperforms SeqKD, and student performance is positively correlated with teacher size — demonstrating the method's potential for compressing models with massive parameters.

Exposure Bias. Models trained with forward KLD suffer from exposure bias due to the mismatch between teacher-forcing training and free-run generation. MINILLM, which samples from the student during training, alleviates this mismatch. Figure 6 uses the ExAccErr metric (defined in Appendix B.5) to measure the excess accumulated error caused by exposure bias (GPT-2-125M as student, GPT-2-1.5B as teacher, Dolly test set, 10 samples per prompt). The ExAccErr of the baselines continuously grows with generation length, while MINILLM's ExAccErr is much lower and stops accumulating beyond ${\sim}150$ tokens.

Calibration. Prior work has shown that policy-optimization trained models tend to be poorly calibrated. We test calibration of MINILLM and KD baselines on SST2 and BoolQ (LLaMA-7B), using zero-shot classification prompts and label-word probabilities to compute ECE. Table 2 shows that KD and SeqKD are much worse calibrated than the teacher (ECE: Teacher 0.025 / 0.356; KD 0.191 / 0.682; SeqKD 0.243 / 0.681), potentially explaining their poor canonical-benchmark performance. We suspect forward KLD pushes high probabilities to void regions, causing large distribution gaps. MINILLM (ECE 0.099 / 0.502) focuses on the major modes and narrows the calibration gap to the teacher.

Performance on Different Response Length. Figure 7 plots the Rouge-L scores of distilled models against SFT on three S-NI subsets split by ground-truth length. For short responses (${\le}5$ tokens) all methods score low — most training responses are long sentences (training-evaluation distribution shift), and the output space is small enough that the student can cover most teacher modes (so forward / reverse KLD perform similarly). For longer responses (${\ge}6$ tokens), the teacher distribution has many more modes than the student, where MINILLM's advantage over forward-KLD methods becomes substantial. Similar results on UnNI are in Appendix C.3.

Generation Diversity. Prior work argued that minimizing reverse KLD risks losing modes and hurting diversity. We discuss diversity along three axes: (i) producing multiple distinct responses to one prompt; (ii) producing linguistically complex responses; (iii) covering the real data distribution. For (i) we argue that for many practical NLP scenarios (those demanding truthfulness and reliability) one correct response suffices. For (ii) and (iii), Table 3 reports Distinct-4 and test-set LM-loss on the LLaMA family: Teacher 99.3 / 3.55; SFT 99.5 / 3.89; MINILLM 99.0 / 3.95 (DollyEval) — MINILLM preserves distinct-4 proportion and LM-loss, with negligible loss of diversity.

We evaluate the three stabilization strategies in §2.2 by distilling GPT-2-125M from GPT-2-1.5B. Table 4 reports the best validation Rouge-L and corresponding test scores. Removing Length Normalization drops valid R-L from 27.4 → 17.4 (catastrophic) — the model collapses to empty responses. Removing Teacher-Mixed Sampling drops to 22.3 — the model exhibits classic reward hacking (repeating short / meaningless strings that teachers score high). Removing Single-Step Decomposition only drops to 27.0/23.7, but Figure 8 shows the training trajectory becomes much noisier (higher variance). Conclusion: all three are needed — teacher-mixed sampling and length-norm to survive, single-step decomposition to converge faster.

Large Language Models. LLMs solve various NLP tasks generatively. Recent works apply instruction tuning or learning from human feedback to further align LLMs with humans and create general AI assistants. There are also efforts to build open-source LLMs (OPT, LLaMA, Pythia) to facilitate research and industry. However, LLMs' broad capacities emerge only at large parameter scales requiring massive compute — making model compression critical for practical deployment and further research.

Knowledge Distillation. KD is a widely used compression technique, training a student under a teacher's guidance. In NLP, KD has been applied to text classification by mimicking the teacher's output distribution, hidden states, or attention scores. For text generation, the standard KD recipe is to approximately minimize the forward KLD — either using the teacher's per-token output as supervision (word-level KD) or directly training on teacher-generated texts (SeqKD). In this paper we instead minimize the reverse KLD, which is more suitable for LLMs when the teacher distribution is available. Concurrent works (GKD; f-divergence KD) also explore alternative distribution discrepancy metrics for KD.

Distribution Discrepancy Metrics in Text Generation. Discrepancy metrics play a major role in training text generators. Forward KLD is widely used because it reduces to MLE. However, minimizing forward KLD induces zero-forcing (more precisely, zero-avoiding) behavior — models try to cover all modes and sacrifice accuracy on major modes. Prior remedies include reverse KLD, Total Variation Distance, and Optimal Transport. This paper tackles the problem specifically under the KD-for-LLMs scenario.

In this work, we investigate the problem of distilling the knowledge of LLMs into small language models. We find that the standard distillation objective of minimizing the forward KLD is sub-optimal in language generation scenarios because the teacher's output distribution contains many more modes than the student's, and forward KLD forces the student distribution to overestimate the low-probability regions of the teacher distribution. Therefore, we propose MINILLM, which minimizes the reverse KLD between teacher and student distributions and designs an on-policy algorithm to optimize this objective. Extensive experiments show that MINILLM produces more precise responses with higher overall quality than standard KD models. We further find that MINILLM has lower exposure bias, better calibration, and higher long-response performance, while preserving generation diversity.