Pruning & Distillation
Pruning removes weights, and only speeds things up when it removes them in shapes the hardware can skip. Distillation trains a small model on a large model's outputs, and inherits everything the large model believed.
The problem, the obvious approach, and why it breaks
Every lesson starts where the work starts: someone has a problem, and the first model that comes to mind looks fine offline.
The pruned model is ninety percent sparse and no faster; the distilled model is fast and wrong where the teacher was uncertain. What did each technique actually do?
A search team pruned their ranking network to a high sparsity to cut serving cost and saw no latency change. In parallel they distilled it into a model a quarter the size for a low-latency tier, and the small model ranks well on popular queries and poorly on the long tail — the queries the search team was most worried about.
Pruning removes weights, so the model has fewer operations and runs faster in proportion to sparsity. Distillation trains a small model on the big model's outputs, so the small model learns what the big one knows.
The pruned matrices are the same shape with zeros in them. A dense kernel does the same number of multiply-adds whether an operand is zero or not, so serving cost is unchanged; only the compressed file got smaller.
- The pruned matrices are the same shape with zeros in them. A dense kernel does the same number of multiply-adds whether an operand is zero or not, so serving cost is unchanged; only the compressed file got smaller.
- Sparse kernels exist and need either very high sparsity or structured sparsity — whole rows, blocks, or a fixed pattern the hardware supports — to beat the dense kernel. Scattered zeros at moderate sparsity lose to dense arithmetic on nearly every device.
- The student learned what the teacher *output* on the corpus it was shown. On the tail, the teacher's scores were noisy and the corpus was thin, so the student learned little and generalises worse than the teacher did — the capacity gap shows exactly where the teacher's signal was weakest.
- Both were signed off on aggregate ranking metrics, where head queries dominate, and both looked fine.
What is being predicted, and from what data
This domain leads with these two. A target nobody defined precisely is a label nobody can trust, and a dataset nobody can describe is a model nobody can debug.
- The surrounding model scores query–document pairs for ranking; the target of pruning is lower serving cost at equal ranking quality, and the target of distillation is a much faster model whose ranking quality on every query segment stays within a stated tolerance of the teacher.
- Both decisions are per-segment comparisons against the original, and the long-tail segment is the one that matters and the one with the fewest examples.
- The pruned model has most of its weights set to zero by magnitude, scattered across every matrix; the runtime multiplies dense matrices and does not know the zeros are there.
- The student was trained on the teacher's scores for a large sample of logged queries, which is dominated by popular queries; the long tail is a small share of the sample and the teacher is least confident there.
- The segment evaluation — head, torso, tail queries — exists from the teacher's launch.
How it actually works
Precisely enough to predict its behaviour — not a framework API.
- Pruning chooses weights to remove by a criterion — magnitude is the simplest — and typically fine-tunes afterwards to recover. Unstructured pruning removes individual weights anywhere, reaching high sparsity at little aggregate loss; structured pruning removes whole units — channels, attention heads, layers, or blocks in a fixed pattern — reaching lower sparsity at more loss but producing smaller dense matrices, which is what actually reduces arithmetic on real hardware.
- Speed comes from operations the kernel does not perform. A dense kernel performs all of them; a sparse kernel skips zeros at a bookkeeping cost that only pays at high or patterned sparsity. That is why unstructured sparsity rarely speeds up real hardware and why some accelerators support only specific patterns (GPU Fundamentals).
- Distillation trains a student to match the teacher's output distribution — the soft probabilities or scores — rather than the hard label alone. The soft outputs carry more information per example: how confident the teacher was, which alternatives it considered plausible. A student with less capacity can learn a smoother function from them than from one-hot labels, and it can learn from unlabelled data the teacher has scored.
- The honest limits: the student cannot exceed the teacher where the teacher is wrong or uncertain, it fits the corpus it is shown, and a capacity gap costs most where the function is most complex — usually the tail. The student also inherits the teacher's version; when the teacher improves, the student does not.
Zeros are multiplied at full speed
A weight matrix with most entries zero is still a matrix of the same shape. The kernel that multiplies it loads every entry and performs every multiply-add; a zero costs the same as any other number. The pruned model is smaller on disk after compression and identical in arithmetic at runtime.
Real speed needs a kernel that skips work, and skipping has overhead: indexing into a sparse representation, irregular memory access, lost vectorisation. Those costs are only paid back at very high sparsity, or when the zeros come in blocks or patterns the hardware supports natively. Structured pruning produces smaller *dense* matrices instead — fewer heads, narrower layers — and dense kernels are fast on every device.
Zero the smallest weights anywhere. High sparsity, small aggregate quality loss, an excellent-looking compression ratio — and the dense kernel does exactly the same work as before.
Remove whole attention heads or channels by an importance criterion, then fine-tune. Lower sparsity, more quality to recover, and a smaller dense model that every kernel runs faster.
Arithmetic time is set by the operations the kernel performs, not by the number of non-zero operands. Removing a head removes its operations; zeroing its weights does not.
Soft labels carry more than hard ones
A hard label says "this document is relevant". The teacher's output says "this document is relevant with this confidence, and these two others are nearly as plausible". The second carries more information per example, and a student with limited capacity learns a smoother and better-generalising function from it — which is the entire mechanism of distillation.
The mechanism also states its limit. The student learns the teacher's function on the shown corpus; where the teacher was uncertain, the signal is weak, and where the corpus was thin, the student sees little. The tail is both.
1import numpy as np2 3def softmax(z, T=1.0):4 z = z / T5 e = np.exp(z - z.max(axis=1, keepdims=True))6 return e / e.sum(axis=1, keepdims=True)7 8def distill_loss(student_logits, teacher_logits, hard_labels, T=4.0, alpha=0.7):9 # soft targets at temperature T expose the teacher's alternatives10 p_teacher = softmax(teacher_logits, T)11 log_p_student = np.log(softmax(student_logits, T) + 1e-12)12 soft = -(p_teacher * log_p_student).sum(axis=1).mean() * (T * T)13 # hard-label term keeps the student anchored where labels exist14 log_p_hard = np.log(softmax(student_logits) + 1e-12)15 hard = -log_p_hard[np.arange(len(hard_labels)), hard_labels].mean()16 return alpha * soft + (1 - alpha) * hardThe temperature is what makes the teacher's second and third choices visible; at T=1 the soft target is nearly one-hot and the student learns little more than from the label. The hard term is what stops the student from faithfully reproducing the teacher's mistakes on data where the truth is known.
The student is a derived artifact
A distilled model has a parent, and the parent changes. When the teacher is retrained on fresh data the student still encodes the old teacher; when the teacher's tail quality is fixed, the student's is not. Treating the student as a model in its own right — with its own retraining decision — loses the lineage that says why it behaves as it does.
The assumption the low-latency tier rests on is therefore two-part: the student is within tolerance of the *current* teacher on every segment, and the pruned model's speed-up still holds on the current hardware.
The student's ranking quality on each query segment, including the tail, is within the stated tolerance of the teacher that is currently in production.
holds when The distillation corpus covered the tail with weight; the student was evaluated per segment against this teacher version; the teacher has not changed since.
breaks when The teacher is retrained; the student is fine-tuned without the teacher; the query mix shifts toward the tail; a segment appears that the corpus never contained.
respond Re-distil from the current teacher; if the tail gap persists, weight the corpus toward it or accept a larger student for the tail tier.
How to build it
Most important first.
- Prune structurally, in the shapes the target kernel can skip — heads, channels, blocks — and benchmark on the target device before trusting a sparsity number. If the kernel is dense, sparsity is a file-size feature.
- Distil on a corpus that over-represents the segments the product cares about, weighted so the tail contributes signal, and include unlabelled data the teacher can score; check the student per segment against the teacher (Evaluation Slices).
- Match on the teacher's soft outputs with a temperature that exposes the alternatives, and keep a term on the hard labels where they exist so the student is not bound to the teacher's mistakes on labelled data.
- Record the teacher version in the student's lineage and re-distil when the teacher changes; treat the student as a derived artifact that goes stale (Model Lineage).
What to measure
Which number actually maps to the decision — and which numbers look relevant and are not.
- On-device latency and throughput of the pruned model against the dense one on production inputs — the number pruning was for; sparsity percentage is not a proxy for it.
- Per-segment ranking quality of the student against the teacher, tail first; the tail is where the capacity gap appears and where the aggregate hides it.
- Aggregate quality after pruning is dominated by head traffic and says little about whether the removed weights mattered on the tail.
What must stay true after deployment
The field this whole domain exists for. A model is a set of assumptions with weights attached; these are the ones a monitor or a test should be checking.
- The pruned model's speed-up was measured on the production device and kernel, and the sparsity pattern remains one the kernel skips.
- The student's per-segment quality stays within tolerance of the teacher on held-out data, including the tail, and the student is re-distilled or re-evaluated when the teacher version changes.
- The distillation corpus's segment mix matches what production sends, or is deliberately weighted toward the segments where the quality budget is tightest.
- Offline: latency and throughput benchmark for the pruned model on target hardware; per-segment student-versus-teacher report, tail first.
- Online: shadow the student against the teacher on live traffic and log disagreements by segment (Shadow Deployment).
- Over time: re-benchmark the pruned model after device or kernel changes; re-run the student report after any teacher change and alert on tail-segment quality drift.
What can go wrong
- Structured pruning removes the attention heads that were least used on the calibration set — which was head-heavy — and those heads were the ones the tail queries relied on.
- The student is fine-tuned later on new hard labels without the teacher, drifts from the teacher's function, and the "distilled from" claim in the registry becomes false while still recorded.
- A sparse kernel that was faster on the benchmark device is slower on the production device after a driver update, and the pruned model is now the slow one.
- Structured pruning gives real speed at more quality loss per removed parameter than unstructured; unstructured gives a smaller file and, without sparse hardware, nothing else.
- Distillation produces the best small model available and costs a teacher pass over a large corpus, a training run, and a permanent dependency on the teacher's version.
- Weighting the distillation corpus toward the tail improves the tail and spends student capacity there, which can cost a little on the head where most traffic is.
- "Ninety percent sparse means ninety percent fewer operations." It means ninety percent of the operands are zero. A dense kernel multiplies them anyway; only a kernel that skips them, at a sparsity and pattern it supports, does fewer operations.
- "The student learns everything the teacher knows." The student learns the teacher's outputs on the corpus it was shown, with less capacity. It is worst exactly where the teacher was least sure and the corpus was thinnest.
- "Distillation is free labels." Teacher outputs are soft labels with the teacher's errors baked in. Where hard labels exist, keep them in the loss, or the student faithfully learns the teacher's mistakes.
Where this applies
ML advice is stated as universal far more often than it is. These labels say what each claim is specific to — and where CONTESTED appears, the note gives the strongest form of the opposing view.
- MODEL-SPECIFICStructured pruning of heads and channels applies to transformers and convolutional networks; some accelerators support a fixed two-of-four sparsity pattern that makes semi-structured pruning pay, while a CPU without sparse kernels gains nothing from any of it. Distillation applies wherever a soft output exists — probabilities, scores, embeddings.
- DATA-SPECIFICDistillation's advantage is largest when a large unlabelled corpus exists for the teacher to score; with only the labelled set, the student's gain over training on hard labels shrinks, and the tail problem is set by the labelled set's tail.
- SIMPLIFIEDThe distillation loss shown uses a single temperature and a fixed mix with the hard-label term; production recipes add intermediate-layer matching and schedule the mix, and any sparsity or quality figures here are illustrative rather than measured.
Where the depth lives
This domain teaches the model and hands the rest off by name.