VAE Derivation: Variational Inference

Our primary goal in generative modeling is to learn the underlying probability distribution $p(x)$ of some observed data $x$. Variational Autoencoders approach this by introducing latent variables $z$, which are unobserved variables that capture underlying structure or features of the data. We assume data $x$ is generated from these latent variables $z$ via a generative process $p_\theta(x|z)$, where $z$ itself is drawn from a prior distribution $p(z)$. The marginal likelihood of an observation $x$ is then given by:

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

The parameters $\theta$ of this generative model (often a neural network, the "decoder") are what we aim to learn. Maximizing $\log p_\theta(x)$ directly is challenging because this integral is usually intractable, especially when $p_\theta(x|z)$ is complex (like a deep neural network) and $z$ is high-dimensional.

Furthermore, to perform tasks like encoding data into its latent representation, we often need access to the posterior distribution $p_\theta(z|x) = \frac{p_\theta(x|z)p(z)}{p_\theta(x)}$. This is also intractable because computing its denominator, $p_\theta(x)$, is intractable.

This is where variational inference comes into play. Instead of computing the true posterior $p_\theta(z|x)$, we introduce a simpler, tractable distribution $q_\phi(z|x)$ to approximate it. This distribution $q_\phi(z|x)$, often called the variational posterior or inference network (and in VAEs, the "encoder"), is typically parameterized by $\phi$ (e.g., parameters of another neural network). Our goal is to make $q_\phi(z|x)$ as close as possible to the true posterior $p_\theta(z|x)$.

Let's start with the log-likelihood of a data point $x$, $\log p_\theta(x)$. We can rewrite it using our approximate posterior $q_\phi(z|x)$:

$$\log p_\theta(x) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x)]$$

This is valid because $\log p_\theta(x)$ does not depend on $z$, and $\int q_\phi(z|x) dz = 1$. Now, we use the definition of conditional probability, $p_\theta(x) = p_\theta(x,z) / p_\theta(z|x)$:

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

Next, we multiply and divide by $q_\phi(z|x)$ inside the logarithm, a common trick:

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

Using the property $\log(ab) = \log a + \log b$:

$$\log p_\theta(x) = \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{p_\theta(x,z)}{q_\phi(z|x)}\right] + \mathbb{E}_{q_\phi(z|x)}\left[\log \frac{q_\phi(z|x)}{p_\theta(z|x)}\right]$$

Let's examine these two terms. The second term is the Kullback-Leibler (KL) divergence between $q_\phi(z|x)$ and the true posterior $p_\theta(z|x)$:

$$\mathbb{E}_{q_\phi(z|x)}\left[\log \frac{q_\phi(z|x)}{p_\theta(z|x)}\right] = D_{KL}(q_\phi(z|x) || p_\theta(z|x))$$

The KL divergence is always non-negative ($D_{KL} \ge 0$), and it is zero if and only if $q_\phi(z|x) = p_\theta(z|x)$.

The first term is defined as the Evidence Lower Bound (ELBO), denoted $\mathcal{L}(\theta, \phi; x)$:

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

So, we have the fundamental identity:

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

Decomposition of the marginal log-likelihood into the ELBO and the KL divergence between the approximate and true posteriors. Maximizing the ELBO is our tractable objective.

Since $D_{KL}(q_\phi(z|x) || p_\theta(z|x)) \ge 0$, it follows that:

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

This is why $\mathcal{L}(\theta, \phi; x)$ is called the "Evidence Lower Bound": it provides a lower bound on the log-likelihood of the data (the "evidence"). By maximizing the ELBO with respect to both the generative model parameters $\theta$ and the variational parameters $\phi$, we are effectively:

1. Increasing the lower bound on $\log p_\theta(x)$, pushing it higher.
2. Minimizing the KL divergence $D_{KL}(q_\phi(z|x) || p_\theta(z|x))$, making our approximation $q_\phi(z|x)$ closer to the true posterior $p_\theta(z|x)$. If $q_\phi(z|x)$ becomes a perfect approximation, the KL divergence becomes zero, and the ELBO equals the true log-likelihood.

The ELBO can be rewritten in a form that is often more intuitive for VAEs. Starting from $\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x,z) - \log q_\phi(z|x)]$ and using $p_\theta(x,z) = p_\theta(x|z)p(z)$:

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

The second term is again a KL divergence: $D_{KL}(q_\phi(z|x) || p(z))$. This gives us the most common form of the ELBO used in VAEs:

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

Let's break down these two components:

• Reconstruction Term: $\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]$. This term measures how well the decoder $p_\theta(x|z)$ can reconstruct the input data $x$ given latent samples $z$ drawn from the encoder's approximation $q_\phi(z|x)$. Maximizing this term encourages the VAE to learn meaningful latent representations $z$ that retain sufficient information to reconstruct $x$. For instance, if $p_\theta(x|z)$ is a Gaussian distribution, this term becomes a squared error loss. If $p_\theta(x|z)$ is a Bernoulli distribution (for binary data), this becomes a binary cross-entropy loss.
• Regularization Term: $D_{KL}(q_\phi(z|x) || p(z))$. This term acts as a regularizer. It measures the divergence between the approximate posterior distribution $q_\phi(z|x)$ (output by the encoder for a given $x$) and the prior distribution $p(z)$ over the latent variables. The prior $p(z)$ is typically chosen to be a simple distribution, like a standard multivariate Gaussian $\mathcal{N}(0, I)$. This KL term encourages the encoder to distribute the latent representations $z$ similarly to the prior. This regularization is essential for ensuring that the latent space is well-structured and can be sampled from to generate new data. Without it, the encoder might learn to produce $q_\phi(z|x)$ distributions that are far apart for different $x$, leading to a non-smooth or "gappy" latent space.

The VAE, therefore, consists of:

1. An encoder network $q_\phi(z|x)$ that takes data $x$ and outputs parameters for the distribution of $z$ (e.g., mean and variance if $q_\phi(z|x)$ is Gaussian). Its parameters are $\phi$.
2. A decoder network $p_\theta(x|z)$ that takes a latent sample $z$ and outputs parameters for the distribution of $x$. Its parameters are $\theta$.
3. A chosen prior $p(z)$ over latent variables, usually $\mathcal{N}(0,I)$.

The training process involves optimizing the ELBO $\mathcal{L}(\theta, \phi; x)$ with respect to both $\theta$ and $\phi$ using techniques like stochastic gradient ascent. This derivation provides the theoretical justification for the VAE objective function, balancing data reconstruction with regularization of the latent space. This balance is what allows VAEs to learn rich, structured representations and generate novel data samples. The term "variational" in Variational Autoencoder and Variational Inference refers to methods from the calculus of variations, where we optimize functionals (functions of functions), in this case, finding the best $q_\phi(z|x)$ within a family of distributions parameterized by $\phi$.

In the subsequent sections, we will explore each component of this objective, the practicalities of optimizing it (like the reparameterization trick), and the implications of this formulation.