Variational Inference

Overview

Bayesian inference begins with the same identity in every textbook:

$p(\theta \mid x) \;=\; \frac{p(x \mid \theta)\,p(\theta)}{p(x)}, \qquad p(x) \;=\; \int p(x \mid \theta)\,p(\theta)\,d\theta.$

The identity is conceptually clean: pick a prior, write a likelihood, divide. The catch is that for almost any model worth fitting, the integral defining $p(x)$ has no closed form, and the right-hand side cannot even be evaluated without it.

The two scalable responses to this obstacle split the modern Bayesian computation literature in half. Markov chain Monte Carlo, the subject of Bayesian Computation and MCMC on formalstatistics, builds a Markov chain whose stationary distribution is the posterior and samples from it without ever computing $p(x)$ — asymptotically exact at the cost of long chains and post-hoc convergence diagnostics. Variational inference, the subject of this topic, takes a different tack: pick a tractable family $\mathcal{Q}$ of distributions over $\theta$, find the member $q^* \in \mathcal{Q}$ closest to the posterior under some divergence, and report $q^*$ as the approximate posterior. VI is fast and deterministic. It is also approximate in a way MCMC isn’t, and the approximation error has a structure worth understanding before any algebra is on the page.

This topic develops the machinery in six sections. §1 frames the geometry — VI is a projection. §2 derives the evidence lower bound (ELBO), the algebraic device that makes reverse-KL projection tractable. §3 develops mean-field VI and CAVI on a Bayesian Gaussian mixture. §4 lifts both restrictions of §3 — full-batch and conjugate — through stochastic VI, the score-function gradient, the reparameterization gradient, and ADVI. §5 attacks the structural error of mean-field with full-rank Gaussian VI and normalizing flows. §6 closes with connections to the rest of the curriculum and the limits a working Bayesian needs to know about.

1. The intractable posterior

1.1 Why the evidence is intractable

A concrete example fixes the obstacle. Bayesian logistic regression takes data $\mathcal{D} = \{(x_i, y_i)\}_{i=1}^n$ with covariates $x_i \in \mathbb{R}^d$ and binary responses $y_i \in \{0, 1\}$, a parameter vector $\theta \in \mathbb{R}^d$, and the likelihood

$p(y_i \mid x_i, \theta) \;=\; \sigma(x_i^\top \theta)^{y_i}\bigl(1 - \sigma(x_i^\top \theta)\bigr)^{1 - y_i}, \qquad \sigma(z) = \frac{1}{1 + e^{-z}},$

with a Gaussian prior $\theta \sim \mathcal{N}(0, \tau^2 I_d)$ on the regression coefficients. Up to the unknown evidence, the posterior is

$p(\theta \mid \mathcal{D}) \;\propto\; \exp\!\Bigl\{\,\sum_{i=1}^n y_i\,x_i^\top \theta - \log\bigl(1 + e^{x_i^\top \theta}\bigr)\,\Bigr\} \cdot \exp\!\Bigl\{-\tfrac{1}{2\tau^2}\|\theta\|^2\Bigr\},$

which is log-concave in $\theta$ but not Gaussian: the $\log(1 + e^{x_i^\top \theta})$ term is not quadratic in $\theta$, so the posterior is not in the same exponential family as the prior. Conjugacy fails. The marginal likelihood

$p(\mathcal{D}) \;=\; \int_{\mathbb{R}^d} p(\mathcal{D} \mid \theta)\,p(\theta)\,d\theta$

has no analytic expression — the integrand is a product of $n$ logistic terms times a Gaussian, and no antiderivative reduces the $d$-dimensional integral to closed form. Even evaluating $p(\theta \mid \mathcal{D})$ at a single $\theta$ requires that integral.

The pattern generalizes. Conjugacy — the property that lets us write $p(\theta \mid x)$ in closed form — is the exception, not the rule. Most models that lift a likelihood from a textbook example into a real application break conjugacy at some point: a logit link, a non-Gaussian noise model, a hierarchical structure with cross-level couplings, or a deep neural network whose weights play the role of $\theta$. For all such models, the integral defining $p(x)$ is the bottleneck.

1.2 Two responses: sampling versus optimization

MCMC sidesteps the intractable integral by exploiting a structural fact: ratios of the unnormalized posterior $\tilde p(\theta) = p(x \mid \theta)\,p(\theta)$ at two parameter values cancel the unknown $p(x)$ exactly, so a Markov chain whose transition probabilities depend only on those ratios has the right stationary distribution. The chain runs, eventually mixes, and the post-mixing samples are drawn from the posterior. The cost is wall-clock time: a chain on a large model can take hours or days, and convergence diagnostics are part of every analysis.

Variational inference replaces the sampling problem with an optimization problem. We choose a variational family $\mathcal{Q}$ — a parametric set of distributions over $\theta$ — and search for the member of $\mathcal{Q}$ closest to the true posterior under a divergence $D\bigl(q,\, p(\cdot \mid x)\bigr)$:

$q^* \;=\; \arg\min_{q \in \mathcal{Q}}\, D\bigl(q,\, p(\cdot \mid x)\bigr).$

