3
$\begingroup$

In a Variational Autoencoder (VAE), given some data $x$ and latent variables $t$ with prior distribution $p(t) = \mathcal{N}(t \mid 0, I)$, the encoder aims to learn a distribution $q_{\phi}(t)$ that approximates the true posterior $p(t \mid x)$ and the decoder aims to learn a distribution $p_\theta(x\mid t)$ that approximates the true underlying distribution $p^*(x\mid t)$.

These models are then trained jointly to maximize an objective $L(\phi, \theta)$, which is a lower bound for the log-likelihood of the training set:

$$L(\varphi, \theta) = \sum_i \mathbb{E}_{q_\varphi} \log \frac{p_\theta(x_i\mid t) p(t)}{q_\varphi(t)} \leq \sum_i \log \int p_\theta(x_i\mid t)p(t) \, dt$$

According to section C.2 in the original paper from Kingma and Welling (https://arxiv.org/pdf/1312.6114.pdf), when we model $p_{\theta}(x|t)$ as a family of gaussians, the decoder should output both the mean $\mu(t)$ and the (diagonal) covariance $\sigma^2(t) I$ for the gaussian distribution.

My question is: isn't this optimization problem ill-posed (just like maximum likelihood training in GMMs)? Having an output for the variance (or log-variance, as is most common), if the decoder can produce a perfect reconstruction for a single image in the training set (i.e. $\mu(t_i)=x_i$) then it can set the corresponding variance $\sigma^2(t_i)$ to something arbitrarily close to zero and therefore the likelihood goes to infinity regardless of what happens with the remaining training examples.

I know that most gaussian VAE implementations have a simplified decoder that outputs the mean only, replacing the term $\mathbb{E}_{q_\varphi} \log p_\theta (x_i\mid t)$ by the squared error between the original image and the reconstruction (which is equivalent to setting the covariance to be always the identity matrix). Is this because of the ill-posedness of the original formulation?

$\endgroup$
1
$\begingroup$

I think the KL divergence term keeps the problem well-defined. Intuitively, you can think of it as a "coding cost", where specifying a very narrow posterior distribution is expensive.

Consider the case where you have a single datapoint $x = 0$, and your latent space is 1D with the standard VAE prior $\mathcal{N}(0,1)$. One possible variational posterior would then have $\mu(x) = 0$ and $\sigma(x) = \sigma^*$ for some unknown optimal value of $\sigma^*$. Then the decoder would simply be the identity function.

The loss would be $$ \begin{align*} E_{z \sim \mathcal{N}(0, \sigma^*)} [\log P(0\mid z)] - \mathcal{D}_{KL}(Q(z\mid X)\parallel P(z)) &= \lambda (\sigma^*)^2 - \left( \log \frac{1}{\sigma^*} + \frac{(\sigma^*)^2}{2} - \frac{1}{2} \right) +c\\ &= \lambda'(\sigma^*)^2 - \log \sigma^* + c' \\ 2\lambda' \sigma^*-\frac{1}{\sigma^*} &= 0 \\ \sigma^* &= (2\lambda')^{-\frac{1}{2}} \end{align*} $$

where $\lambda$ is proportional to the precision of $P(X|z)$ and $\lambda' = \lambda - \frac{1}{2}$. As long as $\lambda > \frac{1}{2}$, then the KL term prevents the posterior from collapsing.

$\endgroup$
  • $\begingroup$ Maybe you're right, but there are two things I did not get in your answer: 1- How did you get the result $E_{z \sim \mathcal{N}(0, \sigma^*)} [\log P(0\mid z)] = \lambda (\sigma^*)^2$? 2- Why do you differentiate the result with respect to $\sigma^*$? This only proves that, for every fixed $\lambda$, there exists a critical point in the objective function at $\sigma^* = (2\lambda')^{-\frac{1}{2}}$. Assuming that your result for in 1. is correct, the problem still seems ill-posed to me. Fix $\sigma^* > 0$ and take $\lambda$ arbitrarily large, then the maximizing objective is unbounded. $\endgroup$ – D... Dec 11 '18 at 18:09

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.