All posts
AI & ML

Knowledge Distillation Is a Data Problem Wearing a Loss Function's Clothes

Most explainers open with temperature and KL divergence, as if distillation were a clever loss trick. After building generative models where a small student imitates a stronger source, I think the loss is the least interesting part. The real work is choosing what the teacher generates, catching where its confident wrongness gets copied verbatim, and knowing the three failure modes tutorials never mention. Here is the version I wish I had read first.

Dhileep Kumar7 min read
Knowledge Distillation Is a Data Problem Wearing a Loss Function's Clothes

Open any distillation write-up and you get the same three beats: a big teacher, a small student, and a softmax with a temperature knob that softens the teacher's probabilities so the student can learn from the 'dark knowledge' in its near-misses. All of that is true, and I will explain it. But it buries the lede. In practice, the loss function is the part you write once and never touch again. The part that decides whether your student is good or garbage is upstream of the loss entirely: it is the data the teacher produces and the flaws it smuggles into that data. Distillation is a data problem cosplaying as a math problem.

That framing changes what you spend your time on. If you believe distillation is a loss trick, you tune temperature and alpha for a week. If you believe it is a data pipeline, you spend that week auditing what the teacher actually said, and you get a far better student — because every real problem in this game hides in the data, not the math.

The mental model: teacher as a labeler that hands you its uncertainty

Here is the cleanest way to hold it in your head. A normal training label is a single finger pointing at the right answer: this token, this class, done. A teacher model is a labeler that instead hands you a full opinion — 'I am 70 percent sure it is A, 25 percent B, and I would not rule out C. ' That distribution is the product. The single right answer was always available for free; what you are paying the teacher for is the shape of its doubt.

Why does the shape matter? Because it encodes similarity structure the label cannot. A teacher shown a photo of a husky says mostly-dog, a little wolf, definitely-not-teacup. Those relative weights tell the student that dogs and wolves live near each other and teacups do not — a geometry the student would otherwise have to rediscover from scratch. That is the famous 'dark knowledge,' and temperature exists purely to keep the small probabilities from being crushed to zero so the student can actually see them.

You are not paying the teacher for the right answer. You already had that. You are paying it for a well-calibrated account of everything it almost said.

The uncomfortable corollary: if the teacher's uncertainty is wrong — overconfident, biased, hallucinated — you are faithfully teaching the student to be wrong in exactly the same shape. The loss does its job perfectly and the student gets worse. Hold that thought; it is the source of two of the three gotchas below.

How LLM distillation is actually done today (and how the textbook lies)

The textbook version matches token-level distributions directly: you need the teacher's logits, you minimize a divergence between student and teacher probabilities, and it transfers the most signal per example. It is beautiful and it is mostly not what people do, because it requires owning both models and running the teacher in lockstep with training. Most of the recent open-model boom is the cruder cousin: sequence-level distillation. You run the teacher over a big pile of prompts, save its generated outputs as plain text, and fine-tune the student on that synthetic dataset like any other supervised data. No logits, no lockstep, no shared vocabulary required.

That distinction is not academic — it decides your whole engineering plan. Logit distillation is a training-time coupling problem. Sequence-level distillation is a dataset-generation problem, which means every hard-won lesson about data quality, dedup, and filtering from ordinary fine-tuning applies directly. The teacher is just a very expensive, very good synthetic-data vendor.

The loss, and why it is the boring part

When you do have logits, the core is genuinely small: blend two objectives so the student copies the teacher's soft distribution and still gets the real answer right. Written out, it is about six lines that you will paste once and forget:

python
import torch.nn.functional as F

def distill_loss(student_logits, teacher_logits, labels, T=2.0, alpha=0.5):
    # Soft target: match the teacher's softened distribution.
    s_log = F.log_softmax(student_logits / T, dim=-1)
    t_prob = F.softmax(teacher_logits / T, dim=-1)
    soft = F.kl_div(s_log, t_prob, reduction='batchmean') * (T * T)

    # Hard target: still get the real answer right.
    hard = F.cross_entropy(student_logits, labels)

    return alpha * soft + (1.0 - alpha) * hard

The one line people forget is the T-squared multiplier on the soft term. Softening by temperature shrinks the gradients by roughly one over T-squared, so without that rescale your soft loss quietly stops mattering as you raise the temperature, and you spend an afternoon wondering why the teacher signal 'does nothing. ' That is the single most common footgun in the math — and notice it is still just bookkeeping. Everything that determines quality happens before this function is ever called.

