Knowledge distillation compresses a large, accurate model into a smaller, cheaper one by training the small model to copy the large one. This page covers the loss, the role of temperature, the main variants, a runnable PyTorch example and the mistakes that quietly ruin results.
Knowledge distillation trains a small student model to imitate a large teacher model. Instead of learning only from the dataset's hard labels ("this image is a cat"), the student also learns from the teacher's full output distribution ("86% cat, 11% dog, 3% fox"). Those extra numbers carry what Hinton and colleagues called dark knowledge: which wrong answers are nearly right, and therefore how the teacher generalises.
The idea goes back to Buciluă, Caruana and Niculescu-Mizil's 2006 model compression work and was popularised by Hinton, Vinyals and Dean's 2015 paper Distilling the Knowledge in a Neural Network. Well-known results include DistilBERT, which keeps most of BERT's accuracy with about 40% fewer parameters, and many of today's small open LLMs, which are trained partly on outputs of larger models.
Large models are accurate but expensive: memory, latency and cost per request grow with parameter count. Many production settings — phones, browsers, high-traffic APIs, edge devices — cannot afford the largest model. Training a small model from scratch on the same labels usually lands well short of the teacher. Distillation narrows that gap because soft targets are a richer training signal than one-hot labels: each example tells the student about every class, not just one.
Distillation also composes with other compression techniques. A common pipeline is: distil to a smaller architecture, then quantize the student to 8 or 4 bits, then serve it with an optimised runtime.
Both models see the same input. The teacher's logits zt and the student's logits zs are each divided by a temperature T before the softmax:
pi = exp(zi / T) / Σj exp(zj / T)
The distillation loss is the Kullback-Leibler divergence KL(pteacher ‖ pstudent), multiplied by T2 so that its gradients stay comparable in size as T changes. It is usually mixed with ordinary cross-entropy on the true labels using a weight α. Only the student's weights are updated; the teacher is frozen.
Temperature matters because a confident teacher puts almost all its mass on one class, so its soft targets look nearly identical to hard labels. Raising T flattens the distribution and exposes the relative sizes of the small probabilities — "dog" is far more plausible than "car" for a cat photo. Typical values are between 2 and 10, tuned on validation data.
Variants differ in what the student imitates. Response-based distillation matches final outputs, as above. Feature-based distillation also matches intermediate activations or attention maps, often through a small projection layer when widths differ (FitNets, TinyBERT). Sequence-level distillation, common for language models, simply fine-tunes the student on text generated by the teacher; it needs only the teacher's outputs, not its logits, which is why it works with API-only models (subject to the provider's terms of use).
This self-contained script distils a wide MLP into a narrow one on synthetic data. It runs on CPU in seconds.
import torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0)
# synthetic 10-class problem
X = torch.randn(6000, 20)
W = torch.randn(20, 10)
y = (X @ W + 0.5 * torch.randn(6000, 10)).argmax(1)
Xtr, ytr, Xte, yte = X[:5000], y[:5000], X[5000:], y[5000:]
def mlp(h): return nn.Sequential(nn.Linear(20, h), nn.ReLU(), nn.Linear(h, 10))
def train(model, loss_fn, epochs=30):
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for _ in range(epochs):
for i in range(0, 5000, 100):
xb, yb = Xtr[i:i+100], ytr[i:i+100]
opt.zero_grad(); loss_fn(model, xb, yb).backward(); opt.step()
return (model(Xte).argmax(1) == yte).float().mean().item()
teacher = mlp(512)
print("teacher acc", train(teacher, lambda m, x, y: F.cross_entropy(m(x), y)))
teacher.eval()
T, alpha = 4.0, 0.3
def kd_loss(student, x, y):
s = student(x)
with torch.no_grad():
t = teacher(x)
soft = F.kl_div(F.log_softmax(s / T, 1), F.softmax(t / T, 1),
reduction="batchmean") * T * T
return alpha * F.cross_entropy(s, y) + (1 - alpha) * soft
torch.manual_seed(1); print("student alone ", train(mlp(16), lambda m, x, y: F.cross_entropy(m(x), y)))
torch.manual_seed(1); print("student + KD ", train(mlp(16), kd_loss))Compare the two student lines. On small synthetic problems the gain can be modest and varies with the seed; the gap is usually larger on real data where the teacher has learned useful similarity structure between classes. Note the reduction="batchmean" argument: it is the mathematically correct KL reduction in PyTorch, and the default "mean" silently scales the loss down by the number of classes.
For LLMs the recipe usually looks different from the classroom version. Matching full next-token distributions requires the teacher and student to share a tokenizer and requires access to the teacher's logits, which rules out most hosted APIs. Teams therefore choose between two routes. When both models are local and share a vocabulary, they compute a token-level KL loss over the teacher's top-k probabilities at each position, which keeps memory manageable for vocabularies of 100,000 or more tokens. When the teacher is only reachable through an API, they generate a large set of prompts, collect the teacher's answers (sometimes several per prompt, filtered by a grader or by checking final answers), and fine-tune the student on that synthetic corpus with ordinary cross-entropy.
Two practical lessons recur. First, prompt diversity matters more than volume: thousands of near-duplicate prompts teach the student one narrow skill. Second, the student must be evaluated on held-out tasks the teacher did not generate, otherwise you only measure how well it copies the teacher's style.
| Method | What changes | Needs retraining? | Typical use |
|---|---|---|---|
| Knowledge distillation | Architecture (smaller model) | Yes, full student training | Big latency or memory reductions; new hardware targets |
| Quantization | Number precision (FP16 to INT8/INT4) | No (post-training) or light (QAT) | Fast win for serving; see vLLM and quantization |
| Pruning | Removes weights, heads or layers | Usually fine-tuning afterwards | Structured pruning for speed, unstructured for size |
| LoRA / adapters | Adds small trainable matrices | Fine-tuning only | Task adaptation, not compression |
| Smaller model trained from scratch | Architecture | Yes | Baseline that distillation should beat |
kl_div expects log-probabilities for the student and probabilities for the teacher; swapping them trains the wrong objective without an error.