Chapter 67 — Knowledge distillation: teaching a small model to imitate a big one

In one sentence

Distillation trains a small student network to reproduce a large teacher network's outputs, which transfers far more of the teacher's learned behaviour into a small model than training that same small model on raw labels ever could.

The problem

A big, accurate model is often too expensive to run everywhere you want it: too slow for a browser, too large for a phone, too costly per request at scale (Chapter 39's cost model). Making a model smaller by just training a smaller architecture on the original labels usually underperforms badly — a tiny network has limited capacity and struggles to independently rediscover everything the big one learned. Distillation is the fix: don't train the small model on the same raw labels — train it to copy the big model.

The idea, in plain language

A trained classifier doesn't just output "cat" — internally, before the final decision, it has a distribution over all classes: cat: 0.82, dog: 0.11, fox: 0.05, car: 0.0001, .... That distribution carries information a single hard label throws away: the model is saying "this looks quite a lot like a dog too, and a little like a fox, but nothing like a car." Hinton et al. called this dark knowledge — the relative sizes of the wrong answers' probabilities encode which mistakes are reasonable, and that signal is exactly what's missing when a student trains on hard labels alone.

Distillation trains the student to match the teacher's full output distribution, not just its final answer.

How it actually works: the classic KD loss

For classification, "matching the distribution" is made precise with a temperature-scaled softmax (Chapter 4's softmax, with one extra knob):

softmax_T(z)_i = exp(z_i / T) / sum_j exp(z_j / T)

At T = 1 this is ordinary softmax. As T increases, the distribution softens — a confident [10, 1, 1] logit vector becomes something closer to uniform, which is exactly what exposes the dark knowledge: at T = 1 the wrong-answer probabilities are so close to zero they carry almost no gradient signal; at T = 3–5 they become large enough for the student to actually learn from the relative ranking of wrong answers.

The full training loss blends two terms:

L = alpha * CE(y_true, student_logits)
    + (1 - alpha) * T^2 * KL( softmax_T(teacher_logits) || softmax_T(student_logits) )

CE(y_true, student_logits) — ordinary cross-entropy against the real label (Chapter 9's loss), so the student doesn't drift from ground truth.
KL(...) — how different the student's softened distribution is from the teacher's; this is the distillation term, and it's what carries the dark knowledge.
T^2 — a scaling correction, because the gradient of the KL term naturally shrinks by a factor of ~1/T^2 as temperature rises; without it, raising T would silently weaken the distillation signal.
alpha — how much to trust ground truth vs. the teacher; typical values are 0.1–0.5, i.e. mostly trust the teacher.

Regression distillation (what we'll actually use in Chapter 68)

Classification needs the temperature trick because softmax outputs are already squashed into (0, 1) and the useful signal (dark knowledge) hides in tiny probabilities. A model that predicts continuous numbers — a temperature, a price, a coordinate — has no such squashing to undo. Its output already is the full, information-rich signal. So regression distillation drops the temperature and the KL divergence entirely, and just trains the student on the teacher's raw output with an ordinary loss:

L = MSE( student(x), teacher(x) )

No true labels are needed at this stage at all — you can run the teacher on any input, including inputs that were never in the original training set, and use its output as a free label. This is the version of distillation Chapter 68's browser project uses.

The three families of distillation

Kind	What the student matches	Example
Response-based	The teacher's final output (logits, or predictions)	What Hinton's 2015 paper does; what Chapter 68 does
Feature-based	The teacher's intermediate hidden layers/attention maps	TinyBERT, MiniLM — the student is trained to reproduce internal representations, not just the final answer
Relation-based	The relationships between the teacher's representations of different examples	Used when the pattern that matters is "how examples relate to each other," e.g. in some embedding-model distillation

Response-based is simplest to implement and is where you should start; the other two exist for cases where matching only the final answer leaves too much of the teacher's structure behind.

Real examples worth knowing

DistilBERT (2019) — 40% fewer parameters than BERT-base, 60% faster, retains ~97% of BERT's language-understanding benchmark performance. Trained with a combination of response-based distillation and matching BERT's embedding layer.
TinyBERT and MiniLM — go further with feature-based distillation, matching intermediate attention patterns, not just the final layer.
DistilGPT2 — the same idea applied to a generative model.
Many of today's small "edge" LLMs (in the 1B–4B range, e.g. the small variants in the Llama, Gemma, and Phi families) are trained partly by distilling from a larger sibling model in the same family — this is why a 3B model released alongside a 70B model from the same lab often outperforms a 3B model trained independently from scratch.

Worked example: dark knowledge, by hand

A teacher classifying an image into {cat, dog, car} outputs logits [4.0, 2.0, -3.0]. Compare hard label vs. soft label:

Hard label (one-hot):     cat=1.0,   dog=0.0,   car=0.0
Softmax at T=1:           cat=0.87,  dog=0.12,  car=0.0003
Softmax at T=4:           cat=0.51,  dog=0.33,  car=0.16

The hard label says nothing about dog vs. car. T=1 softmax already leans toward "somewhat dog-like, not at all car-like" but the gap is small enough to be a weak training signal. T=4 makes that relationship large and learnable: the student is explicitly told "this is much more dog-like than car-like," which is exactly the kind of fine-grained similarity structure a tiny model benefits most from being handed directly, instead of having to rediscover it from millions of raw examples.

Practice (20 min) — feel the effect of temperature

import math

def softmax_t(logits, T=1.0):
    scaled = [z / T for z in logits]
    m = max(scaled)
    exps = [math.exp(z - m) for z in scaled]
    s = sum(exps)
    return [e / s for e in exps]

logits = [4.0, 2.0, -3.0]
for T in [1, 2, 4, 8]:
    probs = softmax_t(logits, T)
    print(f"T={T}: " + ", ".join(f"{p:.3f}" for p in probs))

Run it and watch the car probability climb from ~0.0003 toward something the student can actually get a gradient from. Then push T to 20 — notice the distribution eventually becomes too flat, discarding the teacher's real preference for cat. Picking T is a real tradeoff, not a knob you max out.

Common confusions

Distillation is not quantization (Chapter 66/Week 23's precision reduction — same architecture, fewer bits per weight) and not pruning (removing individual weights/neurons from an existing network). All three shrink a model, by different mechanisms, and are routinely combined: distill first for a smaller architecture, then quantize the result for a smaller file.
The student never needs to see the original training data. In response- and feature-based distillation, the teacher's outputs on any input are the label. This is why distillation datasets can be much larger than the original labelled set — you only need inputs, not annotations.
A distilled student can occasionally beat its teacher on held-out data. This isn't a paradox: if the teacher was trained on noisy raw labels, its smooth output function has effectively denoised them, and the student learns that smooth function directly. You will see exactly this in Chapter 68's numbers.

Check yourself

What does a soft label carry that a hard label doesn't?
Why does the KD loss for classification multiply by T^2?
Why doesn't regression distillation need a temperature?
Name one thing distillation does not do that quantization does.