1
$\begingroup$

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?

$\endgroup$

2 Answers 2

1
$\begingroup$

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.

$\endgroup$
1
$\begingroup$

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.

$\endgroup$
5
  • $\begingroup$ Thank you for your helpful answer. Does this mean that p(x) is the entire model? i.e. it is the function that includes both encoder and decoder and takes an x as input and as output produces another x', which could be different each time due to the sampling from the latent space? And we are trying to make sure that p(x) over all the x in our dataset produces a distribution of x's that are as similar as possible to our original dataset? $\endgroup$ Mar 10, 2023 at 9:13
  • $\begingroup$ @TitusBuckworth I would say that the "entire model" is $p_\theta(x, z)$. (Again, I encourage you to write the optimization parameters like $\theta$ explicitly.) You observe only $x$ and you want to simultaneously calculate a good maximum likelihood estimation of $\theta$ and infer the latent variables $z$, which you don't observe. This is a hard problem due to intractability of the integral for the evidence. Variational approximation $q_\phi(z|x)$ is used to deal with this intractability. $\endgroup$
    – Kostya
    Mar 10, 2023 at 16:11
  • $\begingroup$ @TitusBuckworth See this answer ai.stackexchange.com/a/27991/20538 for in-depth math on that. $\endgroup$
    – Kostya
    Mar 10, 2023 at 16:12
  • $\begingroup$ Thanks. So you say that $p_\theta(x)$ is a family of functions. What is the input and output of those functions? What is the input and output for $p_\theta(x)$ with optimal $\theta$? $\endgroup$ Mar 14, 2023 at 11:54
  • $\begingroup$ @TitusBuckworth It is a family of probability distributions. I'd really recommend reading Wikipedia link on Maximum Likelihood Estimation. $\endgroup$
    – Kostya
    Mar 14, 2023 at 12:13

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .