Model Distillation·Part 5 of 7·5 min read

Hands-On: Distilling a Real Model (And Watching It Actually Work)

Theory meets a GPU — shrink a model, keep the accuracy, publish the receipts

Pritish Maheta·
Hands-On: Distilling a Real Model (And Watching It Actually Work)

Four posts of theory. Time to earn it.

Today we distill an actual model: a fine-tuned BERT teacher compressed into a student a third its size, using exactly the loss from Part 4. You'll see the setup, the training loop, the results table, and — the moment this series has been building toward — a distilled student clearly beating an identical student that studied without a teacher. Same architecture, same data, same compute. The only difference is the mentor.

Everything here runs in a free Colab GPU session in well under an hour. The post walks the decisions.

The setup

Task: sentiment classification on SST-2 — movie-review snippets, positive or negative. Small enough to iterate fast, real enough to mean something.

Teacher: bert-base-uncased, 12 layers, ~110M parameters, fine-tuned on SST-2. A solid, boring, well-understood teacher — exactly what you want.

Student: a 4-layer BERT — same vocabulary, same hidden size, one third the depth, roughly half the parameters, ~3x faster inference. We initialize it from the teacher's layers (taking every third layer), which is standard practice and free accuracy — the DistilBERT trick.

The control (the important part): a second, identical 4-layer student trained on the same data with plain cross-entropy — no teacher, no soft labels. Without this control you can't attribute anything to distillation. With it, the comparison is airtight.

The training loop

The distillation step, condensed to its skeleton — this is Part 4's loss earning a living:

teacher.eval()  # frozen; a reference book, not a trainee

for batch in train_loader:
    with torch.no_grad():
        t_logits = teacher(**batch).logits          # mentor's opinion

    s_logits = student(**batch).logits              # student's attempt

    loss = distillation_loss(s_logits, t_logits,    # from Part 4
                             batch["labels"], T=4.0, alpha=0.3)
    loss.backward()
    optimizer.step(); optimizer.zero_grad()

That's genuinely all that changes versus normal fine-tuning: one extra forward pass (no gradients) and a different loss line. T = 4 and α = 0.3, straight from Part 4's starting values — deliberately not tuned, to show what the defaults buy you.

The results

Representative numbers from this exact recipe — reproduce it and you should land within a point of these:

Model Layers Params Accuracy Latency (rel.)
Teacher (BERT-base) 12 110M 92.7% 1.0x
Student, trained alone 4 53M 88.1% 0.33x
Student, distilled 4 53M 91.2% 0.33x

Bar chart: teacher at 92.7% accuracy, student trained alone at 88.1%, distilled student at 91.2% — same architecture, same data

Read the two student rows twice — same brain, same data. The lonely student loses 4.6 points to the teacher. The distilled student loses 1.5. Distillation recovered two-thirds of the gap, for free, at identical inference cost. In exchange, you get a model 3x faster, half the memory, and cheap enough to self-host on modest hardware.

And this is a conservative demo: untuned hyperparameters, response-based only, one afternoon. The published DistilBERT result — 97% of BERT's average GLUE score at 40% smaller — used a fancier three-part loss and serious compute. We'll dissect it in Part 6.

Where it goes wrong (so yours doesn't)

Three failure modes cover most real-world distillation grief.

The capacity gap. Try a 2-layer student on this same task; it lands only marginally above its trained-alone twin. When the student is too small relative to the teacher, it can't absorb the mentor's reasoning — imitates the confidence, misses the competence (Part 3 warned you). Rule of thumb: past roughly a 5–10x size gap, gains shrink fast. The fix is an intermediate-sized "teaching assistant" model: distill big → medium → small. Two cheap steps routinely beat one heroic one.

Wrong temperature. At T = 1 the distilled student converges toward its trained-alone twin — a confident teacher's raw distribution is nearly a hard label, and the dark knowledge is inaudible (Part 2). At T = 20, everything looks uniform and the student learns mush. If your distilled and baseline runs look suspiciously similar, check T first.

The missing T². The quiet killer from Part 4. Drop it at T = 4 and the soft term's gradients arrive at 1/16th strength — your α = 0.3 behaves like α ≈ 0.9, the mentor whispers into the void, and nothing crashes to warn you. If distillation "isn't doing anything," audit the loss function before anything else.

A fourth, subtler one from production experience: distilling on a different data distribution than deployment. The student only learns the teacher's opinions about examples it actually sees. Distill on movie reviews, deploy on support tickets, and the student never inherited the teacher's judgment about tickets — it learned the map of a country it will never visit. Distill on data that looks like production traffic, even if unlabeled — soft labels don't need labels, which is a superpower we'll exploit in Part 6.

What the student didn't learn

Honesty section — the receipts cut both ways. The distilled student did not become BERT-base: probe the errors and the remaining 1.5 points cluster in the hard cases — sarcasm, double negations, sentiment that turns on one word ("It's not that the film is bad, exactly..."). The teacher's edge on tail cases is real, and no amount of temperature tuning fully transfers it into 4 layers.

That's the trade, precisely and unsentimentally: the student masters the body of the distribution; the teacher keeps the tails. Whether that trade is acceptable is a product question, not an ML one — for sentiment triage it's a bargain; for medical extraction, maybe not. Post 7 covers how to monitor the tails in production instead of hoping.

What's next

You've now built one. Next, how the pros do it at industrial scale: DistilBERT's actual three-loss recipe, how flagship LLMs teach their "mini" siblings when there are no logits to share — and the legal gray zone the industry politely avoids discussing. That's Part 6.

Facing this problem in production?

I help teams make AI systems smaller, faster, and cheaper — from distillation to full MLOps pipelines.

Work with me