Vol. XII · Mid-Training Plate 7 of 10knowledge distillation, soft targets, teacher, student, logits, Gemma 2, Llama 3.2, KL divergenceDE

Learning From a Larger Model.

Instead of “the next token was X”, tell the student everything the teacher believed about the next token.

The ordinary objective gives one bit of feedback per position: the token that actually came next. A trained model knows more than that — at each position it holds a whole distribution over plausible continuations, with near-misses ranked. Knowledge distillation (Hinton and colleagues, 2015) trains a small student to match a larger teacher’s distribution, minimising the KL divergence from the teacher’s probabilities instead of the cross-entropy with a one-hot label.

Distillation has moved into pretraining. Gemma 2’s 2B and 9B models were trained by distillation from the 27B rather than by next-token prediction, and Meta used logits from Llama 3.1 8B and 70B as token-level targets when pretraining Llama 3.2 1B and 3B. The workshop uses micro as the teacher on held-out text and shows what its soft targets contain that the hard label does not.

Below: left, micro’s distribution after the chosen prefix — the soft target — beside the one-hot hard label. Right, along a passage, the entropy of the teacher’s distribution at each position, and whether its top choice matched the token that came.

What you will learn
Workshop
micro as teacher on a held-out passage. Left: its ten most likely next tokens after the prefix ending at the chosen position (soft target), and the token that actually followed (hard label). Right: the entropy of micro’s distribution at every position (bits), coloured by whether its top choice was the actual next token.

More signal per position

Along this passage micro’s top choice is the actual next token at 32% of positions, and its distribution has a mean entropy of 5.52 bits: at a typical position it spreads its belief over a few dozen plausible tokens, and the actual token is among its top ten at 53% of positions. A hard label says only which one occurred; the soft target says how plausible each alternative was — that “a” and “the” were both reasonable, that a synonym was nearly as likely. A student trained on it receives a gradient for every token in the vocabulary at every position, not just one.

That is why distillation trains small models more efficiently per token: Gemma 2’s report found a 2B model distilled from a larger teacher clearly better than the same model trained from scratch on the same tokens. The costs are a teacher forward pass for every training token and, if targets are stored, their size — the top 64 logits for 15 trillion tokens at 6 bytes each would take about 5,760 TB — so targets are usually truncated to the top few and computed on the fly. Distillation also appears after pretraining, in fine-tuning on a teacher’s outputs (plate XVII.9) and in compressing models for deployment (plate XXIV.6).

The student learns the teacher’s mistakes too. Where the teacher is confidently wrong, the soft target pulls the student towards the same error; distillation transfers a model’s beliefs, not the truth.

Terms

Knowledge distillation
Training a student to match a teacher’s predicted distribution.
Soft target
The teacher’s full probability distribution over the next token.
Hard label
The single token that actually occurred.
Teacher and student
The larger model whose outputs are matched, and the smaller model trained to match them.
Distillation gives the student the teacher’s whole distribution at every position — far more signal than one correct token — at the cost of running the teacher.

Sources

  1. Hinton, Vinyals and Dean, “Distilling the Knowledge in a Neural Network”, arXiv:1503.02531, 2015.
  2. Gemma Team, “Gemma 2”, arXiv:2408.00118, 2024.
  3. Meta, Llama 3.2 model card, September 2024.