This note covers KL divergence specifically as an implemented loss function — the mathematical concept was already introduced in KL Divergence; here the focus is on how it's actually wired into training a real model, especially the input-format gotcha that trips up most beginners.
Formula, Restated
As a loss, \(p\) is typically a fixed target distribution and \(q\) is the model's output distribution being trained to match \(p\).
PyTorch's Input Convention — The Critical Detail
nn.KLDivLoss expects its first argument to already be log-probabilities (i.e. log(q), not \(q\) itself), and its second argument to be plain probabilities (\(p\)). This is a genuinely easy detail to get backward or forget entirely:
import torch
import torch.nn as nn
import torch.nn.functional as F
q_logits = torch.tensor([[1.0, 2.0, 0.5]]) # model's raw output
p_target = torch.tensor([[0.2, 0.5, 0.3]]) # target distribution (already valid probabilities)
log_q = F.log_softmax(q_logits, dim=1) # convert model output to LOG-probabilities first
loss_fn = nn.KLDivLoss(reduction='batchmean')
print(loss_fn(log_q, p_target)) # log_q first, p_target second -- this order mattersNumerical Example
\(p=[0.7,0.3]\), \(q=[0.5,0.5]\) — reusing the exact example from KL Divergence:
Where This Is Used as an Actual Training Loss
| Use Case | What \(p\) and \(q\) Represent |
|---|---|
| VAE regularization term | \(q\) = the encoder's learned latent distribution; \(p\) = a fixed prior (usually standard normal) |
| Knowledge distillation | \(q\) = a small "student" model's output distribution; \(p\) = a larger "teacher" model's output distribution |
| RLHF / policy regularization | \(q\) = the fine-tuned model's output distribution; \(p\) = the original pretrained model's distribution (keeps updates from drifting too far) |
In every case, KL divergence is used as an additional loss term, typically summed with a primary task loss (e.g. reconstruction loss for a VAE, or the standard training loss for the student model in distillation) — rarely the sole loss on its own.
Common Mistakes
- Passing raw probabilities (instead of log-probabilities) as the first argument to
nn.KLDivLoss— this silently computes a wrong value, since the function expects to receive \(\log q\) directly, not \(q\). - Using KL divergence as a standalone classification loss instead of cross-entropy — as shown in KL Divergence, they differ only by a constant when the target is a fixed one-hot label, so cross-entropy is simpler and standard for that specific case; KL divergence earns its keep when both distributions being compared are non-trivial (like a VAE's learned latent distribution).
Interview Relevance
Q: "In knowledge distillation, why is KL divergence used between the student and teacher outputs instead of just cross-entropy against the hard labels?" KL divergence lets the student learn from the teacher's full output distribution — including the relative probabilities assigned to incorrect classes ("dark knowledge"), which often encodes useful information about class similarity that a single hard label discards entirely. Cross-entropy against only the true label would ignore this richer signal.
Practice Question
In a VAE, why does the KL divergence term push the encoder's learned latent distribution toward matching the prior, rather than letting it be arbitrary?
