Module 9: Variational Inference

This module develops variational inference (VI) as an optimization-based approximation to Bayesian inference. We treat VI as a problem of projecting the posterior distribution onto a restricted family of distributions by minimizing a divergence functional. The exposition is measure-theoretic and emphasizes:

the functional-analytic properties of the Kullback–Leibler divergence,
the ELBO (evidence lower bound) as a variational characterization of the marginal likelihood,
the structure of mean-field families and coordinate ascent variational inference (CAVI), and
the relationship between VI and MCMC in terms of bias, variance, and asymptotics.

1. KL Divergence and the ELBO

We work on a measurable parameter space \(\bigl(\Theta, \mathcal{T}\bigr)\) with a \(\sigma\)-finite reference measure \(\nu\) (typically Lebesgue measure on \(\mathbb{R}^d\)). All densities below are understood with respect to \(\nu\).

1.1 KL Divergence: Measure-Theoretic Definition

Let \(P\) and \(Q\) be probability measures on \(\bigl(\Theta, \mathcal{T}\bigr)\). Assume that

\[ Q \ll P, \]

so that the Radon–Nikodym derivative \(\frac{dQ}{dP}\) exists. The Kullback–Leibler divergence from \(Q\) to \(P\) is

\[ \operatorname{KL}(Q \Vert P) := \int_\Theta \log\Bigl( \frac{dQ}{dP}(\theta) \Bigr) \, Q(d\theta) = \int_\Theta \frac{dQ}{dP}(\theta) \log\Bigl( \frac{dQ}{dP}(\theta) \Bigr) \, P(d\theta), \]

with the convention that \(a \log a = 0\) at \(a=0\) and that \(\operatorname{KL}(Q \Vert P) = +\infty\) if \(Q \not\ll P\).

In the common case where \(P\) and \(Q\) admit densities \(p\) and \(q\) with respect to \(\nu\), we write

\[ \operatorname{KL}(q \Vert p) = \int_\Theta q(\theta) \log \frac{q(\theta)}{p(\theta)} \, \nu(d\theta). \]

Basic properties. Under mild integrability assumptions:

\(\operatorname{KL}(Q \Vert P) \ge 0\),
\(\operatorname{KL}(Q \Vert P) = 0\) if and only if \(Q = P\) (as measures),
\(Q \mapsto \operatorname{KL}(Q \Vert P)\) is convex.

These follow from Jensen’s inequality applied to the convex function \(x \mapsto x \log x\).

1.2 Posterior Approximation and Definition of the ELBO

Let \(x\) denote observed data. Write the joint density

\[ p(\theta, x) = p(x \mid \theta)\, p(\theta), \]

and suppose the posterior density exists,

\[ p(\theta \mid x) = \frac{p(\theta, x)}{p(x)}, \quad p(x) = \int p(\theta, x)\, d\theta. \]

Let \(\mathcal{Q}\) be a family of candidate densities \(q(\theta)\) (with respect to the same reference measure). Variational inference chooses

\[ q^* \in \arg\min_{q \in \mathcal{Q}} \operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr). \]

Direct optimization of \(\operatorname{KL}(q \Vert p(\cdot \mid x))\) is often inconvenient because it involves the (intractable) posterior density. We therefore introduce the evidence lower bound (ELBO),

\[ \mathcal{L}(q) := \int q(\theta) \log \frac{p(\theta, x)}{q(\theta)} \, d\theta = \mathbb{E}_q\bigl[ \log p(\theta, x) \bigr] - \mathbb{E}_q\bigl[ \log q(\theta) \bigr]. \]

We will show that maximizing \(\mathcal{L}(q)\) is equivalent to minimizing \(\operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr)\).

Theorem 1 Theorem 9.1 (ELBO–KL Decomposition).

For any density \(q\) such that \(\operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr) < \infty\), we have

\[ \log p(x) = \mathcal{L}(q) + \operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr). \]

In particular, \(\mathcal{L}(q) \le \log p(x)\), with equality if and only if \(q = p(\cdot \mid x)\) almost everywhere.

