Deriving the Evidence Lower Bound (ELBO)

In the previous sections, we established that Variational Autoencoders (VAEs) approach generative modeling from a probabilistic perspective. Instead of mapping an input $x$ to a single point $z$ in the latent space, the VAE encoder, parameterized by $\phi$, outputs parameters for a probability distribution $q_\phi(z|x)$ over the latent space. The decoder, parameterized by $\theta$, then defines a likelihood $p_\theta(x|z)$, representing the probability of generating $x$ from a given latent variable $z$.

Our ultimate goal in training a generative model is to maximize the probability (or likelihood) of observing the actual data $x$ under our model. This marginal likelihood is given by integrating over all possible latent variables $z$:

$$p_\theta(x) = \int p_\theta(x|z) p(z) dz$$

Here, $p(z)$ is a chosen prior distribution for the latent variables, often a standard multivariate Gaussian ($N(0, I)$). Unfortunately, computing this integral is generally intractable for complex models like neural networks used as decoders, because it requires integrating over all possible values of $z$, which can be high-dimensional.

Introducing the Variational Approach

Since we cannot directly optimize $p_\theta(x)$, VAEs use an approach derived from variational inference. The core idea is to introduce an approximate distribution $q_\phi(z|x)$ (our encoder) that aims to estimate the true but intractable posterior distribution $p_\theta(z|x) = p_\theta(x|z)p(z) / p_\theta(x)$. We want to make $q_\phi(z|x)$ as close as possible to $p_\theta(z|x)$.

A standard way to measure the "closeness" between two distributions is the Kullback-Leibler (KL) divergence. Let's examine the KL divergence between our approximate posterior $q_\phi(z|x)$ and the true posterior $p_\theta(z|x)$:

$$KL(q_\phi(z|x) || p_\theta(z|x)) = \int q_\phi(z|x) \log \frac{q_\phi(z|x)}{p_\theta(z|x)} dz$$

This represents the information lost when using $q_\phi$ to approximate $p_\theta$. We want to minimize this KL divergence. Let's expand the definition using the definition of the posterior $p_\theta(z|x)$:

$$KL(q_\phi(z|x) || p_\theta(z|x)) = E_{z \sim q_\phi(z|x)} [\log q_\phi(z|x) - \log p_\theta(z|x)]$$

Using Bayes' theorem, $p_\theta(z|x) = \frac{p_\theta(x|z) p(z)}{p_\theta(x)}$, we can substitute $\log p_\theta(z|x) = \log p_\theta(x|z) + \log p(z) - \log p_\theta(x)$:

$$KL(q_\phi(z|x) || p_\theta(z|x)) = E_{z \sim q_\phi(z|x)} [\log q_\phi(z|x) - (\log p_\theta(x|z) + \log p(z) - \log p_\theta(x))]$$

Rearranging the terms inside the expectation:

$$KL(q_\phi(z|x) || p_\theta(z|x)) = E_{z \sim q_\phi(z|x)} [\log q_\phi(z|x) - \log p(z) - \log p_\theta(x|z) + \log p_\theta(x)]$$

Since $\log p_\theta(x)$ does not depend on $z$, the expectation $E_{z \sim q_\phi(z|x)}[\log p_\theta(x)]$ is simply $\log p_\theta(x)$. We can separate the terms:

$$KL(q_\phi(z|x) || p_\theta(z|x)) = E_{z \sim q_\phi(z|x)} [\log q_\phi(z|x) - \log p(z)] - E_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)] + \log p_\theta(x)$$

The first term is the KL divergence between the approximate posterior $q_\phi(z|x)$ and the prior $p(z)$: $KL(q_\phi(z|x) || p(z))$. So, we have:

$$KL(q_\phi(z|x) || p_\theta(z|x)) = KL(q_\phi(z|x) || p(z)) - E_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)] + \log p_\theta(x)$$

The Evidence Lower Bound (ELBO)

Now, let's rearrange this equation to isolate the intractable log-likelihood $\log p_\theta(x)$:

$$\log p_\theta(x) = E_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)] - KL(q_\phi(z|x) || p(z)) + KL(q_\phi(z|x) || p_\theta(z|x))$$

This equation is significant. It decomposes the log-likelihood of the data into three terms. Notice that the KL divergence $KL(q_\phi(z|x) || p_\theta(z|x))$ is always non-negative ($KL \ge 0$). Therefore, if we drop this term, the remaining expression forms a lower bound on the log-likelihood. This lower bound is called the Evidence Lower Bound (ELBO), often denoted as $\mathcal{L}(\theta, \phi; x)$:

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

Since $KL(q_\phi(z|x) || p_\theta(z|x)) \ge 0$, we have:

$$\log p_\theta(x) \ge \mathcal{L}(\theta, \phi; x)$$

This relationship is fundamental to VAE training. Instead of maximizing the intractable $\log p_\theta(x)$ directly, we maximize its tractable lower bound, the ELBO $\mathcal{L}(\theta, \phi; x)$. Maximizing the ELBO serves two simultaneous objectives:

1. It pushes the lower bound $\mathcal{L}$ closer to the actual log-likelihood $\log p_\theta(x)$.
2. Since $\log p_\theta(x) - \mathcal{L}(\theta, \phi; x) = KL(q_\phi(z|x) || p_\theta(z|x))$, maximizing the ELBO is equivalent to minimizing the KL divergence between the approximate posterior $q_\phi(z|x)$ and the true posterior $p_\theta(z|x)$. This means our encoder $q_\phi(z|x)$ becomes a better approximation of the true posterior as the ELBO increases.

Relationship between the total log likelihood, the ELBO, and the KL divergence gap. Maximizing the ELBO pushes it upwards towards the total log likelihood, effectively minimizing the KL gap.

Components of the ELBO

Let's analyze the two terms comprising the ELBO objective function:

1. Reconstruction Term: $E_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)]$ This term measures the expected log-likelihood of the original input data $x$ being generated by the decoder $p_\theta(x|z)$, where the latent code $z$ is sampled from the distribution proposed by the encoder $q_\phi(z|x)$. Maximizing this term encourages the decoder to learn to reconstruct the input accurately from its latent representation. In practice, for inputs like images, this often translates to minimizing a reconstruction loss like Mean Squared Error (MSE) or Binary Cross-Entropy (BCE), depending on the data distribution assumptions.

2. Regularization Term (KL Divergence): $- KL(q_\phi(z|x) || p(z))$ This term measures the KL divergence between the approximate posterior distribution $q_\phi(z|x)$ produced by the encoder for a given input $x$, and the prior distribution $p(z)$ (e.g., a standard Gaussian). Maximizing this term (or minimizing the positive KL divergence) encourages the encoder to produce distributions $q_\phi(z|x)$ that are close to the prior $p(z)$. This acts as a regularizer, forcing the latent space to adopt the structure of the prior (e.g., a smooth, centered Gaussian cloud). This regularization is what enables the VAE to generate novel samples by drawing $z$ directly from the prior $p(z)$ and feeding it to the decoder.

The Optimization Objective

The VAE is trained by maximizing the ELBO with respect to both the encoder parameters $\phi$ and the decoder parameters $\theta$. This objective creates a balance: the reconstruction term pushes for accurate data representation, while the KL divergence term enforces a regular structure on the latent space suitable for generation.

In practice, we often minimize the negative ELBO:

$$\text{Loss}_{VAE} = - \mathcal{L}(\theta, \phi; x) = - E_{z \sim q_\phi(z|x)} [\log p_\theta(x|z)] + KL(q_\phi(z|x) || p(z))$$

Minimizing this loss function using gradient descent (enabled by the reparameterization trick, discussed next) trains the VAE to find parameters $\phi$ and $\theta$ that yield both good reconstructions and a well-structured latent space aligned with the prior $p(z)$. Understanding the ELBO and its components is central to understanding how VAEs learn to generate data.