Deriving the ELBO for VAEs

Hello! Let's continue our exploration of generative models.

Introduction

In our last lesson, you successfully implemented a Variational Autoencoder (VAE) from scratch. We focused on the practical "how": the architecture, the crucial reparameterization trick, and the two-part loss function combining reconstruction and KL divergence. However, we took the form of this loss function as a given.

Today, we transition from the "how" to the "why." This lesson is dedicated to deriving that loss function from first principles. Your learning outcome is to derive the Evidence Lower Bound (ELBO) objective for VAEs. We will unpack the mathematics to understand:

1. Why we need an alternative to maximizing the data likelihood directly.
2. How the principles of variational inference allow us to define a tractable objective.
3. How this objective, the ELBO, naturally decomposes into the reconstruction and regularization terms you implemented previously.

This derivation is a cornerstone of modern generative modeling and will provide a solid theoretical foundation for your understanding.

The Generative Goal and Its Intractability

The fundamental goal of any generative model is to learn the true underlying data distribution, $p(x)$. A good model should assign a high probability to data points similar to those in our training set and a low probability to everything else. Therefore, the most natural objective is to maximize the log-likelihood of our training data: $\log p(x)$.

In a VAE, we model the data $x$ as being generated from a latent variable $z$. The probability of observing a single data point $x$ is obtained by considering all possible latent variables that could have generated it. This is done by marginalizing over $z$:

$$p(x)=\int p(x,z)dz=\int p(x|z)p(z)dz$$

Here, $p(z)$ is the prior distribution over the latent space (e.g., a standard normal distribution), and $p(x|z)$ is the likelihood of the data given a latent variable, which is defined by our decoder network.

This integral presents a major problem. The decoder network defines a highly complex, non-linear function from $z$ to $x$. This makes the integral analytically intractable—there's no closed-form solution. Furthermore, the latent space is often high-dimensional, making numerical approximation via sampling (e.g., Monte Carlo methods) computationally infeasible due to the curse of dimensionality.

Since we cannot optimize $\log p(x)$ directly, we need a different approach. This is where variational inference comes in.

Variational Inference and the ELBO Derivation

The core idea of variational inference is to approximate the true, intractable posterior distribution $p(z|x)$ with a simpler, tractable distribution $q_{\varphi }(z|x)$. In a VAE, this is our encoder, parameterized by $\varphi$. We want to make $q_{\varphi }(z|x)$ as close as possible to the true posterior $p(z|x)$.

To measure the "closeness" of these two distributions, we use the Kullback-Leibler (KL) divergence. Our goal is to minimize $D_{KL}(q_{\varphi }(z|x)||p(z|x))$.

Let's see how this minimization goal leads us to the ELBO. The following video provides a clear, step-by-step walkthrough of the mathematical derivation.

Let's formalize and review the steps from the video. The derivation starts with the log-likelihood of a single data point, $\log p(x)$, and cleverly manipulates it.

1. Start with $\log p(x)$. We can express this as an expectation with respect to our approximate posterior $q_{\varphi }(z|x)$, since $\log p(x)$ does not depend on $z$:
   $\log p(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p(x)]$

2. Apply Bayes' Rule. We can write $p(x)=\frac{p(x,z)}{p(z|x)}$. Substituting this in:
   $\log p(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{p(z|x)}\right]$

3. Multiply and Divide by $q_{\varphi }(z|x)$. This is the key trick. We introduce our approximation into the equation:
   $\log p(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \left(\frac{p(x,z)}{q_{\varphi }(z|x)}\cdot \frac{q_{\varphi }(z|x)}{p(z|x)}\right)\right]$

4. Split the Logarithm. Using the property $\log (ab)=\log (a)+\log (b)$:
   $\log p(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \frac{p(x,z)}{q_{\varphi }(z|x)}\right]+{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \frac{q_{\varphi }(z|x)}{p(z|x)}\right]$

5. Identify the Terms. We now have two important terms:

   • The first term is what we will call the Evidence Lower Bound (ELBO), or $\mathcal{L}(x)$.
   • The second term is precisely the definition of the KL divergence between our approximation and the true posterior, $D_{KL}(q_{\varphi }(z|x)||p(z|x))$.

This gives us the fundamental identity:

$$\log p(x)=\mathcal{L}(x)+D_{KL}(q_{\varphi }(z|x)||p(z|x))$$

Since the KL divergence is always non-negative ($D_{KL}\ge 0$), we have the inequality:

$$\log p(x)\ge \mathcal{L}(x)$$

This is why $\mathcal{L}(x)$ is called a lower bound on the log-evidence (another term for log-likelihood). By maximizing the ELBO, we are pushing up this lower bound, which in turn pushes up the log-likelihood of our data. We have successfully found a tractable proxy objective!

Test your understanding!

