Knowledge Distillation

Knowledge distillation is a model compression technique in which a small "student" neural network is trained to reproduce the behavior of a larger, more accurate "teacher" model. The student learns from the teacher's output probabilities or internal features in addition to, or instead of, the ground-truth labels, and these richer targets let it reach higher accuracy than the same architecture trained on labels alone. The result is a model fast and small enough for phones, embedded boards, and other edge AI hardware that keeps much of the teacher's quality. Buciluǎ, Caruana, and Niculescu-Mizil showed in 2006 that an ensemble's knowledge could be compressed into one model, and Hinton, Vinyals, and Dean named the approach knowledge distillation in 2015 and introduced the temperature-softened targets most methods still use.

How Knowledge Distillation Works#

A label says only that an image shows a wolf. A trained teacher's softmax output says more: the image looks somewhat like a husky, a little like a cat, and nothing like a car. These relative probabilities of wrong classes, often called dark knowledge, reach the student as soft targets. Dividing the logits by a temperature T greater than 1 before the softmax flattens the distribution so the small probabilities carry enough weight to learn from. The distillation loss is the Kullback-Leibler divergence between the teacher's and the student's softened distributions, multiplied by T² so its gradients stay comparable in size to those of the ordinary label loss:

import torch
import torch.nn.functional as F

# Logits for one wolf photo over the classes [wolf, husky, cat, car]
teacher_logits = torch.tensor([[6.0, 4.0, 1.0, -2.0]])
student_logits = torch.tensor([[4.0, 1.0, 1.5, 0.0]])

# Dividing logits by a temperature T > 1 softens the teacher's distribution
for T in (1.0, 4.0):
    probs = F.softmax(teacher_logits / T, dim=1)
    print(f"T={T}: teacher probabilities {[round(p, 3) for p in probs[0].tolist()]}")

# Distillation loss: KL divergence between the softened distributions, scaled by T^2
T = 4.0
soft_targets = F.softmax(teacher_logits / T, dim=1)
student_log_probs = F.log_softmax(student_logits / T, dim=1)
loss = F.kl_div(student_log_probs, soft_targets, reduction="batchmean") * T**2
print(f"Distillation loss at T={T}: {loss.item():.3f}")

At T=1 the teacher gives the wolf class 0.875 and leaves only 0.006 for cat and 0.0 for car. At T=4 the same logits become 0.493, 0.299, 0.141, and 0.067, which tells the student that a wolf resembles a husky more than a cat, and a cat more than a car. The student minimizes this loss, here 1.244, usually as a weighted sum with its regular task loss on the ground-truth labels. The teacher stays frozen throughout and runs only forward passes:

Types of Distillation#

• Response-based: The student matches the teacher's final outputs, as in the soft-target method above, and needs nothing from the teacher beyond its logits. DistilBERT combined it with language modeling and hidden-state losses during pretraining to make BERT 40% smaller and 60% faster while retaining 97% of its language understanding.
• Feature-based: The student matches the teacher's intermediate feature maps, with a small learned projector bridging any difference in channel count. FitNets introduced these "hints" and trained a student with about 10.4 times fewer parameters than its teacher that outperformed it on CIFAR-10.
• Sequence-level: A language model teacher generates text, and the student is fine-tuned on those outputs. DeepSeek fine-tuned its R1 distilled models, based on Qwen2.5 and Llama 3, on 800,000 samples curated with DeepSeek-R1.

Feature-based distillation suits detectors, whose outputs are thousands of boxes and scores rather than one class distribution. Ultralytics YOLO uses this form when training receives a distill_model teacher: a frozen larger YOLO26 model supplies the three neck feature maps that feed its detection head, and the student matches them through 1×1 convolution projectors under a teacher-confidence-weighted L2 loss scaled by dis (6.0 by default). On COCO val2017 the distilled YOLO26 nano model reaches 41.5 mAP_50-95, against 40.9 for the same model trained without a teacher, at identical inference cost.

from ultralytics import YOLO

# Student: the compact YOLO26 nano model that will be deployed
model = YOLO("yolo26n.pt")

# Teacher: a larger YOLO26 checkpoint that guides training and is discarded afterward
results = model.train(data="coco8.yaml", epochs=5, distill_model="yolo26s.pt")

Distillation vs. Related Terms#

• Transfer learning: Reuses a model's own weights as the starting point for a new task. Distillation trains a different, usually smaller architecture to imitate a model on the same task.
• Model pruning: Removes weights, channels, or layers from a trained network. Pruned models are often fine-tuned with distillation from the original to recover accuracy.
• Model quantization: Stores weights and activations at lower precision, such as INT8. It shrinks the same architecture and combines with distillation, which chooses a smaller one.
• Model ensemble: Averages several models at inference time. Distilling an ensemble into one student was the original motivation for the technique.
• Dataset distillation: Despite the similar name, it compresses a dataset instead of a model. Wang et al. compressed MNIST's 60,000 training images into 10 synthetic images, one per class, that bring a network with a fixed initialization close to its original accuracy in a few gradient steps.