What is the effect of KL divergence between two Gaussian distributions as a loss function in neural networks?

Question

In many deep neural networks, especially those based on VAE architecture, a KL divergence term is added to the loss function. The divergence is computed between the estimated Gaussian distribution and prior. Since Gaussian distribution is completely specified by mean and co-variance, only those two parameters are estimated by the neural network. For Gaussian distributions, KL divergence has a closed form solution. By minimizing KL divergence, we bring the estimated distribution closer to the prior.

My question is, since Gaussian distribution is completely specified by mean and co-variance, why don't we just take MSE between estimated parameters and prior parameters? Minimizing MSE between mean and co-variance also brings the two distributions closer. Does taking KL divergence have any significance?

Answer 1

It's too strong of an assumption (I am answering generally, I am sure you know. Coming to VAE later in post), that they are Gaussian.

You can not claim that distribution is X if Moments are certain values. I can bring them all to the same values using this.

Hence if you can not make this assumption it is cheaper to estimate KL metric

BUT with VAE you do have information about distributions, encoders distribution is $q(z|x)=\mathcal{N}(z|\mu(x),\Sigma(x))$ where $\Sigma=\text{diag}(\sigma_1,\ldots,\sigma_n)$, while the latent prior is given by $p(z)=\mathcal{N}(0,I)$. Both are multivariate Gaussians of dimension $n$, for which in general the KL divergence is: $$ \mathfrak{D}_\text{KL}[p_1\mid\mid p_2] = \frac{1}{2}\left[\log\frac{|\Sigma_2|}{|\Sigma_1|} - n + \text{tr} \{ \Sigma_2^{-1}\Sigma_1 \} + (\mu_2 - \mu_1)^T \Sigma_2^{-1}(\mu_2 - \mu_1)\right] $$ where $p_1 = \mathcal{N}(\mu_1,\Sigma_1)$ and $p_2 = \mathcal{N}(\mu_2,\Sigma_2)$.

In the VAE case, $p_1 = q(z|x)$ and $p_2=p(z)$, so $\mu_1=\mu$, $\Sigma_1 = \Sigma$, $\mu_2=\vec{0}$, $\Sigma_2=I$. Thus: \begin{align} \mathfrak{D}_\text{KL}[q(z|x)\mid\mid p(z)] &= \frac{1}{2}\left[\log\frac{|\Sigma_2|}{|\Sigma_1|} - n + \text{tr} \{ \Sigma_2^{-1}\Sigma_1 \} + (\mu_2 - \mu_1)^T \Sigma_2^{-1}(\mu_2 - \mu_1)\right]\\ &= \frac{1}{2}\left[\log\frac{|I|}{|\Sigma|} - n + \text{tr} \{ I^{-1}\Sigma \} + (\vec{0} - \mu)^T I^{-1}(\vec{0} - \mu)\right]\\ &= \frac{1}{2}\left[-\log{|\Sigma|} - n + \text{tr} \{ \Sigma \} + \mu^T \mu\right]\\ &= \frac{1}{2}\left[-\log\prod_i\sigma_i^2 - n + \sum_i\sigma_i^2 + \sum_i\mu^2_i\right]\\ &= \frac{1}{2}\left[-\sum_i\log\sigma_i^2 - n + \sum_i\sigma_i^2 + \sum_i\mu^2_i\right]\\ &= \frac{1}{2}\left[-\sum_i\left(\log\sigma_i^2 + 1\right) + \sum_i\sigma_i^2 + \sum_i\mu^2_i\right]\\ \end{align}

You see that the mean minimisation is the same so the only cost factor is between these stds and the covariance matrix values. We can see that its more expensive to evaluate these integrals in covariance matrix and then minimise them, then just to minimise these std.

TL;DR cheaper, but you are right it could be done. Good question!