The choice of $\mathcal{Q}$ controls expressiveness; the choice of $D$ controls what closest means. Once $q^*$ is in hand, posterior expectations $\mathbb{E}_{p(\cdot \mid x)}[f(\theta)]$ are approximated by $\mathbb{E}_{q^*}[f(\theta)]$, which by construction is something we can compute. Both choices are deferred — $\mathcal{Q}$ to §3 and §5, $D$ to §1.3 below — but the bird’s-eye view is already informative: VI trades the exact-but-slow of MCMC for the approximate-but-fast of optimization, and the quality of the trade depends entirely on whether $\mathcal{Q}$ is rich enough to host a good approximation to the true posterior.

1.3 The geometric picture

Reading the optimization above out loud makes the geometry visible. The space of probability distributions on $\Theta$ is some infinite-dimensional set; the true posterior is a point in that set; the family $\mathcal{Q}$ is a much smaller submanifold parametrized by some finite-dimensional vector $\phi$; and $q^*$ is the point on $\mathcal{Q}$ closest to the posterior under $D$. Variational inference is a projection. When $\mathcal{Q}$ is too restrictive to contain the posterior, which is almost always the case, the projection introduces a bias, and the structure of that bias depends on what $D$ punishes most.

Two natural choices of $D$ give qualitatively different projections.

Forward KL, $\mathrm{KL}\bigl(p(\cdot \mid x)\,\|\,q\bigr)$, takes the expectation under the true posterior:

$\mathrm{KL}\bigl(p(\cdot \mid x)\,\|\,q\bigr) \;=\; \mathbb{E}_{p(\cdot \mid x)}\!\left[\log \frac{p(\theta \mid x)}{q(\theta)}\right].$

Computing this expectation requires sampling from the posterior, which is exactly what VI is trying to avoid. So forward KL is a non-starter as a stand-alone VI objective — though it does appear in the literature, paired with sampling-based corrections (importance-weighted bounds, reweighted wake-sleep) that lie outside this topic.

Reverse KL, $\mathrm{KL}\bigl(q\,\|\,p(\cdot \mid x)\bigr)$, takes the expectation under $q$:

$\mathrm{KL}\bigl(q\,\|\,p(\cdot \mid x)\bigr) \;=\; \mathbb{E}_{q}\!\left[\log \frac{q(\theta)}{p(\theta \mid x)}\right].$

This is tractable in the sense that matters: $q$ is the distribution we chose, so we can sample from it (or compute expectations under it analytically when $q$ is in a simple family). The integrand involves $\log p(\theta \mid x) = \log p(x, \theta) - \log p(x)$, and the unknown constant $\log p(x)$ separates cleanly from the parts that depend on $q$ — which is the entire point and the subject of §2.

Reverse KL is the standard VI objective, and it comes with a famous shape: the mode-seeking behavior previewed in KL divergence. A reverse-KL minimizer punishes mass placed where the target has none — $q(\theta) > 0$ where $p(\theta \mid x) \approx 0$ blows up the integrand — but pays no penalty for failing to cover regions where the target has mass and $q$ does not. When the target is multimodal, and the family is unimodal, the minimizer locks onto one mode rather than spreading across all of them. This is a feature in some applications (sharp predictions, decisive point estimates) and a pathology in others (genuine multimodality, identifiability problems where the modes correspond to distinct scientific hypotheses).

The figure below illustrates the geometry on two synthetic 2D targets. Panel (a) shows a posterior with off-diagonal correlation; panel (b) overlays the reverse-KL projection onto the mean-field family of diagonal Gaussians — the projection captures the center but cannot represent orientation, and (we will see in §3 and again in §5) systematically under-estimates the marginal variance along correlated directions. Panel (c) shows the mode-seeking pathology: a bimodal target on the left, the reverse-KL projection onto a single Gaussian on the right — locked onto one mode, with a symmetric solution at the other.

The §2 derivation makes the unknown $\log p(x)$ disappear from the reverse-KL objective via an algebraic rearrangement called the evidence lower bound. With the ELBO in hand, §3 develops the workhorse mean-field algorithm, §4 extends VI to large-scale and non-conjugate problems, §5 lifts the mean-field assumption with structured families and normalizing flows, and §6 closes with connections and limits.

2. The ELBO

Reverse KL is tractable because $q$ is a distribution we chose; we can sample from it or compute expectations under it. But the integrand

$\mathrm{KL}\bigl(q\,\|\,p(\cdot \mid x)\bigr) \;=\; \mathbb{E}_q\!\left[\log q(\theta) - \log p(\theta \mid x)\right]$

still requires $\log p(\theta \mid x) = \log p(x, \theta) - \log p(x)$, and $\log p(x)$ is the unknown evidence we cannot compute. The way out is an algebraic rearrangement that isolates $\log p(x)$ as a constant — a constant from $q$‘s point of view, irrelevant to optimization. The rearranged objective is the evidence lower bound, or ELBO. Three properties make it the central object of variational inference: it is computable from the model and a sample from $q$ alone; it lower-bounds $\log p(x)$ uniformly in $q$; and maximizing it over $q \in \mathcal{Q}$ is equivalent to minimizing reverse KL to the posterior.

2.1 From Jensen to the ELBO

The derivation is a single application of Jensen’s inequality. Recall: a function $\varphi$ is concave if $\varphi(\mathbb{E}[Z]) \ge \mathbb{E}[\varphi(Z)]$ for every integrable random variable $Z$. The natural logarithm is concave on $(0, \infty)$, so for any positive random variable $Z$,

