In VAE, why use MSE loss between input x and decoded sample x' from latent distribution?

Variational Autoencoders (VAEs) are based on the concept of Variational Inference (VI) and use two Neural Networks similar to Vanilla Autoencoders (AEs) for function approximation. I understood the derivation of the Evidence-Lower-Bound (ELBO) and the role of the two ELBO terms that make up the objective for training a VAE:

1. Expectation of P(x|z) with z ~ Q(z|x) -> "reconstruction loss"
2. KL-Divergence of Q(z|x) and P(z) -> "regularization term"

The second regularization term is very clear to me, as it basically forces the latent representation to be distributed like the prior on z.

However, the first one is not as intuitive to me. Let's assume we speak about the standard VAE that assumes Gaussians as Q(z|x) and P(x|z), where the encoder produces means and variances for the latent dimensions of z and the decoder produces the means for P. For this specific VAE, the first loss term is equal to the MSE loss between the predicted mean of the decoder and the input x, isn't it?

My intuition ends here: Why can we use this reconstruction loss to judge the pixel-wise differences between the input x and the decoded latent sample z? I mean, x gets encoded and specifies some latent distribution over z's. Now we sample from it to get a specific z. But this z does not have to be meant to be the latent code for the input, doesn't it? So the decoded version of z does not have to show the same image content? Let's say we are using face images. The input could be male. The sampled z could be producing a female person? So using MSE between these two seems to be wrong?

In VAEs the conditional distribution $$p(x|z)$$ is assumed to be a Gaussian distribution, i.e. $$p(x|z)=\mathcal N(x;f_{dec}(z), \sigma^2 I)$$ where $$\sigma^2$$ is a hyperparameter. Hence the first term of ELBO, the logarithm of $$p(x|z)$$, would be just a MSE loss between $$f_{dec}(z)$$ and $$x$$.

I'll go through some of your questions:

Why can we use this reconstruction loss to judge the pixel-wise differences between the input x and the decoded latent sample z?

Implicitly, we assume that pixel values are Normally distributed with uniform diagonal covariance.

I mean, x gets encoded and specifies some latent distribution over z's. Now we sample from it to get a specific z. But this z does not have to be meant to be the latent code for the input, doesn't it?

I think you are mixing two generative strategies here: VAEs and GANs (generative adversarial networks). Most (not all) GANs don't try to replicate inputs, instead they simply try to create realistic looking images to fool the discriminator. Since in GANs we do not have a mapping from image to latent code, there is no correspondence (BiGANs/ALI, for example, has this).

VAEs come from a complete probabilistic point of view. In VAEs we want to retrieve the posterior of the distribution that generates the images.

$$p(z|x) = \frac{p(x|z)p(z)}{p(x)}$$

This is very expensive (and actually intractable) due to $$p(x)$$, and we instead approximate it by $$q(z|x)$$, an amortized model of the actual posterior (this the VAE encoder).

So the decoded version of z does not have to show the same image content? Let's say we are using face images. The input could be male. The sampled z could be producing a female person? So using MSE between these two seems to be wrong?

As I explained above, ideally, the reconstructed image should look like the input as best as well as the latent code can afford it.

Maybe you are not considering the fact that for a given input you get its corresponding latent vector (mu and sigmas in VAE that assumes Gaussian distribution for z).

This extends to a mini-batch of inputs. For e.g., if you have 100 images in your mini-batch, the encoder is outputting 100 latent vectors (each vector is of dimension D). See the below image to get better visualization. The batch is shown as ? and the dimension of the latent vector is 2.

You then generate a sample (according to VAE paper, 1 sample is sufficient thanks to re-parameterization trick low variance properties) using the mu and sigma that corresponds to its corresponding input image. This means for a batch of 100 inputs you will have 100 samples.

So if your first input image is that of Bob (a male) and the second input image is that of Alice (a female) you will have latent vectors [mu_bob, sigma_bob] & [mu_alice, sigma_alice]. You would then have sample_bob (using mu_bob & sigma_bob) and sample_alice (using mu_alice & sigma_alice).

Now, at the onset of your training regime, these mus and sigmas will not be good but the idea is that the encoder network will learn to generate "relevant/appropriate" mus and sigmas for your inputs. As such this is no different than any representation learning-based setup.

• Thanks for the answer. But it still does not solve my problem of understanding: mu_bob and sigma_bob define a gaussian just like mu_alice and sigma_alice. Usually, a N(0,1) prior on he latent space is used and encouraged by the regularization term of the ELBO objective. Hence, the latent distributions for bob's and alice's will become more similar, the longer training has been, won't they? So, when sampling from the gaussian defined by [mu_bob, sigma_bob] it is also likely to get a latent vector z that gets decoded by the decoder network into an image that looks like alice -> mse loss bad? Sep 21, 2020 at 21:20
• @JonasG. Just to be clear - mu and sigma are high dimensional. A combination of high dim latent space may/should help. That said if you are only concerned by MSE loss, indeed it is not the best loss function for "face" reconstruction albeit still works for low dim images (e.g. celeba) Some of the directions in which VAE are being improved are - Using sophisticated prior and posterior (e.g. normalizing flows) and not just normal, using more stochastic nodes/layers, balancing KL & Reconstruction loss, and of course better loss functions (SSIM, perceptual, etc). Sep 23, 2020 at 13:00