Distillation in practice

Distillation is training a small model to behave like a large one, either by matching its full output distribution or, for language models, by generating outputs from the large model on your real inputs, filtering them, and fine-tuning the small one on what survives. This section does both, then insists on the comparison that decides whether it was worth doing: against a smaller model with a better prompt, and against a mid-sized model, either of which may reach the same quality with no training at all. Skipping the filter is the other common way these projects fail, because the student copies the teacher's mistakes as faithfully as its successes.

Soft targets#

A teacher model's full output distribution carries more information than the correct label alone. For an image of a cat, the teacher might put a small probability on 'dog' and almost none on 'truck'. Training a student to match that distribution teaches it how classes relate. Raising the softmax temperature spreads the distribution out so this information is visible.

import torch.nn.functional as F

def distill_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.5):
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.softmax(teacher_logits / T, dim=-1),
        reduction="batchmean",
    ) * T * T                                   # keeps gradient size comparable across T
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1 - alpha) * hard

The T squared factor is the detail that gets dropped and then quietly breaks the recipe. Dividing logits by T shrinks the gradient of the soft term by roughly a factor of T squared, so without it the soft loss fades out as you raise the temperature and alpha stops meaning what you think it means. Put it in.

How much extra information is really in those soft targets is measurable, and the answer decides whether the technique is worth setting up. Train a width-512 teacher on a six-class problem, then look at what its output distributions contain at different temperatures.

import torch, torch.nn as nn, numpy as np
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler

X, y = make_classification(n_samples=20000, n_features=40, n_informative=20, n_classes=6,
                           n_clusters_per_class=2, class_sep=1.0, flip_y=0.02, random_state=1)
scaler = StandardScaler().fit(X[:8000])
as_t = lambda a: torch.tensor(scaler.transform(a), dtype=torch.float32)
Xtr, ytr = as_t(X[:8000]), torch.tensor(y[:8000])       # the teacher's training data
Xun = as_t(X[8000:14000])                               # unlabelled inputs nobody trained on
Xte, yte = as_t(X[14000:18000]), torch.tensor(y[14000:18000])

def build(width):
    return nn.Sequential(nn.Linear(40, width), nn.ReLU(),
                         nn.Linear(width, width), nn.ReLU(), nn.Linear(width, 6))

def new(width, seed=0):
    """Fixed initialisation, so two runs differ only in what we changed."""
    torch.manual_seed(seed)
    return build(width)

def fit(m, Xa, ya, epochs=25, lr=3e-3, seed=0, teacher=None, T=4.0, alpha=0.7):
    torch.manual_seed(seed)
    opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=1e-2)
    g = torch.Generator().manual_seed(seed)
    if teacher is not None:
        with torch.no_grad():
            tlogits = teacher(Xa)
    for _ in range(epochs):
        m.train()
        perm = torch.randperm(len(Xa), generator=g)
        for i in range(0, len(Xa), 128):
            idx = perm[i:i + 128]
            opt.zero_grad()
            out = m(Xa[idx])
            loss = nn.functional.cross_entropy(out, ya[idx])
            if teacher is not None:
                loss = alpha * distill_loss_only(out, tlogits[idx], T) + (1 - alpha) * loss
            loss.backward()
            opt.step()
    m.eval()
    return m

def distill_loss_only(student_logits, teacher_logits, T):
    return F.kl_div(F.log_softmax(student_logits / T, dim=-1),
                    F.softmax(teacher_logits / T, dim=-1),
                    reduction="batchmean") * T * T

def accuracy(m):
    with torch.no_grad():
        return (m(Xte).argmax(1) == yte).float().mean().item()

teacher = fit(new(512), Xtr, ytr, epochs=30)
print(round(accuracy(teacher), 4))                      # 0.8727

with torch.no_grad():
    tlogits = teacher(Xte)
for T in (1.0, 2.0, 4.0):
    p = torch.softmax(tlogits / T, dim=-1)
    H = -(p * p.clamp_min(1e-12).log()).sum(-1).mean()
    print(T, round(H.item(), 4), round(p.max(-1).values.mean().item(), 4))
# 1.0  0.0993  0.9601      almost all the mass is on one class
# 2.0  0.2228  0.9165
# 4.0  0.5304  0.8150      now there is a distribution to learn from
# for reference, the most entropy six classes can carry is ln(6) = 1.7918

This is also the fastest way to know whether distillation is worth trying at all. A binary classifier's teacher has almost nothing to say beyond its top choice, and distilling one usually buys nothing measurable. The more classes, and the more genuinely confusable they are, the more there is in the distribution.

Does it help when you already have the labels#

Mostly not, and it is better to find that out in ten lines than after a fortnight. Train a small student on the same 8,000 labelled examples the teacher saw, with and without the soft targets, three seeds each.

