On this page
A large model that works well can cost too much to run. Knowledge Distillation (KD) trains a small model, the student, to copy a large trained model, the teacher. The big model you cannot afford to serve teaches the small model you can.
The problem is deployment. An ensemble of averaged models is the most reliable accuracy boost, but too expensive to serve. Bucila et al. showed an ensemble's knowledge could be compressed into one model; distillation is the simplest version of that idea.
The teacher's soft outputs carry information the hard labels never will. A teacher that says 90% cat, 8% dog, 2% car is telling you which mistakes are reasonable. That ranking, the "dark knowledge", is what the student learns, from the teacher's softmax (the step that turns raw scores into probabilities) softened with a temperature, alongside the usual hard-label loss.
Distillation acts on every training step of the student. Each batch goes through the frozen teacher, with no gradients, and through the student, and the loss mixes both outputs; the companion run recomputes the teacher's outputs batch by batch. Nothing carries over between steps except the student's own weights, and the price per step is one extra teacher forward pass.
It works at scale. DistilBERT (Sanh et al. 2019) is BERT-base (a pretrained language model that reads text in both directions) with 40% of the parameters removed, and it keeps 97% of BERT's performance on the GLUE language benchmark; Hinton et al. (2015) showed the effect first on digits and speech.
Besides the hard label ("this is a cat"), the student trains on the teacher's soft predictions, its probability over all classes. To make them more useful, the teacher's logits, its raw score per class, are softened with temperature scaling. The logit for class i is zi, and the temperature is T.
With a higher temperature the student picks up the teacher's inductive biases (tendencies its architecture and training baked in) and gets a much better signal, especially when labels are sparse or noisy.
The student learns from two targets: the true labels, through cross-entropy (the loss that punishes a confident wrong answer hardest), and the teacher's soft predictions, through KL divergence (one number for how far apart two probability distributions sit). A mixing weight α (alpha) blends them into the total loss L; the true label is y.
Figure 4. Soft targets and the true label mix at 0.7 and 0.3; only the student trains.
Figure 5. The student after ten steps; the slider sets T.
The listing is the loss card line for line, with the card's names; the defaults are the run's T = 2.0 and α = 0.7.
import torch.nn.functional as F
def kd_loss(z_s, z_t, y, T=2.0, alpha=0.7):
p_t = F.softmax(z_t / T, dim=1)
log_p_s = F.log_softmax(z_s / T, dim=1)
kl = F.kl_div(log_p_s, p_t, reduction="batchmean")
ce = F.cross_entropy(z_s, y)
soft = alpha * T**2 * kl
return (1 - alpha) * ce + softlog_p_s, both batch by classes, softened at temperature T and compared by kl (the loss card's second line). F.kl_div takes the student's log-probabilities first and the teacher's probabilities second, which gives KL(pt ‖ ps); reduction="batchmean" sums over classes and averages over the batch.Every training step the teacher sets soft targets, the student guesses, the loss mixes both targets, and the student updates.
Figure 6. One step; the label and the frozen teacher both teach.
The student learns the teacher's ranking, not just its answer; a hard label would teach only "cat".
The companion notebook runs the whole recipe (Figure 7). The teacher is a ResNet-18 pretrained on ImageNet (a large labeled image dataset) with its classifier fine-tuned on STL-10 (96×96 color images in 10 classes). The student is DeeperStudentCNN, a plain CNN (convolutional neural network) with no residual connections, far shallower than the teacher but not smaller: almost all of its parameters sit in one wide fully connected layer, so this run shrinks depth, not size. Training uses Adam (an optimizer that adapts each weight's step size) for 10 epochs (full passes over the data).
The listing is that loop, with teacher, student and the STL-10 train_loader already built. The teacher stays in eval mode and runs under torch.no_grad(), so it is frozen and builds no gradient graph; only the student's weights reach the optimizer.
import torch
teacher.eval()
student.train()
params = student.parameters()
opt = torch.optim.Adam(params, lr=1e-3)
for epoch in range(10):
for x, y in train_loader:
with torch.no_grad():
z_t = teacher(x)
z_s = student(x)
loss = kd_loss(z_s, z_t, y)
opt.zero_grad()
loss.backward()
opt.step()The repo also trains a different, shallower student, StudentCNNWithFeatures, to mimic the teacher's intermediate features: an MSE (mean squared error) term between the two networks' feature maps joins the CE and KL terms, with weights CE 0.5, KL 0.3, MSE 0.2, temperature 2.0, 10 epochs (Figure 7).
53% is a long way from 93%, but this student is shallow and has no residual connections. Its training loss levelled off near 0.24 from epoch 8 of its 10, so the run had stopped improving before it ended. The repo trains no student without a teacher, so the run shows where the distilled student landed, not how much distillation added.
The two students differ in architecture, so 42.48% against 53.46% says nothing about feature matching against logit matching. What carries over is the practical point: matching raw activations (each layer's outputs) across very different architectures needs care, because feature scales differ.
The distillation loss matches outputs. Other kinds of distillation match other things the teacher computes.
| Kind | Matches | Loss | Helps when |
|---|---|---|---|
| Response (the original method) | The softened outputs | KL between the soft targets and the student's guess | The default; works across architectures |
| Feature (FitNets, Romero et al. 2015) | A middle layer | Squared distance from a projected student feature | Deep, thin students |
| Relation (RKD, Park et al. 2019) | Distances and angles between examples | Gap between the two networks' distances and angles | Embedding tasks, where the layout of examples matters |
| Self (Born-Again, Furlanello et al. 2018) | A trained model of its own architecture | The same KL loss, one generation to the next | Same-size retraining; soft targets regularize |
Distillation changes two lines of an ordinary training loop: one extra teacher forward pass per batch, and kd_loss in place of cross-entropy. The ResNet-18 loop with the kd_loss function from The distillation loss is a complete implementation.
torch.no_grad() in eval mode, so it stores no activations for backward and has no optimizer state, but its weights must sit in memory.teacher.eval() leaves dropout and batch-norm updates on, so the soft targets drift from step to step.F.kl_div expects log-probabilities for the student and plain probabilities for the teacher; swapping them or leaving out reduction="batchmean" gives a wrong scale with no error.A distilled student can then be shrunk further: Pruning removes the weights it barely uses, and quantization stores the rest in fewer bits (Post-Training Quantization and Quantization-Aware Training). The non-contrastive self-supervised methods use a form of self-distillation without labels, and the KL and cross-entropy terms are the ones in Cross-Entropy Loss.