Knowledge distillation

Problem

The model that wins on accuracy is rarely the one you can afford to serve. Ensembles, very deep networks and large language models cost memory, latency and money on every request, while the device or the budget only allows something several times smaller. Training the small model from scratch on the original labels usually leaves a gap.

The large model already knows more than the labels say. For a photo of a cat, the label is just "cat", but a good classifier also says it looks a little like a dog and almost nothing like a car. Distillation asks how to move that knowledge into the small model: train it to match the teacher's outputs, not only the one-hot targets. Bucilă, Caruana and Niculescu-Mizil did this in 2006 to compress large ensembles into small neural nets; Hinton, Vinyals and Dean named it distillation in 2015 and built it around a softmax temperature, which this page is mostly about.

Intuition

A one-hot label carries at most log₂ K bits per example and says nothing about which wrong answers are nearly right. The teacher's probabilities say how similar this input is to every class, which is information about the shape of the function the teacher learned. A student that matches those probabilities on an input learns, from that one input, roughly what the teacher would have learned from many nearby labelled examples.

The catch is that a confident teacher hides most of this in tiny numbers: 0.0003 for "car" looks the same as 0 to a cross-entropy loss. Dividing the logits by a temperature T > 1 before the softmax spreads the probability out, so the ratios between the wrong classes, the "dark knowledge", become large enough to matter. The ranking never changes; only how loudly the teacher states it.

Mechanics

Temperature softmax

Given logits z, the teacher's softened distribution and the student's (logits v) are

p_i(T) = exp(z_i/T) / Σ_j exp(z_j/T),   q_i(T) = exp(v_i/T) / Σ_j exp(v_j/T)

T = 1 is the ordinary softmax; T → ∞ tends to the uniform distribution and T → 0 to the one-hot argmax. Dividing by T does not change which logit is largest, so the teacher's prediction is the same at every temperature.

The distillation loss

The student is trained on a weighted sum of the usual cross-entropy with the label y, at T = 1, and a KL divergence between the softened distributions:

L = α · CE(y, q(1)) + (1 − α) · T² · KL(p(T) ‖ q(T))

The teacher is frozen; only the student's logits v receive gradients. KL(p ‖ q) and the cross-entropy H(p, q) differ by H(p), which does not depend on the student, so either can be used; KL is zero when the student matches, which makes it easier to read. At test time the student uses its ordinary softmax, T = 1.

Why multiply by T²

The gradient of the KL term with respect to a student logit is short:

∂ KL(p(T) ‖ q(T)) / ∂v_i = (q_i(T) − p_i(T)) / T

One factor 1/T comes from the chain rule through v/T. The other comes from the softening itself: as T grows both distributions approach uniform, and for logits with mean zero q_i − p_i ≈ (v_i − z_i) / (N T) over N classes. Together the gradient falls like 1/T², so without correction raising T would quietly turn the soft term off. Multiplying the loss by T², as Hinton, Vinyals and Dean recommend, keeps the soft and hard terms in roughly the same proportion whatever T is.

The same approximation shows what very high temperatures do: T² times the gradient tends to (v_i − z_i)/N, the gradient of matching the logits themselves with a squared error. Ba and Caruana trained shallow students exactly that way, regressing the teacher's logits with an L2 loss, so logit matching is the high-temperature limit of distillation.

Matching more than outputs

The teacher's hidden layers are knowledge too. FitNets (Romero et al.) trains a deeper, thinner student using the teacher's intermediate representations as hints: because the student's hidden layer is narrower, a small regressor maps it onto the teacher's, and the two are matched with a squared error. DistilBERT combines three terms during pretraining: masked language modelling, distillation on the output distribution, and a cosine loss aligning the student's and teacher's hidden states. Any such feature loss needs a projection when widths differ, and a choice of which layers to pair.

Distilling language models

A language model's output at each position is already a distribution over the vocabulary, so token-level (word-level) distillation is the loss above summed over positions, with teacher-forced inputs from a fixed dataset. It requires teacher and student to share a tokenizer, otherwise the two distributions are over different sets.

Sequence-level distillation (Kim and Rush) instead runs the teacher's beam search and trains the student with ordinary cross-entropy on the teacher's output sequences, approximating the teacher's distribution over whole sequences by its mode. Fine-tuning a small model on text a larger model generated is this method, and it works across tokenizers.

Both train the student on sequences it did not produce, while at inference it conditions on its own earlier tokens; small errors lead it into prefixes it never saw in training. On-policy distillation samples sequences from the student and has the teacher score every token of them: GKD (Agarwal et al.) frames this as training on self-generated mistakes, and MiniLLM (Gu et al.) minimizes the reverse KL, KL(q ‖ p), with an on-policy optimization.

