The idea: a small model learns from a big one
Knowledge distillation trains a small, fast student model to copy a big, slow teacher model that is already trained. The student is not trained only on the right answers (the hard labels). It is trained on the teacher's full list of probabilities for every class (the soft targets).
Why does this help? A label says only "cat". The teacher says more: "cat, but a bit like a dog and a fox, and not at all like a car or a truck". Those small numbers for the wrong classes tell the student which classes look alike. Hinton, Vinyals and Dean (2015) called this dark knowledge.
The page has one photo, a cat, and five classes: cat, dog, fox, car, truck. The teacher is frozen with logits 9.0, 5.5, 4.5, 0.5, 0.0. The student's five logits start at 0, so it gives every class 0.20. Press Train Step to see one step with every hop, or Train 20 Steps to see many.
Temperature
A normal softmax turns logits z into probabilities. With a temperature T, the logits are divided by T first:
q_i = exp(z_i / T) / Σ_j exp(z_j / T)
At T = 1 the teacher is very sure: cat 0.96, dog 0.03, fox 0.01, car 0.0002, truck 0.0001. The useful tail is almost zero. At T = 4 the same logits give cat 0.51, dog 0.21, fox 0.17, car 0.06, truck 0.05. The ranking does not change, but now it is easy to see. A larger T makes the bars flatter (Demo 2).
The loss and its gradient
The student is trained on a mix of two losses. q means softmax at temperature T, p means the normal softmax at T = 1:
L = α · T² · KL(q_t ‖ q_s) + (1 − α) · CE(y, p_s)
- The soft loss KL(q_t ‖ q_s) is the KL divergence from the teacher's soft targets to the student's soft probabilities. It is 0 when they are equal.
- The hard loss CE(y, p_s) = −ln p_s(cat) is the normal cross-entropy with the label.
- α sets the mix. The page uses α = 0.9, so most of the signal comes from the teacher.
The gradient on each student logit is simple:
∂L/∂z_s,i = α · T · (q_s,i − q_t,i) + (1 − α) · (p_s,i − y_i)
If the student gives a class too much, the gradient is positive (red on the canvas) and the logit goes down. If it gives too little, the gradient is negative (green) and the logit goes up. Each step does z_s ← z_s − η · ∂L/∂z_s with η = 0.5 (Demo 1).
Why the T²? Without it, the gradient of KL(q_t ‖ q_s) on the logits shrinks about like 1/T² as T grows (Hinton et al., 2015). Multiplying by T² keeps the soft loss about the same size whatever T is, so changing T does not also change how much the soft loss counts against the hard loss.
Hard labels vs soft targets
With hard labels only (α = 0), the label is 1 for cat and 0 for everything else. The gradient pushes the cat logit up and pushes the four wrong classes down by exactly the same amount. After 40 steps dog, fox, car and truck all have logit −0.89: the student cannot tell that a cat looks more like a dog than a truck (Demo 3).
With distillation at T = 4 the student's logits after 40 steps are 4.87, 1.23, 0.21, −3.01, −3.29. That is the teacher's logits minus a constant, nearly: the student has copied the ranking dog > fox > car > truck (Demo 4). At T = 1 the soft targets are almost the hard label, the tail gradients are tiny, and after 40 steps car and truck are still equal (Demo 5).
| training, 40 steps | student logits (cat, dog, fox, car, truck) | ranking of the wrong classes |
|---|---|---|
| hard label only, α = 0 | 3.57, −0.89, −0.89, −0.89, −0.89 | none ✗ |
| distillation, T = 1, α = 0.9 | 3.25, −0.55, −0.80, −0.95, −0.95 | dog > fox > car ≈ truck |
| distillation, T = 4, α = 0.9 | 4.87, 1.23, 0.21, −3.01, −3.29 | dog > fox > car > truck ✓ |
In a real model the student has shared weights, not one logit per photo. What it learns about "cat looks like dog" from many photos helps it on photos it has never seen. That is why a distilled student usually beats the same student trained on labels alone.
Where it is used
- DistilBERT (Sanh et al., 2019) has 6 layers instead of BERT's 12, about 40% fewer parameters, runs about 60% faster and keeps about 97% of BERT's scores on the GLUE benchmark.
- DeiT trains a vision transformer with a CNN teacher, using an extra distillation token.
- Large language models: small models are trained on a big model's next-token probabilities (logit distillation) or just on text the big model wrote (sequence-level distillation). A distilled draft model is also what Speculative Decoding needs.
Distillation is one of the ways to make a model cheaper, next to quantization, which keeps the model but stores its weights in fewer bits. See Softmax & Cross-Entropy for the softmax and its gradient p − y, Backpropagation for how a gradient on the logits reaches the weights, and How a Large Language Model Works for the next-token probabilities.
What the page leaves out
The student's weights and back-propagation through its layers (the page trains its five output logits directly); many training examples and generalization to new ones; matching hidden layers or attention maps (feature distillation, as in TinyBERT); self-distillation and born-again networks; online distillation where teacher and student train together; and choosing T and α by validation.