PaperlayerView exact revision on arXiv

Source-linked AI summary

MiniLLM: On-Policy Distillation of Large Language Models

Yuxian Gu, Li Dong, Furu Wei, Minlie Huang

arXiv:2306.08543v6cs.CLcs.AI

TL;DR

White-box distillation for generative LLMs is under-explored, motivating a method that transfers teacher distributions into smaller models. MiniLLM replaces forward KLD with reverse KLD and optimizes it on-policy. It consistently improves response quality and precision over KD baselines across instruction-following tasks, scales from 120M to 13B parameters, and shows better calibration, lower exposure bias, and stronger long-text generation, while addressing optimization challenges with stabilizing strategies.

  • Problem

    White-box KD for generative LLMs is under-explored, even though open-source models make teacher output distributions available for reducing computational demand.

  • Method

    MiniLLM minimizes reverse KLD instead of forward KLD and derives an on-policy optimization procedure with variance, reward-hacking, and length-bias controls.

  • Results

    MiniLLM consistently outperforms standard KD baselines across five instruction-following datasets and scales across model families from 120M to 13B parameters.

  • Takeaways & Limitations

    The method produces more precise, higher-quality responses with lower exposure bias, better calibration, and stronger long-response generation than the baselines.

  • Takeaways & Limitations

    On-policy optimization suffers from high variance and reward hacking, while the unnormalized reward favors short responses; the method adds mitigation strategies for these issues.

Abstract

from arXiv · show

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.

1 Introduction

White-box distillation of generative LLMs remains under-explored, despite open-source models making access to teacher distributions increasingly valuable. MiniLLM replaces forward KLD with reverse KLD and uses on-policy optimization, yielding stronger instruction-following results across model scales.

  • White-box KD for LLMs is under-explored, while open-source LLMs increase its practical value for reducing computational demand.
  • Forward KLD is sub-optimal for generative LLMs because limited-capacity students overestimate the teacher’s low-probability regions.The teacher distribution can contain more modes than the student can express, causing unlikely free-run generations.
  • MiniLLM minimizes reverse KLD with an on-policy optimization approach that targets the teacher’s major modes and avoids void regions.The method adds single-step decomposition, teacher-mixed sampling, and length normalization to stabilize and improve optimization.
  • Across five instruction-following datasets, MiniLLM consistently outperforms standard KD baselines and scales from 120M to 13B parameters.It also shows lower exposure bias, better calibration, stronger long-response generation, and negligible diversity loss.

2 Method

MiniLLM distills a teacher’s generative distribution by minimizing reverse KLD through on-policy student sampling. Its training combines variance reduction, teacher-mixed sampling, length normalization, and language-modeling preservation.

  • MINILLM: Knowledge Distillation with Reverse KLD: Reverse KLD makes the student seek major teacher modes rather than fit every sampled sequence or assign mass to void regions.This mode-seeking behavior is intended to improve correctness and faithfulness in text generation.
  • On-Policy Distillation: Policy-gradient optimization estimates the reverse-KLD gradient from responses sampled on-policy from the student or a teacher-student mixture.The teacher distribution scores generated tokens, while the student distribution supplies the generation policy.
  • On-Policy Distillation: Single-step decomposition computes token-level generation quality directly over the vocabulary, reducing estimator variance and accelerating convergence.The decomposition isolates each step’s quality because early-token errors accumulate across the sentence.
  • On-Policy Distillation: Teacher-mixed sampling uses ep(yt|y<t,x) = α · p(yt|y<t,x) + (1 −α) · qθ(yt|y<t,x) to suppress degenerate high-scoring samples and alleviate reward hacking.Importance sampling rewrites the resulting gradients as an unbiased estimator, although multiplying token weights can create high variance.
  • On-Policy Distillation: Length normalization counteracts the tendency of the reverse-KLD reward to favor short responses.The training pipeline also preserves canonical language-modeling performance with an auxiliary pretraining loss.
  • Training Algorithm: Training first fine-tunes the student on task data, then performs clipped on-policy updates combining single-step, length-normalized, and pretraining gradients.The procedure starts from a student pretrained on a long-document corpus and selects the lowest-validation-loss checkpoint for subsequent optimization.

