From EM to the ELBO
In my previous post, I described the EM algorithm and showed how it can be used to perform maximum likelihood estimation when the likelihood is intractable but the complete-data likelihood (when we include the latent variables) and its posterior expectation are tractable. Therefore, it seems like it could be applied to any generative modeling problem. In this post, I will show why naive EM sometimes (and maybe often actually) cannot be used and how it naturally leads to variational ideas and the omnipresent ELBO.
Generative modeling
In generative modeling, we are interested in generating new data from the underlying distribution (that we assume exists) of an existing dataset. Formally, given a dataset X_1,\dots,X_n of i.i.d. samples from an unknown distribution p_{\mathrm{data}}, we want to learn a model p_\theta parameterized by \theta that approximates p_{\mathrm{data}} and from which we can sample. One way to do this is to directly parameterize p_\theta(x_0) and then fit \theta by maximizing the likelihood of the data. The issue is that p_\theta needs to be a true probability distribution, so it must be non-negative and integrate (or sum) to one. Therefore, we need to choose a parameterization that satisfies these constraints. There are multiple approaches, for instance using the autoregressive factorization (as in LLMs) or parameterizing the unnormalized log-probability, i.e. p_\theta(x_0)=\frac{e^{f_\theta(x_0)}}{Z_\theta}, where Z_\theta is the normalizing constant. Models of the latter form are called energy-based models (EBMs), but the usually intractable normalizing constant causes problems for maximum-likelihood training, and sampling is also generally difficult.
Another way to parameterize p_\theta is to introduce a latent variable, and assume that conditionally on this latent variable, the distribution of the data is simple. Formally, we introduce a latent variable Z and write p_\theta(x_0)=\int p_\theta(x_0\mid z)p_\theta(z)\,dz, where p_\theta(z) is a prior distribution (typically a Gaussian) over the latent variable. Intuitively, we can think of Z as a hidden representation of the data that makes decoding x (sampling from p_\theta(x_0\mid z)) easy. This is the approach that we will focus on in this post. Again, we want to fit \theta by maximizing the likelihood of the data, but now the likelihood is intractable because of the integral over z, so we are typically in the setting where we would like to use EM.
Why naive EM is not always applicable
Recall that in EM, we first compute the posterior distribution p_{\theta^t}(z\mid x_0) of the latent variable given the data and the current parameter \theta^t. We then maximize the expected complete-data log-likelihood, \theta^{t+1} = \arg\max_\theta \mathbb{E}_{p_{\theta^t}(z\mid x_0)}[\log p_\theta(x_0, z)]. For the full dataset, we sum this expectation over the data points.
So we need to be able to compute p_\theta(z\mid x_0) = \frac{p_\theta(x_0\mid z)p_\theta(z)}{\int p_\theta(x_0\mid z)p_\theta(z)\,dz}. The normalizing constant in the denominator is exactly the likelihood of the data that we are trying to maximize. For a generic nonlinear decoder, this integral, and therefore the expectations needed in the E-step, are not tractable.
Approximating the E-step leads to the ELBO
Since the issue is to compute the posterior distribution, a natural idea is to approximate it with a variational distribution q_\phi(z\mid x_0) parameterized by \phi. Introducing this variational distribution and applying Jensen’s inequality, we have \begin{aligned} \log p_\theta(x_0) &= \log \int p_\theta(x_0, z)\,dz \\ &= \log \int q_\phi(z\mid x_0)\frac{p_\theta(x_0, z)}{q_\phi(z\mid x_0)}\,dz \\ &\geq \int \log\left(\frac{p_\theta(x_0, z)}{q_\phi(z\mid x_0)}\right) q_\phi(z\mid x_0)\,dz. \end{aligned} And we have the identity \log p_\theta(x_0) = \mathbb{E}_{q_\phi(z\mid x_0)}\left[\log \frac{p_\theta(x_0, z)}{q_\phi(z\mid x_0)}\right] + \mathrm{KL}(q_\phi(z\mid x_0)\|p_\theta(z\mid x_0)).
The first term is the expected complete-data log-likelihood plus the entropy of q_\phi, and the second term is the KL divergence between the variational posterior and the true posterior. Let us denote the first term by \mathrm{ELBO}(\theta, \phi) = \mathbb{E}_{q_\phi(z\mid x_0)}\left[\log \frac{p_\theta(x_0, z)}{q_\phi(z\mid x_0)}\right]. Then we have \log p_\theta(x_0) = \mathrm{ELBO}(\theta, \phi) + \mathrm{KL}(q_\phi(z\mid x_0)\|p_\theta(z\mid x_0)).
This is the key identity. Since p_\theta(x_0) does not depend on \phi, maximizing the ELBO with respect to \phi is equivalent to minimizing the KL divergence between the variational posterior and the true posterior, so this is an approximate E-step. For a fixed q_\phi, its entropy does not depend on \theta, so maximizing the ELBO with respect to \theta is exactly the variational M-step. Therefore optimizing the ELBO can be seen as coordinate ascent on variational EM, and provided that \phi and \theta are expressive enough, this converges to p_\theta = p_{\mathrm{data}}.
Something that should be emphasized is that in the EM update, we are optimizing the likelihood of the data. This coordinate-ascent variational EM reduces to the EM update when the variational family is expressive enough to reach q_\phi^\star(z\mid x_0) = p_\theta(z\mid x_0) in which case we are also guaranteed to increase the likelihood of the data. However, when the variational family is not expressive enough, we are optimizing a different objective which may not increase the lielihood or reach the maximum likelihood estimator, even at optimum, this is actually the point of the next section. Another detail worth noting is that in a VAE, we usually optimize \phi and \theta simultaneously using stochastic gradient descent. This is not exactly coordinate-ascent EM, but it optimizes the same dataset-level ELBO so the interpretation stays valid.
Variational auto-encoders and variational EM
Until now, we didn’t say anything about the variational family chosen for q_\phi or the model family chosen for p_\theta. In practice, we usually choose a simple family for q_\phi (e.g. Gaussian), but this choice is not innocent. With a fixed prior p(z), the negative VAE objective averaged over the data can be written as -\mathbb{E}_{p_{\mathrm{data}}}[\mathrm{ELBO}(\theta,\phi)] = H(p_{\mathrm{data}}) + \mathrm{KL}(p_{\mathrm{data}}\|p_\theta) + \mathbb{E}_{p_{\mathrm{data}}}[\mathrm{KL}(q_\phi(z\mid x_0)\|p_\theta(z\mid x_0))]. There are two effects here:
- Is p_{\mathrm{data}} in the model family of p_\theta?
- Is p_\theta(z\mid x_0) in the variational family of q_\phi(z\mid x_0)?
The two effects are actually very different and competing. For 1, the model p_\theta(x_0) = \int \mathcal{N}(x_0; f_\theta(z), \sigma^2 I)\mathcal{N}(z; 0, I)\,dz is a continuous mixture of Gaussians whose centers are given by f_\theta(z), so it can represent many continuous distributions when f_\theta is expressive and \sigma can be small. A fixed \sigma, however, limits the scale of the details that the model can represent. For 2, the situation is different. Since f_\theta is usually a highly nonlinear neural network, the posterior p_\theta(z\mid x_0) \propto \mathcal{N}(x_0; f_\theta(z), \sigma^2 I)\mathcal{N}(z; 0, I) is generally not Gaussian. Therefore, q_\phi(z\mid x_0) = \mathcal{N}(z; \mu_\phi(x_0), \Sigma_\phi(x_0)) is usually not able to represent the true posterior exactly. This leads to an interesting effect where even if there exists a \theta^\star such that p_{\theta^\star} = p_{\mathrm{data}}, optimizing the ELBO might not lead to p_\theta = p_{\mathrm{data}} since it also has to balance \mathrm{KL}(q_\phi(z\mid x_0)\|p_\theta(z\mid x_0)). Therefore, we look for a model that is close to the data distribution while having a posterior that is simple for q_\phi to approximate. This does not by itself imply that the latent variables are meaningful.
As a simple example, suppose that the data are spread uniformly around the unit circle, and consider the family of models Z\sim\mathcal N(0,1),\qquad X\mid Z=z\sim\mathcal N((\cos(az),\sin(az)),\tau^2I). When a is small, the values of z with high prior probability only cover part of the circle. The posterior is then simple, but the fit to the data is poor. As a increases, the latent line wraps around the circle several times, improving the data fit but creating several latent explanations for the same observation. For x_0=(1,0), p_a(z\mid x_0)\propto \exp\left(-\frac{z^2}{2}-\frac{1-\cos(az)}{\tau^2}\right). In the figure, q is the local Gaussian approximation around the latent explanation with the smallest absolute value. The data-fit penalty is the average negative log-likelihood relative to its best value over the displayed range. Its sum with the inference gap is therefore the negative ELBO up to an additive constant.