Knowledge Distillation: Teacher-Student Model Compression

Deploying a 50M-parameter vision transformer on a Cortex-M7 is not an option — but the accuracy that large model represents is often exactly what the product needs. Knowledge distillation (KD) is the technique that lets a small model trained to mimic a large one recover much of that accuracy without carrying the parameter count. Unlike pruning or quantization, which shrink an existing model, distillation trains a new, architecturally independent student from scratch, using the teacher's output distribution as a richer training signal than ground-truth labels alone. For embedded AI pipelines this is often the difference between an 89%-accurate 200 KB model and an 84%-accurate one — a gap that matters when the model gates a safety or UX decision.

Why Soft Labels Beat Hard Labels

A conventional classifier is trained against one-hot labels: correct class = 1, everything else = 0. This throws away information the teacher network naturally encodes in its output — the relative confidences across wrong classes.

Consider a teacher classifying a spoken digit "3." Its softmax output might be:

class:   3      8      5      others
prob:   0.82   0.11   0.05   0.02 (split)

The teacher is saying "3, but this could plausibly be an 8." That structure — dark knowledge — reflects learned similarity between classes and is far more informative than the label "3" alone. A student trained to reproduce this whole distribution learns decision boundaries closer to the teacher's, even with far fewer parameters.

The Distillation Loss

Hinton's original formulation (2015) softens the softmax with a temperature T:

softmax_T(z_i) = exp(z_i / T) / Σⱼ exp(z_j / T)

• T = 1 recovers the ordinary softmax.
• T > 1 flattens the distribution, amplifying the small probabilities assigned to non-target classes — exactly the signal you want the student to learn from.

The total training loss combines two terms:

L = α · L_hard(y, softmax(student_logits))
  + (1-α) · T² · L_soft(softmax_T(teacher_logits), softmax_T(student_logits))

• L_hard — standard cross-entropy against the true label.
• L_soft — cross-entropy (or KL-divergence) between teacher and student soft outputs, both computed at temperature T.
• T² scaling factor — compensates for the gradient magnitude shrinking as T grows (gradients of the soft loss scale as 1/T², so this restores comparable gradient magnitude to the hard-label term).
• α — typically 0.1–0.5; weights how much the student trusts ground truth vs. teacher output. Pure distillation (α = 0) is used when labels are noisy or teacher accuracy exceeds label accuracy.

Typical T values are 3–20; higher T is used when the teacher is very confident (near-saturated softmax) and you need to expose the finer-grained class relationships hidden under that confidence.

Worked Example: Sizing a KWS Student Model

Assume a keyword-spotting teacher (DS-CNN, ~500K params, 94.5% top-1 accuracy) and a target student that must fit in 64 KB of flash on a Cortex-M4 at 8-bit quantization.

1. Budget check: 64 KB flash ÷ 1 byte/weight (INT8) ≈ 64,000 weights max, leaving headroom for code and buffers → target ≈ 40K–50K parameters for the student.
2. Baseline (no distillation): training a 45K-parameter DS-CNN directly on labels yields ~89.0% accuracy — a 5.5-point gap versus the teacher.
3. With distillation: train the same 45K-parameter architecture using teacher soft labels, T = 4, α = 0.3, on the same dataset (no additional data needed — this is the key practical advantage: KD needs no new labels, only teacher inference on existing training data).
4. Typical result: reported KD gains in the KWS/DS-CNN literature are 1.5–3 points of accuracy recovered at equal student size — landing the student around 91–92%, closing roughly half the gap to the teacher.
5. Verify the deployment budget: 45K params × 1 byte (INT8) = 45,000 bytes ≈ 44 KB flash for weights, plus ~8–12 KB for activation buffers and runtime → fits within the 64 KB budget with margin. Confirmed feasible.

The number to sanity-check is always parameter count × bytes/weight against the flash budget before investing in a distillation training run — KD changes accuracy, not the compression ratio itself.

Distillation Variants

Feature-based distillation adds an auxiliary loss matching intermediate tensors (often via an L2 loss on projected feature maps, since teacher and student channel counts usually differ and need a learned adapter layer). This helps when the student is too shallow for output-only matching to transfer enough structure — common when compressing detection or segmentation backbones rather than simple classifiers.

Practical Design Decisions

• Student architecture choice matters more than KD hyperparameters. Distillation cannot make an architecturally-inadequate student (too few layers to represent the decision boundary at all) match the teacher — search for the smallest architecture that gets within ~5 points of the teacher without KD first, then use KD to close the remainder.
• Teacher quality caps student quality. If the teacher itself is only 85% accurate, the student will asymptotically approach 85%, not exceed it (barring regularization side-effects). Always distill from your best available teacher, not the deployed one.
• Combine with quantization, not instead of it. KD produces a full-precision (FP32) student; that student still needs post-training quantization or QAT to hit the target flash/RAM footprint. The two techniques are complementary, applied in sequence: distill → quantize.
• Data-free / limited-data distillation: when the original training set isn't available at deployment time, synthetic inputs or a small proxy dataset can still transfer meaningful signal, though with reduced fidelity — relevant for federated or privacy-constrained pipelines.
• Inference cost of the teacher is irrelevant at deployment — it runs only during training (often once, offline, on a workstation/server), so a teacher can be arbitrarily large without affecting the embedded target.
• Loss balance drift: monitor L_hard and L_soft separately during training; if L_soft dominates and α is too low, the student can overfit to teacher mistakes on out-of-distribution inputs the teacher itself handles poorly.

Key Takeaways

• Knowledge distillation trains a compact student to match a large teacher's soft output distribution, transferring inter-class similarity information that hard labels discard.
• The distillation loss blends hard-label cross-entropy and temperature-softened teacher-student cross-entropy, with T² scaling to balance gradient magnitudes.
• Typical embedded gains are 1.5–5 accuracy points recovered at a fixed parameter budget, using the same training data as a normal run — no extra labeling cost.
• Student architecture capacity sets a ceiling KD cannot exceed; teacher accuracy sets a ceiling the student cannot exceed either.
• KD and quantization are complementary and applied in sequence: distill in full precision first, then quantize the resulting student for flash/RAM targets.