Model Distillation (Knowledge Distillation)

Model distillation trains a small model — the student — to reproduce the outputs of a large one — the teacher — instead of training it on the labels alone. Hinton, Vinyals and Dean introduced the method in 2015 as a way to move the behavior of an expensive ensemble into a single deployable model, crediting earlier model compression work by Caruana and his collaborators.

The mechanism is the teacher's probability vector. A hard label says cat. A teacher output of [0.7 cat, 0.2 dog, 0.1 rabbit] also says that this image is much more dog-like than rabbit-like, and that extra structure is what the student learns from.

Soft targets against hard labels

Same student architecture, same training rows, one difference: what it is asked to match.

import numpy as np
from sklearn.datasets import make_classification
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

x, y = make_classification(n_samples=6000, n_features=12, n_informative=8,
                           n_redundant=0, flip_y=0.01, class_sep=1.1, random_state=0)
xtr, ytr, xte, yte = x[:4000], y[:4000], x[4000:], y[4000:]

teacher = KNeighborsClassifier(n_neighbors=25).fit(xtr, ytr)
soft = teacher.predict_proba(xtr)[:, 1]
"""The teacher's soft output is a probability per row, not a 0/1 label.
   Distillation trains the student to match that probability."""

hard_student = DecisionTreeClassifier(max_depth=4, random_state=0).fit(xtr, ytr)
soft_student = DecisionTreeRegressor(max_depth=4, random_state=0).fit(xtr, soft)
soft_pred = (soft_student.predict(xte) >= 0.5).astype(int)

print(f"{'model':<34}{'test accuracy':>14}{'agrees with teacher':>21}")
for name, pred in (("teacher: 25-neighbor vote", teacher.predict(xte)),
                   ("student on hard 0/1 labels", hard_student.predict(xte)),
                   ("student on teacher probabilities", soft_pred)):
    print(f"{name:<34}{(pred == yte).mean():>14.4f}{(pred == teacher.predict(xte)).mean():>21.4f}")

print(f"\nteacher stores {xtr.size:,} feature values to answer one query")
print(f"each student stores {hard_student.tree_.node_count} tree nodes")
# model                              test accuracy  agrees with teacher
# teacher: 25-neighbor vote                 0.9815               1.0000
# student on hard 0/1 labels                0.9060               0.9095
# student on teacher probabilities          0.9105               0.9170
#
# teacher stores 48,000 feature values to answer one query
# each student stores 31 tree nodes

In this run the distilled student is slightly better than the label-trained student, 0.9105 against 0.9060, and agrees with the teacher more often, 0.9170 against 0.9095. The architectures have the same depth limit, but the target and split objective differ. This is one comparison, not a universal size of distillation gain.

The teacher scores 0.9815 and both fitted students remain near 0.91 here. Their limited function class can restrict how closely they reproduce the teacher, but this single fit does not prove the best achievable student accuracy. Matching the teacher’s entire function is also different from matching or exceeding its accuracy against truth. Students can outperform teachers, as Born Again Networks demonstrate. Keep the hard-label student as a baseline and measure the actual accuracy–cost trade-off rather than treating teacher accuracy as an unreachable ceiling.

The last two lines show why anyone bothers. The teacher holds all 4,000 training rows — 48,000 stored feature values — to answer a single query; each student is 31 nodes and a handful of comparisons.

Temperature, and the factor of T squared

A well-trained teacher is often too confident to be informative. If its output is [0.999, 0.0006, 0.0004] the ranking among the wrong classes is present but numerically invisible. The 2015 paper's fix is to raise the temperature in the softmax, \(q_i = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}\), when producing targets and when training the student, which spreads the mass and makes the small differences carry gradient.

The loss is a weighted average of two objectives: one against the teacher's soft targets, one against the true labels. The paper adds an implementation detail that is easy to omit and easy to be bitten by — the soft-target term is multiplied by \(T^2\), because gradients from the softened distribution scale down as \(T\) rises, and without that factor the balance between the two terms would silently change every time you tuned the temperature.

Both \(T\) and the weighting are choices you have to search over, which is a real cost. Distillation also does not remove the teacher from the budget: you still pay to train or obtain it, and you pay to run it over the transfer set. What it removes is the cost at inference time, repeated for every request, which is usually where the money is. If the goal is serving cost, compare it against quantization and pruning on the same footing rather than assuming distillation wins.

What the published results do and do not promise

DistilBERT is the most cited example. Sanh and colleagues report reducing BERT's size by 40% “while retaining 97% of its language understanding capabilities and being 60% faster”. Those three numbers come as a set from one model distilled during pre-training, evaluated on that paper's benchmarks. They are a demonstration that the tradeoff can be favorable, not a rate you should expect to reproduce on a different teacher, task, or student architecture.

The same caution applies to the retained 3%. Aggregate benchmark scores hide where the loss landed, and a student that matches the teacher on average can be systematically worse on a rare class that matters to you. Compare teacher and student on the slices you care about, not on a single headline number.

One property is often overlooked: the transfer set does not need labels. The teacher supplies a target for any input you can feed it, so unlabeled data — often abundant — can be used to train the student. That, rather than the accuracy bump, is frequently the reason distillation is worth the setup.

References

• Hinton, G., Vinyals, O., and Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531. arxiv.org
• Sanh, V., Debut, L., Chaumond, J., and Wolf, T. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv:1910.01108. arxiv.org
• scikit-learn documentation, sklearn.tree.DecisionTreeRegressor and sklearn.neighbors.KNeighborsClassifier. scikit-learn.org
• Furlanello et al. (2018), Born Again Neural Networks