Knowledge Distillation for Large Language Models: From Soft Targets to On-Policy Reverse KL

Knowledge Distillation for Large Language Models: From Soft Targets to On-Policy Reverse KL Knowledge distillation (KD) has become the primary mechanism for transferring capabilities from massive proprietary models to smaller, deployable open-weight models. The technique originated in classification, but applying it to auto-regressive language models exposed fundamental mismatches: token-level forward KL forces students to cover the teacher's full output distribution, while supervised training

7 min
Knowledge Distillation for Large Language Models

Knowledge Distillation for Large Language Models: From Soft Targets to On-Policy Reverse KL

Knowledge distillation (KD) has become the primary mechanism for transferring capabilities from massive proprietary models to smaller, deployable open-weight models. The technique originated in classification, but applying it to auto-regressive language models exposed fundamental mismatches: token-level forward KL forces students to cover the teacher's full output distribution, while supervised training on fixed teacher-generated data creates train-inference distribution shift. Modern LLM distillation — MiniLLM, GKD, and the DeepSeek-R1 distilled series — solves these problems by switching to reverse KL, sampling from the student during training, and unifying distillation with reinforcement learning.

The Classical Foundation: Temperature-Scaled Soft Targets

Hinton, Vinyals, and Dean (2015) introduced distillation as a way to compress an ensemble or large model into a smaller deployable one [^hinton2015]. The key insight: a trained model assigns probabilities to incorrect classes, and the relative magnitudes of these "dark knowledge" probabilities encode a rich similarity structure. By raising the softmax temperature T>1T > 1, the teacher produces a softened distribution qi=exp(zi/T)jexp(zj/T)q_i = \frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)} that reveals this structure. The student is trained to match these soft targets using the same high temperature, then runs at T=1T=1 for inference.

For classification, the loss is cross-entropy between teacher and student distributions at temperature TT, scaled by T2T^2 to preserve gradient magnitudes when combining with hard labels. Matching logits directly is a special case as TT \to \infty.

DistilBERT (Sanh et al., 2019) showed this works for BERT-style encoders, halving parameters while retaining 97% of performance [^sanh2019]. But auto-regressive generation changes the game.

Sequence-Level Distillation: The First Adaptation to Generation

Kim and Rush (2016) observed that NMT models predict sequences, not independent tokens [^kimrush2016]. Standard word-level KD minimizes token-wise cross-entropy:

LWORD-KD=j=1Jk=1Vq(tj=ks,t<j)logp(tj=ks,t<j)\mathcal{L}_{\text{WORD-KD}} = -\sum_{j=1}^J \sum_{k=1}^{|V|} q(t_j=k | \mathbf{s}, \mathbf{t}_{<j}) \log p(t_j=k | \mathbf{s}, \mathbf{t}_{<j})

They proposed two sequence-level variants. SeqKD generates a new training dataset by running beam search with the teacher, then trains the student on these teacher-generated sequences with standard MLE. Sequence-level KD minimizes the KL between the full sequence distributions q(ts)q(\mathbf{t}|\mathbf{s}) and p(ts)p(\mathbf{t}|\mathbf{s}), approximated by sampling from the teacher. Both eliminated the need for beam search at inference — the student could use greedy decoding and match the teacher's beam-search quality, running 10× faster.

This established a pattern: for generation, distillation often means data generation from the teacher, not just logit matching.

The Core Problem: Forward KL and Exposure Bias

When the teacher is a large LLM and the student is much smaller, two problems emerge.

Model under-specification. Forward KL KL(pteacherqstudent)\text{KL}(p_{\text{teacher}} \| q_{\text{student}}) is mean-seeking: it forces the student to put mass wherever the teacher has mass. A smaller student cannot cover all modes of a larger teacher's distribution. It spreads probability thinly over low-probability regions, producing degenerate samples [^huszar2015]. Reverse KL KL(qstudentpteacher)\text{KL}(q_{\text{student}} \| p_{\text{teacher}}) is mode-seeking: it forces the student to concentrate on high-probability teacher modes, avoiding hallucinated low-probability tokens.

Distribution mismatch (exposure bias). Supervised KD trains on a fixed dataset of teacher-generated (or ground-truth) sequences. At inference, the student generates auto-regressively, conditioning on its own previous tokens — which may drift from the teacher's distribution. The student never sees its own mistakes during training. This is the classic imitation learning problem identified by Ross et al. (2011) [^ross2011].

MiniLLM: Reverse KL via Policy Gradients

Gu et al. (2023) reformulated LLM distillation as an RL problem [^gu2023]. The reverse KL objective at the sequence level:

LrevKL=Eyqθ(x)[logqθ(yx)p(yx)]\mathcal{L}_{\text{revKL}} = \mathbb{E}_{y \sim q_\theta(\cdot|x)} \left[ \log \frac{q_\theta(y|x)}{p(y|x)} \right]

