What does the approximate posterior on latent variables, $q_\phi(z|x)$, tend to when optimising VAE's

Question

The ELBO objective is described as follows

$$ ELBO(\phi,\theta) = E_{q_\phi(z|x)}[log p_\theta (x|z)] - KL[q_\phi (z|x)||p(z)] $$

This form of ELBO includes a regularisation term in the form of the KL divergence which drives $q_\phi(z|x) \rightarrow p(z)$ when optimising ELBO.

However we also have the overall expression for the loglikelihood which is defined as follows (proof provided here)

$$ p_\theta(x) = ELBO(\phi,\theta) + KL[q_\phi(z|x)||p_\theta(z|x)] $$

Rearranging the above equation as follows

$$ \max\limits_\phi ELBO(\phi,\theta) = \max\limits_\phi p_\theta(x) - KL[q_\phi(z|x)||p_\theta(z|x)] $$

We can see that maximising ELBO w.r.t $\phi$ in this form causes $q_\phi(z|x) \rightarrow p_\theta(z|x)$

These two ways of describing how VAEs learn conflicts my understanding of what happens to the approximate distribution during training.

Is it simply just trying to match both the prior $p(z)$ and the posterior $p_\theta(z|x)$ or am I missing something

Answer 1

Practically, when optimizing VAE, you assume that prior $p(z)\sim N(0,1)$; i.e. the unit Gaussian distribution. However, in testime you sample z from $p(z|x)$; the encoder model. Why is that?

Let's go back to the start. We have a model $p_{\theta}(x)$ and the data $\{x_1, ..., x_N\}$. Solving the maximum log-likelihood problem, we have \begin{equation} \begin{split} \theta &= argmax_{\theta}\frac{1}{N}\sum_i log p_{\theta}(x_i) \\ &= argmax_{\theta}\frac{1}{N}\sum_i \left( \int p_{\theta}(x_i|z)p(z)dz \right) \end{split} \end{equation}

which is intractable to calculate. So what to do now?

Here it comes the Variational Inference: "Use the expected log-likelihood instead." \begin{equation} \begin{split} \theta &= argmax_{\theta}\sum_i E_{z \sim p_{\theta}(z|x)}\left[ logp_{\theta}(x,z) \right] \end{split} \end{equation}

We approximate $q(z) \approx p_{\theta}(z|x)$. Thus, we unfold $logp(x)$: \begin{equation} \begin{split} logp(x) &= log \int p_{\theta}(x_i|z)p(z)dz \\ &= log \int p_{\theta}(x_i|z)p(z) \frac{q(z)}{q(z)}dz \\ &= log E_{z \sim q(z)}\left[ \frac{p_{\theta}(x|z)p(z)}{q(z)} \right] \\ &\geq E_{z \sim q(z)}\left[ log \frac{p_{\theta}(x|z)p(z)}{q(z)} \right] \\ &= E_{z \sim q(z)}\left[ logp_{\theta}(x|z) + logp(z) \right] - E_{z \sim q(z)}\left[ logq(z)\right] \\ &= E_{z \sim q(z)}\left[ logp_{\theta}(x|z) + logp(z) \right] + H(q) \\ &= ELBO(p,q) \end{split} \end{equation}

where $H(q)$ is the entropy of q. As you highlighted we can also write $$ ELBO(p,q) = logp(x) - D_{KL}\left( q(z) || p(z|x) \right) $$

This means that (a) maximizing $ELBO(p,q)$ w.r.t. to $q$ then KL-Divergence is minimized and (b) maximizing $ELBO(p,q)$ w.r.t to $p$ then the model is improved as the log-likelihood is improved. So this point of yours is true.

However, how can we actually train this model?

The answer is by using Amortized Variational Inference! Practically, we use 2 Neural Networks, $\phi$ and $\theta$, so that we have 2 models: $q_{\phi}(z|x)$ (encoder) and $p_{\theta}(x|z)$ (decoder). Thus, we replace $q(z)$ with $q_{\phi}(z|x)$ and $ELBO(p,q)$ with $ELBO(\theta,\phi)$ .

To fully answer your question, I should reach your first ELBO formula. I will unfold my first equation about ELBO:

\begin{equation} \begin{split} ELBO(\theta,\phi) &= E_{z \sim q_{\phi}(z)}\left[ logp_{\theta}(x|z)\right] + E_{z \sim q_{\phi}(z)}\left[ logp(z) \right] + H\left(q_{\phi}(z|x)\right) \\ &= E_{z \sim q_{\phi}(z)}\left[ logp_{\theta}(x|z)\right] - D_{KL}\left( q_{\phi}(z|x) || p_{\theta}(z) \right) \end{split} \end{equation}

Therefore, using the Amortized Variational Inference, we maximize the above objective (which is the same as the one above).