$\log \mathbb{E}[Z] \;\ge\; \mathbb{E}[\log Z].$

Now the importance-sampling identity: for any density $q$ on $\Theta$ that puts positive mass everywhere $p(x, \theta)$ does,

$p(x) \;=\; \int p(x, \theta)\,d\theta \;=\; \int q(\theta) \cdot \frac{p(x, \theta)}{q(\theta)}\,d\theta \;=\; \mathbb{E}_q\!\left[\frac{p(x, \theta)}{q(\theta)}\right].$

The marginal likelihood is the expectation under $q$ of the ratio between the joint and $q$. Taking logs and applying Jensen’s inequality to $\log$:

$\log p(x) \;=\; \log \mathbb{E}_q\!\left[\frac{p(x, \theta)}{q(\theta)}\right] \;\ge\; \mathbb{E}_q\!\left[\log \frac{p(x, \theta)}{q(\theta)}\right] \;=\; \mathbb{E}_q[\log p(x, \theta)] - \mathbb{E}_q[\log q(\theta)].$

The right-hand side is the evidence lower bound. We give it a name and a notation.

Two properties are immediate from the derivation:

1. $\mathrm{ELBO}(q) \le \log p(x)$ for every admissible $q$ — Jensen’s inequality.
2. $\mathrm{ELBO}(q)$ is computable whenever the model’s joint density $p(x, \theta)$ is computable and we can take expectations under $q$. The unknown $\log p(x)$ does not appear.

The first property explains the name: the ELBO is a lower bound on the evidence. The second explains why it matters: every quantity on the right-hand side is something we built and chose, so we can evaluate it.

2.2 The evidence decomposition

Jensen’s inequality tells us $\mathrm{ELBO}(q) \le \log p(x)$ but not by how much. The gap turns out to have an exact, illuminating form.

The decomposition is the operational identity of variational inference. Three consequences come immediately.

Maximizing the ELBO is equivalent to minimizing the reverse KL. Because $\log p(x)$ does not depend on $q$, it is a constant in the optimization over $q \in \mathcal{Q}$. So

$\arg\max_{q \in \mathcal{Q}} \mathrm{ELBO}(q) \;=\; \arg\min_{q \in \mathcal{Q}} \mathrm{KL}\bigl(q\,\|\,p(\cdot \mid x)\bigr).$

The intractable optimization (reverse KL involves $\log p(\theta \mid x)$, which we cannot evaluate) and the tractable optimization (the ELBO involves $\log p(x, \theta)$, which we can) have the same argmax.

The bound is tight when $\mathcal{Q}$ contains the posterior. Reverse KL is non-negative and equals zero iff $q = p(\cdot \mid x)$ almost everywhere. So if the family is rich enough that $p(\cdot \mid x) \in \mathcal{Q}$, the optimum $q^* \in \mathcal{Q}$ equals the posterior and $\mathrm{ELBO}(q^*) = \log p(x)$. Outside the conjugate cases this almost never happens; the §3 mean-field GMM and §4 Bayesian logistic regression both have $\mathrm{ELBO}(q^*) < \log p(x)$, and the gap measures the structural error of the family.

The ELBO is a free lower bound on the evidence. For model comparison or hyperparameter selection, $\log p(x)$ is exactly the quantity we want to compare across models — a higher evidence is a better fit. Computing $\log p(x)$ requires the intractable integral; the ELBO gives a lower bound that costs only a forward pass under $q$. This is the substrate of Variational Bayes for Model Selection.

2.3 Three equivalent forms

Definition 1 gives one expression for the ELBO. Two algebraic rearrangements produce equivalent expressions, each useful in a different context.

Form 1: Energy plus entropy. The entropy of $q$ is $H(q) = -\mathbb{E}_q[\log q(\theta)]$. Substituting into Definition 1:

$\mathrm{ELBO}(q) \;=\; \mathbb{E}_q[\log p(x, \theta)] + H(q).$

Reading: the ELBO rewards $q$-mass in regions of high joint log-density (the energy term, by analogy with the Boltzmann distribution’s $-\beta E$) and rewards $q$ for being spread out (the entropy term). Maximizing the ELBO trades off between concentrating on the mode of $p(x, \theta)$ and avoiding collapse to a delta. This form is canonical in the statistical-mechanics framing of VI as variational free-energy minimization — Feynman’s variational principle for the partition function, transcribed into Bayesian language.

Form 2: Reconstruction minus prior-KL. Decompose $\log p(x, \theta) = \log p(x \mid \theta) + \log p(\theta)$:

$\mathrm{ELBO}(q) \;=\; \mathbb{E}_q[\log p(x \mid \theta)] + \mathbb{E}_q[\log p(\theta)] - \mathbb{E}_q[\log q(\theta)] \;=\; \mathbb{E}_q[\log p(x \mid \theta)] - \mathrm{KL}\bigl(q(\theta)\,\|\,p(\theta)\bigr).$

Reading: the ELBO rewards $q$ for placing mass on $\theta$ values that explain the data (the reconstruction term, $\mathbb{E}_q[\log p(x \mid \theta)]$) but penalizes $q$ for moving away from the prior (the prior-KL term). This form is the workhorse of latent-variable models: in a variational autoencoder, the reconstruction term is the expected log-likelihood of the observed image given a latent code drawn from $q$, and the prior-KL term keeps the encoder distribution close to the standard-Gaussian prior. The two terms encode the model’s bias-variance tradeoff at the level of $q$.

