Compress Neural Networks with Knowledge Distillation

A model can meet your quality target in a notebook and still be too expensive to serve. A large network may consume too much memory, add unacceptable latency, or make high request volume costly. Knowledge distillation addresses this problem by using a stronger model, called the teacher, to guide the training of a smaller student model.

The key idea is richer than copying the teacher’s final answer. The teacher produces a distribution across possible outputs, and that distribution can reveal useful relationships between alternatives. A student can train against those soft targets while also using the original labels.

This article builds the idea from a small classification example, explains temperature and the distillation loss, and covers the checks that determine whether a smaller student is actually a good deployment trade-off.

Start with more information than a hard label

Consider an image classifier with three classes:

cat
dog
car

For one training image, the ground-truth label is cat. A conventional target encodes that fact as:

cat    1.00
dog    0.00
car    0.00

That target says which class is correct, but nothing about the relationship among the incorrect classes.

Now suppose a well-performing teacher produces:

cat    0.82
dog    0.17
car    0.01

The teacher agrees that cat is the answer, but it also considers dog much more plausible than car. That extra structure can help a student. The student isn’t merely asked to put all probability on cat; it is also encouraged to reproduce the teacher’s relative preferences.

These probabilities are often called soft targets. They can carry information that a one-hot label cannot express.

This does not make the teacher correct by definition. If the teacher has systematic errors, the student can inherit them. The original labels remain valuable because they provide a target independent of the teacher’s preferences.

Temperature exposes relationships in the teacher output

Neural classifiers commonly produce logits before converting them to probabilities with softmax. For logits z_i, ordinary softmax is:

p_i = exp(z_i) / sum_j exp(z_j)

Knowledge distillation often introduces a positive temperature T:

p_i(T) = exp(z_i / T) / sum_j exp(z_j / T)

At T = 1, this is ordinary softmax. A temperature above 1 reduces the gaps between logits before softmax, producing a softer probability distribution.

Take teacher logits:

cat    5
dog    2
car    0

At ordinary temperature, the distribution is strongly concentrated on cat. At a higher temperature, dog and car receive more probability mass. Their relative positions become easier for the student loss to represent.

Temperature does not add information that was absent from the logits. It changes how strongly differences in those logits appear in the probability distribution used for distillation.

The teacher and student should use the same distillation temperature when their softened distributions are compared. At deployment, the student normally returns to the inference behavior required by the application; the training temperature is not automatically a serving setting.

Combine teacher guidance with ground truth

A common distillation objective combines two terms:

total_loss =
    alpha * hard_target_loss
    + (1 - alpha) * T^2 * soft_target_loss

Here:

• hard_target_loss compares the student’s ordinary prediction with the ground-truth label.
• soft_target_loss compares the teacher and student distributions computed with temperature T.
• alpha controls the balance between direct supervision and teacher guidance.
• T^2 compensates for the way temperature reduces gradient magnitudes in the softened objective.

For classification, cross-entropy or KL divergence can be used to compare softened distributions, depending on the implementation. These objectives differ by a term that is constant with respect to the student when the teacher distribution is fixed, so they produce the same student optimum under that condition, although reported loss values differ.

The exact weighting is a hyperparameter choice, not a universal constant. A student that is much smaller than its teacher may need a different balance from a student with nearly the same capacity.

A useful mental model is:

ground truth says: "put probability on the correct answer"
teacher says:      "preserve these relationships among answers"
student objective: "satisfy both as well as your capacity permits"

Walk through one training step

Suppose the teacher sees an example and emits these softened probabilities:

class    teacher
cat      0.60
dog      0.30
car      0.10

The student currently emits:

class    student
cat      0.45
dog      0.15
car      0.40

The hard label cat pushes the student toward the correct class. The distillation term adds a more specific signal: the student’s car probability is far too high relative to the teacher, while its dog probability is too low.

After optimization, a possible student distribution might move toward:

class    student
cat      0.56
dog      0.27
car      0.17

This is only a teaching example; an actual update depends on logits, optimizer state, loss weights, and model parameters. The important point is that the teacher supplies directional information across all classes rather than supervision for only the winning class.

That extra signal is especially relevant when many classes have meaningful similarities. A teacher may assign related objects similar scores, giving the student information about the teacher’s decision surface.

Distillation is not the same as copying predictions

It is tempting to generate labels with a large model and train a small model on those labels. That can be useful, but it discards information if only the teacher’s top choice is retained.

Consider two teacher outputs:

example A
cat    0.51
dog    0.48
car    0.01

example B
cat    0.98
dog    0.01
car    0.01

Both have the same top class: cat. Hard pseudo-labels make the two examples look identical from the teacher’s perspective. Soft targets preserve the distinction between an uncertain boundary case and a confident case.

Knowledge distillation can also use signals beyond final output probabilities. Some methods align hidden representations, attention maps, or intermediate features. Those variants can be effective, but they add architectural assumptions and extra loss design. Output-level distillation is the cleanest place to start because teacher and student do not need matching internal shapes.

Choose a student for the deployment constraint

Distillation does not decide the student architecture for you. Start from the resource limit that matters in production.

If memory is the main constraint, parameter count and numerical precision may dominate. If latency is the issue, hardware utilization, operator efficiency, sequence length, batch size, and memory movement can matter as much as raw parameter count. If cost per request is the target, measure the complete serving path rather than relying on model size alone.