3 Experiments

MiniLLM is evaluated against SFT, KD, and SeqKD across instruction-following datasets, model families, scales, and analyses of quality, calibration, exposure bias, and generation behavior. It generally outperforms baselines, scales from 120M to 13B parameters, and benefits from stabilization strategies that prevent reward hacking and reduce training variance.

  • 3.2 Results: MINILLM outperforms baselines in almost all cases across model families, evaluation sets, and Rouge-L and GPT-4 feedback, with stronger performance beyond Dolly.The comparison covers SFT, KD, and SeqKD on multiple instruction-following datasets.
  • 3.2 Results: MINILLM produces precise responses, sometimes exceeding teacher Rouge-L scores, while on-policy training samples student responses and alleviates exposure bias.The paper attributes this mechanism to reducing the mismatch between teacher-forcing training and free-run generation.
  • 3.2 Results: MINILLM’s improvement remains consistent from 120M to 13B parameters across GPT-2, OPT, and LLaMA model families.The experiments use GPT-2, OPT, and LLaMA students with larger corresponding teachers.
  • 3.2 Results: MINILLM achieves better human preference than all baselines on SelfInst with LLaMA, performing comparably to the teacher.The human evaluation compares responses using Win, Tie, and Loss judgments.
  • 3.3 Analysis: MINILLM performs better than standard KD models across response-length subsets, although all methods score low on prompts requiring five or fewer tokens.The short-response weakness is attributed to a training–evaluation distribution shift toward longer responses.
  • 3.4 Ablation Studies on Optimization Strategies: Teacher-mixed sampling and length normalization stabilize training, while single-step decomposition reduces variance; removing these strategies can cause reward hacking and poor generation.Without the stabilizing strategies, students may generate repeated, short, or meaningless strings with high teacher scores.

4 Related Work

Related work frames knowledge distillation as teacher-guided student training, with established applications in classification and emerging use in text generation. Standard generative KD approximately minimizes forward KLD, whereas this paper studies reverse KLD for white-box LLM distillation.

  • Knowledge distillation trains a small student model with guidance from a larger teacher model [HVD15].
  • Text-classification KD mimics teacher outputs, hidden states, or attention scores, while generative KD uses token-level teacher supervision or teacher-generated texts.
  • Standard generative KD approximately minimizes forward KLD, which can cover all target modes while sacrificing accuracy on major modes [Hus15].
  • This paper instead minimizes reverse KLD when the teacher distribution is available, arguing that it better suits LLM generation.

5 Conclusion

The paper introduces MiniLLM for distilling white-box LLMs into smaller language models through reverse-KLD optimization. Experiments report more precise, higher-quality responses alongside improved calibration, reduced exposure bias, stronger long-text generation, and good diversity.

  • MiniLLM minimizes reverse KLD because forward KLD overestimates low-probability teacher regions when students cannot represent the teacher distribution.
  • MiniLLM uses an optimization algorithm designed to train small language models against the teacher distribution.
  • MiniLLM produces more precise, higher-quality responses than standard KD models, with lower exposure bias, better calibration, stronger long-text generation, and good diversity.

A Derivations

The derivations recast reverse-KLD distillation as on-policy reinforcement learning, where the teacher supplies a reward and the student is optimized as a policy. They connect teacher logits to a soft Bellman formulation and derive policy-gradient estimators.

  • A Perspective of MINILLM from Inverse Reinforcement Learning: The derivation treats token generation as a Markov Decision Process and casts the teacher-guided student objective as reinforcement learning.
  • A Perspective of MINILLM from Inverse Reinforcement Learning: The student is trained to maximize a reward induced from the teacher model rather than merely clone teacher behavior.
  • A Perspective of MINILLM from Inverse Reinforcement Learning: The teacher output logits parameterize the Q-function with γ = 1, yielding the token reward used by the optimization.
  • A Perspective of MINILLM from Inverse Reinforcement Learning: Maximum-entropy reinforcement learning adds the student distribution entropy to the reward-maximization objective.
  • Derivation of Equation 3: The derivation establishes approximate objective equivalence and computes reverse-KLD gradients with the Policy Gradient Theorem using future rewards.

