What is knowledge distillation in deep neural networks?
Knowledge distillation (KD) transfers useful predictive behavior from a trained teacher model to a student model, often a smaller and cheaper neural network. Instead of learning solely from one-hot ground-truth labels, the student also learns from the teacher's soft probability distribution over possible answers. That distribution reveals relationships among classes that a single correct-label target does not express.
For example, a teacher might be highly confident that an image contains a cat but assign a moderate probability to dog and almost none to airplane. The small probability differences can teach the student that cats and dogs share more visual characteristics than cats and airplanes.
The classic method was described by Hinton, Vinyals and Dean in *Distilling the Knowledge in a Neural Network* (2015). This lesson concerns teacher–student learning, softened outputs, temperature, KL divergence and the deployment of a trained student, not pruning or quantization.
1. What do the teacher and student do?
- Teacher: A well-trained model, large architecture or ensemble that provides informative outputs. It is usually frozen during conventional offline distillation.
- Student: A smaller or otherwise more deployment-friendly network, trained to predict the ground-truth classes *and* imitate the teacher's behavior.
- Hard labels: Ordinary targets such as
[1,0,0]for cat, dog and airplane give the correct class but say nothing about similarities among the other classes. - Soft targets: A teacher probability distribution such as
[0.867,0.117,0.016]encodes the teacher's confidence across all three classes. These values are *teacher predictions*, not verified true probabilities of the scene. - Deployment: After training, use the student alone for inference in a standard offline KD setup. The teacher is normally not needed for every prediction.
Distillation is not simply copying teacher weights. Teacher and student may have different architectures or parameter counts. The knowledge transfer occurs through loss terms built from predictions or intermediate representations.
2. Why use a temperature in softmax?
Given a vector of class logits z=[z₁,…,z_C], compute a softened class distribution:
p_i(T) = exp(z_i/T) / Σ_j exp(z_j/T), for temperature T > 0.
- T=1: The usual predictive softmax distribution.
- T>1: Logit differences are divided by a larger number, producing a generally softer, less peaked distribution. Smaller-probability alternatives become easier for the student to observe.
- T<1: Produces a generally sharper distribution, emphasizing the top prediction.
- T very large: For fixed finite logits, softmax approaches a near-uniform distribution, which may dilute useful information unless the objective is designed appropriately.
- T→0 positive: For a unique maximum logit, the output becomes almost one-hot; numerical implementations should not literally divide by zero.
Temperature is a training hyperparameter, not an inherent confidence-calibration guarantee. Choosing a high T does not necessarily improve the student's generalization.
3. Beginner-friendly numerical example: Three classes and two networks
Consider class order [cat, dog, airplane]. Suppose one image has the true label cat (index 0).
For illustration, the teacher and student currently output these logits (raw, unnormalized scores):
- Teacher logits:
z_teacher = [4,2,0]. - Student logits:
z_student = [2,1,0].
These are hypothetical outputs, not probabilities. We will use exactly the same logits for all the calculations below.
Step A — Teacher probabilities at T=1
p_teacher(1) = softmax([4,2,0]).
The exponentials are approximately [54.5982,7.3891,1], and their sum is approximately 62.9873.
- Teacher probabilities at T=1: `[0.866813, 0.117310, 0.015876]`.
The teacher favors cat, assigns a smaller but meaningful probability to dog, and assigns very little probability to airplane.
Step B — Student probabilities at T=1
p_student(1) = softmax([2,1,0]).
Exponentials are approximately [7.3891,2.7183,1]; their sum is approximately 11.1073.
- Student probabilities at T=1: `[0.665241, 0.244728, 0.090031]`.
The student predicts cat but is less confident than the teacher.
Step C — Increase temperature to T=2
Divide both sets of logits by 2 before softmax:
- Teacher softened logits:
[4,2,0]/2 = [2,1,0]. - Student softened logits:
[2,1,0]/2 = [1,0.5,0].
Apply softmax again:
- Teacher at T=2: `P_T = [0.665241, 0.244728, 0.090031]`.
- Student at T=2: `Q_T = [0.506480, 0.307196, 0.186324]`.
Now the teacher's non-cat predictions are more visible than at T=1: dog carries about 24.5% and airplane about 9.0% of the *softened teacher distribution*. Those are not claims about real-world class truth; they are training signals describing the teacher's comparative judgments.
Step D — Calculate ordinary hard-label cross-entropy
The true class is cat, so hard-label cross-entropy of the student is computed at T=1:
L_hard = −log(p_student(1)[cat]) = −ln(0.665240956) ≈ 0.407606.
Hard-label CE looks only at the student's assigned probability to the correct class. It does not directly target the complete teacher distribution.
Step E — Calculate the teacher-to-student KL divergence
Use the standard direction KL(P_T || Q_T):
KL(P_T || Q_T) = Σ_i P_T(i) ln(P_T(i)/Q_T(i)).
For our example, the three summands are approximately:
- Cat:
0.665241 × ln(0.665241/0.506480) ≈ +0.181387. - Dog:
0.244728 × ln(0.244728/0.307196) ≈ −0.055636. - Airplane:
0.090031 × ln(0.090031/0.186324) ≈ −0.065483.
Sum: `KL(P_T||Q_T) ≈ 0.060269`.
Individual KL terms can be negative; the total KL divergence remains nonnegative for valid probability distributions.
Step F — Apply temperature-squared scaling
A common student-training distillation term is:
L_soft = T² × KL(P_T || Q_T).
At T=2, T²=4:
L_soft ≈ 4 × 0.0602688975 = 0.241076.
Why T²? For fixed logits and suitably high temperatures, the softmax gradients shrink approximately as 1/T². The T² factor compensates for this tendency and helps preserve the relative strength of the soft-target learning signal. This is a training convention motivated by gradient behavior, not a probability normalization operation or a rule that every distillation objective must use.
Step G — Combine hard-label and distillation losses
Define the explicit weighting convention:
L_total = (1−α) × L_hard + α × T² × KL(P_T || Q_T).
Choose α=0.7, so the soft-target term receives weight 0.7 and hard-label CE receives weight 0.3:
- Hard contribution:
0.3 × 0.4076059644 ≈ 0.122282. - Soft contribution:
0.7 × 0.2410755900 ≈ 0.168753. - Total: `L_total ≈ 0.122282 + 0.168753 = 0.291035`.
Meaning: The student is encouraged both to predict the correct class and to approximate the teacher's probability pattern. The number 0.291035 is a single-example toy training loss—not accuracy, runtime, or evidence of a successful production student.
4. What does "dark knowledge" add beyond one-hot labels?
A one-hot cat label is [1,0,0], giving identical zero targets for dog and airplane. It does not say which wrong answer is less wrong under the teacher's learned representation.
Our teacher's softened targets at T=2 are:
- Cat: 0.665241.
- Dog: 0.244728.
- Airplane: 0.090031.
The teacher assigns considerably more weight to dog than airplane. That extra structure can guide a smaller student about similarities and decision boundaries.
Important limitation: This is the teacher's belief, which can be wrong, overconfident, biased, or distorted by out-of-distribution inputs. A weak or poorly calibrated teacher can teach a student undesirable behavior. Distillation does not make inaccurate soft targets true.
5. Why exactly KL, and which distribution goes first?
The standard distillation KL objective is KL(teacher || student):
KL(P_teacher || Q_student) = Σ_i P_teacher(i) log(P_teacher(i)/Q_student(i)).
- The teacher provides a fixed target in offline distillation, and the student is optimized to match it.
- With a fixed teacher, minimizing KL(teacher||student) is equivalent for student gradients to minimizing teacher-target cross-entropy, because teacher entropy is constant with respect to student weights.
- KL is asymmetric: KL(student||teacher) is a different objective and can emphasize errors differently.
- PyTorch's
F.kl_divexpects the student log-probabilities as its first argument and the teacher probabilities as the second argument. This API order is easy to misread; its computed mathematical divergence here is still KL(teacher||student). - Use
reduction="batchmean"for the correct per-example KL sum over classes, rather thanreduction="mean", which divides by all elements and changes the loss scale. - For extreme logits, a log-space teacher target with
log_target=Truecan improve numerical robustness.
6. Temperature, alpha, and the learning trade-offs
- T too low: The teacher distribution can be very peaked; small differences between non-target classes are mostly hidden.
- T moderately high: Softer teacher targets may reveal relationships among classes, making imitation more informative.
- T too high: Probabilities can become almost uniform, possibly losing useful differentiation; a large T² multiplier does not magically restore all lost information.
- α close to 0: Training relies mostly on hard-label ground truth, limiting the benefit of transfer from the teacher.
- α close to 1: Training relies mostly on the teacher, risking stronger inheritance of teacher errors or biases.
- No universal optimum: Tune T and α on a validation set, monitoring both accuracy and deployment characteristics. Papers and libraries may define α as the hard loss weight rather than the soft weight; check the equation, not merely the parameter name.
- Class distribution and imbalance matter: Minority or rare classes can be harmed if teacher errors dominate, so assess per-class and tail metrics.
Additional nuance: Changing temperature affects teacher *and* student probability distributions. Use the same T for both branches in the KL term, while computing conventional hard-label CE on the unsoftened student logits.
7. Practical PyTorch training example
This example uses a fixed pretrained teacher, a trainable student, and the exact combined loss convention from the numerical exercise.
PYTHONimport torch import torch.nn.functional as F T = 2.0 alpha = 0.7 # Weight on teacher/student soft-target KL term teacher.eval() # Freeze teacher behavior (dropout, batch norm) teacher.requires_grad_(False) student.train() # Optimize student weights for images, labels in train_loader: images = images.to(device) labels = labels.to(device) optimizer.zero_grad(set_to_none=True) # Offline KD: teacher serves as a fixed target. with torch.no_grad(): teacher_logits = teacher(images) student_logits = student(images) hard_ce = F.cross_entropy(student_logits, labels) teacher_probs = F.softmax(teacher_logits / T, dim=-1) student_log_probs = F.log_softmax(student_logits / T, dim=-1) # First input = student log probabilities, # second target = teacher probabilities. soft_kl = F.kl_div( student_log_probs, teacher_probs, reduction="batchmean" ) loss = (1 - alpha) * hard_ce + alpha * (T * T) * soft_kl loss.backward() optimizer.step() # For deployment, evaluate and use student.eval() without teacher.
Code assumptions: teacher, student, train_loader, device, and the student's optimizer must be provided by the application. Teacher and student must output aligned class logits in matching order and be evaluated on the same inputs. Both outputs have shape [batch_size, class_count].
Good practices:
- No teacher gradient updates:
teacher.eval()alone does not prevent autograd;torch.no_grad()and/or freezing parameters avoid unnecessary teacher-gradient work. - Do not detach student logits: The student requires gradients from both loss terms.
- Use log-softmax: More numerically stable than taking the log of an already-rounded softmax.
- Correct KL reduction:
batchmeansums over classes and averages over the batch. - Separate training and inference cost: Teacher evaluation makes distillation training more expensive, but the deployed student does not require the teacher.
8. Offline, online, and self-distillation
Offline distillation
- Train or obtain a strong teacher.
- Freeze the teacher.
- Train the student using teacher soft targets, ground truth, or selected intermediate teacher features.
- Deploy the student alone.
Advantages: Teacher targets remain stable; simpler optimization; can precompute teacher logits for a fixed dataset if the data pipeline permits.
Limitations: Requires access to a suitable teacher; storing teacher logits has a cost; teacher mistakes may propagate.
Online distillation
- Multiple networks train simultaneously, exchanging prediction distributions or learning from a changing ensemble/peer target.
- Deep Mutual Learning is one example of peer networks teaching each other during training.
- A fixed, already-trained large teacher is not required in every online scheme.
Cautions: Targets change as networks learn; there is additional training compute and risk of mutually reinforcing mistakes. Keep the teacher/peer-gradient policy explicit.
Self-distillation
- Uses a model's deeper outputs, an earlier checkpoint, an exponential-moving-average teacher, or another part/version of the same architecture to supervise student branches or subsequent iterations.
- Does not necessarily shrink the architecture; any speed benefit depends on the model ultimately deployed.
9. Intermediate feature distillation: Learning more than logits
The student can also imitate hidden representations, attention patterns, feature maps, or other internal signals.
- Feature matching: Choose teacher and student hidden layers that represent comparable information and add a loss such as
||h_student − h_teacher||². - Shape mismatch: If the teacher's hidden feature has 64 channels and the student's has 32, a learned projection or adapter can map between dimensions before comparing them.
- FitNets: Romero et al. proposed using intermediate teacher "hints" to help train thin, deep student networks.
- Feature normalization/alignment: Match spatial resolution, masking, channel order and representation statistics appropriately; otherwise an MSE may reward arbitrary alignment rather than useful knowledge.
- Total objective: A system can combine label CE, logits-based distillation and a feature loss with separately chosen weights.
- Inference: A training-only projection head may be discarded when not required by the student architecture; check that it truly is not needed for deployment.
Do not assume teacher and student must have the same architecture. Logits-based KD needs aligned output labels; feature-based KD usually requires additional alignment.
10. Does distillation actually make a model faster or smaller?
It can, if the student architecture is designed to be cheaper.
- Smaller parameter count: A student with fewer layers/channels or smaller embeddings generally needs less raw parameter storage, all else equal.
- Lower inference computation: A genuinely smaller student may have fewer MACs/FLOPs, but this must be profiled; shapes, operator kernels and memory transfer matter.
- Accuracy versus latency trade-off: A student may approach teacher quality at much lower cost, but can also lose accuracy, calibration, robustness or rare-case performance.
- No automatic speed gain: Distilling into an unchanged architecture changes weights, not necessarily compute or model size.
- Training overhead: Teacher pretraining and extra teacher passes/soft-target loss consume resources; evaluate that cost against expected inference savings.
- Student only at inference: In the conventional offline scenario, removing the teacher from the deployment path is essential for the hoped-for cost reduction.
Illustrative deployment budget, not a benchmark: If a teacher has 40 million FP32 parameters and a student has 10 million FP32 parameters, raw parameter storage is approximately 160 MB versus 40 MB (decimal). The student has 75% fewer parameters and 4× less raw parameter storage. This does not prove 4× lower latency, memory or power for the full running system.
11. How should KD be evaluated?
- Establish three baselines: Teacher accuracy; student trained only on hard labels; and student trained with KD under the same task and student architecture.
- Keep comparisons fair: Match train/validation splits, augmentations, student compute budget and optimizer setup when comparing student training methods.
- Sweep T and α: Test a modest grid and retain the values that actually improve held-out metrics.
- Check classwise behavior: Evaluate rare labels, calibration, confusion patterns, out-of-distribution data and distribution shifts—not just average accuracy.
- Profile deployment: Measure student parameter count, serialized size, peak memory, latency and throughput on the actual target device and batch size.
- Record teacher-training cost: Decide whether extra teacher inference/training is justified by repeated inference savings.
- Compare against alternatives: A hard-label-only small student, architecture redesign, and suitable inference optimization can sometimes meet goals without KD.
Failure signs: A poor teacher, incorrect class ordering, mismatched preprocessing, badly weighted KL, excessively high temperature, or student capacity too small for the teacher behavior can all produce inferior students.
12. Common mistakes and limitations
- Using one-hot labels as the teacher distribution: This removes the relative class information that makes classical soft-target distillation useful.
- Softmaxing at T=1 for the KD term but multiplying by T² for T>1: Teacher and student KL branches must consistently use the intended temperature.
- Reversing KL:
F.kl_div(log_student, p_teacher)is the correct argument order for KL(teacher||student). - Using `reduction="mean"` unknowingly: That scales KL by class count and changes the effective α.
- Applying temperature to the hard-label CE: Conventional hard CE uses unscaled student logits; temperature is for the soft loss.
- Forgetting the T² convention: May distort the relative gradient magnitudes of hard/soft components when tuning T.
- Using `teacher.eval()` without disabling gradients: Eval controls stochastic layers, not whether the framework stores autograd graphs.
- Updating the teacher unintentionally: Offline KD should keep teacher targets fixed; set no-grad/freeze clearly.
- Missing label alignment: Student and teacher class indices must represent the same classes in the same order.
- Assuming a smaller architecture always retains accuracy: The student may lack capacity or be unable to learn teacher behavior.
- Assuming KD equals network compression: A same-size student can learn through KD without becoming smaller.
- Assuming the teacher is correct: Distillation can transfer biases, errors and overconfidence.
- Ignoring inference measurements: Parameter count alone does not establish the deployed speed.
- Confusing learning and structural removal: KD learns the student from outputs/features; it does not itself prune channels or simulate INT8 rounding.
Curriculum boundaries: Question 297 explained quantization-aware training; Question 298 compared structured and unstructured pruning. Question 299 covers teacher–student soft targets, temperature and distillation losses, not pruning or quantization. Question 300 will explain Grad-CAM, a method for visualizing which CNN regions influence predictions; it is an interpretability lesson, not a teacher–student training technique.
Key takeaway: Knowledge distillation trains a student to learn from both ground truth and a teacher's softened predictions. Temperature exposes information among classes, KL divergence measures the teacher/student distribution mismatch, and the common T²-weighted soft loss complements hard-label cross-entropy. A smaller deployed student can reduce inference cost, but accuracy and latency gains must be measured.
Research and technical references: Hinton, Vinyals & Dean, *Distilling the Knowledge in a Neural Network*; PyTorch, *Knowledge Distillation Tutorial*; PyTorch, *torch.nn.functional.kl_div*; Romero et al., *FitNets: Hints for Thin Deep Nets*; Zhang et al., *Deep Mutual Learning*.
Sources & Research References
Learning Sequence
Deep Learning Curriculum
What is Grad-CAM and how does it explain CNN predictions?
Explore Related Questions
Deepen your understanding across related Deep Learning topics.