scratch = [accuracy(fit(new(48, s), Xtr, ytr, seed=s)) for s in range(3)]
distil = [accuracy(fit(new(48, s), Xtr, ytr, seed=s, teacher=teacher)) for s in range(3)]
print([round(v, 4) for v in scratch], round(float(np.mean(scratch)), 4))
print([round(v, 4) for v in distil], round(float(np.mean(distil)), 4))
# [0.819,  0.8255, 0.8215]  0.8220     student trained on the labels
# [0.8255, 0.8298, 0.833 ]  0.8294     same student, plus the teacher

The case it is actually for#

Distillation is a way of turning a teacher plus unlabelled inputs into labelled data. Give the student 500 labels and 6,000 inputs nobody has labelled, and let the teacher fill them in.

Xs, ys = Xtr[:500], ytr[:500]
with torch.no_grad():
    yun = teacher(Xun).argmax(1)                # the teacher's guesses, right 86.9% of the time
Xboth, yboth = torch.cat([Xs, Xun]), torch.cat([ys, yun])

few = [accuracy(fit(new(48, s), Xs, ys, epochs=60, seed=s)) for s in range(3)]
hard = [accuracy(fit(new(48, s), Xboth, yboth, seed=s)) for s in range(3)]
soft = [accuracy(fit(new(48, s), Xboth, yboth, seed=s, teacher=teacher)) for s in range(3)]
for r in (few, hard, soft):
    print(round(float(np.mean(r)), 4))

500 real labels only                         0.4854
500 + 6,000 of the teacher's hard guesses    0.7949
500 + 6,000 of the teacher's distributions   0.8205
the same student with all 8,000 real labels  0.8220

Two things in that table are the reason this technique is used at all. The teacher's labels were 86.9% correct and the student trained on them reached 79.5%, which is far better than the 48.5% it could manage from the labels it had. And the soft targets were worth another 2.6 points over the same guesses as hard labels, because a distribution tells the student which mistakes are near misses. The student ends up almost as good as if somebody had labelled 8,000 examples by hand, which is exactly the trade an LLM distillation project is making: the teacher's time is cheap, and human labelling is not.

Distilling language models with generated data#

For LLMs, the common form is simpler: have a large model produce high-quality outputs for many inputs from your task, filter them, and fine-tune a small model on the result. The small model learns to do your task the way the large one does, at a fraction of the serving cost.

• Use real inputs from your task, with the variety of production traffic.
• Filter the teacher's outputs with your eval checks. Distilling mistakes teaches mistakes.
• Hold out a test set the teacher never saw, and compare student and teacher on it.

Choosing the student#

The student's size is the decision that determines whether the project is worth doing, and it has to be swept rather than guessed. Running the same comparison at six widths on the problem above shows something the single-width experiment could not.

width   parameters   from labels   distilled   difference
   12          726      0.6563       0.6498      -0.0065
   24        1,734      0.7688       0.7617      -0.0071
   48        4,614      0.8220       0.8294      +0.0074
   96       13,830      0.8307       0.8465      +0.0158
  192       46,086      0.8512       0.8475      -0.0037
  384      165,894      0.8642       0.8575      -0.0067

teacher, width 512, 286,726 parameters:            0.8727

There is a band where it pays and it is narrower than the technique's reputation suggests. A student too small cannot represent what the teacher is telling it, so the soft targets are a distraction from the labels. A student large enough to learn the task well does better from the real labels than from an imperfect teacher, because the teacher is only 87% right and the labels are 100% right. In between, where the student has the capacity to use extra information but not enough to find it alone, distillation is worth one and a half points. Sweep the widths, and note that the sign of the effect changes, so a single run at an arbitrary width can tell you that distillation helps, or that it hurts, depending on which width you picked.

For language models the same decision is made between released models rather than by choosing a width, and the comparison to run is fixed: the large model prompted, the small model prompted, the small model distilled, and a mid-sized model prompted. The fourth is the one people skip and it frequently wins, because it needs no training run and no maintenance. A distilled model is a training pipeline you now own: it has to be rebuilt when the task changes, re-evaluated when the teacher is upgraded, and hosted by somebody.

Also decide what the student is allowed to be worse at. Distillation transfers the teacher's behaviour on the distribution you sampled, so the student inherits the common cases and not the general ability that handled the rare ones. Route what the student has not seen to the teacher rather than pretending the gap does not exist: a confidence threshold on the student, or a simple check on input length or type, recovers most of the loss for a small fraction of the traffic and therefore a small fraction of the cost.

Check the terms#

Some model providers' terms restrict using their outputs to train other models, and open model licences vary. Check before building a product on distilled weights.

The comparison that decides it#

On your task eval:              quality   cost per 1k requests
large model, prompted             0.91          high
small model, prompted             0.78          low
small model, distilled            0.88          low      <- worth it only if this beats the row above clearly