Machine Learning from Scratch

The Big Question

Imagine a 400B-parameter teacher model that answers questions well, but needs to run on phones, browsers, embedded devices, or edge hardware. Deploying the teacher everywhere is too slow and expensive.

Why train an expensive model for every device when a smaller model can learn to imitate it?

Distillation trains a student model on the behaviour of a stronger teacher. The student is not compressed by copying parameters; it is compressed by learning the teacher's responses, probabilities, and sometimes internal representations.

Core Intuition

A professor does not merely say “the answer is C.” They also explain why A is close, why B is wrong, and why D is implausible. That extra information teaches more than a hard label.

If the teacher predicts dog 0.71, wolf 0.22, fox 0.05, and cat 0.02, the hard label says only dog. The soft distribution reveals dark knowledge: wolf is much more plausible than cat.

Interactive Demo

Distillation lab

Teacher distribution

Dog50.7%

Wolf27.8%

Fox13.1%

Cat8.4%

Temperature2.0

Hard label versus soft label

Hard label

Dog = 1, all others = 0

Soft label

Dog is most likely, but wolf is still plausible.

Student learning

Distillation steps55%Match hidden features too

Dog44.1%

Wolf23.4%

Fox18.9%

Cat13.6%

Student size

Accuracy

82%

KL gap

0.03

Distillation Loss

Ordinary cross entropy with hard labels is:

L=-\sum_i y_i\log p_i

For distillation, the teacher produces probabilities $q_i$ and the student produces $p_i$ . A common loss uses KL divergence:

L_{\text{KD}}=T^2\sum_i q_i^{(T)}\log\frac{q_i^{(T)}}{p_i^{(T)}}

The softened probabilities come from temperature-scaled softmax:

p_i(T)=\frac{e^{z_i/T}}{\sum_j e^{z_j/T}}

Larger $T$ flattens the distribution and reveals secondary preferences. The $T^2$ factor is commonly used to keep gradient scales comparable as temperature changes.

KL appears naturally because:

KL(q\Vert p)=\sum_i q_i\log\frac{q_i}{p_i}=H(q,p)-H(q)

The teacher entropy $H(q)$ is constant with respect to the student, so minimising cross entropy with soft labels also minimises KL.

Types Of Distillation

Output distillation

Student matches teacher probabilities or logits.

Feature distillation

Student matches hidden states: ||h_T - h_S||^2.

Attention distillation

Student matches teacher attention patterns.

Sequence distillation

Teacher generates full target sequences for student training.

Self-distillation

A model or later checkpoint teaches another version of itself.

Online distillation

Teacher and student are trained together.

Implementation

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.7):
    hard = F.cross_entropy(student_logits, labels)

    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    soft = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") * temperature**2

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

with torch.no_grad():
    teacher_logits = teacher(inputs)

student_logits = student(inputs)
loss = distillation_loss(student_logits, teacher_logits, labels)

For LLMs, the teacher may generate synthetic instruction examples or full response sequences. The student then trains on those examples, often with ordinary supervised fine-tuning plus optional soft-logit matching when logits are available.

Interview Discussion

What are soft labels?

The full teacher probability distribution, not just the correct class.

What is dark knowledge?

Information in the teacher's relative probabilities for incorrect classes.

Why use temperature?

It softens the distribution so secondary preferences are easier for the student to learn.

What can go wrong?

A poor teacher, capacity mismatch, bad temperature, distribution mismatch, and inherited hallucinations or bias.

Active Recall

1. Why do soft labels contain more information than hard labels?

2. Write the temperature-scaled softmax.

3. Explain why minimising soft-label cross entropy minimises KL up to a constant.

4. Name three forms of distillation beyond output matching.