Model Compression via Knowledge Distillation

A large model may be accurate but awkward to deploy. It can require too much memory, take too long to answer, or cost too much to serve at the scale an application needs. Knowledge distillation offers one way to trade some of that capacity for a smaller model.

The basic idea is simple: train a student not only from the original labels, but also from the behavior of a stronger teacher. The teacher’s output contains more information than the winning class alone. A picture labelled “cat,” for example, may still receive small probabilities for “fox” and “dog.” Those relative values reveal which alternatives the teacher considers similar.

Distillation is a training method, not a guarantee. A student with too little capacity cannot reproduce everything the teacher knows, and a biased or badly calibrated teacher can pass those problems on.

The Classic Distillation Objective

Let the teacher and student produce logits $z_t$ and $z_s$ over $K$ classes. A temperature $T>0$ softens each distribution:

At $T=1$, these are ordinary softmax probabilities. A larger temperature spreads probability across more classes, exposing relationships that would be almost invisible in a sharply peaked distribution.

A common objective combines the teacher target with the ground-truth label:

Here:

• $D_{\mathrm{KL}}(p_t\parallel p_s)$ asks the student to match the teacher distribution;
• cross-entropy keeps the student tied to the actual labels;
• $\alpha$ controls the balance;
• the $T^2$ factor compensates for the way temperature shrinks logit gradients.

The direction of the KL divergence matters. Standard response distillation uses the teacher as the target distribution. Reversing it changes the behavior and should be an explicit design choice, not an unnoticed notation swap.

Why Soft Targets Can Help

Suppose a teacher gives a handwritten digit the following probabilities:

The hard label only says “3.” The soft target also says that this example resembles an 8 more than a 1. A student can use that extra structure, particularly when labels are sparse or noisy.

Soft targets are not automatically truthful. A teacher may be confidently wrong, overfit a spurious feature, or perform unevenly across groups. Retaining the hard-label term gives the student another source of supervision, but it does not erase teacher bias.

A Correct PyTorch Training Step

import torch
import torch.nn.functional as F

def distillation_loss(
    student_logits,
    teacher_logits,
    labels,
    temperature=4.0,
    alpha=0.7,
):
    soft_loss = F.kl_div(
        F.log_softmax(student_logits / temperature, dim=-1),
        F.softmax(teacher_logits / temperature, dim=-1),
        reduction="batchmean",
    ) * (temperature ** 2)

    hard_loss = F.cross_entropy(student_logits, labels)
    return alpha * soft_loss + (1.0 - alpha) * hard_loss

teacher.eval()
student.train()

for inputs, labels in train_loader:
    optimizer.zero_grad()

    with torch.no_grad():
        teacher_logits = teacher(inputs)

    student_logits = student(inputs)
    loss = distillation_loss(student_logits, teacher_logits, labels)
    loss.backward()
    optimizer.step()

Putting the teacher in evaluation mode disables training-time dropout and uses the intended normalization behavior. torch.no_grad() also avoids storing a teacher backward graph. The teacher and student still need to agree on the class vocabulary and preprocessing.

What Can Be Distilled?

Output or response distillation

The student matches logits or probabilities. This is the simplest form and works even when teacher and student architectures differ substantially.

For multilabel classification, the output is a set of independent sigmoid logits rather than one softmax distribution. The distillation loss should match that output semantics instead of forcing a categorical KL objective onto it.

Feature distillation

The student matches one or more internal teacher representations:

$P$ is a learned projection when dimensions differ. Feature matching can offer a richer signal, but choosing layers and normalization is part of the experiment. Forcing a narrow student to copy every teacher activation may make optimization harder rather than easier.

Relational distillation

Instead of matching individual activations, relational methods preserve relationships among examples, tokens, or layers. A method might match pairwise distances:

Published methods differ in how distances and angles are normalized. There is no single equation that represents every form of relational knowledge distillation.

Attention distillation

Transformer students can be trained to match attention maps, hidden states, value relations, or combinations of them. Attention matrices have different heads and dimensions across architectures, so a mapping or projection is often needed. Matching attention is an auxiliary target; it is not proof that the student uses the same reasoning process.