Form 3: Evidence minus posterior-KL. This is Theorem 1 rearranged:

$\mathrm{ELBO}(q) \;=\; \log p(x) - \mathrm{KL}\bigl(q\,\|\,p(\cdot \mid x)\bigr).$

Reading: the ELBO equals the log-evidence (a constant) minus the reverse KL of $q$ to the posterior. This form makes the bound’s tightness transparent — the ELBO is closest to $\log p(x)$ exactly when $q$ is closest to the posterior under reverse KL — and is the cleanest form for theoretical arguments. It is not directly useful for optimization, since the right-hand side contains the unknown $\log p(x)$ and the unknown posterior.

The three forms are the same identity rearranged. Which one we reach for depends on what we want: Form 1 for thermodynamic intuition, Form 2 for implementation in latent-variable models, Form 3 for proofs.

2.4 Worked example: a conjugate Gaussian

A model with every quantity having a closed form lets us numerically verify the identity and visualize the ELBO landscape. We take the simplest such model.

Model. Observe $x_1, \dots, x_n \stackrel{\mathrm{iid}}{\sim} \mathcal{N}(\theta, \sigma^2)$ with $\sigma^2 > 0$ known. Place a Gaussian prior $\theta \sim \mathcal{N}(0, \tau^2)$ on the unknown mean. The joint density is

$p(x, \theta) \;=\; \prod_{i=1}^n \mathcal{N}(x_i \mid \theta, \sigma^2) \cdot \mathcal{N}(\theta \mid 0, \tau^2).$

Exact posterior. Completing the square in $\theta$ inside $\log p(x, \theta)$ gives the posterior in closed form:

$p(\theta \mid x) \;=\; \mathcal{N}\bigl(\theta \mid \mu_{\mathrm{post}},\, \sigma_{\mathrm{post}}^2\bigr), \qquad \sigma_{\mathrm{post}}^2 = \left(\frac{1}{\tau^2} + \frac{n}{\sigma^2}\right)^{-1}, \qquad \mu_{\mathrm{post}} = \sigma_{\mathrm{post}}^2 \cdot \frac{n \bar x}{\sigma^2},$

where $\bar x = n^{-1} \sum_i x_i$ is the sample mean. The posterior precision is the prior precision plus $n$ times the data precision, and the posterior mean is the posterior precision times the data’s natural-parameter contribution. This is the standard Gaussian–Gaussian conjugate result.

Variational family. Take $\mathcal{Q} = \{\mathcal{N}(\theta \mid \mu, s^2) : \mu \in \mathbb{R},\, s^2 > 0\}$, the family of all 1D Gaussians. Crucially, the true posterior is in this family — the conjugate setup is the rare case where $\mathcal{Q}$ contains $p(\cdot \mid x)$.

Closed-form ELBO. With $q(\theta) = \mathcal{N}(\theta \mid \mu, s^2)$, every expectation in Definition 1 has a closed form. We use the standard identity $\mathbb{E}_q[(a - \theta)^2] = (a - \mu)^2 + s^2$ for any constant $a$, plus $H(q) = \tfrac{1}{2} \log(2\pi e s^2)$:

$\mathrm{ELBO}(\mu, s^2) = -\frac{n}{2} \log(2\pi \sigma^2) - \frac{1}{2\sigma^2}\!\left[\sum_{i=1}^n (x_i - \mu)^2 + n s^2\right] - \frac{1}{2} \log(2\pi \tau^2) - \frac{\mu^2 + s^2}{2\tau^2} + \frac{1}{2}\log(2\pi e s^2).$

Setting partial derivatives to zero recovers the exact posterior:

$\frac{\partial \mathrm{ELBO}}{\partial \mu} = \frac{1}{\sigma^2}\sum_i(x_i - \mu) - \frac{\mu}{\tau^2} = 0 \;\implies\; \mu^* = \mu_{\mathrm{post}},$

$\frac{\partial \mathrm{ELBO}}{\partial s^2} = -\frac{n}{2\sigma^2} - \frac{1}{2\tau^2} + \frac{1}{2 s^2} = 0 \;\implies\; (s^*)^2 = \sigma_{\mathrm{post}}^2.$

The variational optimum exactly equals the true posterior, as Theorem 1 predicted: when $\mathcal{Q}$ contains $p(\cdot \mid x)$, the ELBO maximum saturates at $\log p(x)$ and $\mathrm{KL}(q^* \,\|\, p(\cdot \mid x)) = 0$.

The figure below makes this concrete on a simulated dataset ($n = 20$, $\sigma^2 = 1$, $\tau^2 = 100$, true $\theta = 1.5$): panel (a) plots the ELBO surface over $(\mu_q, \log s)$ with the optimum marked, panel (b) overlays three candidate $q$‘s on the exact posterior (the optimum, an off-center version, and a too-narrow version), and panel (c) shows the evidence decomposition numerically — for each candidate, $\mathrm{ELBO}(q)$ sits below $\log p(x)$ by exactly $\mathrm{KL}(q \,\|\, p(\cdot \mid x))$, with the gap collapsing to zero at $q^*$.

