# Is the optimization of the Gaussian VAE well-posed?

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?

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.

• 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. – D... Dec 11 '18 at 18:09