Loss Functions › KL Divergence

Kullback-Leibler (KL) divergence measures how much one probability distribution $q$ differs from a reference distribution $p$. In neural network training, it appears when we want the model's output distribution to match a target distribution, not just predict a single class.

$$D_{\text{KL}}(p \| q) = \sum_{k=1}^{K} p_k \log \frac{p_k}{q_k}$$

where:

• $p_k$: the probability of event $k$ under the reference distribution $p$
• $q_k$: the probability of event $k$ under the approximate distribution $q$
• $\log \frac{p_k}{q_k}$: the log ratio, measuring how much $q$ deviates from $p$ for event $k$

KL divergence is non-negative (by Jensen's inequality) and equals zero if and only if $p = q$. Importantly, it is asymmetric: $D_{\text{KL}}(p \| q) \neq D_{\text{KL}}(q \| p)$. When $p$ is the true distribution, $D_{\text{KL}}(p \| q)$ penalizes cases where $q$ assigns low probability to events that $p$ considers likely. This is the relevant direction for model training.

The relationship between cross-entropy and KL divergence is:

$$\mathcal{L}_{\text{CE}}(p, q) = H(p) + D_{\text{KL}}(p \| q)$$

where:

• $H(p) = -\sum_k p_k \log p_k$: the entropy of the true distribution $p$, a constant with respect to model parameters
• $D_{\text{KL}}(p \| q)$: the KL divergence between $p$ and $q$

Since $H(p)$ is constant during training, minimizing cross-entropy is equivalent to minimizing KL divergence. This connection explains why cross-entropy is the correct loss for maximum likelihood estimation: it minimizes the information gap between the model's distribution and the true distribution.

KL divergence appears explicitly in variational autoencoders (VAEs) as a regularization term, and in knowledge distillation, where it is used to train a student model to match a teacher model's output distribution. In distillation, the student is trained to minimize $D_{\text{KL}}(p_{\text{teacher}} \| p_{\text{student}})$: the teacher's soft probability distribution provides richer training signal than one-hot labels, because it encodes the teacher's beliefs about how similar various classes are to each other.

The asymmetry of KL divergence has practical consequences. $D_{\text{KL}}(p \| q)$ is large when $q$ assigns low probability to events that $p$ considers common (mode-seeking behavior). $D_{\text{KL}}(q \| p)$ is large when $q$ spreads probability mass over regions where $p$ has low probability (mean-seeking behavior). For model training, you generally want mode-seeking behavior: learn to correctly handle the common cases well. This is why cross-entropy (which minimizes $D_{\text{KL}}(p_{\text{data}} \| p_{\text{model}})$) is preferred over the reverse direction.