The conjugate case is not VI’s typical regime — most real models break conjugacy somewhere, and the §3 mean-field GMM is the first example where $\mathrm{ELBO}(q^*) < \log p(x)$ at the optimum. But the identity from Theorem 1 is universal: every ELBO maximization is a reverse-KL projection, and the gap to $\log p(x)$ is always the KL of $q^*$ to the true posterior, whether that gap is zero or not.

3. Mean-field VI and CAVI

The §2 conjugate Gaussian example is misleading, because conjugacy is rare. For most Bayesian models, the posterior is not in the same family as the prior, the ELBO has no closed-form maximum even on simple parametric families, and Definition 1’s sole virtue is that it is an expression we can compute — not yet an algorithm we can execute. The mean-field assumption is the algorithmic move that turns the ELBO into something we can optimize coordinate-by-coordinate, and coordinate-ascent variational inference (CAVI) is the algorithm that results. CAVI is the workhorse of variational inference in conjugate exponential-family models — Bayesian Gaussian mixtures, latent Dirichlet allocation, factorial hidden Markov models, mixed-effects models, sparse linear regression with conjugate priors. The trade-off it makes is structural: mean-field cannot represent posterior correlations between variables in different blocks. §5 returns to that limit.

3.1 The mean-field family

Partition the latent variables into $m$ blocks: $\theta = (\theta_1, \ldots, \theta_m)$ with $\theta_j \in \Theta_j$. The blocks are arbitrary and chosen by the modeler — usually one block per parameter (fully factorized) or one block per logical group (e.g., all component means together, all assignments together).

The factorization is across blocks; within a block, $q_j$ is unconstrained — any density on $\Theta_j$ is admissible. So the family is simultaneously flexible (each $q_j$ can be any shape) and restrictive (the joint $q$ assumes block-independence).

What we gain is computational tractability. The ELBO under mean-field separates:

$\mathrm{ELBO}(q) \;=\; \mathbb{E}_q[\log p(x, \theta)] \;-\; \sum_{j=1}^m \mathbb{E}_{q_j}[\log q_j(\theta_j)],$

where the entropy term is a simple sum over blocks because $\log q(\theta) = \sum_j \log q_j(\theta_j)$ under the factorization. The first term doesn’t separate — it depends jointly on all $q_j$ — but conditional on fixing every factor except one, it does, which is the door §3.2 walks through.

What we lose is correlation. If the true posterior couples $\theta_j$ and $\theta_k$ in different blocks (which it generically does — even simple regression posteriors couple the coefficient with the intercept), no $q \in \mathcal{Q}_{\mathrm{MF}}$ can represent that coupling. The §1 panel (b) figure already showed this on a 2D Gaussian: the diagonal-Gaussian projection captures the center but cannot tilt to match the posterior’s correlated shape. §5 quantifies this loss as systematic under-estimation of marginal posterior variance along correlated directions.

3.2 The CAVI update

The optimization strategy for mean-field VI is coordinate ascent: cycle through the blocks $j = 1, \ldots, m$, holding all factors except $q_j$ fixed and updating $q_j$ to its conditional optimum. The update has a clean closed form.

The reading: the optimal $q_j$ given the others is proportional to the exponentiated expected complete log-joint, with the expectation taken over the other factors. This is also the form of the complete conditional $p(\theta_j \mid \theta_{-j}, x)$ as a function of $\theta_j$ — except with $\theta_{-j}$ replaced by an expectation under $q_{-j}$. CAVI replaces conditioning on a value with averaging over a distribution.

Algorithm (CAVI). Initialize $q_1, \ldots, q_m$. Repeat until convergence:

• For $j = 1, \ldots, m$: set $q_j(\theta_j) \leftarrow Z_j^{-1} \exp \mathbb{E}_{q_{-j}}[\log p(x, \theta)]$.

Convergence is monitored by tracking the ELBO; halt when the ELBO change between sweeps falls below a tolerance.

The monotonicity of the ELBO under CAVI follows immediately from Theorem 2: each block update either increases the ELBO (if $q_j$ wasn’t already optimal) or leaves it unchanged, so the ELBO is non-decreasing across iterations. Like EM, CAVI is guaranteed to converge to a local maximum of the ELBO; the global maximum is generally unattainable because the ELBO is non-concave in $q$ for non-trivial models — particularly for mixtures, where the symmetry across component labels generates exponentially many local maxima.

3.3 Conjugate exponential families

CAVI’s update has a closed form whenever the model is conditionally conjugate — meaning each complete conditional $p(\theta_j \mid \theta_{-j}, x)$ is in an exponential family. The derivation is short, and the result is operationally important.

Suppose

$p(\theta_j \mid \theta_{-j}, x) \;=\; h_j(\theta_j) \exp\!\Bigl\{\, \eta_j(\theta_{-j}, x)^\top T_j(\theta_j) \;-\; A_j\bigl(\eta_j(\theta_{-j}, x)\bigr) \,\Bigr\},$

with base measure $h_j$, sufficient statistics $T_j$, natural parameter $\eta_j(\theta_{-j}, x)$, and log-partition $A_j$. The exponential-family form means the dependence on $\theta_j$ enters only through $T_j(\theta_j)$ and $h_j(\theta_j)$, while the dependence on the rest enters only through $\eta_j$ and $A_j$.

Writing $\log p(x, \theta) = \log p(\theta_j \mid \theta_{-j}, x) + \log p(x, \theta_{-j})$ and noting that the second term is constant in $\theta_j$:

