← All notes

A student gains from a teacher only when the teacher knows more than the labels

By · · 5 min read

Claim. Distilling a teacher into a small student helped in my experiment when the teacher was better than the student could get from the labels alone, and hurt when it wasn't. Imitating a weak teacher left the student below both the teacher and the labels-only baseline.

Knowledge distillation is often presented as a free improvement: train a small student to match a large teacher's softened predictions as well as the labels, and the student comes out better than it would on its own. Hinton et al. (2015)1 explained why that can work. A teacher's full distribution over classes says how similar the wrong answers are to the right one, information that a one-hot label throws away. I wanted to see the other side as well: what happens when the teacher has less to offer than the labels, and how much of the teacher a student actually absorbs.

The setup

The task is 20-class classification of noisy synthetic images built from fixed random templates. The training set is deliberately small, 2,000 images, so that a network can memorise it; validation uses 5,000 fresh images. I trained two teachers of the same size: a strong one, trained for 4,000 steps, and a weak one, stopped after 300. Their validation accuracies were 66.8% and 62.3%.

The student is a network a quarter as wide. I trained it for 4,000 steps five ways, three seeds each: on labels alone, and on a mix of the labels and the teacher's predictions softened at temperature 4. The mix gave the teacher half or all of the weight, once for each teacher. The distillation term is the usual one, the KL divergence to the teacher's softened distribution, scaled by the temperature squared so that its gradient keeps a comparable size.

The student's validation accuracy while it trains — From the runs behind the note's table, three seeds averaged. (interactive figure)
Student trained onValidation accuracyTraining accuracy
Labels only62.4%100%
Labels and strong teacher, half and half64.0%100%
Strong teacher only63.9%100%
Labels and weak teacher, half and half61.1%100%
Weak teacher only60.3%98.9%

Means over three seeds.

What the student gained, and what it couldn't

Every student memorised the small training set: 100% training accuracy, against about 62% to 64% on new images. The strong teacher improved the student's validation accuracy by 1.6 points half-and-half and 1.5 alone. That is the regularising effect: the teacher's soft targets carry more information per image than the labels and pull the student towards a smoother decision function, even though both still fit the training set perfectly.

The student did not become the teacher. Trained on the strong teacher alone, it reached 63.9% against the teacher's 66.8%. A narrower network has less capacity, but that is only part of the story. Stanton et al. (2021)2 found that students often fail to match their teachers' predictions even when they have the capacity to, and traced much of the gap to optimisation: matching a distribution is a harder target than it looks.

The weak teacher

With the weak teacher, distillation made things worse. Half-and-half lost 1.3 points against labels alone, and pure imitation lost 2.1, landing at 60.3%, below the weak teacher's own 62.3%. The student inherited the weak teacher's mistakes and then lost a little more in the imitation, the same gap as with the strong teacher. When the teacher knows less than the labels can teach, its soft targets aren't a gentler version of the truth; they are a competing, worse signal.

That gives a simple test before using distillation: compare the teacher with a student trained on labels alone, on held-out data. If the teacher isn't clearly better, there is nothing to transfer, and imitation can only cap the student. The cap is also lower than the teacher, by the student's own imitation gap.

Distillation as a regulariser, and its cost

In my runs the benefit came almost entirely as regularisation of a student that memorises a small dataset. Two other things decide whether it is worth having:

  • its cost: the teacher runs a forward pass on every training batch, which can dominate the training time when the teacher is large;
  • whether simpler regularisation, such as augmentation, weight decay or label smoothing, would buy the same 1.6 points without a teacher, which I haven't tested here and would want to know before paying that cost.

Evaluating on a teacher's labels has the mirror-image problem: it rewards imitation, in a pseudo-labelled test set measures agreement with the labeller. Distillation is also one of the softer ways to protect old behaviour while adding new outputs, compared with the exact ways in adding an output without moving the old ones.

What I still don't know

I only varied the teacher's strength by stopping it early, which makes it worse in a particular way: less confident and less accurate everywhere. A teacher trained on different data, better than the labels on some inputs and worse on others, is the more common case, and a per-input weighting of the two targets might then do better than either. I also kept the temperature at 4 throughout. The temperature controls how much of the teacher's information about wrong answers the student sees, and the right value probably depends on how good those wrong answers are.

References

  1. Distilling the Knowledge in a Neural Network — Geoffrey Hinton, Oriol Vinyals and Jeff Dean, arXiv 2015
  2. Does Knowledge Distillation Really Work? — Samuel Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A. Alemi and Andrew Gordon Wilson, NeurIPS 2021