Knowledge Distillation & Model Compression

Follow the teacher-to-student path

The teacher helps create the training signal, but the student normally serves requests by itself. Distillation changes training; it does not automatically add teacher work to student inference.

1. 1

   Baseline

   Qualify the teacher

   Freeze a version whose task quality, calibration, failure slices, and data lineage are understood.

2. 2

   Inputs

   Build transfer examples

   Use representative, permitted inputs that cover the workload the student must handle.

3. 3

   Signals

   Train the student

   Combine ground-truth loss with teacher outputs or compatible intermediate features.

4. 4

   Evidence

   Release independently

   Benchmark the exported student on target hardware and block promotion when any required gate fails.

Name the four moving parts

See what temperature and loss weight actually change

For response distillation, divide teacher and student logits by temperature T before applying softmax. Larger T spreads probability across more classes. The soft loss is commonly multiplied by T² so its gradient scale remains comparable as temperature changes.

A clear objective names the convention instead of relying on an ambiguous alpha:

total_loss = teacher_weight × soft_loss + (1 - teacher_weight) × label_loss

The right temperature and weight are empirical hyperparameters. They must be selected on validation evidence, not copied from a universal range.

Choose a transfer signal the architectures can support

No signal is automatically superior. Feature or attention matching can add memory, implementation, and architecture-coupling costs without improving the deployed student.

Implement one correct response-distillation step

The teacher runs in evaluation mode under torch.no_grad(). The student remains trainable. PyTorch's KLDivLoss contract expects the model input in log space, so the student branch uses log_softmax; reduction="batchmean" preserves the mathematical KL definition.

The function returns the differentiable total loss plus detached metrics. The caller still owns mixed precision, gradient accumulation, optimizer stepping, checkpointing, distributed reduction, and evaluation.

Build a transfer set that represents deployment

Distillation cannot transfer behavior the teacher never produces on the transfer inputs. Treat the transfer set as a versioned dataset with the same rigor as ordinary training data.

• Cover important classes, cohorts, languages, sequence lengths, devices, and difficult examples.
• Keep validation and test examples out of training, teacher-output caches, and hyperparameter selection.
• Record the teacher artifact, preprocessing version, label mapping, and output schema with every cached signal.
• Recompute teacher outputs when the teacher or preprocessing contract changes.
• Confirm that data licenses, privacy rules, and retention policies allow teacher inference and output storage.

For an online teacher, training spends teacher compute on every batch. For cached logits, training is cheaper and reproducible, but storage grows with examples and output width. Neither mode changes the need for an untouched evaluation set.

Select a student using measured deployment evidence

Parameter count and artifact size are useful inventory fields, but they do not determine end-to-end latency or peak memory. Kernel support, sequence shape, batching, preprocessing, data movement, and the target runtime all matter.

The next lab uses an illustrative benchmark snapshot, not a universal model ranking. Change the target profile and candidate to see which constraint becomes the release blocker.

Encode the release contract beside the training recipe

The release gate should compare the exact exported artifact against the approved baseline on identical inputs and target hardware. Aggregate quality alone is not enough when a required slice can regress.

Promotion should require all of the following:

• Overall task quality meets the declared floor.
• Every required slice meets its own floor or maximum regression budget.
• Calibration, robustness, safety, and fairness checks match the product risk.
• Target-runtime p50 and p95 latency, peak memory, throughput, and energy or cost are measured.
• A bounded canary can roll back to the previous artifact.
• Training lineage identifies the teacher, student initialization, transfer data, loss convention, and hyperparameters.

Diagnose failures before adding complexity

Combine compression techniques deliberately

Distillation, quantization, and pruning change different parts of the system:

• Distillation trains a separate student to reproduce selected behavior.
• Quantization changes numeric representation and often depends on backend support and calibration.
• Pruning removes weights or structures and only speeds inference when the runtime can exploit the resulting sparsity.

They can be combined, but each additional transform creates a new artifact that needs the same quality and target-hardware evaluation. Learn the numeric trade-offs in model quantization and the broader release discipline in model compression techniques.

Use primary references for the implementation contract

• Distilling the Knowledge in a Neural Network introduces soft targets, temperature scaling, and the gradient-scaling rationale.
• PyTorch Knowledge Distillation Tutorial demonstrates a frozen teacher, trainable student, softened logits, label loss, and T² scaling.
• PyTorch KLDivLoss documentation defines the log-space input contract and recommends batchmean for the mathematical KL value.
• Keras Knowledge Distillation example provides an independent official implementation of the same teacher-student objective.