Why optimise log p(x) rather than log p(x|z) in a Variational AutoEncoder?

Question

Background

The loss function in a Variational AutoEncoder is the Evidence Lower Bound (ELBO):

$\mathbb{E}_q[log$ $p(x|z)] - KL[q(z)||p(z)]$

And has this inequality:

$log$ $p(x) \ge \mathbb{E}_q[log$ $p(x|z)] - KL[q(z)||p(z)]$

It is said in the literature that we want to maximise $log$ $p(x)$, but instead, we maximise the ELBO, which we know is less than or equal to $log$ $p(x)$ (because this is an easier objective).

First Question

I am struggling to understand why the objective is to maximise $log$ $p(x)$ in the first place.

Why is it not enough to maximise $\mathbb{E}_q[log$ $p(x|z)]$, which is the reconstruction loss?

I think I don't understand what exactly $p(x)$ means here.

I thought $p(x)$ was just the distribution of $x$, which are the input variables we have observed.
So how can we maximise $p(x)$, if it is the actual distribution of data in the real world?

Main Question

The Variational AutoEncoder is trying to learn an encoding from $x$ to $z$ and decoding from $z$ to $x'$, such that the difference between $x$ and $x'$ is minimized.

How does $log$ $p(x)$ relate to this?
What does it mean and how does maximising it help?

Answer 1

TLDR: We're doing a maximum likelihood fit of our model. The VAE sets this up in a way that doesn't require evaluating the model likelihood, but instead expresses a lower bound in terms of reconstruction error.

Here's an explanation in two parts.

1. Variational inference in general

We have a generative model $z \rightarrow x$ with latent variable $z \sim p(z)$, and observations, $x$, generated by $p(x|z)$. We want to get the posterior $p(z|x)$, but this inference may be computationally infeasible even if $p(z)$ and $p(x|z)$ are known because $p(x)$ is hard to evaluate. So, we instead approximate $p(z|x)$ by $q(z)$.

Denoting the ELBO by $\mathcal{L}(x)$, we have (by expanding $D_{KL}$): $$ \log p(x) = D_{KL}\big(q(z) \big|\big| p(z|x) \big) + \mathcal{L}(x). $$ Note that $p(x)$ is part of the generative model so, given data $x$, $p(x)$ is constant regardless of the approximation $q(z)$. This means that increasing $\mathcal{L}(x)$ by changing $q(z)$ is equivalent to decreasing $D_{KL}\big(q(z) \big|\big| p(z|x) \big)$ and vice versa.

Thus, the max ELBO approximation is the best in the sense that $q(z)$ is closest to $p(z|x)$ in KL terms. This is the classic reasoning for maximizing the ELBO.

2. Variational autoencoder

For the variational autoencoder model, we

1. choose $p(z)$ for the latent space (eg. gaussian)
2. parameterize $p(x | z)$ by the decoder network
3. parameterize $q(z)$ by the encoder network (eg. $z \sim \mathcal{N}( \mu(x), \sigma(x))$)

Now $p(x|z)$ is not fixed, but rather fit. So, as we train, we want to improve both the approximate inference model (encoder) and the generative model (decoder). This means that $p(x)$ is not fixed either as we train, so by increasing $L(x)$, we can either increase $\log p(x)$ or decrease $D_{KL}\big(q(z) \big|\big| p(z|x) \big)$ or both.

Why is increasing $\log p(x)$ a good thing? $\log p(x)$ tells us how likely that data we observe are under our model, this is a max likelihood fit of our (decoder parameterized) model. However, like the true posterior, $p(x)$ is difficult to calculate closed form, hence the ELBO.

That's the key idea of the VAE: it expresses the model max likelihood objective (which is hard to compute) in terms of reconstruction loss (which is easy to compute), giving us a generative model and approximate posterior along the way.

Answer 2

I thought $p(x)$ was just the distribution of $x$, which are the input variables we have observed. So how can we maximize $p(x)$, if it is the actual distribution of data in the real world?

I think the main confusion is coming from what "maximize" means in all that. What you have is a family of functions $p_\theta(x)$, parametrized by some parameters $\theta$. It could be a neural network or some distribution in analytic form or some kind of graphical model. By varying the parameters $\theta$ you are getting different distributions over $x$.

Now, you also have a sample of $x$'es that you observed "in the real world". And we would like to find such parameters $\theta$ that $p_\theta(x)$ reflects the real world distribution as close as possible. The standard way to do that is to find $\theta$ that maximizes the log-likelihood of the observed data $x$: $$\theta_{best} = \arg \max_\theta \log p_\theta(x)$$ This is called the maximum likelihood estimation. There are many extra details about it in the context of variational inference and VAEs, but that's the basic idea of what "maximize" means in all these cases.

The Variational AutoEncoder is trying to learn an encoding from $x$ to $z$ and decoding from $z$ to $x'$, such that the difference between $x$ and $x'$ is minimized.

You've described a general case of autoencoder - "variational" adds some extra details to it.

The "decoder" in VAE is not just a function from $z$ to $x$ - it is a probability distribution $p_\theta(x,z)$. With $\theta$ being parameters of the decoding neural network. Given some values for latent variables $z$ you obtain a distribution over observed variables $x$.

For the "encoder" we would ideally like to just have a conditional probability distribution of $z$ given $x$: $$p_\theta(z|x) = \frac{p_\theta(x,z)}{p_\theta(x)} = \frac{p_\theta(x,z)}{\int p_\theta(x,z) dz}$$ The problem is in the denominator - the integral for marginal density of observations (also called evidence) $p_\theta(x)$ is computationally very hard to evaluate. This is what is called intractability of the evidence and the reason for the whole field of variational inference to exist.

VAE approximates the true posterior $p_\theta(z|x)$ with a variational approximation $q_\phi(z|x)$. Where $\phi$ are parameters of the encoder network, the distribution usually factorizes following mean field approximation approach.

The need for $q_\phi(z|x)$ to approximate $p_\theta(z|x)$ leads to the whole ELBO derivation and the KL divergence term in the VAE loss.