Given the equation $\log p(x)=\mathcal{L}(x)+D_{KL}(q_{\varphi }(z|x)||p(z|x))$, consider a fixed data point $x$. What is the relationship between maximizing the ELBO, $\mathcal{L}(x)$, and minimizing the KL divergence, $D_{KL}(q_{\varphi }(z|x)||p(z|x))$?

Show answer

For a fixed data point $x$, its true log-likelihood $\log p(x)$ is a constant value. Therefore, the equation becomes $\text{Constant}=\mathcal{L}(x)+D_{KL}(...)$. Maximizing $\mathcal{L}(x)$ is thus equivalent to minimizing $D_{KL}(q_{\varphi }(z|x)||p(z|x))$. This beautifully connects our practical objective (maximizing the ELBO) to our theoretical goal (making our approximation $q_{\varphi }(z|x)$ as close as possible to the true posterior $p(z|x)$).

For a detailed textual walk-through of this derivation, the following resource is excellent.

Unpacking the ELBO: From Theory to Loss Function

Now, let's rearrange the ELBO into the form you saw in the last lesson. This will reveal the reconstruction and regularization terms.

The ELBO is:

$$\mathcal{L}(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \frac{p_{\theta }(x,z)}{q_{\varphi }(z|x)}\right]$$

(Here we've added $\theta$ to denote the parameters of our generative model/decoder).

Using the chain rule of probability, $p_{\theta }(x,z)=p_{\theta }(x|z)p(z)$, we get:

$$\mathcal{L}(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}\right]$$

Splitting the logarithm again:

$$\mathcal{L}(x)={\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]+{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}\left[\log \frac{p(z)}{q_{\varphi }(z|x)}\right]$$

Rearranging the second term gives the final, most common form of the ELBO:

$$\mathcal{L}(x)=\underset{\text{Reconstruction Term}}{\underbrace{{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]}}-\underset{\text{Regularization Term}}{\underbrace{D_{KL}(q_{\varphi }(z|x)||p(z))}}$$

This is exactly the objective we want to maximize!

• The Reconstruction Term: ${\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]$ is the expected log-likelihood of the input data $x$ given a latent sample $z$ drawn from the encoder's output. Maximizing this term forces the decoder to become good at reconstructing the original data. In practice, this corresponds to minimizing a reconstruction loss like Binary Cross-Entropy or Mean Squared Error.
• The Regularization Term: $-D_{KL}(q_{\varphi }(z|x)||p(z))$ measures the divergence between the distribution produced by the encoder for a given input, $q_{\varphi }(z|x)$, and the fixed prior over the latent space, $p(z)$. Maximizing this term is equivalent to minimizing the KL divergence, which pushes the encoded distributions to be close to the prior (e.g., a standard normal $\mathcal{N}(0,I)$). This is what prevents the model from "cheating" by encoding each point into a separate, far-flung distribution, thus ensuring the latent space remains smooth and continuous.

Since standard optimizers perform minimization, our final loss function is simply the negative of the ELBO:
$\text{Loss}=-\mathcal{L}(x)=-{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\log p_{\theta }(x|z)]+D_{KL}(q_{\varphi }(z|x)||p(z))$

This is precisely the combination of reconstruction loss and KL divergence loss you implemented in the previous lesson.

An Alternative Derivation with Jensen's Inequality

It's often insightful to see that a fundamental result can be reached from different directions. The ELBO can also be derived using Jensen's Inequality, which applies to concave functions like the logarithm.

The video below offers this alternative perspective, which can deepen your intuition about why the ELBO is a lower bound.

Conclusion

In this lesson, we have bridged the gap between the practical implementation of a VAE and its theoretical underpinnings. You have seen how a seemingly arbitrary loss function arises naturally from the principles of variational inference.

Key Takeaways:

• The goal of a generative model is to maximize the data log-likelihood $\log p(x)$, but this is intractable for VAEs due to a complex integral.
• Variational Inference provides a solution by introducing a tractable approximation to the posterior, $q_{\varphi }(z|x)$ (the encoder), and optimizing it to be close to the true posterior.
• This leads to the Evidence Lower Bound (ELBO), a tractable objective function that we can maximize as a proxy for the true log-likelihood.
• The ELBO elegantly decomposes into two parts: a reconstruction term that ensures the model can regenerate the data, and a regularization term (a KL divergence) that enforces a smooth, structured latent space suitable for generation.
• Maximizing the ELBO is equivalent to minimizing the negative ELBO, which gives us the practical loss function you've already used: Reconstruction Loss + KL Divergence Loss.

Preview of the next lesson:
Now that you have a firm grasp of the probabilistic approach to generation with VAEs, we will explore a completely different paradigm. In the next lesson, you will build a basic Generative Adversarial Network (GAN). Instead of probabilistic inference, GANs frame generation as a two-player game between a generator and a discriminator. This will open up a new set of powerful techniques and conceptual tools.