A worked decision: distill, fine-tune, or just prompt?

Say you run a support tool that classifies incoming tickets into forty categories. A frontier model does it beautifully at, say, a few cents per call and half a second of latency — illustrative numbers, but the shape is realistic. You are doing a couple million calls a month and both the bill and the latency hurt. What do you actually do?

Walk the ladder. First ask whether a better prompt on a mid-tier model closes the gap — often it does, and you write zero training code. If not, ask whether plain fine-tuning on your historical labeled tickets is enough; you may already have thousands of human-labeled examples sitting in your ticket system. You reach for distillation specifically when you have the hard task nailed by an expensive model but not enough human labels to fine-tune a cheap one directly — the teacher manufactures the labels you are missing. That is exactly this scenario: the frontier model is right, humans labeled only a fraction, so you let the teacher label a hundred thousand tickets and fine-tune a small student on those. The student inherits the frontier model's judgment on your forty categories at a fraction of the cost, and — the quiet win — it is now a plain classifier you can quantize and run on a boring CPU box. As a table, since the real question is where your bottleneck sits:

  • Prompt a mid model — use when the capability gap is small and you have no labeled data. Fastest to ship, nothing to train, but you stay on someone's API and per-call pricing.
  • Fine-tune directly — use when you already have thousands of real human labels for the exact task. No teacher needed; you avoid inheriting a teacher's biases. Breaks down when you lack labeled data or the task is genuinely hard.
  • Distill (sequence-level) — use when an expensive model already nails the task but you lack labels to train a cheap one. The teacher becomes your labeler. Watch the licensing and the fact that you copy the teacher's mistakes wholesale.
  • Distill (logit-level) — use when you own both models and want maximum signal per example. Best fidelity, but you pay in training-time coupling and shared-tokenizer constraints.

Three gotchas the tutorials skip

These actually cost time. Two follow directly from the 'data problem' framing; the third is a plumbing trap I hit on a related project, flagged honestly as adjacent rather than dressed up as a distillation benchmark.

1. The teacher's confident errors become ground truth you cannot debug

In ordinary fine-tuning, a wrong label is noise — the model averages over it. In sequence-level distillation, the teacher's wrong answers arrive with high confidence and internal consistency, so the student does not average them out; it learns them as rules. Worse, they are invisible: the loss goes down, the eval on the teacher's own style of question looks great, and the systematic error only surfaces in production on the slice where the teacher was quietly bad. The fix is unglamorous and lives entirely in the data layer — sample the teacher's outputs, have a human or a second stronger model spot-check a few hundred, and filter before you train. You are QA-ing a dataset, not tuning a loss.

2. A student cannot exceed its teacher — so measure the right ceiling

Distillation transfers capability; it does not create it. The student is bounded above by what the teacher knew, which means benchmarking your student against the teacher is the wrong test — of course it loses. The right test is: does the student match the teacher on the narrow slice you actually deployed for, at the cost and latency you needed? A tightly specialized student often ties a giant general teacher inside its domain and craters everywhere else, and that is a success, not a regression. If your eval set is broad and general, you will conclude distillation 'failed' when it did exactly what it should.

3. 'Loaded successfully' can be a lie — verify the student actually inherited weights

This one I hit for real, on a generative voice model rather than a distilled classifier, so treat it as a transferable plumbing lesson, not a distillation result I measured. When you copy weights between a source and a target model — teacher checkpoint into student init, adapter into base, whatever — a key-name mismatch (a stray module prefix from a compiled or wrapped model) makes a permissive load silently skip the mismatched tensors. Nothing errors. The model loads. It also produces confident nonsense, because half its weights are still random. The general defenses: never load with mismatched keys silently, log how many tensors actually matched versus how many the checkpoint held, and run one fixed input through the model right after loading to confirm the output is sane before you trust a single training step on top of it.

The bottom line

Distillation is how a model too big to deploy becomes one you can run — pair it with quantization and it is the one-two punch behind nearly every small model that feels smarter than it should. But do not spend your week on temperature and alpha. Spend it deciding what the teacher generates, spot-checking that pile for confident lies, evaluating the student against the narrow job you actually gave it, and verifying it truly inherited the weights you think it did. The teacher thinks hard once. Your job is to make sure the student learned the right lessons and not the confident wrong ones — and that job lives in the data, not the divergence.

Share

Enjoyed this?

Get the next deep dive in your inbox. No spam — just the stories worth reading.

Subscribe to the newsletter

Comments