KL Loss with a unit Gaussian

Question

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!

Answer 1

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.

Answer 2

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