The direction of the KL matters when the student is too small to match everywhere. Forward KL(p ‖ q), the classical choice, is large wherever the teacher has mass and the student does not, so the student spreads to cover every mode, including putting mass between them. Reverse KL(q ‖ p) is large wherever the student puts mass the teacher does not, so the student is pushed onto the teacher's main modes and may ignore small ones: mode-seeking. MiniLLM uses it to stop the student overestimating regions where the teacher assigns low probability.

The pieces in PyTorch

PyTorch has no distillation module; it is built from log_softmax and KLDivLoss. KLDivLoss takes the student's log-probabilities as input and the teacher's probabilities as target (or log-probabilities with log_target=True). Its default reduction "mean" averages over every element, which is not the KL; kl_div warns about it and recommends "batchmean", which sums and divides by the batch size. CrossEntropyLoss also accepts class probabilities as the target, which gives H(p, q): the same gradient as KL.

Worked example

Implementation

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.5):
    hard = F.cross_entropy(student_logits, labels)
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.log_softmax(teacher_logits / T, dim=-1),
        reduction="batchmean",          # "mean" would also divide by the class count
        log_target=True,
    )
    return alpha * hard + (1 - alpha) * T * T * soft

teacher.eval()
for x, y in loader:
    with torch.no_grad():               # the teacher is never updated
        teacher_logits = teacher(x)
    loss = distillation_loss(student(x), teacher_logits, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Token-level distillation for a language model: flatten (batch, seq, vocab)
# to (batch * seq, vocab) and drop padded positions before calling the same loss.

This is the formula the worked example evaluates in plain Python. Passing log-probabilities with log_target=True avoids taking the log of a teacher probability that underflowed to zero. When the teacher is expensive, its logits are often computed once and stored; for large vocabularies only the top few per position are kept, which truncates the soft target and should be renormalized.

Tradeoffs

• Teacher cost. Every training example needs a teacher forward pass, either online or precomputed and stored. Precomputed logits fix any data augmentation in place; online teachers cost compute at every step.
• Temperature and α. Low T ignores dark knowledge, high T approaches logit matching and asks the student to reproduce scores that do not affect predictions. The worked example's best setting was the lowest T; on other problems it differs.
• Fidelity is not accuracy. Stanton et al. find that students often fail to match their teachers' predictive distributions even when they have the capacity to, largely because of optimization, and that closer matching does not always improve the student's generalization.
• The transfer set. Soft targets can label unlabeled or synthetic inputs, as Bucilă et al. did by generating extra points; the student only learns the teacher's function where the transfer set has inputs.
• Labels still matter. The hard term corrects the teacher where it is wrong; with α = 0 the student inherits every teacher mistake and bias.
• Language models. Token-level KD needs a shared tokenizer and full (or top-k) distributions; sequence-level KD needs only generated text but transfers less per example; on-policy methods cost a student sampling step and a teacher scoring pass per update.
• Alternatives. Quantization and pruning shrink the same model instead of training a new one, and they combine with distillation.

Connections

References

• Bucilă, Caruana and Niculescu-Mizil. Model Compression. KDD 2006.
• Ba and Caruana. Do Deep Nets Really Need to be Deep? NIPS 2014. arXiv:1312.6184
• Hinton, Vinyals and Dean. Distilling the Knowledge in a Neural Network. NIPS 2014 Deep Learning Workshop. arXiv:1503.02531
• Romero, Ballas, Kahou, Chassang, Gatta and Bengio. FitNets: Hints for Thin Deep Nets. ICLR 2015. arXiv:1412.6550
• Kim and Rush. Sequence-Level Knowledge Distillation. EMNLP 2016. arXiv:1606.07947
• Sanh, Debut, Chaumond and Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. 2019. arXiv:1910.01108
• Stanton, Izmailov, Kirichenko, Alemi and Wilson. Does Knowledge Distillation Really Work? NeurIPS 2021. arXiv:2106.05945
• Gu, Dong, Wei and Huang. MiniLLM: Knowledge Distillation of Large Language Models. 2023 (later versions retitled MiniLLM: On-Policy Distillation of Large Language Models). arXiv:2306.08543
• Agarwal, Vieillard, Zhou, Stanczyk, Ramos, Geist and Bachem. On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes. ICLR 2024. arXiv:2306.13649
• PyTorch v2.15.0: torch/nn/modules/loss.py, torch/nn/functional.py.