Knowledge Distillation: teaching a tiny model with a big one's soft answers

The best model is usually too big to ship -- it needs a GPU, costs a lot per query, and adds latency that kills anything real-time or on-device. But you rarely need a model to be big when it is answering; you need it big while it is learning. Knowledge distillation splits those two jobs: train a large, accurate teacher once, then transfer what it knows into a small, cheap student that you actually deploy. The clever part is what gets transferred.

One-hot labels throw away most of the signal

Say you are classifying an image of a cat into five classes: cat, dog, tiger, car, plane. The ground-truth label is one-hot: cat = 1, everything else = 0. That label quietly asserts something false -- that the image is equally not-a-dog, not-a-tiger, and not-a-car.

A trained teacher knows better. Show it the cat and it outputs something like:

cat 0.86   dog 0.09   tiger 0.04   car 0.003   plane 0.001

Look at the wrong answers. The teacher thinks this cat resembles a dog far more than a car, and a tiger more than a plane. Those ratios between the wrong classes encode real structure -- the geometry of the problem, what looks like what. Geoffrey Hinton called it dark knowledge, and a one-hot label destroys every bit of it.

If the student learns to reproduce the whole distribution instead of just the top answer, it inherits that structure -- each example now carries a rich vector of information rather than a single integer.

Temperature: turning up the contrast on the small numbers

There is a catch: a confident teacher's softmax is so peaky that the dark knowledge sits at 0.001 -- too small to drive learning. The fix is temperature. Before softmax, divide the logits by a temperature T:

soft = softmax(logits / T)

• T = 1 is ordinary softmax: peaky, dark knowledge invisible.
• T = 4 flattens the distribution and lifts the small probabilities into view.
• T -> infinity becomes perfectly uniform -- all information gone.

So T dials how much inter-class structure you expose: too low is back to one-hot, too high drowns the signal toward uniform. Values of 2 to 6 are typical, and the same T applies to both teacher and student.

The loss: two teachers, one weighted sum

The student trains against two targets at once.

The soft term measures how far the student's softened distribution q is from the teacher's softened distribution p, using KL divergence:

KL(p || q) = sum p_i * log(p_i / q_i)

which is zero exactly when the student matches the teacher. The hard term is ordinary cross-entropy against the true one-hot label at T = 1 -- it keeps the student anchored to ground truth in case the teacher is wrong. Blend them with a weight alpha:

L = alpha * CE(student, true_label)  +  (1 - alpha) * T^2 * KL(p || q)

Two details worth pausing on:

Why the T^2? The soft-loss gradients shrink in proportion to 1 / T^2 (the 1/T inside each softmax shows up twice via the chain rule). Multiplying by T^2 restores them to the magnitude of the hard-label gradients, so a single learning rate works for both.

Why is alpha usually small? Values of 0.1 to 0.5 are common: the student learns mostly from the teacher's soft targets. Set alpha = 0 for pure teacher-matching; set it to 1 (with T = 1) and you are back to ordinary training.

In PyTorch the whole thing is a few lines:

import torch.nn.functional as F

def distill_loss(s_logits, t_logits, target, T=4.0, alpha=0.3):
    p    = F.softmax(t_logits / T, dim=-1)
    logq = F.log_softmax(s_logits / T, dim=-1)
    soft = (T * T) * F.kl_div(logq, p, reduction="batchmean")
    hard = F.cross_entropy(s_logits, target)
    return alpha * hard + (1 - alpha) * soft

The teacher is loaded, set to eval(), and wrapped in torch.no_grad() -- it produces targets but never updates. Only the small student learns.

The surprise: the student often generalizes better

A distilled student is often better than the same small network trained on the raw labels. Soft targets act like an informed label smoothing: instead of forcing every wrong class to exactly zero (which breeds overconfidence and brittle boundaries), they hand the student a sensible, non-zero target for every class. More bits of supervision per image means it learns more from less data, and the teacher smooths over label noise along the way. The idea stretches, too: in self-distillation the student shares the teacher's architecture and still improves.

Where you have already met it

Distillation is how big models get shipped. DistilBERT keeps about 97% of BERT's performance at 40% of the size and 60% faster. On-device vision and speech models are routinely distilled. And modern LLM pipelines distill reasoning traces: a large model generates step-by-step solutions, a small model is trained to reproduce them.

It is complementary to the other two ways of shrinking a model -- pruning (drop unimportant weights) and quantization (fewer bits per weight). They attack different redundancies, so pipelines often distill, then quantize, then prune. Distillation is the one that transfers learned knowledge into a smaller shape.

Play with the temperature and alpha sliders and watch a tiny student slide onto the teacher's curve:

https://dev48v.infy.uk/dl/day33-knowledge-distillation.html