A student can be:

• the same architecture family with fewer layers or narrower hidden dimensions;
• a different architecture that supports the same task outputs;
• a model designed specifically for the target device or runtime.

The smallest possible student is rarely the right initial target. Severe capacity reduction can create an optimization problem the teacher signal cannot solve. A practical process is to establish a non-distilled student baseline first, then add distillation and measure the gain at the same architecture and serving configuration.

That comparison separates the benefit of distillation from the benefit of simply choosing a different model.

Measure the student against three baselines

A distilled model should not be evaluated only against the teacher. Three comparisons answer different engineering questions.

First, compare the student with the teacher. This shows the quality sacrificed for lower serving cost.

Second, compare the distilled student with the same student architecture trained only on ground truth. This isolates the value of teacher guidance.

Third, compare the distilled student with simpler deployment alternatives. Depending on the system, quantization, pruning, batching, caching, or a smaller pretrained model may reach the required cost target with less training complexity.

Measure task quality and serving behavior separately. A useful evaluation table might contain:

model               task score    p95 latency    memory    cost/request
teacher             ...
student baseline    ...
distilled student   ...

Do not assume fewer parameters imply proportional latency savings. A smaller network can still run inefficiently on a given accelerator, and fixed request overhead can dominate short workloads.

Keep the teacher evaluation-safe

The teacher is part of the training pipeline, not an oracle. Its mistakes and biases can become training signals.

Evaluate the teacher on slices that matter to the application before using it as a source of targets. If it performs poorly on a rare class, language variety, device condition, or other important slice, distillation can reproduce that weakness even when the ground-truth term remains present.

Teacher confidence also deserves scrutiny. Soft targets can be informative without being calibrated probabilities. Distillation asks the student to imitate a distribution; it does not prove that the numerical confidence values correspond to real-world frequencies.

If the application uses probabilities for risk decisions, evaluate calibration on the final student separately. Matching teacher outputs is not a substitute for validating the student’s probability behavior.

Common mistakes that weaken distillation

Using only teacher outputs when reliable labels already exist

Teacher-only supervision makes the teacher’s errors authoritative. When trustworthy ground-truth labels are available, a mixed objective provides an independent corrective signal.

There are settings where teacher-generated data is the only practical source of supervision. That is a different constraint and should be evaluated as such rather than treated as equivalent to distillation with verified labels.

Tuning temperature without tuning the loss balance

Changing temperature changes the softened distributions and the scale of the distillation gradients. Treating T, the hard-target weight, and the soft-target weight as unrelated knobs can produce misleading comparisons.

Keep the objective definition explicit in experiment records. Two runs that both say “temperature 4” are not comparable if their loss scaling differs.

Comparing models at different serving conditions

Latency measured with different batch sizes, hardware, precision, or input lengths does not isolate the model change. Use representative traffic and identical serving conditions when comparing teacher and student.

Expecting the student to reproduce unlimited teacher capacity

Distillation transfers training signal, not parameter capacity. A very small student may be unable to represent the teacher’s decision function closely. If the quality gap remains large after sensible tuning, increasing student capacity can be more productive than adding more elaborate distillation losses.

Distilling a weak or mismatched teacher

A teacher that excels on a benchmark but performs poorly on your target distribution may provide harmful guidance. Teacher selection should follow the deployment task, not model reputation.

Offline targets can reduce training cost

Running the teacher during every student update can be expensive. If the training examples are fixed and teacher outputs do not depend on changing augmentation or context, you can precompute teacher logits or softened targets once.

The trade-off is storage versus repeated teacher compute.

Precomputation is attractive when teacher inference is expensive and the dataset fits a manageable storage budget. Online teacher inference is more flexible when examples change dynamically, augmentation affects the teacher input, or storing full-vocabulary outputs would be excessive.

For large output spaces, storing complete probability vectors can itself be costly. Approximations such as retaining selected logits may reduce storage, but they change the target distribution and require separate validation. Do not treat a truncated target as mathematically identical to the full teacher distribution.

Distillation fits some problems better than others

Knowledge distillation is a strong candidate when a high-quality teacher already exists, deployment needs a smaller model, and enough representative training data is available to expose the student’s behavior to teacher guidance.

It is less attractive when the teacher barely exceeds the student baseline, the target task has little training data, serving cost is already acceptable, or a simpler compression method meets the requirement. It can also be awkward when teacher access is limited to a remote API that exposes only final text or labels rather than stable score distributions.

For generative models, distillation needs additional care. Token distributions depend on the prefix, and student-generated prefixes can drift away from the teacher’s training trajectories. Sequence-level objectives, generated training data, or other techniques may be needed depending on the goal. The simple classification recipe should not be assumed to transfer unchanged to every autoregressive setting.

Treat distillation as a measured compression experiment

The practical value of knowledge distillation is not that a student becomes a miniature copy of its teacher. It is that the teacher can provide richer supervision than hard labels alone, giving a constrained student a better chance of preserving useful behavior.

Start with a clear student baseline. Add output-level soft targets before introducing more complex intermediate losses. Measure task quality, latency, memory, and cost under the same serving conditions. Check important data slices and probability behavior on the final student.

If the distilled student meets the deployment budget with an acceptable quality gap, the technique has done its job. If it does not, the measurements still tell you whether the next move should be a larger student, a different teacher, another compression method, or no compression at all.