$\mathbb{E}_{q_{-j}}[\log p(x, \theta)] \;=\; \log h_j(\theta_j) \;+\; \mathbb{E}_{q_{-j}}[\eta_j(\theta_{-j}, x)]^\top T_j(\theta_j) \;+\; \mathrm{const}.$

Exponentiating gives the CAVI update:

$q_j^*(\theta_j) \;\propto\; h_j(\theta_j) \exp\!\Bigl\{\, \mathbb{E}_{q_{-j}}[\eta_j(\theta_{-j}, x)]^\top T_j(\theta_j) \,\Bigr\}.$

This is in the same exponential family as $p(\theta_j \mid \theta_{-j}, x)$, with natural parameter $\eta_j^* = \mathbb{E}_{q_{-j}}[\eta_j(\theta_{-j}, x)]$.

The result is “expectation inside the natural parameter”: take the natural-parameter expression for the complete conditional, and replace each occurrence of $\theta_{-j}$ with the appropriate expectation under $q_{-j}$. The class of conditionally conjugate models — where every complete conditional is an exponential family with a tractable expectation — is the natural domain of closed-form CAVI: Bayesian GMMs, LDA, hidden Markov models, factorial HMMs, and the broader family of conjugate hierarchical models that defines structured Bayesian-machine-learning workflows.

3.4 Worked example: Bayesian Gaussian mixture model

The canonical worked example is the Bayesian Gaussian mixture. We take the simplest version that exhibits all the structure: known component variance, a Gaussian prior on component means, and uniform mixing weights.

Model. Observe $x_1, \ldots, x_n \in \mathbb{R}^d$. Place a $K$-component mixture with:

• Component means $\mu_k \stackrel{\mathrm{iid}}{\sim} \mathcal{N}(0,\, \tau^2 I_d)$ for $k = 1, \ldots, K$;
• Cluster assignments $z_i \stackrel{\mathrm{iid}}{\sim} \mathrm{Cat}(1/K, \ldots, 1/K)$ (uniform mixing weights);
• Observations $x_i \mid z_i = k, \mu \sim \mathcal{N}(\mu_k,\, \sigma^2 I_d)$ with $\sigma^2$ known.

The latent variables are $\theta = (\mu, z)$ with $\mu = (\mu_1, \ldots, \mu_K)$ and $z = (z_1, \ldots, z_n)$.

Mean-field family. Take a two-block factorization: one block for the means, one for the assignments.

$q(\mu, z) \;=\; q(\mu) \, q(z), \qquad q(\mu) = \prod_{k=1}^K q(\mu_k), \qquad q(z) = \prod_{i=1}^n q(z_i).$

The within-block factorizations across $k$ and $i$ are part of the mean-field choice rather than implied by the generative model — in general, within-block posterior dependencies can remain — but for this conjugate GMM the §3.3 expectation-inside-the-natural-parameter argument makes them optimal at the CAVI fixed point: given the other block, the complete conditionals factor cleanly across $i$ for $z$ and across $k$ for $\mu$, so writing the factorizations down loses nothing here.

The CAVI updates. Both updates follow from the conjugate exponential-family setup of §3.3.

Update for $q(z_i)$. The complete conditional $p(z_i \mid x, \mu, z_{-i})$ depends only on $x_i$ and $\mu$. It’s categorical with $p(z_i = k \mid x_i, \mu) \propto K^{-1} \mathcal{N}(x_i \mid \mu_k, \sigma^2 I)$. The natural parameter has term $-\tfrac{1}{2\sigma^2}\|x_i - \mu_k\|^2$. Taking the expectation under $q(\mu_k) = \mathcal{N}(m_k, s_k^2 I)$ and using $\mathbb{E}_{q(\mu_k)}[\|x_i - \mu_k\|^2] = \|x_i - m_k\|^2 + d \, s_k^2$ gives the responsibility:

$r_{ik} \;:=\; q(z_i = k) \;\propto\; \exp\!\Bigl\{ -\frac{1}{2\sigma^2} \bigl( \|x_i - m_k\|^2 + d \, s_k^2 \bigr) \Bigr\},$

normalized so $\sum_k r_{ik} = 1$. The variance-correction term $d \, s_k^2$ softens the responsibilities relative to the EM hard-assignment limit: posterior uncertainty about $\mu_k$ pulls the soft assignments toward uniformity.

Update for $q(\mu_k)$. The complete conditional $p(\mu_k \mid x, z, \mu_{-k})$ is a Gaussian — conjugate combination of the prior and the points hard-assigned to component $k$. Replacing the hard assignments with expected responsibilities under $q(z)$ gives the §2 conjugate Gaussian–Gaussian update applied to soft-assigned sufficient statistics:

$q(\mu_k) \;=\; \mathcal{N}(m_k,\, s_k^2 I_d), \qquad s_k^2 \;=\; \left( \frac{N_k}{\sigma^2} + \frac{1}{\tau^2} \right)^{-1}, \qquad m_k \;=\; \frac{s_k^2}{\sigma^2} \sum_{i=1}^n r_{ik} \, x_i,$

where $N_k := \sum_i r_{ik}$ is the effective count of points assigned to component $k$. Compare with §2: the only change is that the data is weighted by responsibilities rather than counted indicator-by-indicator. The connection to EM is now visible — CAVI’s $r$-update is the soft E-step, and CAVI’s $\mu$-update is a Bayesian M-step that returns the full posterior $\mathcal{N}(m_k, s_k^2 I)$ rather than just the MAP point.

