# What is the loss function for a probabilistic decoder in the Variational Autoencoder?

In a VAE with Gaussian output the loss function is usually:$$\sum{(\hat x - x)^2} + KL,$$ so the sum of squared errors plus KL divergence. When I also want to predict the variance of the reconstructed input I need 2 outputs for each dimension of x: mean and variance. But in this case how should I calculate the reconstruction error? Should I sample from a normal distribution with the mean and variance I have as current output and calculate sum of squared errors for that sample?

## 1 Answer

If you're predicting a mean and a variance with the decoder, you can use a log likelihood loss for the reconstruction loss. So if you predict $\mu_x$ and $\sigma_x$, then the reconstruction loss is:

$$2(\mu_x - x) - 2\sigma_x^2 + \operatorname{log}2\sigma_x + C$$

• Thank you! Just to be sure, so in TensorFlow it would be: 2 * (mu - input) - 2 * tf.square(sigma) + tf.log(2 * sigma) where input, mu and sigma are rank 2 tensors with shape (number of samples, input dimensions)? – Márton György Apr 9 '18 at 7:35
• Yeah, that looks right. – Kevin Yang Apr 10 '18 at 18:42