Proof. By definition of KL divergence between \(q\) and the posterior, \[ \operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr) = \int q(\theta) \log \frac{q(\theta)}{p(\theta \mid x)}\, d\theta. \] Using Bayes’ rule \(p(\theta \mid x) = p(\theta, x)/p(x)\), \[ \log \frac{q(\theta)}{p(\theta \mid x)} = \log q(\theta) - \log p(\theta,x) + \log p(x). \] Substitute into the KL expression and integrate termwise: \[ \begin{aligned} \operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr) &= \int q(\theta) \log q(\theta)\, d\theta - \int q(\theta) \log p(\theta,x)\, d\theta + \int q(\theta) \log p(x)\, d\theta \\ &= \mathbb{E}_q[\log q(\theta)] - \mathbb{E}_q[\log p(\theta,x)] + \log p(x) \int q(\theta)\, d\theta. \end{aligned} \] Since \(q\) is a density, the last integral equals 1. Rearranging gives \[ \log p(x) = \underbrace{\mathbb{E}_q[\log p(\theta,x)] - \mathbb{E}_q[\log q(\theta)]}_{\mathcal{L}(q)} + \operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr), \] which is the claimed identity.

The inequality \(\mathcal{L}(q) \le \log p(x)\) follows from nonnegativity of the KL divergence. Equality holds if and only if \(\operatorname{KL}\bigl(q \Vert p(\cdot \mid x)\bigr) = 0\), i.e., \(q = p(\cdot \mid x)\) almost everywhere. \(\square\)

Thus, maximizing the ELBO is equivalent to minimizing the KL divergence to the posterior within the chosen family \(\mathcal{Q}\).

2. Mean-Field Variational Families

We now restrict to a tractable family \(\mathcal{Q}\) of candidate densities. A common choice is the mean-field family, which assumes conditional independence between groups of parameters under \(q\).

2.1 Factorization Assumption

Partition \(\theta\) into blocks

\[ \theta = (\theta_1, \dots, \theta_J), \]

where each \(\theta_j\) takes values in some measurable space \((\Theta_j, \mathcal{T}_j)\). The mean-field variational family is

\[ \mathcal{Q}_{\text{MF}} := \Bigl\{ q(\theta) = \prod_{j=1}^J q_j(\theta_j) : q_j \text{ densities on } \Theta_j \Bigr\}. \]

The factorization is an approximation: in general, the true posterior \(p(\theta \mid x)\) does not factorize in this way, so VI introduces dependence-structure bias.

2.2 Coordinate Ascent Variational Inference (CAVI)

We wish to solve the constrained optimization problem

\[ \max_{q_1,\dots,q_J}\, \mathcal{L}(q) \quad \text{subject to}\quad q(\theta) = \prod_{j=1}^J q_j(\theta_j), \; q_j \ge 0, \; \int q_j(\theta_j)\, d\theta_j = 1. \]

Exact joint optimization is often difficult, but we can use coordinate ascent: repeatedly optimize \(\mathcal{L}(q)\) with respect to a single factor \(q_j\), holding the others fixed. The optimal \(q_j\) has a closed form.

Theorem 2 Theorem 9.2 (Optimal Mean-Field Factor Updates).

Fix \(q_{-j}(\theta_{-j}) = \prod_{k \ne j} q_k(\theta_k)\) and consider \(q_j\) varying over densities on \(\Theta_j\). Then, up to multiplicative normalization,

\[ \log q_j^*(\theta_j) = \mathbb{E}_{q_{-j}}[\log p(\theta, x)] + \text{constant}, \]

or equivalently,

\[ q_j^*(\theta_j) \propto \exp\Bigl(\mathbb{E}_{q_{-j}}[\log p(\theta, x)]\Bigr). \]

Proof. Write the ELBO as \[ \mathcal{L}(q) = \mathbb{E}_q[\log p(\theta,x)] - \mathbb{E}_q[\log q(\theta)]. \] Under the mean-field factorization \(q(\theta) = q_j(\theta_j) q_{-j}(\theta_{-j})\), we can group terms depending on \(q_j\) and those independent of \(q_j\).

First term: \[ \mathbb{E}_q[\log p(\theta,x)] = \int q_j(\theta_j)\, \mathbb{E}_{q_{-j}}[\log p(\theta,x)]\, d\theta_j. \]