The closed-form ELBO. Every term of $\mathrm{ELBO}(q) = \mathbb{E}_q[\log p(x, z, \mu)] - \mathbb{E}_q[\log q(z, \mu)]$ has a closed form. Working through $\log p(x, z, \mu) = \log p(x \mid z, \mu) + \log p(z) + \log p(\mu)$ and the entropies:

$\mathrm{ELBO} \;=\; \underbrace{-\sum_{i,k} \frac{r_{ik}}{2\sigma^2}\bigl( \|x_i - m_k\|^2 + d \, s_k^2 \bigr)}_{\text{soft reconstruction}} \;-\; \underbrace{\sum_k \frac{\|m_k\|^2 + d \, s_k^2}{2\tau^2}}_{\text{prior penalty}} \;+\; \underbrace{H(q(z))}_{\text{assignment entropy}} \;+\; \underbrace{H(q(\mu))}_{\text{mean entropy}} \;+\; \mathrm{const},$

where $H(q(z)) = -\sum_{i,k} r_{ik} \log r_{ik}$, $H(q(\mu)) = \sum_k \tfrac{d}{2}\log(2\pi e s_k^2)$, and the constant absorbs the data-independent normalizers. Tracking this expression across CAVI sweeps is the §3.4 figure.

The figure below shows CAVI applied to a synthetic 2D 2-cluster dataset designed to mirror the structure of Old Faithful eruptions (two well-separated clusters with moderate within-cluster spread). Panel (a) shows the data with a random initialization of component means; panel (b) shows the converged variational posterior, with each $m_k$ at its cluster center and the variational covariance $s_k^2 I$ as a 2-sigma ellipse; panel (c) shows the ELBO across iterations.

The convergence pattern is characteristic of CAVI on mixtures: the ELBO climbs rapidly in the first few sweeps as the responsibilities sort the points, then plateaus as the component means refine. Mixture posteriors are multimodal — any permutation of the component labels gives an equivalent posterior, and CAVI converges to one of those modes. §5 returns to the consequence that under mean-field, the component-mean variances $s_k^2$ are systematically smaller than under the full Bayesian posterior, because the mean-field assumption ignores the posterior coupling between $\mu_k$ and $z$.

4. Stochastic and black-box VI

The CAVI setup of §3 — closed-form coordinate updates derived from conjugate-exponential-family structure — is the right tool for a specific class of models, and a wrong tool everywhere else. Two stresses that modern Bayesian workflows routinely encounter break it. The first is scale: when the dataset is large enough that even a single full-batch pass is expensive, the linear-in-$n$ cost of every CAVI sweep makes the method impractical. The second is non-conjugacy: when the likelihood and prior do not align in a way that produces closed-form complete conditionals, §3.3’s “expectation inside the natural parameter” trick has nothing to operate on, and the CAVI update $q_j^* \propto \exp \mathbb{E}_{q_{-j}}[\log p(x, \theta)]$ has no analytic form.

This section assembles the two responses. Stochastic variational inference (Hoffman et al. 2013) addresses scale by computing CAVI updates from minibatches under a Robbins–Monro step-size schedule. Black-box variational inference (Ranganath, Gerrish, and Blei 2014) and the reparameterization gradient (Kingma and Welling 2014; Rezende, Mohamed, and Wierstra 2014) address non-conjugacy by recasting ELBO maximization as stochastic gradient ascent on a parametric variational family — and the reparameterization gradient is now the dominant computational device of modern VI, the substrate of every variational autoencoder, Bayesian neural network, and deep generative model. The §4.5 worked example applies the reparameterization gradient to Bayesian logistic regression on the Iris dataset using JAX.

4.1 Where CAVI breaks

Big data. A CAVI sweep on the §3 worked example — $n = 270$, $K = 2$, $d = 2$ — costs roughly $nKd \approx 10^3$ floating-point operations per iteration. For the corpus-scale topic models that motivated SVI ($n = 10^7$ documents, $K = 100$ topics, $d = 10^5$ vocabulary terms), the same $nKd$ scaling produces $10^{14}$ operations per sweep, and convergence can require hundreds of sweeps. The cost is not in any single update — each update is closed-form — but in the linear-in-$n$ scaling of every CAVI iteration. Wall-clock time runs to hours or days before the first ELBO estimate stabilizes. The standard remedy in optimization is to replace full-batch gradients with unbiased minibatch estimates and to rely on stochastic approximation theory for convergence. The §4.2 SVI machinery transfers that idea to CAVI’s natural-gradient form.

Non-conjugacy. Bayesian logistic regression is the simplest example. With Bernoulli likelihood $y_i \mid x_i, \theta \sim \mathrm{Ber}(\sigma(x_i^\top \theta))$ and Gaussian prior $\theta \sim \mathcal{N}(0, \tau^2 I)$, the complete log-joint is

$\log p(\mathcal{D}, \theta) \;=\; \sum_{i=1}^n \bigl[y_i\,x_i^\top \theta - \log(1 + e^{x_i^\top \theta})\bigr] - \frac{\|\theta\|^2}{2\tau^2} + \mathrm{const},$

