Knowledge distillation copies a teacher model’s wrong guesses

How do you move what a big AI model knows into a model that is cheaper to run? In 2015, three Google researchers, Jeffrey Hinton, Oriol Vinyals and Jeff Dean, published an answer called knowledge distillation. Distillation trains one model to copy another model's guesses, the wrong ones included. The researchers' biggest test was the listening part of an older Android voice search. That part guesses which speech sound it hears. The team trained ten of these listening models and averaged their guesses together. Then the team trained a single model, the same size as each of the ten, to copy the group. The group and the single copy both made ten point seven word mistakes in every hundred words. In 2024, Google said it built its two smallest Gemma 2 models the same way.

Why would copying another model beat learning from the same examples? Normal training grades a model against an answer key. For a handwritten 2, the key says 2, and every other answer counts as equally wrong. A trained model grades the wrong answers too. Take the paper's example, a photo of a BMW. A trained model gives the photo a very small chance of being a garbage truck, and a far smaller chance of being a carrot. Those grades show what the model learned about which things look alike. In distillation, the model being copied is the teacher, and the copy is the student. The student gets the teacher's marked-up key. The right answer stays the same, and every wrong answer gets a note in the margin.

Take one page of the marked-up key. Show the teacher a handwritten 2. The teacher splits its confidence across all ten digits, and almost all of it lands on the 2. The paper's example gives this 2 a one-in-a-million chance of being a 3, and a one-in-a-billion chance of being a 7. Those tiny numbers say which 2s look like 3s and which look like 7s. At that size, the margin notes barely change what the student learns.

So the paper adds a setting called temperature. Raising the temperature flattens the teacher's guesses. The big guess shrinks, the small guesses grow, and every guess keeps its place in line. The paper used a temperature of twenty on its digits. In our illustration at twenty, the 2 drops to about a third of the confidence. The 3 rises to about a sixth, and the 7 to about an eighth. The teacher and the student both use the high temperature while the student learns. Afterwards, the student goes back to the normal setting of one.

The student trains to match the teacher's flattened guesses on every example. The student also gets a little credit for the plain right answer. Each example now carries ten grades instead of one tick. On handwritten digits, a big teacher model made sixty-seven mistakes on the test set. A smaller model, trained alone against the answer key, made a hundred and forty-six mistakes. The same smaller model made seventy-four mistakes when it copied the teacher at a temperature of twenty.

Next, the team removed every 3 from the student's lessons. The teacher had seen plenty of 3s. The student only saw the teacher's guesses about other digits, including how much each digit looked like a 3. Before any fix, the student got a hundred and thirty-three of the test's one thousand and ten 3s wrong. So the team nudged the student toward answering 3. The team picked the size of that nudge using the test itself. After the nudge, the student got ninety-eight point six percent of the test 3s right.

Back on speech, the paper scores a model on tiny slices of sound, a hundred every second. The score is the share of slices the model labels with the right speech sound. The team trained a speech model on only three percent of the recordings. With the plain answer key, that model labelled forty-four point five percent of slices right. A second model saw the same three percent, with margin notes from a teacher trained on all the recordings. The second model labelled fifty-seven percent right. An ordinary model trained on all the recordings labelled fifty-eight point nine percent right.

This picture is roughly right, and here is where the picture breaks. On the sound-slice score, one ordinary speech model got fifty-eight point nine percent. The group of ten got sixty-one point one percent. The single copy got sixty point eight percent, three tenths of a point short of the group. The paper says more than eighty percent of the group's gain carried over to the copy. The copy runs as one model instead of ten.

The second break is the date. The 2015 paper sorted things into fixed lists of answers, ten digits or fourteen thousand speech sounds. A chat model picks its next token, a word or a piece of a word, from a much longer list. Gemma two's list has about two hundred and fifty-six thousand tokens. Google's Gemma 2 report says its small models copied a big model's spread over every possible next token. Normal training would teach only the one token that came next. Google's report names no temperature setting, so the knob in this picture comes from the 2015 paper.

Train a big teacher first. Show the teacher an example and keep every guess, the winner and the wrong ones. Turn up the temperature so the small guesses show. Train the student to copy that whole spread, plus a little of the plain right answer. The student picks up which things look alike, graded the way its teacher grades them. That is how one speech model came within three tenths of a point of ten models averaged together.