Introduction to variational inference

In probabilistic machine learning, our goal is often not just to find a single “best” set of parameters for a model, but to quantify our uncertainty over all possible parameters. We accomplish this using Bayesian inference.

Let $\Theta \subseteq \mathbb{R}^{P}$ denote the parameter space, and let $\theta \in \Theta$ be the vector of all parameters in our model (e.g., the weights and biases of a neural network). Let $\mathcal{D} = \{(x_n, y_n)\}_{n=1}^{N}$ be our observed training data, consisting of $N$ input–target pairs. Bayesian inference updates our prior beliefs about the parameters, $p(\theta)$, after observing the data, yielding the posterior distribution $p(\theta \mid \mathcal{D})$.

This update is governed by Bayes’ theorem:

Here,

• $p(\theta \mid \mathcal{D})$ is the posterior distribution, summarising our updated beliefs about $\theta$ given the data;

• $p(\mathcal{D} \mid \theta)$ is the likelihood, viewed as a function of $\theta$: for a given choice of parameters, it quantifies how likely the observed data are under the model;

• $p(\theta)$ is the prior distribution, which encodes our beliefs about $\theta$ before observing any data; and

• $p(\mathcal{D}) = \int p(\mathcal{D} \mid \theta)p(\theta)\,\mathrm{d}\theta$ is the marginal likelihood, also called the evidence.

The central computational challenge in most modern Bayesian inference problems is evaluating $p(\mathcal{D})$, which requires integrating the joint distribution $p(\mathcal{D},\theta)$ over the parameter space $\Theta$. For models such as Bayesian neural networks, this integral is high-dimensional and generally lacks an analytic closed-form solution. Consequently, the posterior density $p(\theta \mid \mathcal{D})$ usually cannot be evaluated in normalised form, motivating approximate inference methods. This chapter focuses on variational inference, also known as variational Bayes, a widely used family of scalable approximate inference methods Jordan et al., 1999Wainwright & Jordan, 2008Blei et al., 2017.

A Variational Approach: Bayesian Inference as Optimisation¶

The core idea of variational inference (VI) is to reframe Bayesian inference as an optimisation problem. Instead of representing the generally intractable posterior $p(\theta \mid \mathcal{D})$ exactly, we choose a tractable distribution $q(\theta;\varparams)$ from a variational family $\mathcal{Q}$, where $\varparams$ denotes its variational parameters. Concretely, $\mathcal{Q}$ could be a family of Gaussian distributions or a family of factorised distributions, as developed in Chapter Gaussian mean-field variational inference. The goal is to find the member $q(\theta;\varparams^\ast)$ that best approximates the posterior.

The Kullback–Leibler Divergence¶

Following Kullback & Leibler (1951), the Kullback–Leibler (KL) divergence from $q(\theta;\varparams)$ to the posterior $p(\theta \mid \mathcal{D})$ is

Because the KL divergence is asymmetric, its direction must be stated explicitly. VI seeks variational parameters satisfying

Rewriting the KL Objective¶

Expanding the posterior in Equation (2) using Bayes’ theorem gives

Although the KL divergence itself cannot generally be evaluated exactly because $\log p(\mathcal{D})$ is unavailable, the log-evidence term is constant with respect to $\varparams$ and therefore does not affect the optimiser in Equation (3). The remaining expectation does not require evaluating $p(\mathcal{D})$ and can be estimated using samples from $q(\theta;\varparams)$ whenever the joint density $p(\mathcal{D},\theta)$ and variational density $q(\theta;\varparams)$ can be evaluated.

The Evidence Lower Bound (ELBO)¶

Rearranging Equation (4) yields

The first term defines the evidence lower bound (ELBO) $\mathcal{L}(\varparams)$. Because the KL divergence is non-negative,

The gap between the ELBO and the log evidence is exactly $D_{KL}\left(q(\theta;\varparams)\middle\|p(\theta \mid \mathcal{D})\right)$; the closer the variational distribution is to the posterior in this divergence, the smaller the gap. Since $\log p(\mathcal{D})$ is constant with respect to $\varparams$, Equation (5) also gives

Thus, VI replaces evaluation of the posterior in normalised form with optimisation of an objective that can be estimated under suitable choices of likelihood, prior, and variational family.

Interpreting the ELBO¶

Using $p(\mathcal{D},\theta)=p(\mathcal{D} \mid \theta)p(\theta)$, the ELBO can be written as

The alternative form in Equation (8) reveals a trade-off between two terms:

• Expected log-likelihood (data fit): This term rewards variational distributions that place more probability in regions of parameter space where the observed data have high likelihood.

• KL penalty (regularisation): This term penalises departures of $q(\theta;\varparams)$ from the prior $p(\theta)$ and can therefore be viewed as a complexity penalty.

VI selects the member of $\mathcal{Q}$ that balances fit to the observed data against departure from the prior.

Other Lower Bounds on the Log-Evidence¶

The ELBO is one lower bound on $\log p(\mathcal{D})$. More generally, strictly positive unbiased estimators of $p(\mathcal{D})$ can induce lower bounds through Jensen’s inequality; this optional subsection states the construction underlying importance-weighted and sequential Monte Carlo objectives.

Assume that $q(\theta;\varparams)>0$ whenever $p(\mathcal{D},\theta)>0$. Then

For a concave function $\varphi$ and random variable $X$, Jensen’s inequality states that $\varphi(\mathbb{E}\left[X\right])\geq\mathbb{E}\left[\varphi(X)\right]$. Taking $\varphi=\log$ in Equation (9) gives

More generally, let $r(\mathbf{u};\varparams)$ denote the sampling density of all random variables $\mathbf{u}$ used to construct a strictly positive unbiased estimator $g(\mathcal{D};\mathbf{u},\varparams)$ of $p(\mathcal{D})$, and suppose that $\mathbb{E}_{r(\mathbf{u};\varparams)}\left[\lvert\log g(\mathcal{D};\mathbf{u},\varparams)\rvert\right]<\infty$. Then

Jensen’s inequality gives the lower bound

The standard ELBO is recovered by taking $\mathbf{u}=\theta$, $r(\theta;\varparams)=q(\theta;\varparams)$, and $g(\mathcal{D};\theta,\varparams)=p(\mathcal{D},\theta)/q(\theta;\varparams)$. For example, Burda et al. (2016) introduce the importance-weighted autoencoder (IWAE) objective, which uses $s$ independent samples $\theta_j\sim q(\cdot;\varparams)$ and the estimator

Related lower bounds can be constructed using sequential Monte Carlo Maddison et al., 2017. Although $\mathcal{L}_g(\varparams)$ may not be available in closed form, it can often be optimised using Monte Carlo gradient estimators. However, a tighter bound alone does not guarantee better optimisation or predictive performance, because the signal-to-noise ratio and other properties of the associated gradient estimator also matter Rainforth et al., 2018.

In summary, VI approximates $p(\theta \mid \mathcal{D})$ by optimising a distribution from a tractable family, with the ELBO providing an objective that balances data fit against departure from the prior. Chapter Gaussian mean-field variational inference develops the mean-field construction in detail, and subsequent chapters consider further variational families and optimisation methods.