6
$\begingroup$

For a variational autoencoder, we have that:

$$\mathcal{L}(x,\theta,\phi) := \mathbb{E}_{z \sim q_\phi(z|x)}[\log p_{\theta}(x|z)] -KL[q_{\phi}(z|x) ||p(z)] $$

This is called the variational lower bound or evidence lower bound (ELBO). But I think what we're actually trying to maximize is the log-likelihood of our data:

$$\log p_{\theta}(x) = \mathcal{L}(x,\theta,\phi) + KL[q_{\phi}(z|x)||p_{\theta}(z|x)]$$

There are a few things I'm unsure about, in increasing order of difficulty.

  1. For the actual loss function of a VAE, we use $-\mathcal{L}$, more or less. Of course, it's expensive to actually calculate the expectation, which is why we use a single $z$ sample each time, yes?

  2. We are told to treat $p(z)$ as being $\mathcal{N}(0,1)$, but I don't see what causes it to become normal. We're just told to plug in $\mathcal{N}(0,1)$ for $p(z)$ when calculating the loss, and all this seems to do is ensure that $q_\phi(z|x)$ gets closer to it. We know nothing about $p(z)$, right (besides the fact that if all $p_{\theta}(z|x)$ are close to some distribution, then so is $p_{\theta}(z)$)?

  3. How should we think about that KL divergence in the second formula? It's usually pointed out that for fixed $p_{\theta}$, maximizing $\mathcal{L}$ is equivalent to minimizing the KL div. But $p_{\theta}$ is not fixed; it's being learned. We could also improve $\log p_{\theta}(x)$ (our ultimate goal) by making that divergence worse, couldn't we? What justifies optimizing $\mathcal{L}$ alone (other than tractability)?

Basically, I'm confused by the common explanation that what we really want is to minimize that second KL divergence, and also that the best way to do that is to maximize the ELBO.

$\endgroup$

1 Answer 1

3
$\begingroup$
  1. For the actual loss function of a VAE, we use $−\mathcal{L}$, more or less. Of course, it's expensive to actually calculate the expectation, which is why we use a single 𝑧 sample each time, yes?

Yes. It turns out a single sample MC estimate has fairly low variance in this case. However, the Importance Weighted Autoencoder shows that taking multiple samples can be useful.

  1. It's usually explained that we treat 𝑝(𝑧) as being $\mathcal{N}(0,1)$. But simply plugging in $\mathcal{N}(0,1)$ to the KL divergence in the loss (as we do) just ensures that $𝑞_𝜙(𝑧|𝑥)$ gets close to it. We know nothing about $𝑝(𝑧)$, right (besides the fact that if all $𝑝_𝜃(𝑧|𝑥)$ are close, then so is $𝑝_𝜃(𝑧)$?

I prefer to think of it this way. There are in fact two models: a stochastic encoder (inference model) $q_\phi(x,z)=q_\phi(z|x) q(x)$ and a probabilistic decoder (generative model) $p_\theta(x,z)=p_\theta(x|z) p(z)$. Our goal is to enforce that $ q_\phi(x,z) $ and $ p_\theta(x,z) $ are close (small KL divergence).

Note that $q(x)$ and $p(z)$ are the empirical distribution and Bayesian prior, and are fixed ahead of time (not learned). This is for the vanilla VAE. However, these are very different from the aggregate marginals: $$ q_\phi(z) = \int q_\phi(z|x) q(x) \, dx \;\;\;\;\&\;\;\;\; p_\theta(x) = \int p_\theta(x|z) p(z) \, dz $$ In other words, $p(z)$ is something we choose as part of the model. Our job is then to ensure that the inference marginal $q_\phi(z)$ is close to it. In other words, we know $p(z)$, but we learn $q_\phi(z)$ (also called the inferred prior or aggregate posterior), which we would like to match it.

  1. How should we think about that KL divergence in the second formula? It's usually pointed out that for fixed $𝑝_𝜃$, maximizing $\mathcal{L}$ is equivalent to minimizing the KL div. But $𝑝_𝜃$ is not fixed; it's being learned. We could also improve $\log 𝑝_𝜃(𝑥)$ (our ultimate goal) by making that divergence worse, couldn't we? What justifies optimizing $\mathcal{L}$ alone (other than tractability)?

