I've recently encountered different articles that are recommending to use the KL divergence instead of the MSE/RMSE (as the loss function), when trying to learn a probability distribution, but none of the articles are giving a clear reasoning why the first is better than the others.

Could anyone give me a strong argument why the KL divergence is suitable for this?


2 Answers 2


In the context of Variational Inference (VI): the KL allows you to move from the unknown posterior $p(z \mid x)$, to the known joint $p(z,x)=p(x|z)p(z)$ and optimize only the ELBO. You cannot do this with L2.

$p(z|x)$ is the desired posterior, of which you cannot calculate the evidence (i.e., using Bayes formula we can set: $p(z|x) = \frac{p(x|z)p(z)}{\int_z p(x|z)p(z)dz}$, and you can't calculate the integral in the denominator [also denoted by $p(x)$] due to it's intractability).

Now suppose $q$ is a Variational distribution (e.g. a family of Gaussians which you can control), VI tries to approximates $p(z|x)$ by $q$ by minimizing their KL divergence.

$$KL(q(z)||p(z|x)) = \int_z q(z) \log \frac{q(z)}{p(z|x)}dz = \mathbb E_q[\log q(z)]-\mathbb E_q[\log \frac{p(x|z)p(z)}{p(x)}] =$$

$$ -\mathbb E_q[\log p(x|z)p(z)] + \mathbb E_q[\log q(z)] + \mathbb E_q[\log p(x)] = -ELBO(q) + \log p(x) $$

Since you're only optimizing $q$ (it's the only thing you can control), you can discard the unknown and difficult to compute normalizing constant $p(x)$.

If you would use the (squared) L2 norm you would get: $$\int_z [q(z)-p(z|x)]^2dz = \int_z [q^2(z)-2q(z)p(z|x)+p^2(z|x)]dz $$ While the 3rd term doesn't depend on q, the 2nd term does, and it also requires $p(x)$ to compute.

  • $\begingroup$ Maybe you should clarify that $q$ is the variational distribution (e.g. a family of Gaussians), i.e. the distribution that you're using to approximate the unknown one that you want to learn. From this answer, it's not fully clear why the KL divergence is really better than the L2 norm, which you seem to have defined as an expectation with respect to $z$. I recommend that you edit your post in order to clarify it. $\endgroup$
    – nbro
    May 30 at 8:17
  • $\begingroup$ I added a crash course in VI, but some knowledge is needed of VI in order to understand this. As implied by the "In the context of Variational Inference" that starts the answer. $\endgroup$ May 30 at 10:49

KL-divergence is a measure on probability distributions. It essentially captures the information loss between ground truth distribution and predicted.

L2-norm/MSE/RMSE doesn't do well with probabilities, because of the power operations involved in the calculation of loss. Probabilities, being fractions under 1, are significantly affected by any power operations (square or root), and considering we are calculating the squares of differences of probabilities, the values that are summed are abnormally small, essentially barely learning anything as the random initialization itself starts with an abnormally small loss, almost always staying constant.

L1 norm, on the other hand, does not have any power operations, making it relatively acceptable.

Loss functions, such as Kullback-Leibler-divergence or Jensen-Shannon-Divergence, are preferred for probability distributions because of the statistical meaning they hold. KL-Divergence, as mentioned before, is a statistical measure of information loss between distributions, or, in other words, assuming $Q$ is the ground truth distribution, KL-Divergence is a measure of how much $P$ deviates from $Q$. Also, considering probability distributions, convergence is much stronger in measures of Information Loss such KL-Divergence.

More clarity on the motivation behind Kullback-Leibler can be read here.


You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .