Knowledge Distillation from Scratch — Copying a Big Model into a Small One

No label ever says "dogs look a bit like cats." A trained model's output does — and that gap is what knowledge distillation harvests. Soft labels, temperature-scaled softmax, where the T² factor comes from, sequence-level and synthetic-data distillation, and how DeepSeek-R1 copied an entire reasoning procedure into smaller models.

Distilling the Knowledge in a Neural Network

The label doesn't carry what the teacher knows

Picture a four-way image classifier: dog, cat, truck, airplane. A photo of a dog carries the label "dog" — as numbers, $(1, 0, 0, 0)$. That is the entire contents of the dataset for that image.

Now show the same photo to a large, well-trained model. The output might be $(0.90,\ 0.09,\ 0.005,\ 0.005)$: it's a dog, though a cat isn't out of the question, and it's certainly not a vehicle. The ratios buried in there — dog:cat at roughly 10:1, dog:truck at 180:1 — appear nowhere in the label. Yet they encode something real that the model absorbed from millions of examples: which categories resemble each other and which don't.

Hinton and colleagues called this dark knowledge. Knowledge distillation is the business of transferring it from a large model (the teacher) into a small one (the student).

The analogy: copying answers vs. copying hesitation

Two students work through the same problem set.

The first looks only at the answer key: question 3 is C. Nothing more. The second also reads the grader's note: the answer is C; a lot of people were torn between B and C; A and D are nowhere close.

The second student extracts far more from the identical set of problems, because "B and C are easy to confuse" is information that exists independently of the correct answer. That is the whole intuition, and it cashes out in two practical ways.

More information per example. A one-hot label transmits a single class index; a teacher's output transmits a number for every class. Training converges on less data, and unlabeled data becomes usable — just run the teacher over it and you have soft labels.

It acts as a regularizer. Teaching "definitely dog, everything else exactly zero" pushes the student toward overconfidence. "Probably dog, cat is plausible" doesn't. It's label smoothing, except the smoothing comes from the teacher's judgment rather than a constant.

The mechanism: blurring the distribution with temperature

There's a catch in using the teacher's output directly. A well-trained model tends to output something like $(0.99,\ 0.009,\ 0.0009,\ \dots)$, and the very thing we wanted — whether cat or truck is more dog-like — sits in digits so small that they barely register in the gradient.

The fix is a temperature-scaled softmax: divide the logits (the raw scores that go into the softmax) by $T$ before normalizing.

In words, this takes each logit $z_i$, divides it by the temperature $T$, and turns the results into probabilities. At $T = 1$ it is the ordinary softmax; a larger $T$ flattens the distribution, and a $T$ near zero collapses it into a one-hot vector.

Here $z_i$ is the logit for class $i$, $T$ is a positive number called the temperature, and $p_i(T)$ is the resulting probability. At $T = 1$ this is the ordinary softmax. Raising $T$ shrinks the gaps between the exponentiated terms, flattening the distribution; pushing $T$ toward zero collapses it into a one-hot vector where only the argmax survives.

Temperature, then, is a dial for how much uncertainty you make the teacher admit. Turn it up to around $T = 4$ and $(0.99, 0.009, \dots)$ relaxes into something like $(0.6, 0.2, 0.13, 0.07)$ — the buried ranking becomes a training signal with real magnitude.

The loss: mixing two objectives

The student is told two things at once: match the ground-truth label, and match the teacher's distribution.

In words, the student is trained on two things at once: match the ground-truth label (the cross-entropy term at temperature 1), and close the distance to the teacher's blurred distribution (the KL term at temperature $T$). $\alpha$ sets how much weight each side gets.

The first term is ordinary supervised learning: $y$ is the true label, $p^S(1)$ is the student's output at temperature 1, and $\mathcal{L}_{\mathrm{CE}}$ is cross-entropy. The second is the distillation term, where $p^T(T)$ and $p^S(T)$ are the teacher's and student's distributions at temperature $T$, and $\mathrm{KL}$ measures how far apart two distributions are (covered in KL Divergence from Scratch). The mixing weight $\alpha$ is often set high — 0.9 is common — to lean on the teacher.