Self-distillation

In self-distillation, teacher and student may share an architecture or belong to the same training run. Born-Again Networks, for example, train a new model of the same architecture from a previously trained model on the same task. Improvements are empirical, not guaranteed or necessarily monotonic across generations.

Other methods use earlier checkpoints, deeper layers, or an ensemble of branches as the teacher. The phrase covers several distinct setups, so an article or experiment should say which one it uses.

Distilling Language Models

Language-model distillation can target several signals:

• next-token distributions over a shared vocabulary;
• generated sequences used as supervised data;
• hidden states or attention-related representations;
• preference scores or task-specific outputs;
• intermediate reasoning traces, when those traces are available and suitable to use.

A student trained only on final answers is learning to produce those answers. If it is also trained to generate a rationale, the objective must include the rationale tokens. A loss that conditions on an already supplied rationale does not by itself teach the student to create one at inference time.

Sequence-level distillation replaces or augments reference text with teacher-generated sequences. This can make the target distribution easier for a small model, but it can also reduce diversity and repeat teacher mistakes.

The sequence-distribution problem

Classification distillation evaluates teacher and student on the same fixed input. Autoregressive generation is less convenient. At training time, a student may see prefixes from the dataset or teacher; at inference time, it conditions on its own sampled tokens. One early mistake changes every prefix that follows.

This is a distribution-mismatch problem, not merely a temperature problem. Three common choices answer it differently:

• off-policy token distillation evaluates the teacher on fixed dataset or teacher-generated prefixes;
• sequence-level distillation trains on complete teacher outputs as supervised targets;
• on-policy distillation samples prefixes from the student, then asks the teacher how the student should behave in states the student actually visits.

On-policy training is more expensive because teacher evaluation is coupled to student generation. It can also expose mistakes that a clean teacher-generated dataset never contains. Generalized Knowledge Distillation formalizes this idea for autoregressive models and allows divergences other than the standard forward KL.

The divergence direction changes the pressure. Forward KL,

penalizes the student for failing to cover probability mass assigned by the teacher. Reverse KL,

penalizes the student for placing its own mass where the teacher assigns little. The latter can encourage mode-seeking behavior, which may suit a capacity-limited generator but can also reduce diversity. MiniLLM uses a reverse-KL-based objective for generative language-model distillation; that is an empirical design choice, not a universal replacement for forward KL.

Choosing a Temperature

Temperature is a hyperparameter, not a measure of intelligence. Too low and the teacher target resembles a hard label; too high and meaningful differences may be washed out.

If a schedule is useful, one that decays toward 1 can be written as:

This really does approach 1. The simpler expression $T_{\max}e^{-\lambda t}$ approaches zero, which contradicts any claim that it ends at ordinary softmax temperature.

Tune $T$ and $\alpha$ together on a validation set. Their best values depend on teacher calibration, number of classes, dataset size, and student capacity.

Intermediate-Layer Alignment

A twelve-layer teacher and six-layer student need an explicit alignment policy. Common choices include:

• match every second teacher layer;
• learn projections between selected layer pairs;
• compare a weighted combination of teacher layers;
• supervise only the final few representations.

More matching terms are not always better. They increase memory use and may constrain a student that needs to organize its smaller representation differently.

Multiple Teachers

An ensemble can provide a smoother or more accurate target:

Weights may be fixed or learned. The cost is substantial because training must evaluate or precompute several teachers. If the teachers share the same blind spot, averaging does not remove it.

Data-Free and Low-Data Distillation

Ordinary distillation uses the original training data or another representative dataset. When that data is unavailable, a separate generator or optimization procedure may synthesize inputs that produce informative teacher responses. A genuinely data-free objective cannot quietly rely on samples from the unavailable real dataset.

Synthetic inputs can miss important parts of the production distribution. They may also reveal information about the teacher’s training data, so “no original data in the student loop” should not be confused with a privacy proof.

Distillation Is Not the Same as Other Compression Methods

