A Field Guide to Distillation Recipes — logit, feature, attention, self

Every distillation recipe is an answer to one question: which part of the teacher should the student match? This guide lines up output (logit), intermediate feature (FitNet), attention, and self-distillation in one table and four equations, then turns the choice into a decision you make from constraints — is the teacher behind an API, do the layers line up, do the head counts match.

Distilling the Knowledge in a Neural Network

There is more than one place to copy from

When you learn to cook from a master, there are several things you could imitate. You could copy only the taste of the finished dish. You could copy how the vegetables are cut during prep. You could copy where the master's eyes go while the pan is on the heat. Or you could skip the master entirely and use yesterday's version of yourself as today's standard.

That is the whole taxonomy of distillation recipes. Which part of the teacher does the student match? Change the location and you change what you need from the teacher, the shape of the loss, and the conditions under which it works at all.

The foundation — why a trained model's output carries more information than a hard label — is covered in Knowledge Distillation from Scratch. This article is the map that comes next: which recipe do you actually reach for.

The four recipes in one table

The rest of the article opens each row in turn: equation, then something you can play with, then code.

Recipe 1: output distillation — the fewest ingredients

The classic recipe, still the workhorse. You pull the teacher's and the student's final probability distributions together and never touch the teacher's internals.

Put in words, this says: mix "match the hard label" and "match the teacher's distribution after it has been softened by temperature" in proportion $\alpha$. Here $y$ is the ground-truth label, $\mathrm{CE}$ is cross-entropy, $p_S$ is the student's probability, $p_T^{(T)}$ and $p_S^{(T)}$ are the teacher's and student's distributions flattened by temperature $T$, and $\mathrm{KL}$ measures how far apart two distributions are (KL Divergence from Scratch).

Temperature divides the logits by $T$ before the softmax: raise $T$ and the distribution flattens. Flattening surfaces the runners-up — the candidates that were not first but were close — and that is where the teacher's knowledge lives. The $T^2$ in front of equation (1) is a correction. Dividing by the temperature shrinks the gradient by $1/T^2$, so without multiplying it back, changing the temperature silently changes your effective learning rate.

The strength of this recipe is how little it asks for. You do not need the teacher's weights, the architectures can be completely unrelated, and as long as the vocabulary is shared it runs. That is why it is the default first attempt.

The weakness is the mirror image: information enters at a single point, the exit. Deep students receive a faint signal in their lower layers, and whatever intermediate features the teacher had decomposed the input into cannot be recovered once they have collapsed into a probability vector. Read the next three recipes as three answers to that one problem.

Recipe 2: feature distillation — matching the prep work

If the exit alone is not enough, match the middle too. That is the FitNet idea: take the output of one teacher layer (the hint) and pull the output of one student layer (the guided layer) toward it directly.

Stated in words, all this does is move the point of comparison from the exit to the middle: pick one shelf inside each network, and penalize the difference between what sits on them. What is being compared is work in progress, not the final probability vector.

Read it as: take the student's layer-$j$ feature vector $h_S^{(j)}$, reshape it with a regressor $r$, and square the difference against the teacher's layer-$i$ feature $h_T^{(i)}$. The $\|\cdot\|_2^2$ is the squared distance — element-wise differences squared and summed.

You need $r$ because a thin student and a wide teacher have different feature dimensions to begin with. A single dense layer or 1×1 convolution reconciles them. It is scaffolding: throw it away when training ends.

Two things matter in practice. First, layer pairing. With a 12-layer teacher and a 4-layer student, connecting every third layer is the sane default. Second, scale. Feature magnitudes differ by orders of magnitude across layers, so summing raw MSE lets whichever layer happens to be large dominate the total loss. Normalizing per layer before measuring distance is the standard fix (Normalization Layers from Scratch).

That is the real judgment call in feature distillation. MSE goes after direction and magnitude; cosine distance cares only about direction. When teacher and student represent things at fundamentally different scales, insisting on matching magnitude is asking for something the student cannot give.

Recipe 3: attention distillation — copying where the eyes go

Between Transformers you can go a step further and match the attention matrices themselves. An attention matrix is a table of which token attends to which (Attention from Scratch), so copying it is copying the master's gaze.

Stated in words, this grades how the answer is reached rather than the answer itself: lay the teacher's and the student's "which token is looking at which" tables on top of each other, score them cell by cell, and average over the heads. Arrive at the same output by looking somewhere else and you still take a penalty.

Meaning: for each of the $h$ heads, measure the gap between the student's attention matrix $A_S^{(k)}$ and the teacher's $A_T^{(k)}$, then average. Each $A^{(k)}$ is a (sequence length × sequence length) table for head $k$.

On the vision side, the transferred quantity is not attention itself but a spatial "where is the activation strong" map built from the sum of squares of a feature map (attention transfer). Different mechanics, same idea: copy the distribution of what the model responds to, not just what it outputs.