In plain words: keep checking your answers, but at the same time keep closing the distance to the teacher's blurred distribution.

Where the $T^2$ comes from

This is the most commonly botched line in a distillation implementation. Differentiate the soft loss with respect to a student logit and you get

$$\frac{\partial \mathcal{L}_{\mathrm{soft}}}{\partial z_i} = \frac{1}{T}\big(p_i^{S}(T) - p_i^{T}(T)\big)$$

In words, the gradient of the soft loss with respect to a student logit is the gap between the student's and teacher's probabilities at temperature $T$, scaled by $1/T$. Because the probabilities themselves also shrink with $T$, the whole gradient ends up on the order of $1/T^2$, which is why the $T^2$ factor is needed.

and since at high $T$ the probabilities themselves respond to $z_i/T$ to first order, the gradient as a whole shrinks on the order of $1/T^2$. In other words, raising the temperature silently starves the distillation term of gradient. Multiplying by $T^2$ cancels that shrinkage so the balance against the hard-label term stays fixed as you tune $T$. Skip it and every change to $T$ forces you to re-tune $\alpha$.

Writing it in code

In PyTorch the core is a handful of lines.

import torch.nn.functional as F

def distill_loss(student_logits, teacher_logits, labels, T=2.0, alpha=0.9):
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),   # input must be LOG-probs
        F.log_softmax(teacher_logits / T, dim=-1),   # target in log space too
        reduction="batchmean", log_target=True,
    ) * (T ** 2)                                     # undo the gradient shrinkage
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

Three things trip people up. The first argument to F.kl_div is the student's log-probabilities and the second is the teacher — swap them and you have computed a different quantity (reverse KL). The default reduction="mean" divides by the number of elements, not the batch size, so use "batchmean" if you want the textbook KL. And compute the teacher inside torch.no_grad(), or you will quietly train it too.

Three flavors of distillation

Response-based (logit distillation) — what we just wrote. Only the output layer is matched. Cheap to implement, but it carries a hard constraint: the teacher and student must share an identical vocabulary and tokenizer. Logits are a list of numbers indexed by vocabulary entries; different vocabularies leave you with nothing to subtract.

Feature-based — match intermediate hidden states as well. DistilBERT initialized a half-depth student from every other layer of BERT and optimized three losses together: masked-language-modeling, KL against the teacher's output, and cosine similarity between hidden states.

Relation-based — match the relationships between examples (distances, angles) rather than individual outputs. Even when the two representation spaces have different shapes, "A is near B and far from C" survives the transfer.

What changes for language models

Unlike a classifier, a language model emits one token at a time. The straightforward move is a per-token KL (word-level KD), but for generation something cruder often works better: have the teacher generate the text, then train the student on that text as if it were ground truth. Kim and Rush named this sequence-level KD. You stop copying the distribution and start copying samples drawn from it — which is the ancestor of synthetic-data distillation below.

The other language-model-specific problem is exposure bias. At inference the student walks over context it produced itself, but during training it only ever saw context produced by the teacher or the dataset. One deviation and it's standing somewhere it has never seen. The remedy is on-policy distillation: let the student generate, then query the teacher's logits along the sequences the student actually visited (the GKD and MiniLLM line of work).

The direction of the KL matters here too. Forward KL, $\mathrm{KL}(p^T \| p^S)$, forces the student to put mass anywhere the teacher did, so an under-capacity student ends up smeared and noncommittal. Reverse KL chases only the dominant modes: sharper generations, less diversity.

Synthetic-data distillation — no logits required

When the teacher lives behind an API, or the two models use different vocabularies, logits are simply unavailable. What travels instead is text. Prompt the teacher, collect its responses, and train the student on those (prompt, response) pairs as ordinary supervised fine-tuning data. Much of the recent impression that "small models suddenly got smart" comes from exactly this.

