Knowledge Distillation Explained: Logits, Softmax Temperature, and Why the Loss Is Multiplied by T² (Hinton 2015)

Knowledge distillation: logits z, temperature T, and the T² factor

 

Knowledge distillation is a way to transfer what a large model (the teacher) has learned into a smaller model (the student). The method proposed in the 2015 paper "Distilling the Knowledge in a Neural Network" by Hinton, Vinyals, and Dean is still the most basic form of it.

 

This post walks through the three core ideas of that paper for a classifier whose output layer is a softmax, using equations, tables, and charts.

 

📌 Questions this post answers

Q1. What exactly is z in the equations? → the logit, the value right before the softmax
Q2. What does the temperature T do? → it makes the probability distribution softer
Q3. Why is the loss multiplied by T²? → because raising T shrinks the gradient by 1/T²

 

1. Notation first

Before reading the equations, let's agree on what each symbol means.

 

Symbol	Meaning	Note
zi	The student's logit for class i	Output of the last linear layer, before any activation
vi	The teacher's logit	The paper writes the teacher's logits as v
T	Temperature	T = 1 gives the ordinary softmax
qi	The student's softmax probability	Softmax of z divided by T
pi	The teacher's softmax probability = soft target	Softmax of v divided by the same T
y	The true label (one-hot) = hard target	1 for the correct class, 0 elsewhere
N	Number of classes	10 for digit classification

 

2. z is the logit, the raw score before the softmax

The last linear layer of a classifier outputs one real number per class. That number is the logit z. It has no fixed range: it can be 8.0 or −1.0. The softmax turns logits into probabilities, and distillation adds one more ingredient: the temperature T.

 

Softmax with temperature T. With T = 1 this is the ordinary softmax.

 

Dividing the logits by T changes the result a lot. Suppose the teacher produces these logits for an image of a handwritten "2":

 

Class	2	3	7	1	8	0	4	5	6	9
Logit v	8.0	2.0	1.5	0.5	0.0	0.0	−0.5	−0.5	−1.0	−1.0

 

Here is the softmax of those same logits at different temperatures (computed directly):

 

Temperature T	P("2")	P("3")	P("7")	Other 7 classes	Entropy
1 (ordinary softmax)	99.4%	0.2%	0.1%	0.3%	0.05
2	83.4%	4.2%	3.2%	9.2%	0.79
5	34.5%	10.4%	9.4%	45.7%	2.07
10	19.6%	10.7%	10.2%	59.5%	2.26
20	14.1%	10.5%	10.2%	65.2%	2.29

 

As T grows, the same logits give a flatter distribution, and the runner-up classes "3" and "7" become visible.

 

• T = 1: the exponential blows up the differences between logits, so almost all probability goes to "2". The fact that "3" gets 0.2% and "7" gets 0.1% is effectively zero and provides almost no learning signal.
• Larger T: dividing by T shrinks the gaps between logits, so the distribution softens. The maximum entropy is ln 10 ≈ 2.30, so at T = 20 it is already close to uniform.
• T → ∞: every class converges to 1/N (10%).

 

3. Why a soft distribution is useful: soft targets

The true label y only says "this is a 2". The teacher's soft distribution says more: "it's a 2, but if anything it looks more like a 3 than a 7." The paper calls this knowledge, hidden in the relative probabilities of the wrong classes, "dark knowledge."

 

	Hard target y	Soft target p (T > 1)
Form	[1, 0, 0, …, 0]	[0.35, 0.10, 0.09, …] (T = 5 example)
Information	Only the correct class	Correct class + similarity between classes
Information per example	Low	High (higher entropy)
Gradient variance	High	Low → trains well even with less data

 

4. The distillation loss

The student learns two objectives at the same time:

 

1. Soft loss L_soft: cross-entropy that pushes the student to match the teacher's distribution p^(T) at a high temperature T. Teacher and student use the same T.
2. Hard loss L_hard: the ordinary cross-entropy with the true label y, at T = 1.

 

Distillation loss = T²-corrected soft loss + hard loss, as a weighted sum

 

The paper uses a weighted average of the two and reports the best results when the hard loss gets a considerably smaller weight. After training, the student runs inference at T = 1. The remaining question is why the soft loss is multiplied by T².

 

5. Why multiply by T²: the gradient shrinks by 1/T²

5-1. The gradient of the soft loss

Differentiating the cross-entropy C with soft targets with respect to the student logit z_i gives a very simple expression:

 

Gradient of softmax + cross-entropy. The division of z by T contributes one factor of 1/T through the chain rule.

 

That is one factor of 1/T. On its own it does not make 1/T². The second 1/T comes from the fact that (q_i − p_i) itself shrinks as T grows.

 

