In VAEs, we try to maximize the ELBO $\mathbb(E_q log\ p(x|z) + D_{KL}(q(z|x), p(z))$), but I see that many implement the first term as MSE of the image and it's reconstruction. Is this mathematically sound?

  • $\begingroup$ It may be a good idea to provide 1-2 examples where you saw this because the MSE is not always used. For example, here they use the cross-entropy. $\endgroup$
    – nbro
    Apr 15 '21 at 10:09
  • $\begingroup$ I recently read that MSE loss optimization is equivalent to minimizing Pearson $\chi^{2}$ divergence. Kullback–Leibler divergence (and also cross-entropy) has its drawbacks. Here is explanations of Least Squares loss for GAN $\endgroup$ Apr 15 '21 at 14:06
  • $\begingroup$ As you mentioned, MSE is used to measure the difference between the original and generated images. This encourages the model to preserve the original content. MSE loss can be used as an additional term, which is done in CycleGAN, where the authors use LSGAN loss and cycle-consistent loss, which is MSE-like loss. $\endgroup$ Apr 15 '21 at 14:17
  • $\begingroup$ @nbro, it is not clear why they use BCE there... In fact, that implementation doesn't seem to sample between the encoder and decoder, so even more strange. Looks like they treat the distribution parameters as the input to the decoder $\endgroup$
    – IttayD
    Apr 16 '21 at 9:36
  • $\begingroup$ What do you mean by "sample between the encoder and decoder"? Yes, the input to the decoder is a sample from the latent space, so I am not sure what you mean. $\endgroup$
    – nbro
    Apr 16 '21 at 9:39

On page 5 of the VAE paper, it's clearly stated

We let $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$ be a multivariate Gaussian (in case of real-valued data) or Bernoulli (in case of binary data) whose distribution parameters are computed from $\mathbf{z}$ with a MLP (a fully-connected neural network with a single hidden layer, see appendix $\mathrm{C}$ ).


As explained above and in appendix $\mathrm{C}$, the decoding term $\log p_{\boldsymbol{\theta}}\left(\mathbf{x}^{(i)} \mid \mathbf{z}^{(i, l)}\right)$ is a Bernoulli or Gaussian MLP, depending on the type of data we are modelling.

So, if you are trying to predict e.g. floating-point numbers (in the case of images, these can be the RGB values in the range $[0, 1]$), then you can assume $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$ is a Gaussian, then you can equivalently minimise the MSE between the prediction of the decoder and the real image in order to maximise the likelihood. You can easily show this: just replace $p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$ with the Gaussian pdf, then maximise that wrt the parameters, and you should end up with something that resembles the MSE.

So, to answer your question more directly: yes, minimizing the MSE is theoretically founded, provided that you're trying to predict some floating-point number.

  • $\begingroup$ the pdf is $\frac{1}{\sigma_z \sqrt{2\pi} } e^{-\frac{1}{2}\left(\frac{x-\mu_z}{\sigma_z}\right)^2}$. MSE is $\|\hat{x}-x\|^2$. I don't see the connection other than the fact there's an $(x-\mu_z)^2$ at the power of e. Note that in p(x|z), the original x is used, not the reconstruction. $\endgroup$
    – IttayD
    Apr 18 '21 at 5:48
  • $\begingroup$ See section C.2 in the original paper where they calculate p(x|z) with no MSE / BCE $\endgroup$
    – IttayD
    Apr 18 '21 at 5:51
  • $\begingroup$ Note sure why you say "Note that in p(x|z), the original x is used, not the reconstruction.". Both the original and the reconstructed images are used. z is used to compute the reconstructed image (or pixel). $\endgroup$
    – nbro
    Apr 18 '21 at 11:28

If $p(x|z) \sim \mathcal{N}(f(z), I)$, then $log\ p(x|z) \sim log\ exp(-(x-f(z))^2) \sim -(x-f(z))^2 = -(x-\hat{x})^2$ where $\hat{x}$, the reconstructed image, is just the distribution mean $f(z)$.

It also makes sense to use the distribution mean when using the decoder (vs. just when training), as it is the one with the highest pdf value. So the decoder produces a distribution from which we take the mean as our result.


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.