Second term (entropy part): \[ \mathbb{E}_q[\log q(\theta)] = \mathbb{E}_q\bigl[\log q_j(\theta_j)\bigr] + \mathbb{E}_q\bigl[\log q_{-j}(\theta_{-j})\bigr]. \] The second expectation does not depend on \(q_j\) and can be treated as constant. Hence, up to an additive constant independent of \(q_j\), the ELBO as a functional of \(q_j\) is \[ \mathcal{L}_j(q_j) = \int q_j(\theta_j) \mathbb{E}_{q_{-j}}[\log p(\theta,x)]\, d\theta_j - \int q_j(\theta_j) \log q_j(\theta_j)\, d\theta_j. \]

We must maximize \(\mathcal{L}_j(q_j)\) over all densities \(q_j\) on \(\Theta_j\). This is a functional optimization problem with a normalization constraint \(\int q_j(\theta_j) d\theta_j = 1\). Introduce a Lagrange multiplier \(\lambda\) and consider the Lagrangian \[ \mathcal{F}(q_j) = \mathcal{L}_j(q_j) + \lambda\Bigl( \int q_j(\theta_j) d\theta_j - 1 \Bigr). \]

Taking a first variation in the direction of an arbitrary perturbation \(h\) with \(\int h(\theta_j) d\theta_j = 0\) and using Gâteaux derivatives, we obtain \[ \delta \mathcal{F} = \int h(\theta_j) \Bigl( \mathbb{E}_{q_{-j}}[\log p(\theta,x)] - (1 + \log q_j(\theta_j)) + \lambda \Bigr) d\theta_j. \] At an optimum, \(\delta \mathcal{F} = 0\) for all such \(h\), which implies that the term in parentheses must vanish almost everywhere: \[ \mathbb{E}_{q_{-j}}[\log p(\theta,x)] - (1 + \log q_j^*(\theta_j)) + \lambda = 0. \] Solving for \(\log q_j^*(\theta_j)\) gives \[ \log q_j^*(\theta_j) = \mathbb{E}_{q_{-j}}[\log p(\theta,x)] + (\lambda - 1). \] The constant \((\lambda - 1)\) enforces normalization and does not depend on \(\theta_j\), so we write simply \[ \log q_j^*(\theta_j) = \mathbb{E}_{q_{-j}}[\log p(\theta,x)] + C, \] with \(C\) a constant. Exponentiating both sides and renormalizing yields the stated form. \(\square\)

This theorem is the basis of coordinate ascent variational inference (CAVI): at each step, replace \(q_j\) by \(q_j^*\) while holding \(q_{-j}\) fixed.

2.3 Exponential Family VI

Suppose the joint density \(p(\theta,x)\) belongs to an exponential family of the form

\[ \log p(\theta,x) = \eta(x)^\top T(\theta) - A(\eta(x)) + c(x), \]

and suppose each factor \(q_j\) is restricted to an exponential family with natural parameters \(\lambda_j\). Then the expectation \(\mathbb{E}_{q_{-j}}[\log p(\theta,x)]\) is often linear in sufficient statistics of \(\theta_j\), implying that \(q_j^*\) remains in the same exponential family with updated natural parameters determined by expectations under \(q_{-j}\). This yields closed-form CAVI updates.

3. Variational Inference in a Conjugate Normal–Normal Model

We now work out a concrete example to illustrate VI computations and compare them to exact Bayesian inference.

3.1 Model and Exact Posterior

Consider a Normal–Normal model with known variance:

\[ Y_i \mid \theta \;\sim\; N(\theta, \sigma^2), \quad i=1,\dots,n, \quad \theta \sim N(m_0, v_0), \]

with \(\sigma^2 > 0\), \(v_0 > 0\) known. The posterior is exactly Normal,

\[ \theta \mid y_{1:n} \sim N(m_n, v_n), \]

where

\[ v_n^{-1} = v_0^{-1} + \frac{n}{\sigma^2}, \qquad m_n = v_n\Bigl( v_0^{-1} m_0 + \frac{n \bar{y}}{\sigma^2} \Bigr), \qquad \bar{y} = \frac{1}{n} \sum_{i=1}^n y_i. \]

This is a conjugate setting, so VI is not necessary for tractability; however, it provides a clean sanity check.

3.2 Mean-Field Approximation with Redundant Factorization

Introduce an artificial factorization by writing \(\theta = (\theta_1, \theta_2)\) with the constraint \(\theta_1 = \theta_2\) under the true model. Suppose we approximate the posterior with a mean-field family