which is not in the same exponential family as the prior — the $\log(1 + e^{x_i^\top \theta})$ term is not quadratic in $\theta$. Mean-field CAVI on a Gaussian variational family $q(\theta) = \mathcal{N}(\mu, \mathrm{diag}(\sigma^2))$ would require $\mathbb{E}_q[\log(1 + e^{x_i^\top \theta})]$ in closed form, and that expectation has no analytic expression. Model-specific workarounds exist — local-variational bounds (Jaakkola and Jordan 2000), Laplace-style expansions around the variational mean — but each non-conjugate likelihood needs its own analysis, and no general recipe scales across the modeling landscape. The §§4.3–4.4 black-box machinery is general: any model where we can evaluate $\log p(x, \theta)$ at a sampled $\theta$ admits a stochastic gradient estimator for the ELBO.

4.2 Stochastic VI for big data

Hoffman, Blei, Wang, and Paisley (2013) extended CAVI to the minibatch regime through two observations: that CAVI’s coordinate update is a natural gradient step on the ELBO, and that natural gradients in conditionally conjugate models admit unbiased minibatch estimators.

The natural gradient takes seriously that $q(\theta; \eta)$ lives on a manifold of probability distributions rather than in the Euclidean space of natural parameters. The Euclidean gradient $\nabla_\eta \mathrm{ELBO}$ asks how the ELBO changes as $\eta$ moves a unit step in coordinate space; the natural gradient $\tilde\nabla_\eta \mathrm{ELBO} := F(\eta)^{-1} \nabla_\eta \mathrm{ELBO}$, where $F(\eta) = \mathbb{E}_q[\nabla_\eta \log q \cdot \nabla_\eta \log q^\top]$ is the Fisher information of $q(\theta; \eta)$, asks how the ELBO changes as $q(\theta; \eta)$ moves a unit step in KL distance. The natural gradient is the right object whenever the parameter space’s Euclidean geometry is misleading — which it always is for probability distributions.

The result is operationally direct: the natural gradient points from the current natural parameter to the CAVI target. A natural-gradient ascent step with rate $\rho$ is a damped CAVI update,

$\eta_j^{(t+1)} \;=\; (1 - \rho)\,\eta_j^{(t)} \;+\; \rho \cdot \mathbb{E}_{q_{-j}}\!\bigl[\eta_j(\theta_{-j}, x)\bigr],$

and CAVI itself is the special case $\rho = 1$ — a one-shot natural-gradient step that overshoots a damped trajectory but converges at every iteration.

The minibatching step is now mechanical. The CAVI target $\mathbb{E}_{q_{-j}}[\eta_j(\theta_{-j}, x)]$ in conjugate exponential-family models has the form $\eta_j^{(0)} + \sum_{i=1}^n s_i$ where $\eta_j^{(0)}$ is the prior natural parameter and $s_i$ is a per-datapoint sufficient-statistic contribution. Replacing the full sum with a rescaled minibatch sum,

$\widehat{\mathbb{E}_{q_{-j}}[\eta_j]} \;=\; \eta_j^{(0)} \;+\; \frac{n}{|B|}\sum_{i \in B} s_i,$

gives an unbiased estimator of the CAVI target. Stochastic natural-gradient ascent with this estimator and a Robbins–Monro learning-rate schedule ($\sum_t \rho_t = \infty$, $\sum_t \rho_t^2 < \infty$, typically $\rho_t = (t + \tau)^{-\kappa}$ for $\kappa \in (0.5, 1]$) converges to a local maximum of the ELBO. SVI scales the workhorse CAVI machinery from research-scale to corpus-scale models — LDA on Wikipedia-sized corpora, hierarchical Dirichlet processes on millions of documents, factorial hidden Markov models on long time series — and remains the right tool for any conjugate exponential-family model that doesn’t fit in memory. It is not the right tool for non-conjugate models, where the natural gradient itself is not closed-form.

4.3 Black-box VI: the score-function gradient

For non-conjugate models, we need a gradient estimator that requires only the ability to (i) sample from $q_\phi$ and (ii) evaluate $\log p(x, \theta)$ at a sampled $\theta$. Ranganath, Gerrish, and Blei (2014) showed that the log-derivative trick — also known as the REINFORCE estimator in reinforcement learning (Williams 1992) and the likelihood-ratio estimator in stochastic optimization — provides exactly this.

Let $q_\phi$ be a parametric variational family with parameters $\phi \in \mathbb{R}^P$, and write $\mathcal{L}(\phi) := \mathrm{ELBO}(q_\phi) = \mathbb{E}_{q_\phi}[\log p(x, \theta) - \log q_\phi(\theta)]$.

The estimator is universal in a clean sense: it requires only that we can sample from $q_\phi$ and evaluate both $\log p(x, \theta)$ and $\nabla_\phi \log q_\phi(\theta)$. The model can be arbitrarily non-conjugate, hierarchical, deep — the score-function gradient still applies. Replacing the expectation with $S$ Monte Carlo samples gives the unbiased estimator

$\widehat{\nabla_\phi \mathcal{L}}_{\mathrm{SF}}(\phi) \;=\; \frac{1}{S}\sum_{s=1}^S \bigl(\log p(x, \theta^{(s)}) - \log q_\phi(\theta^{(s)})\bigr)\,\nabla_\phi \log q_\phi(\theta^{(s)}), \qquad \theta^{(s)} \stackrel{\mathrm{iid}}{\sim} q_\phi.$