10
$\begingroup$

I've been implementing a VAE and I've noticed two different implementations online of the simplified univariate gaussian KL divergence. The original divergence as per here is $$ KL_{loss}=\log(\frac{\sigma_2}{\sigma_1})+\frac{\sigma_1^2+(\mu_1-\mu_2)^2}{2\sigma^2_2}-\frac{1}{2} $$ If we assume our prior is a unit gaussian i.e. $\mu_2=0$ and $\sigma_2=1$, this simplifies down to $$ KL_{loss}=-\log(\sigma_1)+\frac{\sigma_1^2+\mu_1^2}{2}-\frac{1}{2} $$ $$ KL_{loss}=-\frac{1}{2}(2\log(\sigma_1)-\sigma_1^2-\mu_1^2+1) $$ And here is where my confusion rests. Although I've found a few obscure github repos with the above implementation, what I more commonly find used is:

$$ =-\frac{1}{2}(\log(\sigma_1)-\sigma_1-\mu^2_1+1) $$ For example in the official Keras autoencoder tutorial. My question is then, what am I missing between these two? The main difference is dropping the factor of 2 on the log term and not squaring the variance. Analytically I have used the latter with success, for what its worth. Thanks in advance for any help!

$\endgroup$
7
$\begingroup$

Notice that by replacing $\sigma_1$ with $\sigma_1^2$ in the last equation you recover the previous (i.e. $\log(\sigma_1) - \sigma_1 \rightarrow 2\log(\sigma_1) - \sigma_1^2$). Leading me to think that in the first case the encoder is used to predict the variance, whereas in the second it is used to predict the standard deviation.

Both formulations are equivalent and the objective is unchanged.

$\endgroup$
  • $\begingroup$ I don't think it can be the case that these are equivalent. Yes, they are both minimised when for zero $\mu$ and unit $\sigma$. However, in the original equation (featuring the variance), the penalty for moving $\sigma$ away from unity is far larger than in the second equation (based on the standard deviation). The penalty for variations in $\mu$ is the same for both, and the reconstruction error would be the same, so using the second version dramatically changes the relative importance of departures of $\sigma$ from unity. What am I missing? $\endgroup$ – TheBamf Jan 13 at 0:28
0
$\begingroup$

I believe the answer is simpler. In VAE, people usually use a multivariate normal distribution, which has covariance matrix $\Sigma$ instead of variance $\sigma^2$. That looks confusing in a piece of code but has the desired form.

Here you can find the derivation of a KL divergence for multivariate normal distributions: Deriving the KL divergence loss for VAEs

$\endgroup$

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.