\[ q(\theta_1, \theta_2) = q_1(\theta_1) q_2(\theta_2), \]

even though the true posterior does not factorize in this way (it is supported on the diagonal \(\theta_1 = \theta_2\)). Applying Theorem 9.2 yields coupled updates for \(q_1\) and \(q_2\) which, at the optimum, restore the equality \(\theta_1 = \theta_2\) in distribution, and the product \(q_1 q_2\) recovers the exact posterior density. This example illustrates that mean-field VI can be exact if the factorization is aligned with conditional independence structure implied by the model.

More realistic examples (e.g., multivariate Normal posteriors with non-diagonal covariance) yield strictly biased mean-field approximations, which we discuss next.

4. Comparison with MCMC: Bias, Variance, and Asymptotics

4.1 Bias–Variance Tradeoff

Let \(f(\theta)\) be a posterior integrable functional of interest, and denote

\[ \mu := \mathbb{E}_{p(\cdot \mid x)}[f(\theta)], \qquad \mu_{\text{VI}} := \mathbb{E}_{q^*}[f(\theta)]. \]

For an ergodic MCMC algorithm with invariant distribution \(p(\cdot \mid x)\) and samples \(\theta^{(1)},\dots,\theta^{(N)}\), the Monte Carlo estimator \[ \hat{\mu}_N = \frac{1}{N} \sum_{k=1}^N f(\theta^{(k)}) \] satisfies a CLT under conditions from Module 5, \[ \sqrt{N}(\hat{\mu}_N - \mu) \overset{d}{\to} N(0, \sigma_f^2). \] Thus, it is (asymptotically) unbiased but has variance \(\sigma_f^2 / N\).
For VI, the quantity \(\mu_{\text{VI}}\) is deterministic once \(q^*\) is obtained. There is no sampling variance, but in general \(\mu_{\text{VI}} \ne \mu\); the difference \(\mu_{\text{VI}} - \mu\) is an approximation bias induced by restricting \(q\) to \(\mathcal{Q}\).

In practice, VI trades off bias against computational speed and ease of implementation.

4.2 Failure Modes of VI

Typical issues include:

Underestimation of posterior variance. Mean-field approximations ignore posterior correlations, which often leads to overly concentrated approximations.
Mode-seeking behavior. Minimizing \(\operatorname{KL}(q \Vert p)\) penalizes placing mass where \(p\) is small but is relatively tolerant of missing modes entirely, so VI tends to focus on one mode of a multimodal posterior.
Sensitivity to initialization. CAVI may converge to local optima of the ELBO in nonconvex problems.

These phenomena should be understood and diagnosed when using VI in place of MCMC.

5. Problem Set 9 (Representative Problems)

ELBO Identity (Measure-Theoretic Proof). Prove Theorem Theorem 1 in full generality using probability measures \(P\) and \(Q\) and Radon–Nikodym derivatives, explicitly checking conditions under which termwise integration and changes of measure are justified.
CAVI Derivation with Functional Analysis. Re-derive Theorem Theorem 2 using functional-analytic language (e.g., viewing the space of densities as a convex subset of \(L^1\)), and show that the ELBO is strictly concave in each factor \(q_j\) when \(p(\theta,x)\) is strictly positive.
VI in a Conjugate Model. For the Normal–Normal model with known variance, work out a variational approximation in which \(q\) is restricted to a Normal family. Verify that the variational optimum coincides with the exact posterior, and compute the ELBO at the optimum.
Bias in a Correlated Gaussian Posterior. Consider a bivariate Normal posterior \(N_2(m, \Sigma)\) with nonzero off-diagonal entries. Approximate it by a mean-field product \(q(\theta_1,\theta_2) = q_1(\theta_1) q_2(\theta_2)\) with \(q_j\) univariate Normals, and explicitly compute \(q_1, q_2\) and the resulting approximation error in means and variances.
VI vs MCMC in High Dimensions. Consider a high-dimensional Bayesian logistic regression model. Discuss qualitatively and, where possible, quantitatively how VI and MCMC scale with dimension and sample size. In particular, relate the computational complexity of CAVI updates to that of one MCMC sweep (e.g., random-walk MH or HMC), and discuss the implications for bias and variance.