Modules / Gen AI / Distillation Lab

Knowledge Distillation in LLMs

Train a small student to imitate a large teacher. Raise the temperature to expose the dark knowledge hidden in the teacher's wrong answers — the signal a plain label can never carry.

Teacher — softened at T

LARGE MODEL

Hard label (one-hot)

WHAT A NORMAL DATASET GIVES YOU

Student — current beliefs

SMALL MODEL

Knowledge Distillation: Teaching a Small Model to Think Big

A large model is accurate but expensive. A small model is cheap but weaker. Knowledge distillation narrows that gap by training the small student to imitate the large teacher — and remarkably, the student often beats an identical model trained on the original labels.

Dark Knowledge: The Central Idea

A training label says "this is a dog" and nothing else. But a trained teacher, shown that same photo, might output dog 0.90, wolf 0.07, cat 0.02, car 0.001.

Those small numbers are not noise — they encode similarity structure the label cannot express: a dog resembles a wolf far more than a car. Geoffrey Hinton named this dark knowledge. It is the teacher's entire learned view of how the classes relate, and it is free supervision on every single example.

Why Temperature Is Essential

There is a catch: a confident teacher's non-target probabilities are so tiny that they contribute almost nothing to the gradient. At T = 1 the distribution is nearly one-hot, so the student learns roughly what the plain label already told it.

Raising the temperature in the softmax flattens the distribution and amplifies the ratios between the small values:

pᵢ = exp(zᵢ / T) / Σ⫺ exp(z⫺ / T)

Drag T from 1 to 6 and watch the runner-up bars rise out of nothing. That is the dark knowledge becoming visible — and becoming trainable. Both teacher and student are softened by the same T, and the distillation loss is scaled by to keep gradient magnitudes comparable to the hard-label term.

The Combined Loss

In practice the student learns from both sources at once:

L = (1 − α) · CE(student, true label)   +   α · T² · KL(teacher ‖ student)

  • The cross-entropy term keeps the student anchored to ground truth, so it cannot inherit the teacher's mistakes wholesale.
  • The KL divergence term pulls the student's whole distribution toward the teacher's, transferring the similarity structure.
  • α balances them. Typical values lean toward the teacher, since the soft targets carry strictly more information per example.

Why It Matters for LLMs

Distillation is how most small production models are made. DistilBERT keeps roughly 97% of BERT's performance at 40% of the size and 60% faster inference. The same recipe produced the compact members of most modern model families.

For generative models the idea extends naturally: the teacher's full next-token distribution at every position is a far richer target than the single token that happened to appear in the text. A related variant, sequence-level distillation, simply trains the student on text the teacher generated.

Interactive Exploration Guide

  1. Look at T = 1. The teacher is nearly one-hot and looks almost identical to the hard label — barely any extra signal.
  2. Raise T to 5. The runner-ups emerge. Note which ones rise: semantically related classes, never random ones. That ordering is the knowledge being transferred.
  3. Press "Train student" repeatedly and watch KL divergence fall as the student's bars converge on the teacher's shape — including the small ones.
  4. Set α = 0 and retrain from reset. The student matches only the one-hot label; its non-target probabilities collapse and the similarity structure is lost.
  5. Try the sentiment example. "not bad at all" is genuinely ambiguous, so the teacher stays uncertain even at T = 1 — and that calibrated uncertainty is itself worth teaching.

Key Takeaway

A hard label answers one question. A teacher's full distribution answers every question at once — how likely each alternative is, and therefore how the classes relate. Temperature is the tool that makes that hidden structure large enough to learn from. The student in this lab performs real gradient descent on the combined loss, so the convergence you watch is genuine optimisation, just at a much smaller scale.

Cheat sheet

Knowledge Distillation in LLMs

Train a small student to imitate a large teacher. Raise the temperature to expose the dark knowledge hidden in the teacher's wrong answers — the signal a plain label can never carry.

GEN AI · vizlearn.in/gen_ai/knowledge_distillation_in_llms.html