Understanding Variational Inference and the ELBO in AI

This article explores how variational inference transforms complex probabilistic inference problems into solvable optimization tasks in artificial intelligence. It covers the fundamental challenge of intractable posteriors, the role of Kullback-Leibler divergence, and the mathematical decomposition of the Evidence Lower Bound. Finally, it explains why maximizing this bound successfully approximates the true underlying data distributions in modern machine learning models.

In Bayesian machine learning, the primary goal is often to infer the posterior distribution of latent variables \(z\) given observed data \(x\), denoted as \(p(z|x)\). According to Bayes' theorem, this is calculated as:

\[p(z|x) = \frac{p(x, z)}{p(x)}\]

The denominator, \(p(x) = \int p(x, z) dz\), is known as the marginal likelihood or "evidence." In complex artificial intelligence models, such as deep neural networks, this integral is computationally intractable because it requires summing or integrating over an exponentially large space of latent configurations. Without \(p(x)\), the true posterior \(p(z|x)\) cannot be calculated directly.

Variational inference bypasses this intractable integration by turning it into an optimization problem. Instead of computing \(p(z|x)\) directly, the method introduces a simpler, tractable family of distributions, \(q(z)\), parameterized by variational parameters. The objective is to adjust these parameters so that the surrogate distribution \(q(z)\) becomes as close as possible to the true posterior \(p(z|x)\).

To measure the closeness between the approximated distribution and the true posterior, variational inference uses the Kullback-Leibler (KL) divergence. The goal is to minimize the KL divergence from \(q(z)\) to \(p(z|x)\), which is defined as:

\[KL(q(z) \parallel p(z|x)) = \mathbb{E}_{q(z)} \left[ \log \frac{q(z)}{p(z|x)} \right]\]

However, this equation still cannot be minimized directly because it contains the intractable \(p(z|x)\) term. By substituting Bayes' rule into the KL divergence formula and applying logarithmic identities, the equation can be rewritten as:

\[KL(q(z) \parallel p(z|x)) = \mathbb{E}_{q(z)} [\log q(z)] - \mathbb{E}_{q(z)} [\log p(x, z)] + \log p(x)\]

Notice that \(\log p(x)\) does not depend on \(z\), allowing it to be moved outside the expectation. Rearranging this equation isolates the intractable log-evidence on one side:

\[\log p(x) = \mathbb{E}_{q(z)} [\log p(x, z)] - \mathbb{E}_{q(z)} [\log q(z)] + KL(q(z) \parallel p(z|x))\]

Because the KL divergence is always non-negative (\(KL \geq 0\)), the first two terms on the right-hand side form a strict lower bound on the log-evidence \(\log p(x)\). This combination of terms is defined as the Evidence Lower Bound (ELBO):

\[ELBO(q) = \mathbb{E}_{q(z)} [\log p(x, z)] - \mathbb{E}_{q(z)} [\log q(z)]\]

Substituting the ELBO back into the rearranged equation yields a fundamental identity in variational inference:

\[\log p(x) = ELBO(q) + KL(q(z) \parallel p(z|x))\]

The marginal log-likelihood \(\log p(x)\) is a fixed quantity with respect to the variational distribution \(q(z)\) because it depends only on the observed data and the model priors. Therefore, the sum of the ELBO and the KL divergence is constant. Consequently, maximizing the ELBO directly forces the KL divergence to decrease. Because the ELBO only depends on the tractable joint distribution \(p(x, z)\) and the chosen surrogate distribution \(q(z)\), it can be computed and optimized using standard gradient ascent techniques.

To understand what maximizing the ELBO actually does to the model, the equation can be algebraically factored into a different form using \(p(x, z) = p(x|z)p(z)\):

\[ELBO(q) = \mathbb{E}_{q(z)} [\log p(x|z)] - KL(q(z) \parallel p(z))\]

Maximizing this form of the ELBO balances two competing objectives that guide the artificial intelligence model toward an accurate approximation:

1. Reconstruction Accuracy: The first term, \(\mathbb{E}_{q(z)} [\log p(x|z)]\), is the expected log-likelihood of the data. Maximizing this term encourages the variational distribution \(q(z)\) to place probability mass on latent variable states \(z\) that accurately explain and reconstruct the observed data \(x\).
2. Prior Regularization: The second term, \(-KL(q(z) \parallel p(z))\), acts as a regularizer. Maximizing the negative KL divergence is equivalent to minimizing the divergence between the variational distribution \(q(z)\) and the prior distribution \(p(z)\). This prevents the model from overfitting by ensuring the approximated posterior does not stray too far from initial prior beliefs.

By iteratively adjusting the parameters of \(q(z)\) to maximize these two terms, variational inference squeezes the gap between the lower bound and the true log-evidence, effectively molding the tractable \(q(z)\) into a reliable proxy for the intractable \(p(z|x)\).