• Quantization reduces the precision used to store or compute with weights and activations.
• Pruning removes weights, channels, heads, or blocks according to a sparsity rule.
• Low-rank factorization replaces a large matrix with smaller factors.
• Distillation changes the training signal so a student imitates a teacher.

They can be combined. A common workflow distils a student and then applies quantization, followed by task-specific evaluation on the actual target hardware.

Model size, latency, throughput, peak memory, and energy are different measurements. Parameter count does not determine all of them, and FLOPs are not the same thing as joules consumed. Report the metric that matters for the deployment.

What Published Results Actually Show

DistilBERT reported a model that was 40% smaller and 60% faster than BERT-base while retaining 97% of its language-understanding performance on the paper’s evaluated tasks. Those figures describe that model and experimental setup; they are not a general law that distillation keeps 97% of any teacher.

TinyBERT used both general and task-specific transformer distillation. Its paper is arXiv:1909.10351, not the DistilBERT paper at arXiv:1910.01108.

Progressive distillation also appears in diffusion models, where it can reduce the number of sampling steps. The paper at arXiv:2202.00512 is about diffusion sampling, not LLM inference.

No universal “compression law” says that a particular student-to-teacher ratio preserves a fixed amount of quality. Architecture, data, training budget, task, and evaluation all matter.

A Practical Workflow

1. Define the deployment target. Set limits for latency, memory, throughput, accuracy, and hardware.
2. Choose a student that can meet it. A student too large misses the point; one too small may have an impossible learning task.
3. Establish non-distilled baselines. Train the same student on labels alone and compare it with quantization or pruning baselines.
4. Cache teacher outputs when practical. This saves repeated teacher inference but can consume considerable storage.
5. Tune the soft and hard losses. Include calibration and per-group metrics, not only average accuracy.
6. Measure on target hardware. A smaller checkpoint does not always produce lower latency without optimized kernels.
7. Test failure cases. Check distribution shift, rare classes, adversarial inputs, and teacher-specific mistakes.

An Experiment That Can Falsify the Claim

A distillation result is weak if it compares only one distilled student with its teacher. At minimum, train:

1. the student on hard labels or ordinary language-model targets;
2. the same student with distillation;
3. the distilled student with teacher outputs shuffled or otherwise ablated when a sanity check is possible;
4. a quantized or pruned baseline when deployment efficiency is the motivation.

Keep architecture, data budget, optimizer-search effort, and stopping rule comparable. Then report:

For an autoregressive student, evaluate generated sequences rather than only teacher-forced token loss. A lower KL on reference prefixes can coexist with poor free-running behavior. Record decoding settings too; changing temperature or top-$p$ between models can overwhelm a small training gain.

The useful artifact is not a table with every cell green. It is a result that reveals the boundary: which student size, data regime, and deployment target made distillation worthwhile, and where the teacher signal stopped helping.

Privacy and Federated Settings

Federated distillation can exchange outputs or summary statistics instead of full model updates. That may reduce communication or raw-data movement, but it is not inherently privacy-preserving. Logits, gradients, synthetic data, and updates can leak information. Strong privacy claims require a threat model and mechanisms such as secure aggregation or differential privacy, with their trade-offs measured.

Conclusion

Knowledge distillation is best understood as supervised imitation with a richer target. It can produce a smaller model that behaves more like a large one, but it cannot squeeze unlimited capability through a fixed bottleneck. The honest test is practical: compare the student with a label-only baseline, measure it on the hardware that will serve it, and check which teacher mistakes survived the transfer.

References

• Hinton, Vinyals, and Dean, Distilling the Knowledge in a Neural Network (2015).
• Furlanello et al., Born-Again Neural Networks (2018).
• Sanh et al., DistilBERT (2019).
• Jiao et al., TinyBERT (2019).
• Park et al., Relational Knowledge Distillation (2019).
• Salimans and Ho, Progressive Distillation for Fast Sampling of Diffusion Models (2022).
• Gu et al., MiniLLM: Knowledge Distillation of Large Language Models (2023).
• Agarwal et al., On-Policy Distillation of Language Models (2023).