- 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.
- 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.
- 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.
