In VAEs, we try to maximize the ELBO = $\mathbb{E}_q [\log\ p(x|z)] + D_{KL}(q(z \mid x), p(z))$, but I see that many implement the first term as the MSE of the image and its reconstruction. Here's a paper (section 5) that seems to do that: Don't Blame the ELBO! A Linear VAE Perspective on Posterior Collapse (2019) by James Lucas et al. Is this mathematically sound?
$\begingroup$
$\endgroup$
9
-
$\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$– nbroApr 15, 2021 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$– Aray KarjauvApr 15, 2021 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$– Aray KarjauvApr 15, 2021 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$– IttayDApr 16, 2021 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$– nbroApr 16, 2021 at 9:39
|
Show 4 more comments
2 Answers
$\begingroup$
$\endgroup$
1
If $p(x|z) \sim \mathcal{N}(f(z), I)$, then
\begin{align} \log\ p(x|z) &\sim \log\ \exp(-(x-f(z))^2) \\ &\sim -(x-f(z))^2 \\ &= -(x-\hat{x})^2, \end{align}
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.
-
$\begingroup$ This seems correct to me (and this is what I had in mind when I wrote my answer above, although I didn't write the derivation), but note that the second $\sim$ is exactly $=$. The first one is correct because you're omitting the denominator of the Gaussian density (which becomes a constant if you optimize wrt $\theta$), so you use approximately equal, which is correct. Moreover, it might be a good idea to add a underscript to $f$, i.e. $f_\theta$ to emphasize that you will be optimizing with respect to those parameters, although you jointly optimize them with $\phi$ (the encoder parameters). $\endgroup$– nbroJun 12, 2022 at 8:16
$\begingroup$
$\endgroup$
3
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 real 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.
It turns out that maximising the Gaussian likelihood is equivalent to minimising the MSE between the prediction of the decoder and the real image. 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. G. Hinton shows this in this video lesson. See also this related answer.
So, yes, minimizing the MSE is theoretically founded, provided that you're trying to predict some real number.
When the binary cross-entropy (instead of the MSE) is used (e.g. here), the assumption is that you're maximizing a Bernoulli likelihood (instead of a Gaussian) - this can also be easily shown.
-
$\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$– IttayDApr 18, 2021 at 5:48
-
$\begingroup$ See section C.2 in the original paper where they calculate p(x|z) with no MSE / BCE $\endgroup$– IttayDApr 18, 2021 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$– nbroApr 18, 2021 at 11:28