Knowledge Distillation in deep learning

Module 16 · Advanced AI · 23 concepts

In depth / 15 lines

Knowledge Distillation, in more detail

Distillation trains a small student model to imitate a large teacher.

The aim is most of the capability at a fraction of the cost.

Training the small model directly on the original labels works less well.

The difference is in what the teacher provides.

A hard label says this is a cat and nothing else.

The teacher's full distribution says cat 0.9, dog 0.08, fox 0.02.

Those soft targets carry information about how the classes relate.

The student learns that a cat resembles a dog more than a lorry.

That extra signal is why distillation beats plain training at the same size.

The teacher's outputs are usually softened with a temperature to expose more of that structure.

The student can also be trained to match intermediate representations, not only outputs.

For language models, distillation often means training on the teacher's generated text.

That is how many small capable open models were produced.

It raises licensing questions, since providers commonly forbid training competitors on their output.

DistilBERT is the classic case, at 40 per cent smaller and 60 per cent faster with most of the quality.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Knowledge Distillation

Training on a distribution instead of a label

\mathcal{L}_{\mathrm{hard}} = - \log \hat{p}_{c}

A one-hot label says the answer is cat and everything else is equally wrong, which is most of the information thrown away.

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

The teacher's softened probabilities say the image is a cat, somewhat like a dog, and nothing like a lorry, and that ranking is the dark knowledge being transferred.

T > 1 \implies \text{softer distribution}

Raising the temperature exposes the small probabilities. At $T = 1$ a confident teacher gives 0.99 and two values near zero, which carries almost no more than the label did.

The loss

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

$p^{\mathcal{T}}, p^{\mathcal{S}}$ : teacher and student distributions at temperature $T$ .
$\alpha$ : the mix, often 0.5 to 0.9 in favour of the teacher.

The true labels stay in the loss, because the teacher is wrong sometimes and the labels are the anchor.

T^{2}

That factor is not cosmetic. Softmax gradients scale as $1/T^{2}$ , so without it the distillation term shrinks as the temperature rises and the two terms fall out of balance.

One distillation, in order

Train or obtain the teacher, and freeze it.
Run the teacher over the training data once, storing its logits rather than recomputing them every epoch.
Train the student on the combined loss at temperature $T$ .
Evaluate the student at $T = 1$ , since the temperature is a training device only.
Compare against the same student trained from labels alone, which is the baseline that matters.

T \in [2, 10]

Temperature and $\alpha$ are the two knobs, and the gain over training from labels alone is typically a few points at the same student size.

Why it works

The teacher's errors are informative.

The relative probabilities encode the teacher's learned similarity structure, so the student learns which classes are confusable rather than discovering it from scratch.

Soft targets also work as regularisation.

They also act as a label smoother that varies per example, which is a better regulariser than a constant because it reflects genuine ambiguity.

Distillation can use unlabelled data, since the teacher supplies the target, and that often matters more than the technique itself.

The variants in current use

For language models the usual form is to generate outputs with a large model and fine-tune a small one on them, which is distillation through samples rather than through logits.

\text{feature matching} : \| \mathbf{h}^{\mathcal{S}} - W \mathbf{h}^{\mathcal{T}} \|^{2}

Matching intermediate representations rather than only the output adds a signal per layer, which helps when the student is much shallower than the teacher.

The limits are consistent.

A student around a quarter of the teacher's size usually keeps most of its quality, and beyond that the gap widens quickly, because there is a floor on what the capacity can represent.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

Why does distillation use the teacher's full probability distribution rather than hard labels?

Advanced AI

23 concepts in this module, then a graded test with every answer explained.

Study this module →