KL (Kullback-Leibler) divergence measures the extra loss from predicting with the wrong chances when the target is itself a set of chances. Cross-entropy is the target's own spread plus KL, and which side sits in the first slot decides what a distilled student or an RLHF (reinforcement learning from human feedback) policy gives up.
The teacher on the Knowledge Distillation page looks at a cat photo and says cat 0.645, dog 0.194, car 0.096, bird 0.065, illustrative numbers from that page's toy, not a trained model. These are its soft targets at temperature 2, which flattens its chances. If the label turned out to be bird, the teacher would be the most surprised.
Call the teacher's chances y, with yk its chance for class k. Logs are natural logs, so the unit is the nat (log base 2 would give bits).
class
yk
surprise
yk · surprise
cat
0.645
0.44
0.28
dog
0.194
1.64
0.32
car
0.096
2.34
0.22
bird
0.065
2.73
0.18
average surprise
1.00
Entropy is the average surprise you feel if you believe the true chances.
sk = −log yk
sk is the surprise of class k, the table's surprise column: 0 for a sure class, large for an unlikely one.
H(y) = −Σkyk log yk= 1.00
H(y) is the entropy of y, the table's total row. Σk adds over the classes and weights each surprise by its chance, so H is an expected value. It measures how spread out the chances are: 0 for a one-hot target, log 4, about 1.39, for an even split over four classes, and 1.00 for this teacher.
Figure 1. Surprise rises fast as a chance falls. ① cat, ② dog, ③ car, ④ bird.
The student comes from the Knowledge Distillation page's training tool, two of its ten steps in; its chances pk in the table are illustrative too. Graded against the teacher, each class costs minus the log of the student's chance, weighted by the teacher's.
class
yk
pk
cross-entropy
KL
cat
0.645
0.451
0.51
0.23
dog
0.194
0.209
0.30
−0.01
car
0.096
0.175
0.17
−0.06
bird
0.065
0.165
0.12
−0.06
total
1.10
0.10
Cross-entropy is your average surprise when you believe the wrong chances. It splits into the surprise nobody can avoid and the extra caused by believing p instead of y; that extra is KL divergence.
H(y, p) = −Σkyk log pk= 1.10
H(y, p) is the cross-entropy, the table's total: the entropy line with the student's surprise in place of the teacher's, still weighted by yk, since the teacher sets the grading.
H(y, p) − H(y) =Σkyk log ykpk= 0.10
Subtract the entropy class by class and the two logs combine into the log of a ratio, the table's KL column: positive where the student gives too little chance (cat), negative where it gives too much (dog, car, bird).
The gap is named KL(y ‖ p), read "KL of y from p". The double bar separates two slots, and the first slot holds the chances that do the weighting, here the teacher's. H(y) does not depend on the student, so lowering cross-entropy lowers KL by the same amount, and soft-target cross-entropy (Hinton et al. 2015) and a KL term train the same student.
Figure 2. Cross-entropy is spread ① plus KL ②; a one-hot label has no spread.
A one-hot target has no spread, so its entropy is 0 and the KL keeps only the surprise at the correct class: 0.80 for this student with cat correct, exactly the categorical loss. Maximum likelihood, and with it a language model's next-token loss, minimizes this KL with the data in the first slot (Maximum Likelihood and EM, section 6).
The NumPy snippet checks the split.
import numpy as np
# cat, dog, car, bird:
# teacher at T = 2, then student
y = np.array([0.645, 0.194,
0.096, 0.065])
p = np.array([0.451, 0.209,
0.175, 0.165])
H_y = -np.sum(y * np.log(y))
H_yp = -np.sum(y * np.log(p))
kl = np.sum(y * np.log(y / p))
print(H_yp, H_y + kl) # equal
1The teacher and student from the table.
2The entropy, 1.00.
3The cross-entropy, 1.10, and the KL, 0.10.
4Prints the split; both values are 1.10.
Section 3 swaps the two slots and gets a different number.
3 KL is not symmetric
Take an illustrative student trained on the labels alone, with no teacher term, that ends at cat 0.94 and 0.02 on each other class. Score it against the teacher both ways.
class
yk
pk
forward
reverse
cat
0.645
0.94
−0.24
0.35
dog
0.194
0.02
0.44
−0.05
car
0.096
0.02
0.15
−0.03
bird
0.065
0.02
0.08
−0.02
total
0.43
0.25
KL(y ‖ p), target first, is called forward KL; KL(p ‖ y), model first, is reverse KL. Each averages the same kind of log ratio under a different set of chances.
KL(y ‖ p) = Σkyk log ykpk= 0.43
Forward KL is the KL line of the cross-entropy card. The teacher's yk is the weight, so it charges where the teacher expects something: dog, 0.194 for the teacher and 0.02 for the student, gives 0.44.
KL(p ‖ y) = Σkpk log pkyk= 0.25
Reverse KL swaps the slots, so the student's pk is the weight and it charges where the student bets: cat, 0.94 against the teacher's 0.645, gives 0.35. The same lists score 0.43 one way and 0.25 the other, so KL is a directed score, not a distance.
KL is never negative, and zero only when the two lists match, which is what distillation aims for (Cover and Thomas 2006, Theorem 2.6.3). In Figure 3 both scores reach 0 when the student's bars sit on the teacher's outlines.
student, drag a barteacher
Figure 3. ① and ② mark the largest forward and reverse terms.
Take distilling a large language model (LLM) into a small one. For one prompt the teacher has two equally good kinds of answer, and the small student can produce only one. Picture all answers on one line, as an illustrative sketch: the teacher is two equal bumps at −2 and +2, each with standard deviation 0.6, and the student is one bell curve.
Section 3's two directions carry over with area in place of a sum: each charges the log ratio wherever its first slot puts weight.
KL(y ‖ p) = ∫ y(x) log y(x)p(x) dx
y(x) and p(x) are the two curves' heights at x; chance is area under a curve, and ∫ … dx adds up area. The weight y(x) is large on both bumps, where a low p(x) makes the log ratio large, so the forward fit stays high over both bumps and, being one curve, over the valley between.
KL(p ‖ y) = ∫ p(x) log p(x)y(x) dx
Now the weight is the student's p(x). Where the student puts chance and y(x) is near 0 the log ratio is huge, so the reverse fit avoids the valley and sits on one bump; the missed bump costs almost nothing, since p(x) is near 0 there.
Figure 4. Forward KL fills the valley ①; reverse KL misses one bump ②.
In this sketch the forward fit puts about 19% of its chance in the middle of the valley, between −0.5 and +0.5, where the teacher has under 1%: a student that generates answers the teacher finds unlikely.
The picture follows Bishop (2006, chapter 10). A mode is a peak of a distribution, one bump here, so forward KL is called mode-covering and reverse KL mode-seeking. MiniLLM (Gu et al. 2024) switches LLM distillation to reverse KL for the reason in the forward panel ①.
5 KL in distillation, RLHF and t-SNE
The running student's KL, 0.098 from the rounded chances in Section 2's illustrative table, is the KL inside the soft part of the Knowledge Distillation page's loss, at that page's temperature T of 2 and soft weight α of 0.7.
step
value
KL of teacher and student at temperature 2
0.098
times T2, which is 4
0.39
times α, which is 0.7: the soft part of the loss
0.27
In RLHF the policy π is the language model being tuned, πSFT is the frozen supervised fine-tuned (SFT) model it starts from, called the reference model, and the reward model scores answers the way people would. t-SNE (t-distributed stochastic neighbor embedding) draws high-dimensional data as a 2D map.
All three are the KL of Section 2. Distillation scales it by a weight, RLHF scales its per-answer form by a coefficient, and t-SNE uses it as is.
L = (1 − α) · CE(z, label)+ α · T2 · KL(y ‖ p)
y and p are the teacher's and student's chances at temperature T (the Knowledge Distillation page calls them pt and ps), and CE, the cross-entropy, scores the student's logits z (its raw scores before the softmax), at temperature 1, against the true label. α weighs the soft part: raise it and the student copies the teacher more and the labels less. T2 undoes the shrink that temperature puts on the soft term's gradient (Knowledge Distillation, section 2). With the teacher in the first slot, the student must cover every class the teacher expects.
reward = r− β · log π(a)πSFT(a)
r is the reward model's score for one sampled answer a to one prompt. β is the KL coefficient, 0.02 in InstructGPT (Ouyang et al. 2022, Appendix C.4): raise it and the policy stays closer to the SFT model; lower it and the policy chases the reward model (reward over-optimization). The log ratio is computed per token and summed over the answer. Answers are sampled from π, so averaged over them it is KL(π ‖ πSFT), reverse KL: the policy is charged for answers the SFT model finds unlikely. With an illustrative answer the policy finds 3 times likelier than the SFT model does, the penalty is about 0.022 (β times log 3), so r drops by about 0.02.
t-SNE cost = KL(P ‖ Q)
For each point, P holds the chance it would pick each other point as its neighbor, high for points close in the original space, and Q holds the same chances measured on the map; P fills the first slot (van der Maaten and Hinton 2008). Forward KL charges a map that puts true neighbors far apart.
Figure 5. Which side fills the first slot in each loss.
6 KL in PyTorch
F.kl_div takes the student's log-chances first and the teacher's chances second, the opposite order from the math notation, and returns KL(y ‖ p), forward KL with the teacher first. Swapped arguments raise no error. nn.KLDivLoss is the same function as a module.
Get the log-chances from F.log_softmax on the logits, not from the log of a softmax, which can hit log 0.
Use reduction="batchmean": the default "mean" also divides by the number of classes, and the docs warn it is not the true KL. For an illustrative batch of two, the running student (KL 0.10) and a student that gives every class 0.25 (KL 0.38), batchmean gives 0.24 and mean gives 0.06.
A zero target chance adds 0; a zero model chance where the target is positive makes KL infinite.
The PyTorch snippet returns the NumPy snippet's 0.10.
import torch
import torch.nn.functional as F
# a batch of one: cat, dog, car,
# bird; the teacher's chances
y = torch.tensor([[0.645, 0.194,
0.096, 0.065]])
# the student's logits
z = torch.log(torch.tensor(
[[0.451, 0.209,
0.175, 0.165]]))
log_p = F.log_softmax(z, dim=-1)
# returns 0.10
kl = F.kl_div(log_p, y,
reduction="batchmean")
1The teacher, a batch of one.
2Student logits, here the log of its chances.
3Log-chances (point 2).
4With the student's log-chances first and the teacher's chances second, the call returns KL(y ‖ p); batchmean divides by the batch size (points 1 and 3).
7 References
Bishop, C. M. (2006). Pattern Recognition and Machine Learning. Springer. Chapter 10.
Cover, T. M., and Thomas, J. A. (2006). Elements of Information Theory, 2nd ed. Wiley. Chapter 2.
Gu, Y., et al. (2024). MiniLLM: Knowledge Distillation of Large Language Models. ICLR. arXiv:2306.08543.
Hinton, G., Vinyals, O., and Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531.
Kullback, S., and Leibler, R. A. (1951). On Information and Sufficiency. Annals of Mathematical Statistics 22(1), 79 to 86.
Ouyang, L., et al. (2022). Training Language Models to Follow Instructions with Human Feedback. NeurIPS. arXiv:2203.02155.
PyTorch documentation for torch.nn.functional.kl_div, torch.nn.KLDivLoss and torch.nn.CrossEntropyLoss.
van der Maaten, L., and Hinton, G. (2008). Visualizing Data using t-SNE. JMLR 9, 2579 to 2605.