is equivalent to maximizing expected reward r(y)=logp(yx)logqθ(yx)r(y) = \log p(y|x) - \log q_\theta(y|x) under the student policy. They derive a token-level policy gradient with three critical stabilizers:

  1. Teacher-mixed sampling: instead of pure on-policy rollouts, mix α=0.2\alpha=0.2 teacher-generated tokens into student sequences to reduce variance and prevent reward hacking.
  2. Length normalization: normalize rewards by sequence length to remove generation-length bias.
  3. PPO-style clipping: clip the importance ratio ρt(θ)=qθ(yty<t,x)p~(yty<t,x)\rho_t(\theta) = \frac{q_\theta(y_t|y_{<t},x)}{\tilde{p}(y_t|y_{<t},x)} to [1ϵ,1+ϵ][1-\epsilon, 1+\epsilon] for stable updates.

The combined gradient includes a pre-training loss term to prevent catastrophic forgetting. Across GPT-2, OPT, and LLaMA families (120M–13B), MiniLLM outperforms SeqKD and standard KD on instruction-following benchmarks, with lower exposure bias, better calibration, and higher long-text quality. In some settings the student exceeds the teacher's Rouge-L scores, attributed to reduced exposure bias from on-policy training.

GKD: Unifying Supervised, On-Policy, and Divergence Choice

Agarwal et al. (2023) generalized the framework: distillation is minimizing a token-level divergence D\mathcal{D} between teacher and student distributions on a mixture of fixed data and student-generated data [^agarwal2023].

LGKD=(1β)E(x,y)Dfixed[D(pθ(x,y<t)p(x,y<t))]+βExDx,yqθ(x)[D(pθ(x,y<t)p(x,y<t))]\mathcal{L}_{\text{GKD}} = (1-\beta) \mathbb{E}_{(x,y) \sim \mathcal{D}_{\text{fixed}}} \left[ \mathcal{D}(p_\theta(\cdot|x,y_{<t}) \| p(\cdot|x,y_{<t})) \right] + \beta \mathbb{E}_{x \sim \mathcal{D}_x, y \sim q_\theta(\cdot|x)} \left[ \mathcal{D}(p_\theta(\cdot|x,y_{<t}) \| p(\cdot|x,y_{<t})) \right]

where β[0,1]\beta \in [0,1] controls the student-data fraction. Special cases:

  • β=0\beta=0, forward KL → supervised KD
  • β=1\beta=1, forward KL → on-policy KD (ImitKD)
  • β=1\beta=1, reverse KL → MiniLLM-style (without policy gradient tricks)
  • β(0,1)\beta \in (0,1), JSD → mixed supervised/on-policy with bounded divergence

GKD shows that mode-seeking divergences (reverse KL, JSD with high α\alpha) consistently outperform forward KL on summarization (XSum), translation (WMT), and reasoning (GSM8K). On-policy and mixed data fractions beat purely supervised data. Crucially, GKD can combine distillation with RL fine-tuning in a single objective: L=LRL+λLGKD\mathcal{L} = \mathcal{L}_{\text{RL}} + \lambda \mathcal{L}_{\text{GKD}}, regularizing toward the teacher instead of the initial policy.

DeepSeek-R1 Distillation: Reasoning Transfer at Scale

DeepSeek-R1 (2025) demonstrated that distillation can transfer reasoning capability, not just style or knowledge [^deepseek2025]. They generated 800K training samples from DeepSeek-R1 (671B MoE) — including chain-of-thought traces — and fine-tuned smaller dense models (Qwen-2.5 1.5B–32B, Llama-3.1 8B–70B) via SFT alone. No RL, no reverse KL, no on-policy sampling.

Results: DeepSeek-R1-Distill-Qwen-1.5B outperforms GPT-4o and Claude-3.5-Sonnet on AIME (28.9%) and MATH (83.9%). The 7B, 14B, and 32B variants achieve 55–73% on AIME. The Llama-70B distilled model reaches 70% AIME. This shows that high-quality teacher-generated reasoning traces, at sufficient scale, can transfer complex reasoning to much smaller models through standard SFT — a simpler recipe than policy-gradient methods, but dependent on an exceptionally strong teacher.

Taxonomy of Modern LLM Distillation

| Method | Divergence | Data Source | Optimization | Key Insight | |--------|-----------|-------------|--------------|-------------| | Standard KD | Forward KL (token) | Fixed teacher data | Supervised MLE | Dark knowledge via temperature | | SeqKD | Sequence MLE | Teacher beam-search data | Supervised MLE | Greedy student matches teacher beam | | MiniLLM | Reverse KL (sequence) | On-policy + teacher-mixed | Policy gradient + PPO | Mode-seeking avoids low-prob mass; on-policy fixes exposure bias | | GKD | Forward KL / Reverse KL / JSD | Mixed (tunable β\beta) | Supervised (no backprop through sampling) | Unified framework; combines with RL | | DeepSeek-R1 Distill | N/A (SFT) | 800K teacher CoT traces | Supervised SFT | Reasoning traces transfer via massive SFT |

Distillation approach comparison

Open Questions and Practical Guidance