B.1 Training Details

Training details define the baselines and describe MiniLLM’s two-phase procedure. Students are first supervised-fine-tuned, then continuously trained with the MiniLLM algorithm using selected baseline hyperparameters.

  • Training Details: Baselines include SFT without KD, KD, and SeqKD, with learning rates, batch sizes, and training epochs searched separately by model size.
  • MINILLM: MiniLLM training has two phases, analogous to RLHF: supervised fine-tuning followed by continued MiniLLM optimization.
  • MINILLM: Phase 1 fine-tunes on instruction-response data for three epochs and selects the lowest-validation-loss checkpoint before Phase 2 training.

B.2 Automatic Evaluation Details

Automatic evaluation samples model responses with fixed decoding settings and evaluates GPT-4 feedback using a separate prompt wrapper.

  • Responses are sampled at temperature 1 with a maximum length of 512 across random seeds 10, 20, 30, 40, and 50.
  • The evaluation prompt wrapper converts each instruction-response pair into a sentence for model evaluation.
  • GPT-4 feedback uses a dedicated prompt with temperature 0.7.

B.3 Human Evaluation Details

Human evaluation compares anonymized model responses using randomly sampled prompts and blinded annotator preferences.

  • Human evaluation randomly samples 50 prompts because additional prompts were found not to materially affect the results.
  • Annotators compare baseline responses with MINILLM responses and select a preferred response or indicate no significant difference.
  • The evaluation interface hides the source model from annotators and is shown in Figure 13.

B.4 Details About Generation Diversity Metrics

Generation diversity is evaluated with distinct 4-grams and language-modeling loss, using repeated-seed measurements and teacher-scaling analysis.

  • Distinct-4 measures generation diversity as the fraction of distinct 4-grams among all generated 4-grams.Table 3 reports averages across five random seeds, while Table 5 shows the seed-wise values on Dolly.
  • Loss is the test-set negative log-likelihood and measures coverage of the real-data distribution through forward KLD.
  • Teacher-scaling analysis compares MINILLM and SeqKD with OPT-1.3M students and OPT 2.7B, 6.7B, and 13B teachers.

B.5 Exposure Bias Analysis

The exposure-bias analysis defines a relative error measure and examines performance across teacher sizes, response lengths, and teacher-mixing settings.

  • Exposure-bias error measures the relative generation error caused by model-generated prefixes after separating oracle-context estimation error.The regret consists of oracle-context error and prefix-induced error, with the latter reflecting exposure bias.
  • MINILLM outperforms baselines in most GPT-J teacher evaluations across GPT-2 and GPT-Neo student sizes.
  • MINILLM consistently performs better as teacher models scale in the OPT-family comparison.
  • On short U-NI responses, KD methods perform similarly, whereas MINILLM outperforms other methods on long responses.

C.4 More Ablation Studies

The ablations examine teacher-mix-in sampling, pre-training loss, and qualitative instruction-following behavior. Together, they show that teacher mixing supports detailed, accurate outputs, while pre-training loss preserves canonical NLP abilities without materially changing instruction-following performance.

  • Effect of Teacher-Mix-in Strength α: α = 0.2 is generally suitable across model families and sizes, while larger models are more robust to teacher-mix-in strength.α = 0.0 samples only from the student, whereas α = 1.0 samples entirely from the teacher.
  • Qualitative Analysis: Without teacher-mixed sampling, distilled models can produce short responses or simply repeat the input.These failure patterns appear in Case #1 and Cases #2, respectively.
  • Effect of Adding Pre-Training Loss: Adding the pre-training loss preserves abilities on canonical NLP tasks while keeping instruction-following performance nearly unchanged.This comparison removes the language-modeling loss on the pre-training corpus in the w/o LPT variant.
  • Qualitative Analysis: MINILLM generates more detailed and accurate responses than baselines on instruction-following cases from the SelfInst dataset.The qualitative cases use prompts sampled from SelfInst and are drawn from the LLaMA model family.
Loading 2306.08543v6…