Three implementation details bite. Head count first: if student and teacher differ, one-to-one pairing does not exist. Part of why MiniLM moved toward transferring only the final layer's self-attention relations is that it sidesteps this pairing problem. Then masking: padding positions carry near-zero weight in both models, so including them in the MSE dilutes the loss with regions that already agree. Finally, before or after softmax: scores and probabilities give gradients with different characters, so check which one a paper's code actually uses.

Recipe 4: self-distillation — distilling without a teacher

Everything so far needed a teacher. The framework survives without one.

Born-again training takes a normally trained model as the teacher and distills into a student with the same architecture. The counterintuitive finding was that the student can come out better despite identical capacity — the usual explanation being that soft targets act as a regularizer.

Depth-wise self-distillation puts the teacher inside a single network: deeper layers teach shallower ones. You attach auxiliary exits to the shallow layers and train them against the final output. Drop the exits afterward and you are back to the original architecture; keep them and you have an early-exit mechanism that answers cheaply when it can.

The trap is blunt. The teacher is you, so your mistakes get reinforced too. With no external source of correctness, systematic errors are not corrected — they are handed forward through the soft labels. Calibrate your expectations to "regularization," not to "make a big model small."

Implementation: one loop, all four terms

The recipes are not exclusive. The normal thing is to weight them and add them up.

teacher.eval()
with torch.no_grad():
    t = teacher(x, output_hidden_states=True, output_attentions=True)
s = student(x, output_hidden_states=True, output_attentions=True)

kd = F.kl_div(F.log_softmax(s.logits / T, -1),      # arg 1 must be LOG-probs
              F.softmax(t.logits / T, -1),
              reduction="batchmean") * (T ** 2)      # never drop the T^2
hint = sum(mse(proj[j](s.hidden_states[j]), t.hidden_states[i])
           for j, i in layer_map)                    # e.g. {1:3, 2:6, 3:9}
attn = sum(mse(a_s[mask], a_t[mask])                 # exclude padding
           for a_s, a_t in zip(s.attentions, t.attentions))

loss = (1 - alpha) * ce + alpha * kd + beta * hint + gamma * attn

That kl_div takes log-probabilities first and probabilities second is the most common mix-up in distillation code. Forgetting log_softmax raises no exception: the loss goes down, accuracy does not follow, and nothing points at the cause. Likewise, reduction must be "batchmean" and not "mean", or the KD term gets divided by the vocabulary size and shrinks into irrelevance.

Choosing by constraint, not by taste

• Teacher behind an API, outputs only → output distillation, no alternative. If even the distributions do not line up because the tokenizers differ, move to sequence-level distillation on text the teacher generated
• Same family, differing only in depth or width → output plus feature. Layer pairing falls out naturally; best return on effort
• Compressing a Transformer during pre-training → add attention distillation, which transfers knowledge that is not tied to a downstream task
• No teacher available, or you want one more point of accuracy → self-distillation, with expectations set at "regularizer"

How this shows up on the job

This choice becomes someone's work when the person responsible for inference cost — an ML engineer, or whoever owns the serving stack — is told to shrink a model without giving up accuracy. The first reach is usually quantization (Quantization from Scratch); distillation enters when that is not enough.

The knobs you actually touch are the temperature T, the hard-label/KD mixing ratio alpha, the feature and attention loss weights beta and gamma, and the layer map. In HuggingFace transformers, output_hidden_states=True and output_attentions=True expose the intermediates; in your own code, PyTorch's register_forward_hook does the same job.

The failure modes, roughly in order of how often they happen:

1. Forgetting teacher.eval(). With dropout and BatchNorm still live, the teacher's output jitters every step and the student is shooting at a moving target. Forget torch.no_grad() as well and you also retain the teacher's gradients, inflating memory badly
2. Dropping the T² correction. Your temperature sweep quietly becomes a learning-rate sweep, and the conclusion inverts
3. Not watching loss scales. Ship beta and gamma at their default and the feature MSE can sit orders of magnitude above the CE term, leaving the student with no incentive to learn the task. Logging each term's value at step zero prevents this entirely
4. Preprocessing drift between teacher and student. Normalization statistics for images, tokenizer for text. Once these diverge, the teacher's output is the answer to a different input, and no amount of matching recovers accuracy

The question that comes up in design reviews is "why raise the temperature, and what is the $T^2$ for?" The answer has two halves: you raise it to expose the runners-up, and you multiply by $T^2$ to restore the gradient the temperature shrank, so that temperature and learning rate stay independent knobs. Being able to separate those two shows you understand the equation and not just the recipe.

Summary

• The recipes differ in where you match the teacher: output, intermediate features, attention, or yourself
• Output distillation demands the least and generalizes furthest; feature and attention pay off within an architecture family; self-distillation needs no teacher but delivers gently
• Choose on constraints before accuracy: can you see inside the teacher, do the layers pair up, do the head counts match
• Implementation failures cluster in known places: log_softmax, batchmean, T², eval(), loss scale