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?
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.
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.
