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!


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.


I would like to expand on F. Evlangeli's answer. Let's denote the first function as $f$ and the second function as $g$, i.e.

$$ f(\mu, \sigma) = -\frac{1}{2}(1 + 2\log(\sigma) - {\sigma}^{2} - {\mu}^{2}) $$

$$ g(\mu, \sigma) = -\frac{1}{2}(1 + \log(\sigma) - \sigma - {\mu}^{2}) $$

As far as math is concerned, the first function is the correct one, if you want to calculate the actual KL divergence. Yet, since variational autoencoders use KL divergence as a regulariser on the latent space, both functions serve the purpose well. That being said, it's important to highlight some differences. Although both functions have the same optimum with respect to $\sigma$, $f$ has a steeper gradient around the optimum. To show this, it's enough to consider partial derivatives with respect to $\sigma$:

$$ \frac{\partial f}{\partial \sigma} = -\frac{-{\sigma}^{2} + 1}{\sigma} $$


$$ \frac{\partial g}{\partial \sigma} = \frac{1}{2} - \frac{1}{2\sigma} $$


Therefore, it should be easier to optimise $f$. At the same time, $g$ is a less stringent regulariser, which might be beneficial given that over-regularisation is a common problem in VAE models. Finally, $g$ is cheaper to compute. Most VAE models estimate $\log(\sigma)$ in the latent layer. Consequently, you only need to calculate $\exp(\log(\sigma))$ in $g$, while $f$ requires additional squaring (which might also increase numerical instability).


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.