2
$\begingroup$

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?

$\endgroup$
1
$\begingroup$

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!

|improve this answer|||||
$\endgroup$
  • $\begingroup$ One clarification. You mean MSE between mean and co-variance is cheaper, right? And KL divergence is a bit complex. One more follow up question. For Images, MSE is a bad metric (perceptually), but SSIM is very good. That is images with low MSE might not always be perceptually good, but those with high SSIM look good. Do you think a similar phenomenon might occur here? That is, minimizing KL divergence might lead to visually better images while minimizing MSE (between mean and cov) might lead to not so good images. Or the other way round. Would it happen like that? $\endgroup$ – Nagabhushan S N Dec 23 '19 at 9:06
  • $\begingroup$ Sorry. First part of my doubt above, I got it cleared. For non Gaussian distributions, computing KL metric is cheaper. But for Gaussian distributions, MSE between mean and cov is cheaper. Right? $\endgroup$ – Nagabhushan S N Dec 23 '19 at 9:17
  • 1
    $\begingroup$ Exactly right! Regarding the follow up: Possible, we can talk empirically and I have seen such things happen with different loss functions however without some strong theory we are speculating and it depends on dataset $\endgroup$ – Noah Weber Dec 23 '19 at 9:19

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.