5-2. The high-temperature approximation

When T is large compared with the logits, we can use exp(x) ≈ 1 + x:

 

Substituting exp(z/T) ≈ 1 + z/T

 

If the logits are zero-meaned for each example (∑z_j = ∑v_j = 0), the denominators become N and the expression simplifies:

 

With zero-mean logits, the gradient is 1/(N·T²) × (z − v)

 

So the 1/T² comes from two places:

 

Source	Size	Why
Chain rule	1/T	The softmax input is z/T, not z
Probability gap (q − p)	≈ (z − v)/(N·T)	As T grows, both distributions approach uniform, so their difference shrinks
Total	≈ (z − v)/(N·T²)	The two effects multiply to 1/T²

 

5-3. Checking with numbers

Using the teacher logits above and slightly different student logits, I computed the size (vector norm) of the gradient directly. For comparison, the hard-loss gradient (T = 1) in this example is about 0.0075.

 

Temperature T	‖∂Lsoft/∂z‖	T² × ‖∂Lsoft/∂z‖
5	0.00595	0.149
10	0.00195	0.195
20	0.00054	0.216
50	0.00009	0.227

 

Left: the soft-loss gradient falls quickly as T grows. Right: after multiplying by T², it stays at a similar size.

 

Raising T from 5 to 50 makes the gradient about 65× smaller. Without T², the soft loss would be drowned out by the hard loss as T grows and would barely contribute to training. With T², the values stay within 0.15–0.23, so changing T does not throw off the balance between the two losses. The larger T is, the more accurate the approximation, and the closer the value gets to constant.

 

Multiplying by T² cancels T, so the gradient size no longer depends on T.

 

💡 Why it matters in practice · In distillation you tune T and the weight α together. Without T², every change of T would force you to re-tune α. With T², you can change T in experiments and the relative contribution of the two losses stays roughly the same. It is less a mathematical necessity than an engineering correction that makes hyperparameter search easier.

 

6. Logit matching is a special case of distillation

Before distillation, Caruana and colleagues (Bucila et al. 2006, Ba & Caruana 2014) compressed large models into small ones by training the student to match the teacher's logits themselves, i.e., by minimizing (z − v)². Look again at the result of 5-2:

 

At high temperature, the distillation gradient is proportional to the gradient of the squared logit error.

 

So when T is very large and the logits are zero-meaned, distillation is equivalent to logit matching. On this basis, the paper treats logit matching as a special case of distillation.

 

	Logit matching (Ba & Caruana)	Distillation (Hinton et al.)
Target	The logits v themselves	Softmax probabilities at temperature T, p(T)
Loss	½(z − v)²	Cross-entropy H(p(T), q(T))
Relationship	Equivalent in the limit T → ∞ (with zero-mean logits)
At intermediate T	Copies even very negative logits exactly	Pays less attention to very negative, noisy logits → can help when the student is small

 

7. Implementing it in PyTorch

Here is the loss above, translated directly. KL divergence differs from cross-entropy only by the teacher's entropy, which is a constant, so the gradient for the student is the same.

 

import torch.nn.functional as F

def distillation_loss(student_logits, teacher_logits, labels, T=4.0, alpha=0.9):
    # 1) soft loss: both use the same temperature T
    soft_student = F.log_softmax(student_logits / T, dim=1)
    soft_teacher = F.softmax(teacher_logits / T, dim=1)
    soft_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (T * T)  # T^2 correction

    # 2) hard loss: ordinary cross-entropy at T = 1
    hard_loss = F.cross_entropy(student_logits, labels)

    return alpha * soft_loss + (1 - alpha) * hard_loss

 

• The teacher is not trained, so compute its logits inside torch.no_grad() or call .detach().
• Remove * (T * T) and change T from 2 to 8: you will see the soft-loss value and its gradient drop sharply.

 

Summary

Concept	In one line
Logit z	The raw output of the last linear layer, right before the softmax
Temperature T	Divide logits by T to soften the distribution. T > 1 only during training; inference uses T = 1
Soft targets	The teacher's high-temperature softmax output; carries class similarity (dark knowledge)
Loss	α·T²·H(p(T), q(T)) + (1 − α)·H(y, q(1))
Why T²	Gradient = 1/T from the chain rule × 1/T from the probability gap = 1/T². T² cancels it
Logit matching	A special case: equivalent to distillation as T → ∞ with zero-mean logits

 

Reference · G. Hinton, O. Vinyals, J. Dean, "Distilling the Knowledge in a Neural Network", NIPS 2014 Deep Learning Workshop (arXiv:1503.02531). The numbers in the tables and charts were computed directly from the example logits assumed in this post.