It's important to note that while $p_\theta(x)$ is being learned (well, an approximation is being learned), there is actually an "correct" one that can be computed via Bayes Rule: $$ p_\theta(z|x) = \frac{p_\theta(x|z) p(z) }{p_\theta(x)} $$ when $ p_\theta(x|z) $ is fixed. In other words, we don't have to learn an inference model; there is an optimal one given by Bayes rule. In theory we can just learn the generative model and use this rule to create the encoder. But this is of course intractable to compute.

How should we think about that KL divergence in the second formula?

We are doing variational inference. We cannot compute the true posterior $p_\theta(z|x)$. So instead we compute an approximation $q_\phi(z|x)$. The second KL says the approximation should be good. Notice something important: when the variational posterior approximation is perfect, then the ELBO equals the marginal likelihood. In other words, when our $q_\phi$ is perfect, then we are indeed optimizing the true marginal likelihood, as you noted we want to do.

We could also improve $\log 𝑝_𝜃(𝑥)$ (our ultimate goal) by making that divergence worse, couldn't we?

No, decreasing it should always make the marginal likelihood worse. I don't have a proof, but I conjecture it is so.

What justifies optimizing $\mathcal{L}$ alone (other than tractability)?

Well, it is a lower bound, so maximizing it guarantees we are "pushing up" the true marginal likelihood. As noted above, when the variational approximation is good, we are guaranteed to be doing the right thing. There is a lot of literature on variational inference theory; I suspect one can say more under more constrained conditions.

Basically, I'm confused by the common explanation that what we really want is to minimize that second KL divergence, and also that the best way to do that is to maximize the ELBO.

I don't know if it is the best way. It is certainly an efficient one. But in some cases, it may be better to perform inference more directly, e.g. Hamiltonian Monte Carlo methods.

$\endgroup$
8
  • 1
    $\begingroup$ Thanks. I am going to have to refresh my understanding to make sense of all this. One thing that jumps out is that I may have been confusing $p_{\theta}$ and $p$ (and similarly for $q$). I take it they're different? Also, in what sense are $p$ and $q$ empirical? We don't have any empirical information about latent variables in the real world, do we? $\endgroup$
    – A_P
    Oct 3, 2019 at 19:27
  • $\begingroup$ Ah, "empirical" was only about q(x) (the distribution of samples, which we have). Another question: when you say "In theory we can just learn the generative model," how would we do that? $\endgroup$
    – A_P
    Feb 5, 2020 at 18:28
  • $\begingroup$ @A_P $p_\theta$ and $p$ only denote whether or not there are any parameters to the distribution. $p(z)$ is fixed and parameterless, while the other $p$ distributions are not (and hence have a $\theta$). When I wrote $q(x)$ is the empirical distribution, it is usually defined to mean the distribution of the data (over data space) - we can trivially sample from this via drawing from the uniform categorical distribution over the training set. $\endgroup$ Feb 24, 2020 at 23:38
  • 1
    $\begingroup$ Thanks again! Let me try to summarize. The original problem is to find a $p(x, z)$ where $p(z)$ is $\mathcal{N}(0,1)$ so that we can draw from $p(x|z)$, subject to maximizing $\prod p(x)$ over our dataset. Then for some reason we find ourselves trying to solve for $p(z|x)$, but that's intractable, so we use V.I. I guess the "some reason" is that we notice that if we knew $p(z|x)$ (and the corresponding $q(z|x)$), we could use the formulas in my question to provide a possibly good estimate $p(x)$? $\endgroup$
    – A_P
    Feb 27, 2020 at 21:28
  • 2
    $\begingroup$ @A_P yeah that seems right. Only thing I'd clarify is that if we knew $p(z|x)$ exactly there would be no need for $q(z|x)$ (the variational distribution) because the latter is merely a (usually poor) approximation for the former. We could then compute $p(x)$ directly. In this case, the ELBO is literally the log likelihood $\log p(x)$ because the KL term disappears (if I recall correctly). PS: are you familiar with normalizing flow models? Like GLOW? Those might help you understand. :) $\endgroup$ Feb 28, 2020 at 21:17

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.