Knowledge Distillation

"My teacher was a hundred times my size and twice as sure of everything. I kept only the parts of its doubt that were useful, and now I answer in a tenth of the time."

A Student Model That Fit in One GPU

The previous section reduced a model's cost by removing weights from a fixed architecture (Section 22.3, pruning and sparsity). Distillation takes a different route to the same goal: instead of shrinking one model, it trains a separate, smaller model from scratch to imitate the large one. Quantization and pruning are constrained to stay near the original network; distillation is free to choose any student architecture that fits the deployment budget, then teach it to behave like the expensive teacher. That freedom is why distillation routinely produces the steepest accuracy-per-parameter curve of the three, and why the largest frontier labs ship distilled variants of their flagship models for everyday serving.

Like everything in this chapter, distillation is single-node economics: it changes what one replica costs to run. We treat it here as a labeled per-node prerequisite, because you cannot size or price a serving fleet (Chapter 23) until you know the unit cost of the model each node holds. The figure below frames the technique and the rest of the section unpacks it.

1. Why a Soft Target Teaches More Than a Hard Label Beginner

A labeled training example gives the model one bit of supervision per class: this image is a cat, not a dog, not a car. A trained teacher, shown the same image, returns something richer: it might assign 0.55 to cat, 0.42 to dog, and 0.03 to car. That distribution encodes a similarity structure the hard label throws away. It says a cat is easily confused with a dog and almost never with a car, a fact the teacher learned from millions of examples and that a small student could never infer from a handful of one-hot labels. Geoffrey Hinton and colleagues named this discarded information dark knowledge, and the central claim of distillation is that learning to reproduce it is what lets a small model punch above its parameter count.

The wrinkle is that a well-trained teacher is often overconfident: it places 0.999 on the right class and spreads almost nothing over the rest, so the informative structure is squashed into digits the student barely sees. Distillation fixes this by raising a temperature $T$ inside the softmax, flattening the distribution so the relative weights of the non-target classes become legible. The student is then trained to match this softened distribution rather than the hard label, and the temperature is the dial that controls how much of the teacher's nuance is exposed.

2. The Distillation Loss and Temperature Intermediate

Make the temperature precise. Given a teacher's logits $z^{\mathcal{T}}$ and a student's logits $z^{\mathcal{S}}$ over $C$ classes, the temperature-scaled softmax is

At $T = 1$ this is the ordinary softmax; as $T$ grows, the distribution flattens toward uniform and the relative weights of the small probabilities are amplified. The student is trained to minimize a convex combination of two terms: a distillation term that matches the teacher's softened distribution, and an ordinary term that matches the true hard label when one is available,

Two details earn their place. The mixing weight $\alpha$ trades the cheap, abundant soft targets against the scarce, authoritative hard labels; pure distillation is $\alpha = 1$. The $T^2$ factor rescales the soft-target gradient, which otherwise shrinks like $1/T^2$ under temperature scaling, so that the two loss terms keep comparable magnitudes as $T$ changes. With these in place the student's gradient is dominated, in the high-temperature limit, by the difference $z^{\mathcal{S}} - z^{\mathcal{T}}$ between student and teacher logits, which is why distillation is sometimes described as logit matching with a soft floor.

3. From Scratch: A Distilled Student Beats Its Same-Size Twin Intermediate

The claim that soft targets buy generalization a hard-label model of the same size cannot reach is easy to demonstrate with no deep-learning machinery at all. The program below builds a nonlinear three-class problem, trains a large teacher on a big labeled corpus, then trains two small students of identical size: one on the handful of hard labels we can afford, the other on the teacher's softened predictions over a larger unlabeled transfer set. The student is a tiny multinomial classifier using only a fraction of the teacher's features, so it is genuinely capacity-limited, exactly the deployment situation distillation is meant for.

import numpy as np

rng = np.random.default_rng(3)

def make_data(n):                                  # nonlinear 3-class spiral rule
    Z = rng.standard_normal((n, 2)) * 1.3
    r = np.sqrt((Z**2).sum(1))
    ang = np.arctan2(Z[:, 1], Z[:, 0])
    y = (np.floor((ang + np.pi) / (2 * np.pi) * 3 + 0.9 * r)).astype(int) % 3
    return Z, y

def softmax(Z, T=1.0):                             # temperature-scaled softmax
    Z = (Z / T); Z = Z - Z.max(1, keepdims=True)
    E = np.exp(Z); return E / E.sum(1, keepdims=True)

def onehot(y, C=3):
    O = np.zeros((y.size, C)); O[np.arange(y.size), y] = 1.0; return O

def train(F, Targets, epochs=2500, lr=0.8, l2=1e-3):   # fit to (soft or hard) Targets
    W = np.zeros((F.shape[1], Targets.shape[1])); b = np.zeros(Targets.shape[1])
    n = F.shape[0]
    for _ in range(epochs):
        G = (softmax(F @ W + b) - Targets) / n     # cross-entropy gradient
        W -= lr * (F.T @ G + l2 * W); b -= lr * G.sum(0)
    return W, b

acc = lambda F, W, b, y: (np.argmax(F @ W + b, 1) == y).mean()

D = 600                                            # shared random-feature lift
W1 = rng.standard_normal((2, D)) * 1.6
b1 = rng.standard_normal(D) * 1.6
feat = lambda X: np.tanh(X @ W1 + b1)

# TEACHER: large model trained on a big labeled corpus (the expensive asset)
Xbig, ybig = make_data(8000)
Wte, bte = train(feat(Xbig), onehot(ybig))

d = 16                                             # student sees only d features
Xte, yte = make_data(20000)
Fte = feat(Xte)[:, :d]

Xlab, ylab = make_data(40)                         # all the hard labels we can afford
Xtr = np.vstack([Xlab, make_data(2000)[0]])        # + a large UNLABELED transfer set
Ftr_big = feat(Xtr); Ftr_s = Ftr_big[:, :d]

# Student A: hard labels only, on the 40 labeled points
WhA, bhA = train(feat(Xlab)[:, :d], onehot(ylab))
student_hard = acc(Fte, WhA, bhA, yte)

# Student B: distilled. Teacher softly labels the whole transfer set at temperature T.
T = 4.0
soft = softmax(Ftr_big @ Wte + bte, T=T)
WhB, bhB = train(Ftr_s, soft)
student_soft = acc(Fte, WhB, bhB, yte)

teacher_test = acc(feat(Xte), Wte, bte, yte)
pT = Wte.size + bte.size; pS = WhA.size + bhA.size
print(f"teacher params (W,b)     : {pT:,}")
print(f"student params (W,b)     : {pS:,}  ({pS/pT:.1%} of teacher)")
print(f"distillation temperature : {T}")
print("-" * 46)
print(f"teacher test accuracy    : {teacher_test:.3f}")
print(f"student, hard labels     : {student_hard:.3f}")
print(f"student, distilled (soft): {student_soft:.3f}")
print(f"gap to teacher recovered : {(student_soft-student_hard)/(teacher_test-student_hard):.0%}")

teacher params (W,b)     : 1,803
student params (W,b)     : 51  (2.8% of teacher)
distillation temperature : 4.0
----------------------------------------------
teacher test accuracy    : 0.852
student, hard labels     : 0.585
student, distilled (soft): 0.745
gap to teacher recovered : 60%

The two students are the same model with the same parameter count and the same optimizer; the only difference is the target they were trained against. The hard-label student, starved of data, lands at 0.585. The distilled student, trained on the teacher's softened distribution over a transfer set it never had labels for, reaches 0.745 and closes most of the distance to the teacher. The dark knowledge in those softened probabilities did real work: it taught the small model the similarity structure of the problem that its scarce hard labels could not. This is the entire value proposition of distillation in one experiment, and it scales from this toy classifier up to billion-parameter language models.

4. Variants, Cost, and When Distillation Wins Advanced

The soft-target recipe above is the response-based variant: the student matches the teacher's final output distribution. Two other families extend it. Feature-based distillation adds terms that align the student's intermediate representations with the teacher's, so the student learns not just what the teacher answers but how it internally arrives there; this helps when the student architecture differs enough that matching only the final layer leaves too little signal. Sequence-level distillation is the form that matters for generative models: instead of matching the teacher's per-token distribution, the student is trained on whole sequences the teacher generates, so it learns to reproduce the teacher's outputs as coherent wholes rather than token by token. Modern LLM distillation leans heavily on this last form, training students on text the teacher writes.

That points at the one real cost of distillation, the cost that ties it back to the distributed theme of this book. Generating the teacher's targets is itself a large inference job: to distill on millions of examples you must run the expensive teacher over all of them first. For a frontier LLM teacher this is a serious distributed-inference workload in its own right, the subject of Chapter 23, and it is the reason distillation is an offline data-generation pipeline, not a training trick you bolt on for free. The economics still favor it because the generation cost is paid once while the serving saving recurs on every replica of the fleet, but the bill is front-loaded onto the very distributed-inference machinery the student is meant to lighten.

When does distillation beat quantization and pruning? Quantization keeps the same architecture and shrinks each number; pruning keeps the architecture and removes some numbers; both are bounded by how far you can perturb a fixed network before it breaks. Distillation is unbounded in that sense: the student can be an entirely different, smaller architecture, so when you need a four- or ten-fold size reduction rather than two-fold, distillation is usually the only one of the three that reaches it without collapse. The cost is that it requires a full training run and the teacher's generated data, whereas post-training quantization needs neither. In practice the three compose: distill to a smaller architecture, then quantize and prune the student, stacking the savings.

# pip install textbrewer
import textbrewer
from textbrewer import GeneralDistiller, TrainingConfig, DistillationConfig

distill_cfg = DistillationConfig(
    temperature=4.0,                 # the T from the loss in Section 2
    hard_label_weight=0.0,           # 1 - alpha; here pure soft-target distillation
    kd_loss_type="ce",               # cross-entropy / KL on softened logits
)
train_cfg = TrainingConfig()
distiller = GeneralDistiller(        # handles the teacher forward pass + combined loss
    train_config=train_cfg, distill_config=distill_cfg,
    model_T=teacher, model_S=student,
    adaptor_T=adaptor, adaptor_S=adaptor)

with distiller:                      # one call replaces the hand-written loop above
    distiller.train(optimizer, train_loader, num_epochs=3, callback=None)

The takeaway for the rest of Part V is that distillation changes the single most important number a serving fleet is built around: the size of the model on each node. The sections that follow attack the runtime cost of whatever model the node ends up holding. The next one is the most consequential for generative serving, the memory the model spends remembering the conversation so far.