Build Your Own Distillation — Growing a Small Model in 100 Lines

The distillation loss fits in twenty lines — and almost everyone who writes it trips on the same three things: the direction of the KL, the choice of reduction, and the missing T². We build the whole rig: freezing the teacher, the loss, the training loop, the teacher-free baseline, a temperature sweep, and four sanity checks that prove the implementation isn't quietly broken.

Distilling the Knowledge in a Neural Network

What You See Only When You Write It Yourself

The distillation loss is short. Twenty lines at most, three if you're honest about it. Which is exactly why people bolt it onto an existing trainer and move on — and then have nothing to hold onto when it doesn't work. Raising the temperature changes nothing. The student lands below the baseline. Is it the loss? The way the teacher is being called? The comparison itself? There's no thread to pull.

The theory side is settled in The Math of Distillation. This article picks up from there and builds a working rig you can turn knobs on. About a hundred lines total. What you discover writing it is that the hard part of distillation isn't the few lines of loss — it's everything around them.

Analogy: A Knob You Can Only Trust If You Wired It

Someone hands you a machine with a temperature dial and says turning it changes the flavor. You turn it. Nothing happens. Two explanations fit equally well: the dial is at a setting that doesn't matter, or the dial isn't connected to anything inside.

From the outside you cannot tell these apart. Only the person who did the wiring can say "the dial is definitely connected, so this is a settings problem." The temperature $T$ in distillation is exactly this dial, and as we'll see, forgetting one line leaves you with a dial that turns freely and does nothing. Training still proceeds normally, so nothing announces that it's broken.

Intuition: The Real Work Is Three Runs, Not One Loss

The point of building this yourself isn't to own a loss function. It's to put yourself in a position where you can compare. That means at least three runs.

• A — the teacher: the large model, trained normally. This is your practical ceiling.
• B — the plain student: the small model, trained on hard labels only, no teacher. This is your floor.
• C — the distilled student: the same small model, trained with the teacher's soft labels mixed in.

There is exactly one number worth reporting: where C sits between B and A. The common mistake is looking at C's accuracy alone and declaring distillation a success. Maybe the student architecture was already strong enough to get there without any teacher. Without B, you can't rule that out. And even if C nearly matches A, that means nothing if A and B were close to begin with.

So the most important discipline in the implementation is this: B and C must differ in the loss and nothing else. Same architecture, same init seed, same epochs, same LR schedule, same batch order, same augmentation. Any drift gets silently booked as the effect of distillation.

Mechanism: Pin Down the Loss

First, temperature-scaled softmax gives us the teacher's and the student's distributions.

Here $z_i$ is the logit for class $i$ (the raw output before softmax) and $T$ is a positive number. Equation (1) says only this: divide the logits by $T$, then apply the usual softmax. Larger $T$ flattens the distribution, which is what makes the runner-up classes legible.

The loss we actually implement is this.

$p$ is the teacher's distribution, $q$ the student's, $y$ the true label, $H$ the cross-entropy, and $\alpha$ a mixing weight between 0 and 1. Read equation (2) as: blend a term that matches the teacher's distribution with a term that hits the true answer, in proportion $\alpha$. Only the hard term uses $q(1)$ — temperature reset to 1 — because that term's job is to be right, not to be informative. The leading $T^2$ compensates for the fact that the soft term's gradient shrinks like $1/T^2$ as temperature rises; the derivation is in The Math of Distillation.

Before writing any code, turn the dial by hand. Seeing how the teacher's output flattens changes how you read a temperature sweep later.

Step 1: Set Up the Teacher, and Freeze It

The teacher is a pretrained model used as-is. There are three things to do to it.

teacher.eval()                      # Dropout / BatchNorm into inference mode
for p in teacher.parameters():
    p.requires_grad_(False)         # no gradients through the teacher

with torch.no_grad():
    teacher_logits = teacher(x)     # same inputs as the student, every step

Forget eval() and the teacher answers with Dropout still active: show it the same image twice and get two different distributions. Your targets become noise. Training still converges, so this is hard to spot — and it's the classic cause of "distillation made things worse." The requires_grad_(False) and no_grad are there so the teacher doesn't burn compute and memory on gradients nobody wants.

If teacher inference is expensive, you can precompute the logits once and store them. That saves a teacher forward pass every epoch, but it comes with a condition. If your augmentation is random, the cached logits are answers to a different image, so they're unusable. You either freeze the augmentation or cache per augmented input. In LLM distillation this gets worse: storing full-vocabulary logits for every token is rarely affordable, so in practice you keep only the top-k.