MiniLLM: On-Policy Distillation of Large Language Models
MiniLLM uses reverse-KL on-policy distillation; LLaMA-7B reaches 73.1 GPT-4 and 23.2 Rouge-L on SelfInst.
Key Findings
Methodology
MiniLLM replaces forward KL[p||q] with reverse KL[q||p], encouraging a limited student to model the teacher's major modes rather than low-probability regions. Policy Gradient yields an on-policy objective, augmented by Single-Step Decomposition, Teacher-Mixed Sampling, and Length Normalization. An auxiliary pretraining loss LPT preserves general language-modeling ability.
Key Results
- Across GPT-2, OPT, and LLaMA families, with students from 120M to 7B parameters, MiniLLM generally beats SFT, word-level KD, and SeqKD. OPT-1.3B reaches 47.0 GPT-4 and 14.8 Rouge-L on SelfInst, versus 36.6 and 12.7 for SeqKD.
- LLaMA-7B obtains 64.1 GPT-4 and 20.7 Rouge-L on VicunaEval, and 73.1 and 23.2 on SelfInst, close to the LLaMA-13B teacher's 75.5 and 23.4. GPT-2-760M reaches 37.7 Rouge-L on UnNI.
- MiniLLM's ExAccErr stops accumulating beyond roughly 150 tokens. On SST2, it obtains ECE/accuracy of 0.099/89.7, versus KD's 0.191/84.7. Removing length normalization reduces Dolly validation Rouge-L from 24.6 to 14.7.
Significance
The paper addresses a central open-source-LLM compression problem: exploiting a white-box teacher distribution rather than merely imitating sampled text. It shows why classification-style forward KL is poorly matched to open-ended generation: an under-capacity student may place mass in regions the teacher barely supports. MiniLLM reduces inference cost while improving long-form generation, calibration, and out-of-distribution instruction following, with direct implications for private, edge, and low-latency deployment.
Technical Contribution
The main contribution is a practical reverse-KLD objective for white-box LLM distillation and a policy-gradient derivation that makes the sequence objective trainable on policy. Equation 7 combines vocabulary-level single-step expectations with a sequence-level advantage term. Teacher mixing suppresses reward hacking; approximate importance ratios reduce variance; length normalization removes a short-response bias; LPT protects canonical language ability.
Novelty
Unlike SeqKD, which asks the student to memorize all teacher-sampled responses, MiniLLM lets the student select teacher-preferred modes within its own capacity. Its novelty is therefore both objective-level and algorithmic: reverse KL is paired with a complete low-variance, anti-reward-hacking on-policy procedure, validated across 120M–13B model scales.
Limitations
- Experiments focus on Dolly-based instruction following, where teachers are fine-tuned on the same task distribution. Generalization to difficult reasoning, factuality, safety, multilingual settings, and code is therefore not established.
- Reverse KL is mode-seeking and may discard useful diversity; on-policy sampling, teacher forward passes, and importance weighting are more expensive than SFT or SeqKD.
- Aggregate metrics provide limited diagnosis of factual errors, hallucination, safety failures, or robustness to weak and heterogeneous teachers.
Future Work
Future work should study adaptive KL direction, temperature, and teacher-mixing schedules; combine distillation with preference optimization, safety constraints, retrieval, and multiple teachers; and evaluate mathematics, code, multilingual, long-context, and black-box settings. Sampling caches and more efficient teacher-query strategies could reduce training cost.
AI Executive Summary
Large language models are capable but expensive to run. Knowledge distillation promises smaller systems, yet conventional forward KL asks a limited student to cover every mode of a teacher's open-ended text distribution. This can push probability into regions the teacher scarcely supports, producing unreliable free-running text. MiniLLM studies the underexplored white-box setting in which the teacher's token probabilities are available.
The method reverses the objective to KL[q||p]. Rather than memorizing every teacher sample, the student generates its own responses and receives teacher-based feedback, emphasizing major modes. Policy Gradient turns this sequence objective into on-policy optimization. Single-Step Decomposition reduces variance, Teacher-Mixed Sampling with α=0.2 limits reward hacking from repetition, Length Normalization prevents empty or overly short answers, and LPT preserves general language ability.
Experiments use GPT-2, OPT, and LLaMA models, Dolly-15K training data, and DollyEval, SelfInst, VicunaEval, S-NI, and UnNI evaluation sets. LLaMA-7B reaches 73.1 GPT-4 points and 23.2 Rouge-L on SelfInst, close to the 13B teacher's 75.5 and 23.4. OPT-1.3B reaches 47.0 and 14.8 on the same set. MiniLLM also reduces exposure bias and improves calibration: SST2 ECE is 0.099 versus 0.191 for KD. The broader lesson is that generative distillation should teach a smaller model which responses are reliable, not force it to reproduce every possible response.
Deep Analysis
Background
Knowledge distillation, introduced for model compression by Hinton et al., was first studied mainly in white-box classification. Recent SeqKD and API-based approaches train generators on teacher-produced text. Open-source LLMs make white-box distillation attractive because full output distributions are available, but generation is multimodal, long-tailed, and far more complex than classification.
Core Problem
Forward KL[p||q] penalizes the student for failing to cover every teacher mode. With limited capacity, this can assign excessive mass to low- or zero-probability regions, causing poor free-running text, exposure bias, and miscalibration. The challenge is to exploit teacher probabilities while training on the student's own generation trajectory.
Innovation
- �� Reverse KLD focuses on major teacher modes.
- �� Policy Gradient turns the sequence objective into on-policy learning.
- �� Single-Step Decomposition computes vocabulary expectations and lowers variance.
- �� Teacher-Mixed Sampling uses ep=αp+(1−α)q to suppress reward hacking.
- �� Length Normalization removes short-output bias, while LPT preserves pretrained ability.
Methodology
- �� Supervised fine-tune the student on Dolly instruction-response pairs and select the best validation checkpoint.
- �� Optimize L=KL[qθ||p]=−Eqθ log(p/q), with token reward rt=log[p(yt|·)/qθ(yt|·)].
- �� Sample from the teacher-student mixture ep and apply approximate per-token importance ratios.
- �� Compute the Single component by summing over the vocabulary; compute the Long component from normalized future rewards RNorm.
- �� Update parameters with Single, normalized Long, and LPT gradients, using clipping for stability.
Experiments
Training uses databricks-dolly-15K: about 12.5K train, 1K validation, and 0.5K test examples. OpenWebText supports GPT-2; the RoBERTa corpus supports other families. Students include GPT-2 120M/340M/760M, OPT 1.3B/2.7B/6.7B, and LLaMA-7B; teachers are GPT-2-1.5B, OPT-13B, and LLaMA-13B. Baselines are SFT, word-level KD, and SeqKD. Metrics include Rouge-L, GPT-4 feedback, human preference, ECE, and ExAccErr.
Results
MiniLLM leads across most settings. LLaMA-7B obtains 64.1 GPT-4 and 20.7 Rouge-L on VicunaEval, and 40.2 Rouge-L on UnNI. OPT-2.7B reaches 52.7 GPT-4 and 17.2 Rouge-L on SelfInst. Larger teachers consistently improve MiniLLM students. Ablations show that removing teacher mixing or length normalization lowers Dolly validation Rouge-L to 20.4 and 14.7, respectively.
Applications
The approach suits local assistants, enterprise question answering, edge devices, and low-latency APIs. It is especially relevant where a white-box teacher, task prompts, and additional sampling compute are available, and where long answers, calibration, and out-of-distribution generalization matter more than simple imitation.
Limitations & Outlook
MiniLLM requires access to teacher probabilities, so it does not directly solve black-box API distillation. Repeated on-policy sampling and teacher inference increase training cost. Reverse KL may sacrifice diversity and does not guarantee factuality or safety. Important next tests include reasoning, code, multilingual, long-context, weak-teacher, and multi-teacher scenarios, together with adaptive objectives and cached sampling.
Plain Language Accessible to non-experts
Imagine a master chef and an apprentice. The usual training method tells the apprentice to copy every dish the master has ever made, including rare experiments and occasional mistakes. The apprentice cannot remember everything, so the copied menu becomes confused. MiniLLM changes the lesson: the apprentice cooks first, then the master evaluates which choices resemble excellent dishes.
The apprentice therefore concentrates on reliable recipes instead of trying to cover every strange possibility. The teacher sometimes joins the sampling process, preventing the apprentice from repeating one phrase merely because it earns an easy score. Rewards are also adjusted for dish length, so the apprentice cannot win by serving an empty plate. Ordinary reading practice is retained so the apprentice does not forget basic language skills.
This produces more precise and stable long answers, especially on unfamiliar instructions. The price is a more expensive training kitchen: the teacher must repeatedly inspect the apprentice's work. And a better imitation is not automatically truthful or safe; those properties still require separate testing.
ELI14 Explained like you're 14
Think of a video game where a 13B model is a pro player and a 1.3B model is a beginner character. Normal training is like forcing the beginner to copy every move in the pro's entire replay, including accidental button mashes. The beginner cannot store all of it, so strange moves appear during a real match.
MiniLLM lets the beginner play first. The pro then says which choices were strong and which were risky. The beginner focuses on the pro's best routes instead of memorizing every possible move. The pro sometimes helps during practice so the beginner cannot farm points by repeating one silly action. Rewards are adjusted for match length too, so the beginner does not quit immediately.
Researchers tested this on Dolly, SelfInst, VicunaEval, S-NI, and UnNI with GPT-2, OPT, and LLaMA. LLaMA-7B scored 73.1 GPT-4 points and 23.2 Rouge-L on SelfInst, close to the LLaMA-13B teacher's 75.5. OPT-1.3B reached 47.0.
The big idea is simple: do not make a smaller model memorize the whole giant model; teach it to choose better actions. Cool, right? The catch is that training costs more, the teacher must reveal probability information, and the student can still hallucinate or make unsafe choices.
Glossary
Reverse KLD
A divergence KL[q||p] measuring how the student differs from the teacher. Unlike forward KL, it is mode-seeking and emphasizes teacher-supported regions.
MiniLLM's central distillation objective.
On-Policy Distillation
The student samples its own or a mixed trajectory and is updated using teacher feedback. This better matches training with free-running inference.
The optimization framework for reverse KLD.
Policy Gradient
An optimization method that changes the probability of sampled actions according to their rewards. Here, token-level teacher-to-student log-probability ratios act as rewards.
Used to derive the sequence-level gradient.
Exposure Bias
The mismatch between training on gold histories and inference on model-generated histories. Errors can accumulate as generation becomes longer.
MiniLLM reduces it through on-policy sampling.
Teacher-Mixed Sampling
Sampling from ep=αp+(1−α)q, combining teacher and student distributions. It reduces degenerate samples and reward hacking.
The experiments set α=0.2.
ExAccErr
A metric for excess accumulated error caused by training-decoding discrepancy. Lower values indicate less exposure bias.
Used to analyze long-form generation.
Open Questions Unanswered questions from this research
- 1 It remains unclear whether reverse KL is consistently superior for mathematics, code, and multilingual reasoning. These settings require controlled comparisons across tasks, languages, teacher quality, and context lengths.
- 2 Mode-seeking behavior may reduce creativity or factual coverage. A principled mechanism is still needed to balance accuracy, diversity, safety, and calibration under a fixed student capacity.
Applications
Immediate Applications
Private enterprise assistants
Organizations can use an open-source white-box teacher, internal instruction data, and MiniLLM to train a smaller model for private servers. Expected benefits include lower memory use and latency while retaining workflow answering and summarization.
Edge-device instruction models
Device makers can distill OPT or LLaMA teachers into hundreds-of-millions-parameter students for offline control and personal assistance. Requirements include task prompts, teacher probability access, and extra training-time sampling compute.
Long-term Vision
Verifiable model compression
Combining MiniLLM with preference feedback, retrieval, and safety constraints could create compression pipelines optimized for factuality, compliance, and low-cost inference, enabling large-scale personalized AI.
Abstract
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. Our code, data, and model checkpoints can be found in https://github.com/microsoft/LMOps/tree/main/minillm.