There is a well-documented trap. The False Promise of Imitating Proprietary LLMs showed that imitation models reproduce the teacher's style with uncanny accuracy while factuality and reasoning lag well behind. They score well with human raters and with LLM-as-judge setups, yet the gap on knowledge benchmarks barely moves. Only the style got distilled. If you evaluate synthetic-data distillation, choose metrics that can tell those two apart.

DeepSeek-R1 — distilling the procedure itself

DeepSeek-R1, released in 2025, is the cleanest example of this line. R1 was trained with reinforcement learning until it wrote out long chains of thought before answering. DeepSeek then had R1 solve a large batch of problems, curated the good outputs into a training set of several hundred thousand samples, and fine-tuned various sizes of Qwen2.5 and Llama on it with supervised fine-tuning only — the released DeepSeek-R1-Distill-Qwen and -Llama models involve no RL stage on the student side.

The notable finding reported in the paper is that distilling from a large reasoning model beat running large-scale RL directly on the same small model. The intuition is straightforward: RL has to discover good reasoning trajectories by trial and error, and a small model lacks the search capacity to stumble onto them. Distillation simply hands over trajectories that have already been found.

And what is being copied here isn't a probability distribution — it's the procedure: where to split into cases, where to check the work, where to say "wait, that's wrong" and back up. Not a single line of logit KL is involved. The word "distillation" has stretched from numbers in an output layer to patterns of behavior. R1's own lineage is covered in The DeepSeek Family.

How this shows up on the job

Who, and when: ML engineers and inference-platform teams reach for distillation first when accuracy is already adequate but latency or cost-per-token isn't. Where quantization compresses by making the same model's numbers coarser, distillation compresses by moving the contents into a different, smaller model. They compose — and the usual order is distill, then quantize.

What you actually tune: temperature $T$ (2–10 for image classifiers; usually kept to 1–2 for LLMs, whose token distributions are already broad), the mixing weight $\alpha$, and the student's depth and hidden size. On the tooling side: Hugging Face TRL's GKDTrainer for on-policy distillation and SFTTrainer for the synthetic-data route, vLLM batch inference for bulk teacher generation, and — if you need logits — a check on the logprobs limit your serving stack exposes. For verifiable tasks like math and code, always put rejection sampling in front, filtering teacher outputs by correctness.

Failure modes that bite:

• Letting the teacher grade the exam. A student trained on a teacher's outputs will naturally please that teacher as a judge. Pair every such evaluation with an independent held-out set and either human review or executable checks (does the test suite pass?).
• Leaving the temperature on at inference. $T$ belongs to the training loss only; the deployed student runs at temperature 1. Training-time temperature leaking into the serving config is a real and quiet bug.
• Losing general capability. Long training on narrow synthetic data erodes what the model could already do. Mix in original data, lower the learning rate, or confine updates to a LoRA adapter — see catastrophic forgetting.
• Writing logit distillation across mismatched vocabularies. Across model families the logit dimensions don't line up in size or meaning. Switch to text-level distillation instead of fighting it.
• Licensing and terms of service. Most commercial APIs explicitly forbid using their outputs to train a competing model, and open-weight licenses often govern derivative naming and redistribution. Whether it is technically possible is a separate question from whether you may — check first.

How it gets asked in interviews and design reviews: Why does a soft label carry more information than a one-hot label? What happens as you raise the temperature, and why is the $T^2$ factor needed? Forward KL or reverse KL — which, and why? Each answer is one equation from this article translated back into words.

Summary

• Distillation transfers the inter-class similarity structure that ground-truth labels never contain
• Temperature-scaled softmax blurs the distribution, KL closes the gap, and $T^2$ compensates for the gradient shrinkage that the blurring causes
• For language models, copying the teacher's generated text has largely displaced copying its distribution
• DeepSeek-R1's distillation is the extreme case: what transfers is a reasoning procedure, not a set of probabilities
• But "only the style transferred" is a real outcome. Unless the evaluation is independent of the teacher, you can finish the project convinced of an improvement that isn't there