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

Question

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?

Answer 1

In VAEs the conditional distribution $p(x|z)$ is (usually) 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$.

Other assumptions of $p(x|z)$ can be used. For example, when you are dealing with discrete $x$ that can only take values from $\{0,1\}$, then you would probably assume $p(x|z)$ to be Bernoulli distribution. I guess that will lead to a binary cross entropy loss.

Answer 2

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.

Answer 3

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.