Knowledge Distillation

More than the right answer

A trained classifier asked about a picture of a 2 might put 0.9 on "2", 0.09 on "7" and almost nothing on "cat". The label says only "2". The teacher's distribution also says that this 2 looks a little like a 7, which is information about the task the label doesn't carry.

Hinton, Vinyals and Dean proposed raising the temperature of the teacher's softmax until it produces a suitably soft set of targets, and training the small model with the same high temperature to match them Established. The training loss is the KL divergence (equivalently, cross-entropy) between teacher and student distributions, usually combined with the ordinary loss on the true labels.

Temperature, in numbers

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

Tiny example. Teacher logits $z = [4, 2, 0]$.

• At $T = 1$: $[0.867, 0.117, 0.016]$. The third class is nearly invisible.
• At $T = 4$: $[0.507, 0.307, 0.186]$. The ranking is the same, but the student now sees how the wrong answers compare.

This is the same temperature as in decoding, used for a different purpose.

For language models

DistilBERT applied distillation during pretraining and produced a model 40% smaller and 60% faster than BERT that retained 97% of its language-understanding performance Established. For generative LLMs the idea extends to training a student on the teacher's next-token distributions, or more simply on text the teacher generates.