Knowledge Distillation: Learning from Teachers

Why It Matters for Modern Models

Knowledge distillation trains a smaller student model to match a larger teacher’s outputs. Sanh et al. (2019) reported that DistilBERT was 40% smaller and 60% faster than BERT while retaining 97% of its performance on their language-understanding evaluations.
A teacher's soft outputs carry information that hard labels lack; when the teacher is confident, much of what it has learned lies in the ratios of the very small probabilities it gives to wrong classes (Hinton et al., 2015).
It is used to train small language models: Gemma 2's 2- and 9-billion-parameter models were trained to match a larger teacher model's next-token probabilities instead of only the observed next token (Gemma Team, 2024), and sequence-level distillation trains a student on the teacher's generated outputs (Kim & Rush, 2016).

Questions to Think About

Why does raising the temperature reveal the teacher's small probabilities, and why does matching soft targets at very high temperature become the same as matching the logits (Hinton et al., 2015)?
Why can a student with a different architecture, which could never copy the teacher's weights, still learn from the teacher's outputs?

Interactive Demo

The interactive demo for this concept is in the newer notebook.

Open the notebook demo

Core Math

If you prefer intuition first, try the demo in the newer notebook, then come back for the math.

Key Equation

\mathcal{L} = (1-\alpha)\, \mathcal{L}_{\text{CE}}(p^{S}, y) + \alpha\, \tau^2\, \mathrm{KL}\big(p^{T,\tau} \,\|\, p^{S,\tau}\big)

A small student network $S$ is trained to match the output distribution of a larger teacher $T$ . Both distributions are softened with a temperature $\tau$ :

p_i^{T,\tau} = \frac{\exp(z_i^T / \tau)}{\sum_j \exp(z_j^T / \tau)}

and $p^{S,\tau}$ is defined the same way from the student's logits $z^S$ . The loss mixes the ordinary cross-entropy with the true label $y$ , computed at $\tau = 1$ , and a soft-target term:

\mathcal{L} = (1-\alpha)\, \mathcal{L}_{\text{CE}}\big(p^{S,1}, y\big) + \alpha\, \tau^2\, \mathrm{KL}\big(p^{T,\tau} \,\|\, p^{S,\tau}\big)

The gradients of the soft-target term shrink as $1/\tau^2$ , so the factor $\tau^2$ keeps the balance between the two terms roughly fixed when $\tau$ changes (Hinton et al., 2015).

Raising $\tau$ enlarges the teacher's small probabilities, so the student learns which wrong classes the teacher considers similar to the right one. At very high $\tau$ , with the logits of each example centred to mean zero, the gradient of the soft-target term becomes proportional to that of matching the logits directly, $\tfrac12 \sum_i (z_i^S - z_i^T)^2$ (Hinton et al., 2015, §2.1).