When to use which?

  • Black-box teacher (API only): SeqKD or standard KD on teacher-generated data.
  • White-box teacher, moderate student: GKD with reverse KL or JSD(α0.9\alpha \approx 0.9), β=0.5\beta=0.5 mixed data. Simpler than MiniLLM, no policy gradient tricks.
  • White-box teacher, very small student: MiniLLM if you can implement policy gradients; otherwise GKD reverse KL.
  • Reasoning transfer, strong teacher available: Generate massive CoT dataset, run SFT (DeepSeek-R1 recipe).

Temperature surprises. GKD finds teacher temperature <1< 1 can help — contrary to classical wisdom. Lower temperature sharpens the teacher distribution, making mode-seeking easier for the student.

Evaluation gap. Standard perplexity and Rouge/BLEU miss exposure bias and calibration. MiniLLM introduces ExAccErr (excess accumulated error) to measure exposure bias directly, and ECE (expected calibration error) on classification probes.

Conclusion

LLM distillation has evolved from temperature-scaled logit matching (Hinton et al., 2015) through sequence-level data generation (Kim & Rush, 2016) to modern frameworks that treat distillation as imitation learning with mode-seeking divergences and on-policy sampling (MiniLLM, GKD). The DeepSeek-R1 result adds a pragmatic note: at sufficient scale, plain SFT on teacher-generated reasoning traces can outperform sophisticated objectives. The choice today depends on teacher access (white-box vs. black-box), student capacity, and whether reasoning transfer is the goal. For most practitioners with white-box teachers, GKD with reverse KL or JSD on mixed data offers the best trade-off of simplicity, stability, and performance.


Sources

[^hinton2015]: Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531. https://arxiv.org/abs/1503.02531 [^sanh2019]: Sanh, V., Debut, L., Chaumond, J., & Wolf, T. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv:1910.01108. https://arxiv.org/abs/1910.01108 [^kimrush2016]: Kim, Y., & Rush, A. M. (2016). Sequence-Level Knowledge Distillation. EMNLP 2016. arXiv:1606.07947. https://arxiv.org/abs/1606.07947 [^huszar2015]: Huszár, F. (2015). How (not) to train your generative model: Scheduled sampling, likelihood, adversary? arXiv:1511.05101. https://arxiv.org/abs/1511.05101 [^ross2011]: Ross, S., Gordon, G., & Bagnell, D. (2011). A reduction of imitation learning and structured prediction to no-regret online learning. AISTATS 2011. https://arxiv.org/abs/1011.0686 [^gu2023]: Gu, Y., Dong, L., Wei, F., & Huang, M. (2023). MiniLLM: Knowledge Distillation of Large Language Models. ICLR 2024. arXiv:2306.08543. https://arxiv.org/abs/2306.08543 [^agarwal2023]: Agarwal, R., Vieillard, N., Stanczyk, P., Ramos, S., Geist, M., & Bachem, O. (2023). GKD: Generalized Knowledge Distillation for Auto-regressive Sequence Models. ICLR 2024. arXiv:2306.13649. https://arxiv.org/abs/2306.13649 [^deepseek2025]: DeepSeek-AI et al. (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948. https://arxiv.org/abs/2501.12948 [^xu2024]: Xu, X., Li, M., Tao, C., Shen, T., Cheng, R., Li, J., Xu, C., Tao, D., & Zhou, T. (2024). A Survey on Knowledge Distillation of Large Language Models. arXiv:2402.13116. https://arxiv.org/abs/2402.13116

Written by

More to read

  • Activation Checkpointing in Large Language Models: How Selective Recomputation Eliminates Memory Bottlenecks

    Large language model pre-training and fine-tuning are fundamentally constrained by GPU memory (VRAM). While distributed techniques such as Fully Sharded Data Parallel (FSDP), ZeRO, and Tensor Parallelism successfully shard model parameters, optimizer states, and gradients across hundreds or thousands of GPUs, activation memory presents a distinct scaling bottleneck. During the forward pass of a transformer model, intermediate tensor outputs must be preserved in GPU memory so that backpropagatio

    1 min
  • Anthropic Demonstrates Autonomous De Novo Protein Design and Chemical Analysis with Claude

    Anthropic Demonstrates Autonomous De Novo Protein Design and Chemical Analysis with Claude Anthropic has published experimental results demonstrating Claude's ability to autonomously design de novo protein binders with physical wet-lab validation and automate complex analytical chemistry workflows. The findings show frontier LLMs acting as autonomous agents across computational biology and molecular characterization pipelines. In the primary experiment, Anthropic evaluated Claude Mythos Previe

    1 min
  • Cerebras Unveils CS-4 Rack-Scale System Powered by Three WSE-3 Turbo Chips and Nexus Architecture

    Cerebras Unveils CS-4 Rack-Scale System Powered by Three WSE-3 Turbo Chips and Nexus Architecture Cerebras Systems has announced the CS-4, a rack-scale AI accelerator system designed around three of its next-generation Wafer Scale Engine 3 Turbo (WSE-3 Turbo) chips and a modular hardware architecture dubbed Nexus. Cerebras confirmed that initial customer shipments for the CS-4 are scheduled to begin in the current quarter. The new system marks a structural shift from Cerebras's single-wafer CS

    1 min