8 Variational Inference

Every probabilistic model with latent variables poses the same fundamental question: given observed data $𝒙$, what can we say about the hidden variables $𝒛$? Answering this question means computing the posterior distribution $p(𝒛|𝒙)$. In simple models such as the Gaussian mixture model of Chapter ch:gmm, the posterior factorises into a short list of responsibilities $p(z_i=k|{𝒙}_i;\theta )$, each available in closed form. But in the deep latent-variable models that drive modern generative AI, including variational autoencoders (Chapter ch:vae), hierarchical VAEs, and diffusion models (Chapter ch:diffusion), the posterior is a high-dimensional, intricately correlated distribution that no closed-form expression can capture.

Variational inference (VI) resolves this impasse by reframing inference as optimisation: instead of computing $p(𝒛|𝒙)$ exactly, we search over a tractable family $𝒬$ for the member $q^{\ast }(𝒛)$ closest to the true posterior. The “closeness” is measured by Kullback–Leibler divergence, and the optimisation target is the evidence lower bound (ELBO), the same quantity we met briefly in the EM algorithm (Section sec:gmm:elbo).

This chapter develops variational inference from first principles. We begin with motivation (Section Introduction and Motivation), revisit the asymmetry of KL divergence (Section KL Divergence Revisited), derive the ELBO in depth (Section The Evidence Lower Bound), build the classical mean-field framework (Section Mean-Field Variational Inference), and present the coordinate ascent algorithm (Section Coordinate Ascent Variational Inference). File B of this chapter continues with stochastic variational inference, amortised inference, and the reparameterisation trick that powers VAEs.

Introduction and Motivation

The Central Problem: Posterior Inference

Consider a latent-variable model specified by a prior $p(𝒛)$ and a likelihood $p(𝒙|𝒛)$. The joint density is $p(𝒙,𝒛)=p(𝒛)p(𝒙|𝒛)$, and Bayes' theorem gives the posterior (Bayes)$$\begin{array}{ccc}& p(𝒛|𝒙)=\frac{p(𝒙,𝒛)}{p(𝒙)}=\frac{p(𝒛)p(𝒙|𝒛)}{{\displaystyle \int p({𝒛}^{\prime })p(𝒙|{𝒛}^{\prime })\mathrm{d}{𝒛}^{\prime }}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The denominator $p(𝒙)=\int p(𝒙,𝒛)\mathrm{d}𝒛$ is the marginal likelihood (or evidence). Computing it requires integrating the joint density over the entire latent space, a sum over $K^N$ configurations for a GMM with $N$ data points, or an integral over ${\mathrm{ℝ}}^m$ for a continuous latent space.

The difficulty is not merely practical but fundamental:

This hardness result means that no algorithm can solve arbitrary inference problems efficiently, and we must accept approximations. Two great families of approximate inference have emerged:

1. Markov chain Monte Carlo (MCMC): construct a Markov chain whose stationary distribution is $p(𝒛|𝒙)$ and collect samples. MCMC is asymptotically exact (given infinite time, the samples converge to the true posterior), but convergence diagnostics are difficult, and the method can be prohibitively slow for large datasets.

2. Variational inference (VI): choose a tractable family $𝒬$ and find the member $q^{\ast }\in 𝒬$ closest to the posterior. VI is biased (the optimum $q^{\ast }$ generally differs from the true posterior) but fast: it inherits the scalability of optimisation, including stochastic gradient methods.

The Inference Landscape

Figure fig:vi:landscape provides a bird's-eye view of inference methods, organised by their accuracy–speed trade-off.

The Core Idea: Inference as Optimisation

Road Map

The remainder of this chapter is organised as follows.

1. Section KL Divergence Revisited: KL Divergence Revisited. We examine the crucial asymmetry between forward and reverse KL, which determines why VI uses the reverse direction.

2. Section The Evidence Lower Bound: The Evidence Lower Bound. We derive the ELBO from multiple perspectives and explore its decompositions.

3. Section Mean-Field Variational Inference: Mean-Field VI. We introduce the fully factorised approximation and derive the optimal factors via functional optimisation.

4. Section Coordinate Ascent Variational Inference: Coordinate Ascent VI. We present the classical CAVI algorithm and apply it to Bayesian Gaussian mixture models.

5. File B continues with stochastic VI, amortised inference, and connections to VAEs (Chapter ch:vae) and diffusion models (Chapter ch:diffusion).

Why Posterior Inference Matters

Before diving into the technical machinery, let us emphasise why computing, or approximating, the posterior $p(𝒛|𝒙)$ is so important. The posterior answers three distinct questions:

1. What do the latent variables look like? In a mixture model, the posterior $p(z_i|{𝒙}_i)$ tells us which cluster generated each observation. In a topic model, it tells us which topics are present in each document. In a VAE, it tells us the latent “code” that summarises each data point.

2. How uncertain are we? Bayesian posteriors quantify uncertainty. The posterior $p(𝒘|𝒟)$ of a neural network's weights tells us not just the best-guess parameters but how confident we should be. This is crucial for safety-critical applications (medical diagnosis, autonomous driving) and for active learning (where we query the most uncertain inputs).

3. How do we learn model parameters? In the EM framework (Chapter ch:gmm), the E-step requires the posterior $p(𝒛|𝒙;\theta )$. In VAEs, the approximate posterior $q_{\varphi }(𝒛|𝒙)$ appears in the ELBO gradient. Posterior inference is not just a downstream task; it is integral to training.

KL Divergence Revisited

The KL divergence was introduced in Chapter ch:divergences (Definition def:kl) and appeared naturally in the ELBO decomposition for GMMs (Proposition prop:ll:decompose). In this section we examine a subtlety that is central to variational inference: the two directions of KL divergence behave very differently, and choosing the “wrong” direction would make VI impractical.

Recap: KL Divergence

Recall that for two continuous distributions with densities $p$ and $q$ on ${\mathrm{ℝ}}^d$, the KL divergence from $q$ to $p$ is (DEF)$$\begin{array}{ccc}& {𝖣}_{\mathrm{KL}}(q\Vert p)=\int q(𝒛)\log \frac{q(𝒛)}{p(𝒛)}\mathrm{d}𝒛={𝔼}_q\left[\log \frac{q(𝒛)}{p(𝒛)}\right].& \text{<span class="tml-eqn"></span>}\end{array}$$ By Gibbs' inequality (Theorem thm:gibbs), ${𝖣}_{\mathrm{KL}}(q\Vert p)\ge 0$ with equality if and only if $q=p$ almost everywhere. Crucially, KL divergence is asymmetric: ${𝖣}_{\mathrm{KL}}(q\Vert p)\ne {𝖣}_{\mathrm{KL}}(p\Vert q)$ in general.

Forward vs. Reverse KL

Let $p(𝒛)$ denote the target distribution (in our setting, $p(𝒛|𝒙)$) and $q(𝒛)$ the approximation. The two directions of KL divergence optimise very different objectives.

The terms “inclusive” and “exclusive” hint at the qualitative difference, which we now make precise.

Visualising the Asymmetry

The following figure is perhaps the single most important illustration for understanding variational inference.

Why Reverse KL for Variational Inference

Given the mode-seeking behaviour of the reverse KL, one might ask: isn't the forward KL a better choice, since it covers all modes? The answer is no, for a pragmatic but decisive reason.

Worked Example: Gaussian Approximation to a Mixture

To build intuition, let us compute both KL divergences explicitly for a toy problem.

Properties of KL Divergence

For completeness, we collect the key properties that will be used throughout.

The Evidence Lower Bound

The evidence lower bound (ELBO) is the central object of variational inference. It is simultaneously:

• a lower bound on the log-evidence $\log p(𝒙)$,

• the objective function that VI maximises,

• a bridge between the intractable posterior and the tractable approximation.

We derive it from two complementary perspectives and then explore its multiple decomposition forms, each of which highlights a different aspect of the inference–learning trade-off.

Derivation Route 1: Jensen's Inequality

Derivation Route 2: KL Decomposition

The second derivation is more revealing: it shows exactly where the gap comes from and why maximising the ELBO is equivalent to minimising the reverse KL divergence.

This is the same decomposition we encountered for GMMs in Proposition prop:ll:decompose (equation ), now stated in full generality. The identity has three immediate consequences:

Three Decompositions of the ELBO

The ELBO can be written in several equivalent forms, each emphasising a different interpretation.

Each form provides a different lens:

1. Energy + Entropy (Form 1): This is the classical variational free energy formulation from statistical physics. Maximising $ℒ$ balances fitting the model (low energy = high ${𝔼}_q[\log p(𝒙,𝒛)]$) against spreading out (high entropy $𝖧(q)$).

2. Reconstruction $-$ Regularisation (Form 2): This is the form most commonly used in VAEs (Section sec:vae:elbo). The first term rewards $q$ for concentrating on latent configurations that explain the data well; the KL term penalises $q$ for deviating from the prior, acting as a regulariser.

3. Evidence $-$ Gap (Form 3): This shows the ELBO as the log-evidence minus an approximation error. It is the most conceptually transparent form but the least computationally useful, since $\log p(𝒙)$ is unknown.

Worked Example: ELBO for a Gaussian Model

To solidify understanding, we compute the ELBO in closed form for a simple conjugate model.

Tightness of the ELBO

The ELBO as a Training Objective

In practice, the ELBO is used not only for inference (optimising $q$) but also for learning (optimising model parameters $\theta$). When the model has parameters $\theta$ (e.g. a neural network decoder), we maximise the ELBO jointly over $q$ and $\theta$: (Joint)$$\begin{array}{ccc}& \underset{\theta ,q\in 𝒬}{\max }ℒ(\theta ,q)=\underset{\theta ,q\in 𝒬}{\max }{𝔼}_{q(𝒛)}\left[\log \frac{p_{\theta }(𝒙,𝒛)}{q(𝒛)}\right].& \text{<span class="tml-eqn"></span>}\end{array}$$ This is the principle behind the VAE objective (Section sec:vae:elbo): the encoder parameterises $q_{\varphi }$, the decoder parameterises $p_{\theta }$, and both are optimised jointly via stochastic gradient ascent on the ELBO.

Connection to Statistical Physics

The ELBO and the Exponential Family

When the model $p(𝒙,𝒛)$ belongs to the exponential family, the ELBO and its gradients take particularly clean forms. This is the setting where classical (non-stochastic) variational inference is most naturally derived.

The ELBO in Deep Generative Models

Mean-Field Variational Inference

The ELBO gap identity (Theorem Theorem 2) tells us to maximise $ℒ(q)$ over a variational family $𝒬$. But which family? The choice of $𝒬$ determines both the quality of the approximation and the difficulty of the optimisation. The simplest and historically most important choice is the mean-field family, in which the variational distribution factorises completely over the latent variables.

The Mean-Field Assumption

The mean-field assumption is drastic: by forcing $q(𝒛)={\prod }_j{q}_j(z_j)$, we ignore all correlations between the groups $z_1,\dots ,z_J$. If the true posterior has strong dependencies between latent variables, as it often does, this introduces bias. But the payoff is computational: the optimal factors have a remarkably simple form.

Optimal Mean-Field Factors

Example: Mean-Field for the Ising Model

The Ising model is the prototypical example from statistical physics.

Example: Mean-Field for a Bivariate Gaussian

Before tackling a full model, let us see mean-field VI in action on the simplest possible example: a two-dimensional Gaussian.

This simple example encapsulates the core trade-off of mean-field VI: tractability (independent factors) versus fidelity (capturing correlations).

Example: Mean-Field for Bayesian Linear Regression

We now apply mean-field VI to a workhorse statistical model.

Mean-Field vs. Exact Posterior

The mean-field approximation decouples the latent variables, which means it cannot represent posterior correlations. The following figure illustrates this concretely.

Structured Variational Families

Coordinate Ascent Variational Inference

The optimal mean-field updates of Theorem Theorem 3 give us the best factor $q_j^{\ast }$ when all other factors are held fixed. But each update depends on the moments of the other factors, creating a system of coupled equations. The natural solution is to iterate: update each factor in turn, cycling through all $J$ groups. This is the coordinate ascent variational inference (CAVI) algorithm.

The CAVI Algorithm

Monotonic Convergence

Extended Example: CAVI for Bayesian Gaussian Mixture Models

We now develop a complete CAVI algorithm for the Bayesian Gaussian mixture model (BGMM), the canonical example in variational inference. This extends the EM algorithm of Chapter ch:gmm by placing priors on all parameters and performing fully Bayesian inference.

We now derive each optimal factor using Theorem Theorem 3.

Update for $q(z_i)$: Cluster Assignments

The terms in $\log p(𝒙,𝒛)$ that involve $z_i$ are $\log p(z_i|𝝅)$ and $\log p({𝒙}_i|z_i,{𝝁}_{1:K})$. By Theorem Theorem 3: (Z Update)$$\begin{array}{cccc}& \log q^{\ast }(z_i=k)& ={𝔼}_{q_{-z_i}}[\log {\pi }_k+\log 𝒩({𝒙}_i|{𝝁}_k,{\sigma }^2\mathrm{I})]+\mathrm{const}& \\ & & ={𝔼}_{q(𝝅)}[\log {\pi }_k]+{𝔼}_{q({𝝁}_k)}[\log 𝒩({𝒙}_i|{𝝁}_k,{\sigma }^2\mathrm{I})]+\mathrm{const}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Define the “responsibility” parameters: (RHO)$$\begin{array}{cccc}& {\rho }_{ik}& \propto \exp \{𝔼[\log {\pi }_k]-\frac{1}{2{\sigma }^2}{𝔼}_{q({𝝁}_k)}[\Vert {𝒙}_i-{𝝁}_k{\Vert }^2\left]\right\},& \text{<span class="tml-eqn"></span>}\end{array}$$ normalised so that ${\sum }_k{\rho }_{ik}=1$. These are the variational analogue of the EM responsibilities (Section sec:gmm:em).

The expectations needed are:

• $𝔼[\log {\pi }_k]$: the digamma function, computed from the Dirichlet factor $q(𝝅)$.

• ${𝔼}_{q({𝝁}_k)}[\Vert {𝒙}_i-{𝝁}_k{\Vert }^2]=\Vert {𝒙}_i-{𝒎}_k{\Vert }^2+\mathrm{trace}({𝑺}_k)$, where ${𝒎}_k,{𝑺}_k$ are the mean and covariance of $q({𝝁}_k)$.

Update for $q({𝝁}_k)$: Cluster Means

The terms involving ${𝝁}_k$ are $\log p({𝝁}_k)$ and the likelihood terms for all ${𝒙}_i$ assigned to cluster $k$: (MU LOG)$$\begin{array}{cccc}& \log q^{\ast }({𝝁}_k)& ={𝔼}_{q_{-{𝝁}_k}}[\log p({𝝁}_k)+\sum _{i=1}^N\mathrm{𝟙}[z_i=k]\log 𝒩({𝒙}_i|{𝝁}_k,{\sigma }^2\mathrm{I})]+\mathrm{const}& \\ & & =-\frac{{\beta }_0}{2}\Vert {𝝁}_k-{𝒎}_0{\Vert }^2-\frac{1}{2{\sigma }^2}\sum _{i=1}^N{\rho }_{ik}\Vert {𝒙}_i-{𝝁}_k{\Vert }^2+\mathrm{const}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Completing the square in ${𝝁}_k$, this is Gaussian: (Q MU)$$\begin{array}{ccc}& q^{\ast }({𝝁}_k)=𝒩({𝝁}_k|{𝒎}_k,{𝑺}_k),& \text{<span class="tml-eqn"></span>}\end{array}$$ with (MU Params)$$\begin{array}{ccc}& {𝑺}_k={({\beta }_0\mathrm{I}+\frac{N_k}{{\sigma }^2}\mathrm{I})}^{-1},{𝒎}_k={𝑺}_k({\beta }_0{𝒎}_0+\frac{1}{{\sigma }^2}\sum _{i=1}^N{\rho }_{ik}{𝒙}_i),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $N_k={\sum }_{i=1}^N{\rho }_{ik}$ is the effective number of points assigned to cluster $k$. This is a precision-weighted average of the prior mean ${𝒎}_0$ and the data centroid ${\overline{𝒙}}_k=\frac{1}{N_k}{\sum }_i{\rho }_{ik}{𝒙}_i$.

Update for $q(𝝅)$: Mixing Weights

The terms involving $𝝅$ are $\log p(𝝅)$ (the Dirichlet prior) and ${\sum }_i\log p(z_i|𝝅)$: (PI LOG)$$\begin{array}{cccc}& \log q^{\ast }(𝝅)& ={𝔼}_{q_{-𝝅}}[({\alpha }_0-1)\sum _k\log {\pi }_k+\sum _{i=1}^N\sum _{k=1}^K\mathrm{𝟙}[z_i=k]\log {\pi }_k]+\mathrm{const}& \\ & & =\sum _{k=1}^K({\alpha }_0-1+N_k)\log {\pi }_k+\mathrm{const}.& \text{<span class="tml-eqn"></span>}\end{array}$$ This is the log-density of a Dirichlet distribution: (Q PI)$$\begin{array}{ccc}& q^{\ast }(𝝅)=\mathrm{Dirichlet}({\alpha }_0+N_1,\dots ,{\alpha }_0+N_K).& \text{<span class="tml-eqn"></span>}\end{array}$$ The expected log-mixing weights needed by the $z_i$-update are: (Elogpi)$$\begin{array}{ccc}& 𝔼[\log {\pi }_k]=\psi ({\alpha }_0+N_k)-\psi \left(\sum _{k^{\prime }=1}^K({\alpha }_0+N_{k^{\prime }})\right),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\psi (\cdot )$ is the digamma function $\psi (x)=\frac{\mathrm{d}}{\mathrm{d}x}\log \mathrm{\Gamma }(x)$.

Summary of CAVI Updates

Computing the ELBO for Monitoring

To monitor convergence, we must evaluate the ELBO $ℒ(q)={𝔼}_q[\log p(𝒙,𝒛)]-{𝔼}_q[\log q(𝒛)]$ at each iteration. Using the mean-field factorisation :

(ELBO FULL)$$\begin{array}{cccc}& ℒ& =\underset{\text{Dirichlet terms}}{\underbrace{{𝔼}_q[\log p(𝝅)]-{𝔼}_q[\log q(𝝅)]}}& \\ & & +\sum _{k=1}^K\underset{\text{Gaussian mean terms}}{\underbrace{({𝔼}_q[\log p({𝝁}_k)]-{𝔼}_q[\log q({𝝁}_k)])}}& \\ & & +\sum _{i=1}^N\underset{\text{data terms}}{\underbrace{({𝔼}_q[\log p(z_i|𝝅)]+{𝔼}_q[\log p({𝒙}_i|z_i,{𝝁}_{1:K})]-{𝔼}_q[\log q(z_i)])}}.& \text{<span class="tml-eqn"></span>}\end{array}$$

Each term involves expectations of log-densities under their respective variational factors, all of which are available in closed form for exponential-family models:

• Dirichlet terms: $𝔼[\log p(𝝅)]-𝔼[\log q(𝝅)]$ involves the difference of two Dirichlet log-normalisation constants plus digamma expectations.

• Gaussian mean terms: $𝔼[\log p({𝝁}_k)]-𝔼[\log q({𝝁}_k)]=-{𝖣}_{\mathrm{KL}}(q({𝝁}_k)\Vert p({𝝁}_k))$, computable via the closed-form KL between Gaussians (Proposition prop:kl-gaussian-mv).

• Data terms: $𝔼[\log p(z_i|𝝅)]={\sum }_k{\rho }_{ik}𝔼[\log {\pi }_k]$; the reconstruction term uses $𝔼[\Vert {𝒙}_i-{𝝁}_k{\Vert }^2]=\Vert {𝒙}_i-{𝒎}_k{\Vert }^2+\mathrm{trace}({𝑺}_k)$; the entropy of $q(z_i)$ is $-{\sum }_k{\rho }_{ik}\log {\rho }_{ik}$.

CAVI Iterations: A Visual Example

Comparison with EM

CAVI for the Bayesian GMM bears a striking resemblance to EM for the frequentist GMM (Chapter ch:gmm). The following table makes the correspondence explicit.

Aspect	EM (frequentist GMM)	CAVI (Bayesian GMM)
Parameters	Point estimates $\theta =\{{\pi }_k,{𝝁}_k,{𝚺}_k\}$	Distributions $q(𝝅),q({𝝁}_k)$
E-step	Compute ${\gamma }_{ik}=p(z_i=k|{𝒙}_i;\theta )$	Update $q(z_i)$: ${\rho }_{ik}$ via
M-step	Maximise $Q(\theta ,{\theta }^{(t)})$ over $\theta$	Update $q(𝝅)$, $q({𝝁}_k)$
Objective	Log-likelihood $\ell (\theta )$	ELBO $ℒ(q)$
Uncertainty	None (point estimate)	Full posterior over all parameters
Overfitting	Can overfit with large $K$	Automatic regularisation via prior
$K$ selection	Requires BIC/AIC	Automatic pruning (empty clusters)

Practical Considerations

CAVI and Natural Gradients

There is an elegant connection between the CAVI coordinate updates and the natural gradient of information geometry.

This connection is important because it explains why CAVI converges rapidly in practice: natural gradients account for the curvature of the distribution space and take optimally large steps. It also paves the way for stochastic natural-gradient methods, where we replace the full-data sufficient statistics with mini-batch estimates, the basis of stochastic variational inference (SVI), covered in File B.

Predictive Distribution

After convergence, the variational approximation can be used for prediction. For the Bayesian GMM, the predictive density for a new observation ${𝒙}_{\ast }$ is:

(Predictive)$$\begin{array}{ccc}& p({𝒙}_{\ast }|{𝒙}_{1:N})\approx \int p({𝒙}_{\ast }|𝒛)q^{\ast }(𝒛)\mathrm{d}𝒛=\sum _{k=1}^K{𝔼}_q[{\pi }_k]\int 𝒩({𝒙}_{\ast }|{𝝁}_k,{\sigma }^2\mathrm{I})q^{\ast }({𝝁}_k)\mathrm{d}{𝝁}_k.& \text{<span class="tml-eqn"></span>}\end{array}$$

Since both the likelihood and $q^{\ast }({𝝁}_k)$ are Gaussian, the inner integral is available in closed form:

(Predictive Closed)$$\begin{array}{ccc}& \int 𝒩({𝒙}_{\ast }|{𝝁}_k,{\sigma }^2\mathrm{I})𝒩({𝝁}_k|{𝒎}_k,{𝑺}_k)\mathrm{d}{𝝁}_k=𝒩({𝒙}_{\ast }|{𝒎}_k,{\sigma }^2\mathrm{I}+{𝑺}_k).& \text{<span class="tml-eqn"></span>}\end{array}$$

The predictive distribution is therefore itself a Gaussian mixture ${\sum }_k{\bar{\pi }}_k𝒩({𝒎}_k,{\sigma }^2\mathrm{I}+{𝑺}_k)$ with ${\bar{\pi }}_k={𝔼}_q[{\pi }_k]=({\alpha }_0+N_k)/{\sum }_{k^{\prime }}({\alpha }_0+N_{k^{\prime }})$. Notice that the predictive covariance ${\sigma }^2\mathrm{I}+{𝑺}_k$ is larger than the noise covariance ${\sigma }^2\mathrm{I}$: the additional term ${𝑺}_k$ captures our epistemic uncertainty about the cluster mean. As $N\to \infty$, ${𝑺}_k\to 0$ and the predictive distribution converges to the plug-in estimate, and the Bayesian and frequentist answers agree asymptotically.

Limitations of CAVI

Exercises

Exponential Families and Conjugacy

The coordinate ascent updates derived in the previous section, while general in principle, become especially clean and computationally attractive when the model belongs to the exponential family. This is no accident: exponential families possess a rich algebraic structure that interacts beautifully with the KL divergence and the ELBO. In this section we formalise this structure and show why exponential families occupy a privileged position in variational inference.

The Exponential Family

The exponential family encompasses nearly every distribution encountered in introductory statistics: the Gaussian, Bernoulli, Poisson, exponential, gamma, beta, Dirichlet, multinomial, Wishart, and von Mises–Fisher distributions are all members. Notable exceptions include the uniform distribution on $[0,\theta ]$ (whose support depends on the parameter), the Cauchy distribution (no sufficient statistic of fixed dimension), and mixture distributions (sums of exponentials are not exponentials).

The representation may look abstract, but its power lies in the properties of the log-partition function $A$. The function $A(𝜼)$ acts as a “generating function” for the moments of $T(𝒙)$: each derivative of $A$ produces a higher moment.

The KL divergence between two members of the same exponential family has a particularly elegant form: (KL)$$\begin{array}{ccc}& {𝖣}_{\mathrm{KL}}(p(𝒙;{𝜼}_1)\Vert p(𝒙;{𝜼}_2))=A({𝜼}_2)-A({𝜼}_1)-({𝜼}_2-{𝜼}_1{)}^{𝖳}\nabla A({𝜼}_1).& \text{<span class="tml-eqn"></span>}\end{array}$$ This is the Bregman divergence generated by $A$, a connection that provides geometric insight into variational inference as projection onto the variational family with respect to this divergence.

Conjugate Priors

When building Bayesian models, a natural question is: can we choose a prior so that the posterior has the same functional form as the prior? If so, the posterior update becomes a simple parameter update rather than a change of distributional family.

When the likelihood belongs to an exponential family with natural parameter $𝜽$ and sufficient statistics $T(𝒙)$, a conjugate prior can always be constructed: (FORM)$$\begin{array}{ccc}& p(𝜽;𝝌,\nu )=\frac{1}{Z(𝝌,\nu )}\exp \{{𝝌}^{𝖳}𝜽-\nu A(𝜽)\},& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝝌$ and $\nu >0$ are hyperparameters, and $Z(𝝌,\nu )$ is the normalising constant. After observing data ${𝒙}_1,\dots ,{𝒙}_N$, the posterior is (Posterior)$$\begin{array}{ccc}& p(𝜽|{𝒙}_{1:N};𝝌,\nu )\propto \exp \{(𝝌+\sum _{i=1}^{N}T({𝒙}_i){)}^{𝖳}𝜽-(\nu +N)A(𝜽)\},& \text{<span class="tml-eqn"></span>}\end{array}$$ which is the same family with updated hyperparameters ${𝝌}^{\prime }=𝝌+{\sum }_{i}T({𝒙}_i)$ and ${\nu }^{\prime }=\nu +N$. The posterior simply “adds” the observed sufficient statistics to the prior pseudo-observations.

CAVI in Exponential Family Models

We now connect exponential families to the coordinate ascent variational inference (CAVI) algorithm developed in the previous section. The payoff is remarkable: for models in which every complete conditional belongs to an exponential family, the optimal mean-field factors have closed-form natural parameter updates.

The Variational EM Algorithm

In sec:gmm:em we introduced the Expectation–Maximisation (EM) algorithm for fitting models with latent variables. Recall the two steps:

• E-step: compute the posterior over latent variables, $p(𝒛|𝒙;\theta )$.

• M-step: maximise the expected complete-data log-likelihood $Q(\theta ,{\theta }^{(t)})$ with respect to the model parameters $\theta$.

For GMMs, the E-step is tractable: the posterior over the cluster assignments $z_i$ factorises and each factor is a categorical distribution (the “responsibilities”). But for more complex models (topic models, deep latent variable models, models with continuous high-dimensional latent spaces) the exact posterior $p(𝒛|𝒙;\theta )$ is intractable. What do we do?

The answer is the Variational EM algorithm: replace the exact E-step with a variational E-step that finds an approximate posterior within a tractable family $𝒬$.

To understand why the exact E-step becomes intractable, consider a topic model with $K$ topics and a document of $N$ words. The posterior over topic assignments requires marginalising over all $K^N$ possible configurations, an exponential sum that no closed-form trick can simplify. In a deep latent variable model, the decoder is a neural network, and the posterior involves inverting this nonlinear mapping, and there is no hope of a closed form. Variational EM sidesteps these difficulties by approximating the posterior rather than computing it exactly.

From EM to Variational EM

Recall the ELBO decomposition (from the previous sections): (Decomp)$$\begin{array}{ccc}& \log p(𝒙;\theta )=\underset{ℒ(q,\theta )\text{(ELBO)}}{\underbrace{{𝔼}_{q(𝒛)}\left[\log \frac{p(𝒙,𝒛;\theta )}{q(𝒛)}\right]}}+{𝖣}_{\mathrm{KL}}(q(𝒛)\Vert p(𝒛|𝒙;\theta )).& \text{<span class="tml-eqn"></span>}\end{array}$$ In standard EM, the E-step sets $q(𝒛)=p(𝒛|𝒙;{\theta }^{(t)})$, making the KL term zero and the ELBO equal to the log-evidence. In variational EM, we cannot set $q$ to the exact posterior; instead, we maximise the ELBO over a restricted family $q\in 𝒬$:

Example: Variational EM for Latent Dirichlet Allocation

Latent Dirichlet Allocation (LDA) is the canonical application of Variational EM [8]. It is a topic model that discovers latent “topics” (distributions over words) from a corpus of documents.

The generative model. For a corpus of $D$ documents, each with $N_d$ words, and $K$ latent topics:

1. For each topic $k=1,\dots ,K$: draw a word distribution ${𝜷}_k\sim \text{Dirichlet}(𝜼)$, where ${𝜷}_k\in {\mathrm{\Delta }}^{V-1}$ (the $(V-1)$-simplex over a vocabulary of $V$ words).

2. For each document $d=1,\dots ,D$: enumerate

3. Draw topic proportions ${𝜽}_d\sim \text{Dirichlet}(𝜶)$.

4. For each word position $n=1,\dots ,N_d$: enumerate

5. Draw a topic assignment $z_{dn}\sim \mathrm{Categorical}({𝜽}_d)$.

6. Draw a word $w_{dn}\sim \mathrm{Categorical}({𝜷}_{z_{dn}})$. enumerate enumerate

The posterior $p({𝜽}_{1:D},{𝒛}_{1:D},{𝜷}_{1:K}|{𝒘}_{1:D};𝜶,𝜼)$ is intractable because the Dirichlet–categorical conjugacy is broken by the coupling through the topic assignments $z_{dn}$.

The mean-field approximation. The standard variational approximation uses a fully factorised family: (MF)$$\begin{array}{ccc}& q({𝜽}_{1:D},{𝒛}_{1:D},{𝜷}_{1:K})=\prod _{k=1}^{K}q({𝜷}_k|{𝝀}_k)\prod _{d=1}^D[q({𝜽}_d|{𝜸}_d)\prod _{n=1}^{N_d}q(z_{dn}|{𝝓}_{dn})],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $q({𝜷}_k|{𝝀}_k)=\text{Dirichlet}({𝝀}_k)$, $q({𝜽}_d|{𝜸}_d)=\text{Dirichlet}({𝜸}_d)$, and $q(z_{dn}|{𝝓}_{dn})=\mathrm{Categorical}({𝝓}_{dn})$.

The computational cost of each VE-step is $O(D\cdot \bar{N}\cdot K)$, where $\bar{N}$ is the average document length: we must update ${𝝓}_{dnk}$ for each of the $D\cdot \bar{N}$ word positions and $K$ topics. For a corpus of $D={10}^6$ documents with average length $\bar{N}=200$ and $K=100$ topics, a single VE-step touches $2\times {10}^{10}$ variational parameters.

Moreover, the VE-step must be iterated to convergence (cycling through the $𝝓$ and $𝜸$ updates multiple times) before performing a single VM-step. This makes each iteration of Variational EM extremely expensive for large corpora.

The VM-step aggregates information from all documents to update the global topic parameters. This full-data sweep is the computational bottleneck that motivates the stochastic methods of the next section.

Stochastic Variational Inference

Coordinate ascent variational inference (CAVI) and Variational EM both require a full pass over the dataset in every iteration. For the LDA example, each VE-step touches every word in every document. When the corpus grows from thousands to millions of documents, this becomes a crippling bottleneck: the cost of a single iteration scales linearly in the dataset size $N$, and many iterations are needed for convergence.

The insight of Stochastic Variational Inference (SVI) [20] is both natural and powerful: subsample the data and use noisy but unbiased gradient estimates to update the global variational parameters. This brings the per-iteration cost down from $O(N)$ to $O(1)$ (independent of dataset size), at the price of noisier updates that nonetheless converge to the same optimum.

Local and Global Latent Variables

The key structural assumption in SVI is a distinction between two types of latent variables.

In LDA, the topic–word distributions ${𝜷}_{1:K}$ are global variables (shared across all documents), while the topic proportions ${𝜽}_d$ and topic assignments $z_{dn}$ are local to document $d$. In a Gaussian mixture model, the component parameters $({𝝁}_k,{𝚺}_k,{\pi }_k)$ are global, and the cluster assignments $z_i$ are local.

The mean-field variational distribution mirrors this decomposition: (MF)$$\begin{array}{ccc}& q({𝒛}_{1:N},𝜷)=q(𝜷|𝝀)\prod _{i=1}^{N}q({𝒛}_i|{𝝓}_i),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝝀$ are the global variational parameters and ${𝝓}_i$ are the local variational parameters for data point $i$.

The SVI Algorithm

The ELBO for the model decomposes as: (ELBO)$$\begin{array}{ccc}& ℒ(𝝀,{𝝓}_{1:N})={𝔼}_q[\log p(𝜷)]-{𝔼}_q[\log q(𝜷)]+\sum _{i=1}^N\left({𝔼}_q\right[\log p({𝒙}_i,{𝒛}_i|𝜷)]-{𝔼}_q[\log q({𝒛}_i)]).& \text{<span class="tml-eqn"></span>}\end{array}$$ The sum over $N$ data points is the bottleneck. The SVI strategy proceeds in three steps:

1. Subsample a minibatch $ℬ\subset \{1,\dots ,N\}$ of size $B$.

2. Local step: for each $i\in ℬ$, optimise ${𝝓}_i$ to convergence (using the current global parameters $𝝀$). This is typically cheap, often one CAVI update per local variable.

3. Global step: form the stochastic natural gradient of the ELBO with respect to $𝝀$ and take a step.

The stochastic gradient is constructed by scaling the minibatch contribution: (GRAD)$$\begin{array}{ccc}& {\hat{\nabla }}_{𝝀}ℒ={\nabla }_{𝝀}{𝔼}_q[\log p(𝜷)-\log q(𝜷)]+\frac{N}{B}\sum _{i\in ℬ}{\nabla }_{𝝀}\left({𝔼}_q\right[\log p({𝒙}_i,{𝒛}_i|𝜷)]-{𝔼}_q[\log q({𝒛}_i)]).& \text{<span class="tml-eqn"></span>}\end{array}$$ The factor $N/B$ rescales the minibatch contribution to be an unbiased estimate of the full-data sum.

For conjugate exponential family models, the global update takes a particularly elegant form. After computing the locally optimal ${𝝓}_i^{\ast }$ for each $i\in ℬ$, the intermediate global natural parameter is (Intermediate)$$\begin{array}{ccc}& \widehat{𝝀}={𝝀}_{\text{prior}}+\frac{N}{B}\sum _{i\in ℬ}t({𝒙}_i,{𝝓}_i^{\ast }),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝝀}_{\text{prior}}$ encodes the prior's contribution and $t({𝒙}_i,{𝝓}_i^{\ast })$ are the expected sufficient statistics from data point $i$ under the optimised local distribution. The global update is then a convex combination (weighted average) of the old parameters and this intermediate estimate: (Update)$$\begin{array}{ccc}& {𝝀}^{(t+1)}=(1-{\rho }_t){𝝀}^{(t)}+{\rho }_t\widehat{𝝀},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\rho }_t\in (0,1]$ is the step size (learning rate) at iteration $t$.

Why Natural Gradients?

A crucial but subtle point in the SVI algorithm is that the global update uses the natural gradient rather than the ordinary (Euclidean) gradient. To understand why, consider the geometry of the variational parameter space.

The ordinary gradient ${\nabla }_{𝝀}ℒ$ gives the steepest ascent direction in Euclidean space; it assumes that a unit step in any direction of $𝝀$ produces an equally significant change. But for probability distributions, this is misleading: a change $\mathrm{\Delta }\lambda =0.01$ in the mean of a Gaussian has a very different effect from the same change in the log-variance.

The natural gradient corrects for this by measuring distances in the space of distributions rather than parameters. It is defined as: (Natgrad)$$\begin{array}{ccc}& {\tilde{\nabla }}_{𝝀}ℒ=𝑭{(𝝀)}^{-1}{\nabla }_{𝝀}ℒ,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝑭(𝝀)={𝔼}_{q_{𝝀}}[{\nabla }_{𝝀}\log q_{𝝀}(𝒛)({\nabla }_{𝝀}\log q_{𝝀}(𝒛){)}^{𝖳}]$ is the Fisher information matrix.

For exponential families, a remarkable simplification occurs: when the variational parameters are the natural parameters $𝜼$, the Fisher information matrix is $𝑭={\nabla }^{2}A(𝜼)$, and the natural gradient update in natural parameter space takes the simple form of the convex combination . This is why SVI works in natural parameter space, where the natural gradient is “built in” to the update rule.

Convergence Guarantees

The step size schedule must satisfy the classical Robbins–Monro conditions [9]: (Robbins)$$\begin{array}{ccc}& \sum _{t=1}^{\infty }{\rho }_t=\infty \text{and}\sum _{t=1}^{\infty }{\rho }_t^2<\infty .& \text{<span class="tml-eqn"></span>}\end{array}$$ The first condition ensures that the step sizes are large enough to reach any point in parameter space (no matter how far the initialisation is from the optimum). The second condition ensures that the noise in the stochastic gradient is eventually averaged out. The schedule ${\rho }_t={(\tau +t)}^{-\kappa }$ satisfies both conditions when $\kappa \in (0.5,1]$.

Example: SVI for Online Topic Modelling

Black-Box Variational Inference

Stochastic variational inference (Stochastic Variational Inference) brought scalability to VI, but it required the practitioner to derive model-specific update equations, exploiting conjugate exponential family structure to obtain closed-form natural gradients. What if the model is not conjugate? What if the likelihood involves a neural network, or the prior is non-standard? In such cases, the elegant closed-form updates of SVI simply do not exist.

Black-Box Variational Inference (BBVI) [21] attacks this problem head-on: it provides a gradient estimator for the ELBO that works for any model and any variational family, requiring only the ability to

1. sample from the variational distribution $q_{𝝀}(𝒛)$, and

2. evaluate $\log p(𝒙,𝒛)$ and $\log q_{𝝀}(𝒛)$ at the sampled points.

No conjugacy, no exponential family structure, no model-specific derivations. The price? Variance, which we must work hard to control.

The Score Function Estimator

The foundation of BBVI is the score function estimator, also known as the REINFORCE trick or the log-derivative trick [22]. It provides a way to differentiate an expectation with respect to the parameters of the distribution being averaged over.

The quantity ${\nabla }_{𝝀}\log q_{𝝀}(𝒛)$ is called the score function of $q$; it measures how sensitive the log-density is to changes in the parameters at a given point $𝒛$.

Applying this identity to the ELBO gradient requires care. Write $ℒ(𝝀)={𝔼}_{q_{𝝀}}[\log p(𝒙,𝒛)-\log q_{𝝀}(𝒛)]$ and note that $\log q_{𝝀}(𝒛)$ itself depends on $𝝀$ both through the expectation (the distribution we average over) and through the integrand (the log-density being evaluated). The full gradient is: (Fullgrad)$$\begin{array}{cccc}& {\nabla }_{𝝀}ℒ& ={\nabla }_{𝝀}{𝔼}_{q_{𝝀}}[\log p(𝒙,𝒛)-\log q_{𝝀}(𝒛)]& \\ & & ={𝔼}_{q_{𝝀}}\left[\right(\log p(𝒙,𝒛)-\log q_{𝝀}(𝒛)){\nabla }_{𝝀}\log q_{𝝀}(𝒛)]-\underset{=0}{\underbrace{{𝔼}_{q_{𝝀}}[{\nabla }_{𝝀}\log q_{𝝀}(𝒛)]}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ The first term applies the score function identity (Lemma 1) to the full ELBO integrand. The second term arises from differentiating $\log q_{𝝀}(𝒛)$ as part of the integrand; it vanishes because ${𝔼}_q[{\nabla }_{𝝀}\log q_{𝝀}(𝒛)]=\int {\nabla }_{𝝀}{q}_{𝝀}(𝒛)\mathrm{d}𝒛={\nabla }_{𝝀}\int q_{𝝀}(𝒛)\mathrm{d}𝒛={\nabla }_{𝝀}1=0$. This yields the key result:

(GRAD)$$\begin{array}{ccc}& {\nabla }_{𝝀}ℒ(𝝀)={𝔼}_{q_{𝝀}(𝒛)}[{\nabla }_{𝝀}\log q_{𝝀}(𝒛)\cdot (\log p(𝒙,𝒛)-\log q_{𝝀}(𝒛)\left)\right].& \text{<span class="tml-eqn"></span>}\end{array}$$

This is an expectation under $q_{𝝀}$, so we can estimate it with Monte Carlo: (MC)$$\begin{array}{ccc}& {\hat{\nabla }}_{𝝀}ℒ=\frac{1}{S}\sum _{s=1}^S{\nabla }_{𝝀}\log q_{𝝀}({𝒛}^{(s)})\cdot (\log p(𝒙,{𝒛}^{(s)})-\log q_{𝝀}({𝒛}^{(s)})),{𝒛}^{(s)}\sim q_{𝝀}.& \text{<span class="tml-eqn"></span>}\end{array}$$

This estimator is unbiased: its expectation equals the true gradient ${\nabla }_{𝝀}ℒ$. Crucially, it requires no knowledge of the model structure beyond the ability to evaluate $\log p(𝒙,𝒛)$; the model is treated as a “black box.”

Variance Reduction

Three classical techniques bring the variance of the score function estimator under control.

Control Variates

If $h$ is positively correlated with $f$, choosing $c>0$ reduces the variance. The optimal coefficient that minimises is: (Optimal)$$\begin{array}{ccc}& c^{\ast }=\frac{\mathrm{𝖢ov}[f,h]}{\mathrm{𝖵ar}[h]}.& \text{<span class="tml-eqn"></span>}\end{array}$$ In practice, $c^{\ast }$ is estimated from the samples themselves.

For BBVI, a common choice is the baseline approach: subtract a scalar $b$ (an estimate of the mean reward) from the “reward” signal $\log p(𝒙,𝒛)-\log q_{𝝀}(𝒛)$ in the gradient estimate . This does not change the expected gradient (since ${𝔼}_q[{\nabla }_{𝝀}\log q_{𝝀}(𝒛)]=0$, a standard property of the score function), but can dramatically reduce variance.

Rao–Blackwellization

If some latent variables can be integrated out analytically, doing so always reduces variance (by the Rao–Blackwell theorem). For example, if $𝒛=({𝒛}_1,{𝒛}_2)$ and the expectation over ${𝒛}_2$ can be computed in closed form, we should compute it analytically and only use Monte Carlo for ${𝒛}_1$. This is model-specific but can yield large variance reductions.

Optimal Baseline Derivation

For the BBVI gradient estimator, consider subtracting a baseline $b$: (Baseline)$$\begin{array}{ccc}& {\hat{\nabla }}_{𝝀}ℒ=\frac{1}{S}\sum _{s=1}^S{\nabla }_{𝝀}\log q_{𝝀}({𝒛}^{(s)})\cdot (\log p(𝒙,{𝒛}^{(s)})-\log q_{𝝀}({𝒛}^{(s)})-b).& \text{<span class="tml-eqn"></span>}\end{array}$$ The variance of a single-sample estimator (for each component $j$ of $𝝀$) is $$\begin{array}{c}\mathrm{𝖵ar}\left[{\hat{g}}_j\right]=𝔼\left[\right(s_j(𝒛)\cdot (f(𝒛)-b){)}^2]-(𝔼[s_j(𝒛)\cdot f(𝒛)]{)}^2,\end{array}$$ where $s_j(𝒛)=\partial \log q_{𝝀}(𝒛)/\partial {\lambda }_j$ and $f(𝒛)=\log p(𝒙,𝒛)-\log q_{𝝀}(𝒛)$. Minimising over $b$ yields: (Optbaseline)$$\begin{array}{ccc}& b_j^{\ast }=\frac{𝔼[s_j(𝒛{)}^{2}f(𝒛)]}{𝔼[s_j(𝒛{)}^2]}.& \text{<span class="tml-eqn"></span>}\end{array}$$ This is the optimal per-component baseline. In practice, a single scalar baseline (averaged across components) is often used for simplicity.

The BBVI Algorithm

Gradient Variance: Score Function vs. Reparameterization

The figure reveals the fundamental tension in BBVI: the score function estimator is maximally general (works for any model and any variational family) but has high variance. The reparameterisation trick (next section) has much lower variance but requires the variational distribution to support a differentiable reparameterisation. For models where the reparameterisation trick applies, it is almost always preferred.

To quantify the difference: in a typical 50-dimensional latent space model, the score function estimator can have variance ${10}^3$–${10}^5$ times larger than the reparameterised estimator. With $S=1$ sample, the score function estimator often produces gradient estimates whose signal-to-noise ratio is below 1, making optimisation impractical. The reparameterised estimator typically achieves a signal-to-noise ratio above 10 even with a single sample, enabling practical training with minibatch sizes of 32–128.

This variance gap is the primary reason that the VAE (ch:vae) uses the reparameterisation trick rather than the score function estimator. For discrete latent variables, where reparameterisation is not available, the variance of BBVI remains a major challenge and an active area of research.

The Reparameterization Trick

The score function estimator is general but noisy. In this section we develop a dramatically better gradient estimator for the important special case when the variational distribution admits a differentiable reparameterisation, meaning we can express samples from $q_{𝝓}(𝒛)$ as a deterministic, differentiable transformation of parameter-free noise.

The reparameterisation trick is the technical engine behind the Variational Autoencoder (VAE, ch:vae). It was independently discovered by two groups in 2013–2014, and its impact on the field cannot be overstated: it made end-to-end training of deep latent variable models practical for the first time.

Separating Randomness from Parameters

The fundamental difficulty in computing ${\nabla }_{𝝓}{𝔼}_{q_{𝝓}(𝒛)}[f(𝒛)]$ is that the distribution we are averaging over depends on the parameters $𝝓$. When we draw $𝒛\sim q_{𝝓}$, the randomness is entangled with the parameters, so we cannot take the gradient inside the expectation directly.

The reparameterisation trick untangles this by expressing the random variable $𝒛$ as a deterministic function of $𝝓$ and a parameter-free noise variable $𝝐$.

The reparameterisation converts a stochastic computation (draw from $q_{𝝓}$) into a deterministic computation applied to fixed noise. This seemingly simple change has profound consequences: because $g_{𝝓}$ is deterministic and differentiable, standard backpropagation can compute ${\nabla }_{𝝓}$ through the entire computation, including through the “sampling” step.

The most important example is the Gaussian reparameterisation: (Gauss)$$\begin{array}{ccc}& 𝒛\sim 𝒩(𝝁,\mathrm{diag}({𝝈}^2))⟺𝒛=𝝁+𝝈\odot 𝝐,𝝐\sim 𝒩(0,𝑰),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝝓=(𝝁,𝝈)$ and $\odot$ denotes element-wise multiplication. The noise $𝝐$ is drawn from a standard Gaussian that does not depend on $𝝓$; all the parameter-dependence is in the deterministic transformation $g_{𝝓}(𝝐)=𝝁+𝝈\odot 𝝐$.

The Reparameterised Gradient

With the reparameterisation, the expectation transforms: (Transform)$$\begin{array}{ccc}& {𝔼}_{q_{𝝓}(𝒛)}[f(𝒛)]={𝔼}_{p(𝝐)}\left[f\right(g_{𝝓}(𝝐)\left)\right].& \text{<span class="tml-eqn"></span>}\end{array}$$ Now the distribution $p(𝝐)$ inside the expectation does not depend on $𝝓$, so we can freely move the gradient inside: (GRAD)$$\begin{array}{ccc}& {\nabla }_{𝝓}{𝔼}_{q_{𝝓}}[f(𝒛)]={\nabla }_{𝝓}{𝔼}_{p(𝝐)}[f(g_{𝝓}(𝝐))]={𝔼}_{p(𝝐)}\left[{\nabla }_{𝝓}f\right(g_{𝝓}(𝝐)\left)\right].& \text{<span class="tml-eqn"></span>}\end{array}$$ The gradient now flows through $f$ and through $g_{𝝓}$ via the chain rule: (Chain)$$\begin{array}{ccc}& {\nabla }_{𝝓}f(g_{𝝓}(𝝐))={\nabla }_{𝒛}f(𝒛){|}_{𝒛=g_{𝝓}(𝝐)}\cdot {\nabla }_{𝝓}{g}_{𝝓}(𝝐).& \text{<span class="tml-eqn"></span>}\end{array}$$

For the Gaussian case, the reparameterised ELBO gradient has an especially clean form. With $q_{𝝓}(𝒛)=𝒩(𝝁,\mathrm{diag}({𝝈}^2))$ and the reparameterisation $𝒛=𝝁+𝝈\odot 𝝐$: (Gradmu)$$\begin{array}{cccc}& {\nabla }_{𝝁}ℒ& ={𝔼}_{𝝐}[{\nabla }_{𝒛}f(𝒛)],& \text{<span class="tml-eqn"></span>}\\ & {\nabla }_{𝝈}ℒ& ={𝔼}_{𝝐}[{\nabla }_{𝒛}f(𝒛)\odot 𝝐],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $f(𝒛)=\log p(𝒙,𝒛)-\log q_{𝝓}(𝒛)$. These gradients are straightforward to compute via automatic differentiation, requiring only a forward pass (to evaluate $f$) and a backward pass (to compute ${\nabla }_{𝒛}f$) through the model.

Computational Graph Perspective

The reparameterisation trick has an elegant interpretation in terms of computational graphs, which is how automatic differentiation frameworks (PyTorch, JAX) implement it in practice.

Reparameterisable Distributions

Not every distribution admits a simple reparameterisation. The following table summarises the most commonly used reparameterisable families.

llll@ Distribution	Parameters $𝝓$	Base noise $p(𝝐)$	Transformation $g_{𝝓}(𝝐)$
Gaussian $𝒩(𝝁,\mathrm{diag}({𝝈}^2))$	$𝝁,𝝈$	$𝒩(0,𝑰)$	$𝝁+𝝈\odot 𝝐$
[4pt] Full-cov. Gaussian $𝒩(𝝁,𝚺)$	$𝝁,𝑳$	$𝒩(0,𝑰)$	$𝝁+𝑳𝝐$
[4pt] Log-normal $\text{LogNormal}(\mu ,{\sigma }^2)$	$\mu ,\sigma$	$𝒩(0,1)$	$\exp (\mu +\sigma ϵ)$
[4pt] Exponential $\text{Exp}(\lambda )$	$\lambda$	$\text{Uniform}(0,1)$	$-\log (ϵ)/\lambda$
[4pt] Gamma $\text{Gamma}(\alpha ,\beta )$	$\alpha ,\beta$	via shape augmentation	implicit
[4pt] von Mises–Fisher $\text{vMF}(𝝁,\kappa )$	$𝝁,\kappa$	rejection sampling	Householder transform

Reparameterisation with Full Covariance

Connection to Pathwise Gradients

The reparameterisation trick is an instance of a more general idea known as pathwise gradients (or pathwise derivatives) in the operations research and stochastic simulation literature. The key insight is the same: when we can express a random quantity as a deterministic function of a parameter and independent randomness, we can differentiate through the function.

In the notation of stochastic optimisation, for a performance measure $J(𝝓)=𝔼[f(𝒛)]$ where $𝒛$ depends on $𝝓$:

• The score function approach (likelihood ratio method) differentiates the measure under which we average.

• The pathwise approach differentiates the sample path $g_{𝝓}(𝝐)$ directly.

When both are available, the pathwise approach typically has lower variance because it exploits the local geometry of $f$ (via ${\nabla }_{𝒛}f$), while the score function approach only uses information about the density $q_{𝝓}$.

Extensions: Handling Non-Reparameterisable Distributions

For distributions that do not admit exact reparameterisation, most notably discrete distributions, several approximations have been developed:

• Gumbel–Softmax / Concrete relaxation. Replace a categorical variable with a continuous relaxation using Gumbel noise: (Gumbel)$$\begin{array}{ccc}& z_k=\frac{\exp ((\log {\pi }_k+g_k)/\tau )}{{\sum }_j\exp ((\log {\pi }_j+g_j)/\tau )},g_k\sim \text{Gumbel}(0,1),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\tau >0$ is a temperature parameter. As $\tau \to 0$, the samples approach one-hot vectors (hard categorical); for $\tau >0$, the samples are continuous and differentiable in $\log 𝝅$.

• Straight-through estimator. Use the discrete samples in the forward pass but pretend the rounding operation is the identity in the backward pass. This introduces bias but works surprisingly well in practice.

• RELAX and other hybrid estimators. Combine the score function estimator with a learned control variate based on a continuous relaxation, achieving lower variance than either approach alone.

Amortized Variational Inference

All the variational inference methods we have seen so far (CAVI, SVI, BBVI, reparameterised VI) share a common pattern: for each data point ${𝒙}_i$, we must optimise variational parameters ${𝝓}_i$ to approximate the posterior $p(𝒛|{𝒙}_i)$. When the dataset contains $N$ data points, this means solving $N$ separate optimisation problems, each requiring multiple gradient steps.

For a dataset of modest size, this is feasible. But what happens when $N={10}^6$ or $N={10}^8$? And what about new data points not seen during training; do we need to run optimisation from scratch for each one?

Amortized inference answers both questions with a single idea: instead of optimising ${𝝓}_i$ separately for each ${𝒙}_i$, train a shared function (typically a neural network) that maps any $𝒙$ directly to the variational parameters $𝝓$ in a single forward pass.

From Per-Instance to Amortized

For example, if $q(𝒛;𝝓)=𝒩(𝒛|𝝁,\mathrm{diag}({𝝈}^2))$, then the inference network outputs $f_{\psi }(𝒙)=(𝝁(𝒙),\log 𝝈(𝒙))$, a mean vector and a log-standard-deviation vector, both computed from $𝒙$ via a neural network with shared parameters $\psi$.

The key insight is one of cost amortisation:

Once $f_{\psi }$ is trained, computing the approximate posterior for a new data point ${𝒙}_{\text{new}}$ (never seen during training) requires only a single forward pass through the network, with no iterative optimisation. This is the “amortisation”: the cost of learning to do inference is paid once during training, and then inference for any input is essentially free.

The term “amortisation” comes from accounting: a large upfront cost (training the network) is spread over many future uses (inference for each data point). Just as a factory invests in machinery that produces many units cheaply, amortised inference invests in a neural network that produces many approximate posteriors cheaply.

The Amortised ELBO

The ELBO for amortized inference over a dataset $\{{𝒙}_1,\dots ,{𝒙}_N\}$ is: (ELBO)$$\begin{array}{ccc}& ℒ(\psi ,\theta )=\sum _{i=1}^N{𝔼}_{q_{\psi }(𝒛|{𝒙}_i)}[\log p_{\theta }({𝒙}_i,𝒛)-\log q_{\psi }(𝒛|{𝒙}_i)],& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\theta$ denotes the generative model parameters. We maximise $ℒ$ jointly over $\psi$ (inference network) and $\theta$ (generative model) using stochastic gradient ascent, with the reparameterisation trick (The Reparameterization Trick) to estimate the gradients through the sampling.

The Amortization Gap

Amortised inference trades optimality for speed. A single neural network cannot perfectly approximate the optimal variational parameters for every possible input, so it must generalise across inputs, and this generalisation introduces a systematic approximation error.

The amortization gap has several sources:

• Limited network capacity. The inference network $f_{\psi }$ may not be expressive enough to represent the mapping $𝒙\mapsto {𝝓}^{\ast }(𝒙)$ exactly. For example, if the optimal mapping is highly nonlinear but $f_{\psi }$ is a shallow network, it will approximate poorly in some regions of input space.

• Training data coverage. The network is trained on a finite dataset and must generalise to unseen inputs. In regions of input space poorly represented in the training data, the amortised parameters may be suboptimal.

• Multi-modality. If the true posterior is multi-modal for some inputs, the mean-field Gaussian output of a standard inference network cannot capture this structure, regardless of network capacity.

In practice, the amortization gap is typically small for well-designed models: less than a few nats per data point. It can be diagnosed by comparing the amortised ELBO against a locally refined ELBO (as in semi-amortised inference, below).

Semi-Amortized Inference

A natural compromise between per-instance optimisation and full amortisation is semi-amortized inference: use the amortised output $f_{\psi }(𝒙)$ as an initialisation for a few steps of local iterative refinement. This combines the speed of amortisation (a good initialisation in one forward pass) with the accuracy of per-instance optimisation (a few gradient steps close the amortization gap).

Concretely:

1. Compute ${𝝓}^{(0)}=f_{\psi }(𝒙)$ via the inference network.

2. For $j=1,\dots ,J$ (typically $J=5$–$20$): $$\begin{array}{c}{𝝓}^{(j)}={𝝓}^{(j-1)}+\alpha {\nabla }_{𝝓}ℒ(𝝓,\theta ;𝒙){|}_{𝝓={𝝓}^{(j-1)}}.\end{array}$$

3. Use ${𝝓}^{(J)}$ as the final variational parameters.

This can be seen as “test-time compute” for inference: spending more computation per data point to get a tighter ELBO.

Amortization and VAEs

The connection to Variational Autoencoders (ch:vae) is now immediate. In a VAE:

• The generative model $p_{\theta }(𝒙|𝒛)p(𝒛)$ is a deep latent variable model with a neural network decoder.

• The approximate posterior $q_{\psi }(𝒛|𝒙)$ is amortised via a neural network encoder (the inference network).

• The ELBO is optimised jointly over $\psi$ and $\theta$ using the reparameterisation trick from The Reparameterization Trick.

Thus the VAE encoder is precisely an amortised inference network: it implements $f_{\psi }(𝒙)=({𝝁}_{\psi }(𝒙),\log {𝝈}_{\psi }(𝒙))$, mapping each input $𝒙$ to the parameters of a Gaussian approximate posterior in a single forward pass.

Example: Amortized Inference for a Latent Variable Model

The Road to Deep Generative Models

Let us pause and survey the path we have traversed in this chapter. We began with the fundamental inference problem: computing the posterior $p(𝒛|𝒙)$. We formulated variational inference as an optimisation problem, minimising ${𝖣}_{\mathrm{KL}}(q\Vert p)$ over a tractable family, and derived the ELBO as the objective.

The journey from classical to modern VI can be summarised as a sequence of generalisations, each removing a limitation of the previous method:

The combination of amortised inference, the reparameterisation trick, and stochastic minibatch training gives us the VAE (ch:vae), the first practical framework for training deep latent variable models at scale. In the next sections, we push further: normalizing flows enrich the variational family beyond mean-field Gaussians, importance-weighted bounds tighten the ELBO, and natural gradients exploit the geometry of the variational distribution for faster convergence.

Normalizing Flows for Richer Posteriors

The variational families considered so far (mean-field factorizations, Mean-Field Variational Inference; exponential-family conditionals, Exponential Families and Conjugacy; and even full-covariance Gaussians) impose strong structural assumptions on the approximate posterior. A mean-field Gaussian, for instance, can never capture correlations between latent dimensions, while a full-covariance Gaussian cannot represent multimodality or skewness. When the true posterior $p(𝒛|𝒙)$ exhibits such features, these families introduce an irreducible approximation gap: the ELBO remains strictly below $\log p(𝒙)$ regardless of optimization quality.

Normalizing flows address this limitation by constructing rich, flexible variational families through the composition of simple invertible transformations. The core idea is deceptively simple: start with a tractable base distribution $q_0({𝒛}_0)$ (e.g., a standard Gaussian) and transform it through a sequence of invertible maps to obtain a complex distribution capable of matching intricate posterior geometries.

Change of Variables via Invertible Maps

Let $f:{\mathrm{ℝ}}^d\to {\mathrm{ℝ}}^d$ be a smooth, invertible (i.e., diffeomorphic) mapping. If ${𝒛}_0\sim q_0({𝒛}_0)$ and $𝒛=f({𝒛}_0)$, then the density of $𝒛$ is given by the change of variables formula: (COV)$$\begin{array}{ccc}& q(𝒛)=q_0(f^{-1}(𝒛))\left|\det \frac{\partial f^{-1}}{\partial 𝒛}\right|=q_0({𝒛}_0){\left|\det \frac{\partial f}{\partial {𝒛}_0}\right|}^{-1}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Taking logarithms yields the key identity: (LOG COV)$$\begin{array}{ccc}& \log q(𝒛)=\log q_0({𝒛}_0)-\log \left|\det \frac{\partial f}{\partial {𝒛}_0}\right|.& \text{<span class="tml-eqn"></span>}\end{array}$$ The Jacobian determinant $\left|\det (\partial f/\partial {𝒛}_0)\right|$ measures how $f$ locally stretches or compresses volume, adjusting the density accordingly.

Composing Flows

A single invertible transformation may not be expressive enough. The power of normalizing flows comes from composition: given a sequence of $K$ diffeomorphisms $f_1,f_2,\dots ,f_K$, define (FLOW Chain)$$\begin{array}{ccc}& {𝒛}_K=(f_K\circ f_{K-1}\circ \cdots \circ f_1)({𝒛}_0),{𝒛}_0\sim q_0({𝒛}_0).& \text{<span class="tml-eqn"></span>}\end{array}$$ By iterating the change-of-variables formula, the log-density of the final iterate ${𝒛}_K$ decomposes as: (FLOW LOG Density)$$\begin{array}{ccc}& \log q_K({𝒛}_K)=\log q_0({𝒛}_0)-\sum _{k=1}^K\log \left|\det \frac{\partial f_k}{\partial {𝒛}_{k-1}}\right|& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒛}_k=f_k({𝒛}_{k-1})$ for $k=1,\dots ,K$. Each transformation $f_k$ may have its own learnable parameters.

ELBO with Normalizing Flows

Substituting the flow density into the standard ELBO (The Evidence Lower Bound), we obtain: (FLOW ELBO)$$\begin{array}{cccc}& {ℒ}_{\text{flow}}& ={𝔼}_{q_0({𝒛}_0)}[\log p(𝒙,{𝒛}_K)-\log q_0({𝒛}_0)+\sum _{k=1}^K\log \left|\det \frac{\partial f_k}{\partial {𝒛}_{k-1}}\right|]& \\ & & =\underset{\text{joint likelihood}}{\underbrace{{𝔼}_{q_0}[\log p(𝒙,{𝒛}_K)]}}+\underset{\text{base entropy}}{\underbrace{𝖧[q_0]}}+\underset{\text{volume correction}}{\underbrace{{𝔼}_{q_0}\left[\sum _{k=1}^K\log \left|\det {𝑱}_k\right|\right]}},& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝑱}_k\stackrel{\mathrm{def}}{=}\partial f_k/\partial {𝒛}_{k-1}$ is the Jacobian of the $k$-th transformation. The first term encourages the transformed samples to land in high-probability regions under the model; the second and third terms together represent the entropy of the flow distribution, preventing collapse.

Since ${𝒛}_K$ is a deterministic function of ${𝒛}_0$, all expectations are over $q_0$: we sample ${𝒛}_0\sim q_0$, push it through the flow, and evaluate the objective. Gradients flow back through the entire chain via the reparameterization trick (The Reparameterization Trick).

Planar Flows

The simplest flow architecture, introduced by Rezende and Mohamed [23], uses planar transformations: (Planar)$$\begin{array}{ccc}& f(𝒛)=𝒛+𝒖h({𝒘}^{\top }𝒛+b),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝒖,𝒘\in {\mathrm{ℝ}}^d$, $b\in \mathrm{ℝ}$, and $h$ is a smooth activation function (typically $\tanh$). This transformation adds a perturbation in the direction $𝒖$ whose magnitude depends on the projection of $𝒛$ onto $𝒘$. Geometrically, it contracts or expands the density along hyperplanes $\{𝒛:{𝒘}^{\top }𝒛+b=\text{const}\}$.

The key computational advantage of planar flows is the efficient Jacobian determinant.

For the transformation to be invertible, we require $\det (\partial f/\partial 𝒛)\ne 0$, which imposes the constraint $h^{\prime }({𝒘}^{\top }𝒛+b){𝒖}^{\top }𝒘>-1$. When $h=\tanh$, this is enforced by reparameterizing $𝒖$ so that ${𝒖}^{\top }𝒘\ge -1+\log (1+e^{{𝒖}^{\top }𝒘})$ [23].

Radial Flows

While planar flows warp the density along hyperplanes, radial flows create contractions and expansions around a reference point ${𝒛}_0\in {\mathrm{ℝ}}^d$: (Radial)$$\begin{array}{ccc}& f(𝒛)=𝒛+\frac{\beta }{\alpha +r(𝒛)}(𝒛-{𝒛}_0),r(𝒛)={\Vert 𝒛-{𝒛}_0\Vert }_2,& \text{<span class="tml-eqn"></span>}\end{array}$$ where $\alpha >0$ and $\beta \in \mathrm{ℝ}$ are scalar parameters. The Jacobian determinant can likewise be computed in $O(d)$: (Radial DET)$$\begin{array}{ccc}& \det \frac{\partial f}{\partial 𝒛}={(1+\frac{\beta }{\alpha +r})}^{d-1}(1+\frac{\beta }{\alpha +r}-\frac{\beta r}{{(\alpha +r)}^2}).& \text{<span class="tml-eqn"></span>}\end{array}$$ Radial flows are naturally suited to creating or removing modes, since they can concentrate or disperse density radially.

Beyond Simple Flows

Planar and radial flows have limited per-layer expressiveness (each adds only a single degree of freedom to the density shape). In practice, many layers are needed for complex posteriors. More powerful architectures have been developed that achieve a better balance between expressiveness and computational cost:

1. RealNVP (Dinh et al. [12]): uses affine coupling layers with triangular Jacobians, yielding $O(d)$ determinant computation and exact inversion.

2. Masked Autoregressive Flows (MAF) (Papamakarios et al. [13]): autoregressive transformations with triangular Jacobians. Density evaluation is efficient but sampling requires sequential passes.

3. Inverse Autoregressive Flows (IAF) (Kingma et al. [14]): the inverse of MAF, enabling efficient sampling at the cost of slower density evaluation, which is ideal for variational inference where sampling is the bottleneck.

We defer a comprehensive treatment of these architectures to the normalizing flows chapter; here we focus on flows as a tool for enriching variational posteriors.

The proof relies on the fact that smooth diffeomorphisms are dense in the space of measure-preserving maps (in the appropriate topology) [24]. In practice, the number of layers $K$ required may be large, and optimization may not find the global optimum, so the approximation guarantee is primarily of theoretical interest.

Importance Weighted Bounds

An orthogonal strategy for tightening the ELBO does not change the variational family at all but instead uses multiple samples from $q$ to form a tighter lower bound on $\log p(𝒙)$. The importance weighted autoencoder (IWAE) bound, introduced by Burda et al. [7], achieves this through a simple application of importance sampling.

The IWAE Objective

Draw $k$ independent samples ${𝒛}^{(1)},\dots ,{𝒛}^{(k)}$ from $q(𝒛|𝒙)$ and define the importance weights: (IW Weights)$$\begin{array}{ccc}& w^{(\ell )}=\frac{p(𝒙,{𝒛}^{(\ell )})}{q({𝒛}^{(\ell )}|𝒙)},\ell =1,\dots ,k.& \text{<span class="tml-eqn"></span>}\end{array}$$ The importance weighted ELBO with $k$ samples is: (IWAE Bound)$$\begin{array}{ccc}& {ℒ}_k^{\mathrm{IWAE}}=𝔼\left[\log \frac{1}{k}\sum _{\ell =1}^k{w}^{(\ell )}\right]=𝔼\left[\log \frac{1}{k}\sum _{\ell =1}^k\frac{p(𝒙,{𝒛}^{(\ell )})}{q({𝒛}^{(\ell )}|𝒙)}\right]& \text{<span class="tml-eqn"></span>}\end{array}$$ where the expectation is over ${𝒛}^{(1)},\dots ,{𝒛}^{(k)}\stackrel{\mathrm{iid}}{\sim }q(𝒛|𝒙)$. Note that ${ℒ}_1^{\mathrm{IWAE}}$ recovers the standard ELBO.

The Tightness Hierarchy

The central theoretical result for the IWAE bound is that more importance samples always yield a tighter bound.

Bias Analysis

A Taylor expansion reveals the rate at which the IWAE bound approaches $\log p(𝒙)$: (IWAE BIAS)$$\begin{array}{ccc}& {ℒ}_k^{\mathrm{IWAE}}=\log p(𝒙)-\frac{1}{2k}\frac{{\mathrm{𝖵ar}}_q[w]}{({𝔼}_q[w]{)}^2}+O(k^{-2}).& \text{<span class="tml-eqn"></span>}\end{array}$$ Thus the gap closes at rate $O(1/k)$, and the gap itself is proportional to the relative variance of the importance weights. If $q$ is close to $p(𝒛|𝒙)$, the weights are nearly constant, the variance is small, and even moderate $k$ suffices.

The Signal-to-Noise Ratio Problem

While increasing $k$ tightens the bound, Rainforth et al. [25] discovered a surprising counterpoint: the gradient signal for the inference network (encoder) deteriorates as $k$ grows.

This result reveals a fundamental tension:

1. Tighter bound: larger $k$ reduces the gap to $\log p(𝒙)$, helping the generative model $p_{\theta }$;

2. Noisier encoder gradient: larger $k$ reduces the SNR for ${\nabla }_{\varphi }$, hurting the inference network $q_{\varphi }$.

Gradient Computation

The gradient of the IWAE objective with respect to both model parameters $\theta$ and inference parameters $\varphi$ can be written using self-normalized importance weights: (IWAE GRAD)$$\begin{array}{ccc}& \nabla {ℒ}_k=𝔼[\sum _{\ell =1}^k{\tilde{w}}^{(\ell )}\nabla \log w^{(\ell )}],\text{where }{\tilde{w}}^{(\ell )}=\frac{w^{(\ell )}}{{\sum }_{m=1}^k{w}^{(m)}}& \text{<span class="tml-eqn"></span>}\end{array}$$ are the normalized importance weights. This shows that the IWAE gradient is a weighted combination of per-sample ELBO gradients, with weights proportional to the importance ratios.

Tucker et al. [26] proposed the doubly reparameterized gradient estimator, which modifies the gradient computation to reduce variance without increasing bias. The key insight is to stop gradients through the normalized weights ${\tilde{w}}^{(\ell )}$ while still obtaining an unbiased estimator, achieving better SNR properties for the encoder.

Stein Variational Gradient Descent

All methods discussed so far represent the approximate posterior through a parametric distribution $q_{\lambda }(𝒛)$. An entirely different philosophy is to represent the posterior as a collection of particles ${\{{𝒛}_i\}}_{i=1}^n$ that are iteratively transported toward high-probability regions of $p(𝒛|𝒙)$ while maintaining diversity. Stein Variational Gradient Descent (SVGD), introduced by Liu and Wang [6], achieves this by combining ideas from Stein's method, kernel methods, and variational inference into an elegant, deterministic update rule.

Stein's Identity and the Stein Operator

The foundation of SVGD is Stein's identity, which characterizes a distribution $p$ through an integral condition involving its score function ${\nabla }_{𝒛}\log p(𝒛)$.

The crucial observation is that Stein's identity holds for $𝝓$ drawn from any sufficiently regular function class. If ${𝔼}_q[\mathrm{trace}({𝒜}_p𝝓)]\ne 0$ for some $𝝓$, then $q\ne p$. This motivates using the maximum violation as a discrepancy measure.

Kernelized Stein Discrepancy

To obtain a tractable discrepancy, we restrict $𝝓$ to lie in a reproducing kernel Hilbert space (RKHS) $\mathscr{H}$ with kernel $k:{\mathrm{ℝ}}^d\times {\mathrm{ℝ}}^d\to \mathrm{ℝ}$.

By the representer theorem in RKHS, the optimal test function admits a closed-form solution: (KSD Optimal)$$\begin{array}{ccc}& {𝝓}^{\ast }(\cdot )\propto {𝔼}_q[k(𝒛,\cdot ){\nabla }_{𝒛}\log p(𝒛)+{\nabla }_{𝒛}k(𝒛,\cdot )].& \text{<span class="tml-eqn"></span>}\end{array}$$ This function ${𝝓}^{\ast }$ represents the direction of steepest descent for the KL divergence ${𝖣}_{\mathrm{KL}}(q\Vert p)$ in the space of distributions, as measured by the RKHS norm.

The SVGD Update Rule

Given $n$ particles ${\{{𝒛}_i\}}_{i=1}^n$ representing the current approximation to $p(𝒛|𝒙)$, SVGD updates them simultaneously: (SVGD Update)$$\begin{array}{ccc}& {𝒛}_i\leftarrow {𝒛}_i+ϵ{\widehat{𝝓}}^{\ast }({𝒛}_i),{\widehat{𝝓}}^{\ast }({𝒛}_i)=\frac{1}{n}\sum _{j=1}^n[k({𝒛}_j,{𝒛}_i){\nabla }_{{𝒛}_j}\log p({𝒛}_j|𝒙)+{\nabla }_{{𝒛}_j}k({𝒛}_j,{𝒛}_i)]& \text{<span class="tml-eqn"></span>}\end{array}$$ where $ϵ>0$ is the step size. The update has a beautiful interpretation as the sum of two forces:

1. Attractive force: $\frac{1}{n}{\sum }_{j}k({𝒛}_j,{𝒛}_i){\nabla }_{{𝒛}_j}\log p({𝒛}_j|𝒙)$ drives each particle toward high-probability regions, weighted by kernel similarity. When $k({𝒛}_j,{𝒛}_i)$ is large (particles are close), particle $i$ is strongly influenced by the gradient at particle $j$.

2. Repulsive force: $\frac{1}{n}{\sum }_j{\nabla }_{{𝒛}_j}k({𝒛}_j,{𝒛}_i)$ pushes particles apart, preventing them from collapsing onto a single mode. For an RBF kernel $k(𝒛,{𝒛}^{\prime })=\exp (-{\Vert 𝒛-{𝒛}^{\prime }\Vert }^2/(2h^2))$, this term pushes ${𝒛}_i$ away from nearby particles.

Theoretical Foundation

The SVGD update can be understood as a discretization of a continuous gradient flow on the space of probability distributions.

This theorem connects SVGD to the theory of Wasserstein gradient flows, providing a rigorous justification for the particle updates. In practice, the finite-particle approximation introduces bias, and the quality of the approximation depends on both the number of particles and the choice of kernel.

Comparison with MCMC

SVGD occupies an interesting middle ground between variational inference and Markov chain Monte Carlo (MCMC):

$({}^{\ast })$In the infinite-particle limit. Finite-particle SVGD has variance collapse issues.

Natural Gradient Variational Inference

Standard gradient-based optimization of the ELBO treats the variational parameters $𝝀$ as living in Euclidean space: the steepest-ascent direction is simply ${\nabla }_{𝝀}ℒ$. But $𝝀$ parameterizes a probability distribution $q_{𝝀}$, and Euclidean geometry is a poor match for the space of distributions. Two parameter vectors $𝝀$ and $𝝀+𝜹$ that are close in Euclidean norm may correspond to distributions that are far apart in KL divergence, and vice versa.

The natural gradient [27] resolves this mismatch by equipping the parameter space with the Fisher information metric, the unique Riemannian metric that respects the geometry of probability distributions.

The Fisher Information Matrix

The Fisher matrix $𝑭$ is positive semidefinite and, for regular models, positive definite. It defines a local inner product ${⟨𝒖,𝒗⟩}_{𝑭}={𝒖}^{\top }𝑭𝒗$ on the parameter space, making it a Riemannian manifold.

For exponential families (Exponential Families and Conjugacy) with natural parameters $𝜼$, the Fisher matrix has a particularly elegant form: (Fisher Expfam)$$\begin{array}{ccc}& 𝑭(𝜼)={\nabla }^{2}A(𝜼),& \text{<span class="tml-eqn"></span>}\end{array}$$ the Hessian of the log-partition function. Since $A$ is convex, $𝑭$ is always positive semidefinite, confirming its role as a metric.

The Natural Gradient

The natural gradient is the solution to the constrained optimization problem: find the direction $𝜹$ that maximizes the linear approximation $\nabla ℒ\cdot 𝜹$ subject to a fixed step in KL divergence: (Natgrad Constrained)$$\begin{array}{ccc}& \tilde{\nabla }ℒ={\mathrm{arg\hspace{0.17em}max}}_{𝜹}\nabla ℒ\cdot 𝜹\text{s.t.}{𝜹}^{\top }𝑭𝜹\le {ϵ}^2.& \text{<span class="tml-eqn"></span>}\end{array}$$ This ensures that each step moves a fixed “distance” in the space of distributions, not in the space of parameters.

Natural Gradient for Exponential Families

For variational families in the exponential family, the natural gradient has a remarkably simple form. Let $𝒎={𝔼}_{q_{𝜼}}[T(𝒛)]=\nabla A(𝜼)$ denote the mean parameters. Since $𝑭(𝜼)={\nabla }^{2}A(𝜼)=\partial 𝒎/\partial 𝜼$, the natural gradient update in natural parameters is: (Natgrad Expfam)$$\begin{array}{ccc}& {𝜼}_{t+1}={𝜼}_t+{\rho }_t𝑭{({𝜼}_t)}^{-1}{\nabla }_{𝜼}ℒ={𝜼}_t+{\rho }_t{\nabla }_{𝒎}ℒ,& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${\nabla }_{𝒎}ℒ$ is the gradient with respect to the mean parameters. This is the natural gradient expressed purely in terms of mean-parameter gradients, with no explicit matrix inversion required.

The Bayesian Learning Rule

For conjugate exponential-family models with prior $p(𝒛)\propto \exp ({𝜼}_0^{\top }T(𝒛))$ and likelihood contributions $p({𝒙}_i|𝒛)\propto \exp ({𝜼}_i{(𝒛)}^{\top }T(𝒛))$, the natural gradient update takes the form of the Bayesian learning rule [28]: (BLR)$$\begin{array}{ccc}& {𝜼}_{t+1}=(1-{\rho }_t){𝜼}_t+{\rho }_t({𝜼}_0+N\cdot {𝔼}_{q_t}[{\nabla }_{𝒎}\log p({𝒙}_i|𝒛)])& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒙}_i$ is a randomly selected data point and $N$ is the dataset size. This has a beautiful interpretation: at each step, the natural parameters are a weighted average of the current estimate and a target formed by combining the prior natural parameters with a scaled stochastic estimate of the data contribution.

Practical Methods

Computing the full Fisher matrix is $O(p^2)$ in the number of parameters $p$, which is prohibitive for large models. Several practical approximations have been developed:

1. Diagonal Fisher: Approximate $𝑭\approx \mathrm{diag}(𝑭)$, reducing storage and inversion to $O(p)$. This is equivalent to using per-parameter adaptive step sizes.

2. KFAC (Kronecker-factored approximate curvature): For neural networks, approximate the Fisher as a Kronecker product of smaller matrices, one per layer.

3. VOGN (Variational Online Gauss-Newton): A natural-gradient method for Bayesian neural networks that maintains a diagonal Gaussian posterior $q(𝒘)=𝒩(𝝁,\mathrm{diag}({𝝈}^2))$ over weights [15].

The Bethe Free Energy and Belief Propagation

We now step back from gradient-based optimization and examine variational inference from a broader, unifying perspective. The variational principle for inference, developed systematically by Wainwright and Jordan [29], reveals that essentially all approximate inference methods, including mean field, belief propagation, expectation propagation, and more, can be understood as different approximations to a single optimization problem. This perspective provides deep insight into why methods work (or fail) and suggests principled ways to design new ones.

The Exact Variational Principle

Consider an undirected graphical model (Markov random field) with potential functions ${\psi }_s({𝒛}_s)$ for nodes $s$ and ${\psi }_{st}({𝒛}_s,{𝒛}_t)$ for edges $(s,t)\in E$ of a graph $G=(V,E)$. The joint distribution is: (MRF)$$\begin{array}{ccc}& p(𝒛;𝜽)=\frac{1}{Z(𝜽)}\exp ({𝜽}^{\top }T(𝒛)),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝜽$ collects all natural parameters and $T(𝒛)$ collects the sufficient statistics. The log-partition function $A(𝜽)=\log Z(𝜽)$ is the central quantity of interest, since marginals and expectations follow from its derivatives.

The exact variational principle expresses $A$ as a convex optimization problem via conjugate duality: (Exact Variational)$$\begin{array}{ccc}& A(𝜽)=\underset{𝝁\in ℳ}{\sup }\{{𝜽}^{\top }𝝁+𝖧(𝝁)\}& \text{<span class="tml-eqn"></span>}\end{array}$$ where $ℳ$ is the marginal polytope (the set of all valid marginal vectors realizable by some distribution) and $𝖧(𝝁)$ is the entropy written as a function of the marginals. The supremum is attained at ${𝝁}^{\ast }={𝔼}_p[T(𝒛)]$, the exact marginals.

This formulation is exact but intractable in general: (i) the marginal polytope $ℳ$ has exponentially many constraints, and (ii) the entropy $𝖧(𝝁)$ has no closed-form expression in terms of marginals for general graphs.

Mean Field as Inner Approximation

Mean-field VI (Mean-Field Variational Inference) can be recast in this framework as making two approximations:

1. Inner bound on $ℳ$: restrict to product distributions, which gives a subset ${ℳ}_{\mathrm{MF}}\subseteq ℳ$. Since we optimize over a smaller set, we obtain a lower bound on $A(𝜽)$.

2. Exact entropy: the entropy of a product distribution factorizes as ${𝖧}_{\mathrm{MF}}={\sum }_s𝖧({\tau }_s)$, which is exact for the product family.

The mean-field approximation is thus: (MF Variational)$$\begin{array}{ccc}& A_{\mathrm{MF}}(𝜽)=\underset{𝝉\in {ℳ}_{\mathrm{MF}}}{\sup }\{{𝜽}^{\top }𝝉+\sum _{s\in V}𝖧({\tau }_s)\}\le A(𝜽).& \text{<span class="tml-eqn"></span>}\end{array}$$ The bound is tight only when the true distribution is itself a product (i.e., the graph has no edges).

Belief Propagation and the Bethe Approximation

Belief propagation (BP) takes the opposite approach to mean field, making approximations in a different direction:

1. Outer relaxation of $ℳ$: replace the marginal polytope with the local polytope: (Local Polytope)$$\begin{array}{ccc}& ℒ(G)=\{𝝉\ge 0|{\sum }_{z_s}{\tau }_s(z_s)=1,{\sum }_{z_s}{\tau }_{st}(z_s,z_t)={\tau }_t(z_t)\forall (s,t)\in E\},& \text{<span class="tml-eqn"></span>}\end{array}$$ which enforces only pairwise marginalization consistency. We always have $ℳ(G)\subseteq ℒ(G)$; equality holds if and only if $G$ is a tree.

2. Bethe entropy approximation: replace the exact entropy with the Bethe entropy: (Bethe Entropy)$$\begin{array}{ccc}& {𝖧}_{\mathrm{Bethe}}(𝝉)=\sum _{s\in V}𝖧({\tau }_s)-\sum _{(s,t)\in E}𝖨({\tau }_{st}),& \text{<span class="tml-eqn"></span>}\end{array}$$ where $𝖧({\tau }_s)$ is the marginal entropy of node $s$ and $𝖨({\tau }_{st})=𝖧({\tau }_s)+𝖧({\tau }_t)-𝖧({\tau }_{st})$ is the mutual information of edge $(s,t)$.

Unlike the mean-field bound, the Bethe approximation is neither an upper nor a lower bound on $A(𝜽)$ in general. However, it can be much more accurate than mean field, especially for graphs with weak interactions or high connectivity.

BP Fixed Points and the Bethe Variational Problem

The connection between message-passing algorithms and variational methods is made precise by the following foundational result.

Exactness on Trees

This explains the remarkable success of BP on tree-structured models [16] and provides insight into why “loopy” BP (applied to graphs with cycles) often works well in practice: many real-world graphical models have locally tree-like structure.

Beyond Bethe: Region-Based Methods

The Bethe approximation uses pairwise interactions. The Kikuchi hierarchy [30] generalizes this to larger regions (clusters of variables), leading to tighter entropy approximations at higher computational cost. The corresponding message-passing algorithms are called generalized belief propagation (GBP). The $k$-th level of the Kikuchi hierarchy uses regions of up to $k$ variables and recovers the exact entropy when $k$ equals the treewidth of the graph.

Variational Inference in Modern Generative Models

Variational inference is not merely a classical technique for approximate Bayesian computation; it is the mathematical engine powering the most successful generative models of the modern era. Variational autoencoders, diffusion models, and even aspects of large language model alignment all rest on variational principles developed in this chapter. We now trace these connections explicitly, demonstrating how the abstract machinery of the ELBO, reparameterization, and amortized inference manifests in concrete generative architectures.

Variational Autoencoders

The variational autoencoder (VAE, ch:vae) is perhaps the most direct application of variational inference. A deep latent variable model specifies: (VAE Model)$$\begin{array}{ccc}& p_{\theta }(𝒙,𝒛)=p(𝒛)p_{\theta }(𝒙|𝒛),p(𝒛)=𝒩(0,𝑰),p_{\theta }(𝒙|𝒛)=𝒩({𝝁}_{\theta }(𝒛),{𝝈}_{\theta }^2(𝒛)),& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝝁}_{\theta }$ and ${𝝈}_{\theta }^2$ are neural networks. The true posterior $p_{\theta }(𝒛|𝒙)$ is intractable, so we introduce an amortized variational distribution: (VAE Encoder)$$\begin{array}{ccc}& q_{\varphi }(𝒛|𝒙)=𝒩({𝝁}_{\varphi }(𝒙),\mathrm{diag}({𝝈}_{\varphi }^2(𝒙))).& \text{<span class="tml-eqn"></span>}\end{array}$$

The VAE training objective is precisely the ELBO: (VAE ELBO)$$\begin{array}{ccc}& ℒ(\theta ,\varphi ;𝒙)={𝔼}_{q_{\varphi }(𝒛|𝒙)}[\log p_{\theta }(𝒙|𝒛)]-{𝖣}_{\mathrm{KL}}(q_{\varphi }(𝒛|𝒙)\Vert p(𝒛)).& \text{<span class="tml-eqn"></span>}\end{array}$$ The first term is the expected reconstruction quality; the second is the KL regularizer pushing the approximate posterior toward the prior (see sec:vae:elbo for a detailed derivation).

Key VI innovations in VAEs.

1. Amortized inference: Rather than optimizing separate variational parameters ${𝝀}_i$ for each data point ${𝒙}_i$, the encoder $q_{\varphi }(𝒛|𝒙)$ shares parameters $\varphi$ across all data points. This introduces an amortization gap: the amortized posterior may be suboptimal for any individual ${𝒙}_i$, but this enables $O(1)$ inference at test time.

2. Reparameterization: The reparameterization trick (The Reparameterization Trick) makes the ELBO differentiable with respect to $\varphi$, enabling joint optimization of encoder and decoder.

3. $\beta$-VAE: The modified objective ${ℒ}_{\beta }={𝔼}_q[\log p_{\theta }(𝒙|𝒛)]-\beta {𝖣}_{\mathrm{KL}}(q_{\varphi }\Vert p)$ with $\beta >1$ encourages disentangled representations at the cost of reconstruction quality.

4. VQ-VAE: Replaces the continuous latent space with discrete codes using vector quantization and the straight-through estimator, sidestepping the KL term entirely.

Diffusion Models as Hierarchical VAEs

Diffusion models (ch:diffusion) can be viewed as hierarchical VAEs with $T$ layers of latent variables [31]. The forward (noising) process defines the variational distribution: (DIFF Forward)$$\begin{array}{ccc}& q({𝒙}_{1:T}|{𝒙}_0)=\prod _{t=1}^{T}q({𝒙}_t|{𝒙}_{t-1}),q({𝒙}_t|{𝒙}_{t-1})=𝒩(\sqrt{1-{\beta }_t}{𝒙}_{t-1},{\beta }_t𝑰),& \text{<span class="tml-eqn"></span>}\end{array}$$ while the reverse (generative) process is: (DIFF Reverse)$$\begin{array}{ccc}& p_{\theta }({𝒙}_{0:T})=p({𝒙}_T)\prod _{t=1}^T{p}_{\theta }({𝒙}_{t-1}|{𝒙}_t),p({𝒙}_T)=𝒩(0,𝑰).& \text{<span class="tml-eqn"></span>}\end{array}$$

The variational lower bound (sec:diff:vdm:elbo) decomposes as: (DIFF VLB)$$\begin{array}{ccc}& \log p_{\theta }({𝒙}_0)\ge -\underset{\text{prior matching}}{\underbrace{{𝖣}_{\mathrm{KL}}(q({𝒙}_T|{𝒙}_0)\Vert p({𝒙}_T))}}-\sum _{t=2}^T\underset{\text{consistency terms}}{\underbrace{{𝖣}_{\mathrm{KL}}(q({𝒙}_{t-1}|{𝒙}_t,{𝒙}_0)\Vert p_{\theta }({𝒙}_{t-1}|{𝒙}_t))}}+\underset{\text{reconstruction}}{\underbrace{{𝔼}_q[\log p_{\theta }({𝒙}_0|{𝒙}_1)]}}.& \text{<span class="tml-eqn"></span>}\end{array}$$ Crucially, each consistency term is a KL divergence between two Gaussians: the true posterior $q({𝒙}_{t-1}|{𝒙}_t,{𝒙}_0)$ (available in closed form via Bayes' rule, see thm:diff:elbo-first) and the learned reverse $p_{\theta }({𝒙}_{t-1}|{𝒙}_t)$, making each term analytically tractable.

Bayesian Neural Networks

In a Bayesian neural network, we place a prior $p(𝑾)$ over the weights and seek the posterior $p(𝑾|𝒟)$. Variational inference approximates this with $q_{\varphi }(𝑾)$, typically a factorized Gaussian: (BNN)$$\begin{array}{ccc}& q_{\varphi }(𝑾)=\prod _{l=1}^L\prod _{ij}𝒩(W_{ij}^{(l)};{\mu }_{ij}^{(l)},{({\sigma }_{ij}^{(l)})}^2).& \text{<span class="tml-eqn"></span>}\end{array}$$

Bayes by Backprop (Blundell et al. [32]) optimizes the ELBO: (Bayes Backprop)$$\begin{array}{ccc}& ℒ(\varphi )={𝔼}_{q_{\varphi }(𝑾)}[\log p(𝒟|𝑾)]-{𝖣}_{\mathrm{KL}}(q_{\varphi }(𝑾)\Vert p(𝑾)),& \text{<span class="tml-eqn"></span>}\end{array}$$ using the reparameterization trick: $W_{ij}={\mu }_{ij}+{\sigma }_{ij}\cdot {ϵ}_{ij}$ with ${ϵ}_{ij}\sim 𝒩(0,1)$.

MC-Dropout as implicit VI. Gal and Ghahramani [33] showed that applying dropout at test time (MC-Dropout) is equivalent to approximate variational inference with a specific variational distribution: each weight is either present (with its trained value) or absent (zeroed out), with Bernoulli probability determined by the dropout rate. This provides a practical uncertainty estimate requiring no architectural changes, only keeping dropout active during prediction and averaging over multiple forward passes.

VI and Large Language Models

Variational principles appear in modern large language models in several contexts:

1. KL regularization in RLHF: Reinforcement learning from human feedback [18] maximizes reward subject to a KL penalty ${𝖣}_{\mathrm{KL}}({\pi }_{\theta }\Vert {\pi }_{\mathrm{ref}})$ that prevents the policy ${\pi }_{\theta }$ from deviating too far from a reference model ${\pi }_{\mathrm{ref}}$. This is structurally identical to the ELBO, with the reward playing the role of the log-likelihood and the KL penalty playing the role of the regularizer.

2. Latent reasoning models: Some approaches to improving LLM reasoning introduce latent “thought” variables $𝒛$ and optimize a variational bound on $\log p(𝒚|𝒙)\ge {𝔼}_{q(𝒛|𝒙,𝒚)}[\log p(𝒚|𝒙,𝒛)]-{𝖣}_{\mathrm{KL}}(q\Vert p(𝒛|𝒙))$, using variational EM to alternately improve the thought generator and the output model.

The Generative Model Family Tree

tab:vi:modern-summary summarizes how VI concepts manifest across modern generative model families.

Model	Variational family	ELBO variant	Key innovation
GMM (ch:gmm)	Categorical + Gaussians	Standard ELBO	EM algorithm
VAE (ch:vae)	Amortized Gaussian	Standard ELBO	Reparameterization
$\beta$-VAE	Amortized Gaussian	Weighted KL	Disentanglement
VQ-VAE	Discrete codes	Commitment loss	Straight-through
Diffusion (ch:diffusion)	Fixed forward process	Hierarchical VLB	Noise prediction
BNN	Factorized Gaussian	Standard ELBO	Bayes by Backprop
LLM (RLHF)	Policy ${\pi }_{\theta }$	Reward$-$KL	KL regularization

Practical Guide and Open Problems

We conclude the chapter with practical guidance for applying variational inference and a survey of open research directions.

Choosing the Right VI Method

The choice of VI algorithm depends on the model structure, dataset size, and desired posterior quality. fig:vi:decision-flowchart provides a decision guide.

Diagnostics

Monitoring VI convergence and quality requires attention to several indicators:

1. ELBO monitoring: Plot the ELBO across iterations. A plateau indicates convergence (but possibly to a local optimum). Sudden drops suggest numerical issues.

2. Posterior predictive checks: Sample $𝒛\sim q(𝒛)$, generate synthetic data ${𝒙}^{\ast }\sim p({𝒙}^{\ast }|𝒛)$, and compare distributional properties with the observed data. Systematic discrepancies reveal model or approximation inadequacies.

3. Comparison to MCMC on small problems: On low-dimensional subproblems or subsets of data, run MCMC to obtain ground-truth posterior summaries. Compare VI marginals, means, and variances to MCMC estimates. Large discrepancies indicate the variational family is too restrictive.

4. Gradient variance monitoring: For BBVI (Black-Box Variational Inference), monitor the variance of gradient estimates. Excessively high variance indicates the need for better control variates or larger mini-batches.

Common Pitfalls

1. Mode collapse in mean-field: Mean-field VI minimizes the reverse KL ${𝖣}_{\mathrm{KL}}(q\Vert p)$, which is mode-seeking: the optimizer concentrates $q$ on a single mode of a multimodal posterior, ignoring others entirely.

2. Poor initialization: CAVI and gradient-based methods can converge to poor local optima. Multiple random restarts or initialization from simpler models (e.g., initializing from the prior) can mitigate this.

3. Posterior collapse in VAEs: When the decoder is too powerful, the model learns to ignore the latent variables entirely, with $q_{\varphi }(𝒛|𝒙)\approx p(𝒛)$ for all $𝒙$ (see sec:vae:collapse).

4. Gradient variance in BBVI: The score-function gradient estimator (Black-Box Variational Inference) can have extremely high variance, making optimization impractical without careful variance reduction via control variates and baselines.

5. Underestimated uncertainty: Variational posteriors, especially mean-field ones, systematically underestimate posterior variance. This can lead to overconfident predictions in downstream tasks.

Open Research Directions

Despite decades of development, variational inference remains an active area of research with several important open problems:

1. Tighter bounds without SNR degradation: The IWAE bound (Importance Weighted Bounds) tightens monotonically in $k$ but degrades encoder gradients. Designing bounds that are simultaneously tight and provide high-quality gradients for all model components remains open.

2. Better structured variational families: Beyond mean-field and normalizing flows, what are the “right” variational families that balance expressiveness, computational cost, and optimization difficulty? Structured mean-field approaches that exploit graphical model structure are promising but underexplored for deep models.

3. VI for discrete structures: Most VI methods assume continuous latent variables. Extending reparameterization-style tricks to discrete latent variables (beyond Gumbel-Softmax) remains challenging, with implications for program synthesis, combinatorial optimization, and symbolic reasoning.

4. Scaling VI to trillion-parameter models: As models grow to hundreds of billions or trillions of parameters, maintaining even a diagonal Gaussian posterior requires doubling the parameter count. Efficient representations of uncertainty at this scale are needed.

5. Connections to optimal transport: Wasserstein distances and optimal transport maps offer alternative geometries for measuring the discrepancy between $q$ and $p$. Replacing KL divergence with Wasserstein-type objectives leads to different algorithms with potentially better properties for multimodal posteriors.

Exercises

The exercises above span the full range of variational inference topics covered in this chapter. Exercises 1–5 test foundational understanding of KL divergence and the ELBO\@; Exercises 6–9 develop fluency with mean-field and CAVI\@; Exercises 10–14 cover exponential families, SVI, and stochastic optimization; Exercises 15–19 address gradient estimation and the reparameterization trick; Exercises 20–21 treat normalizing flows and importance weighting; Exercises 22–24 develop SVGD and natural gradient methods; and Exercises 25–30 connect VI to graphical models, diffusion models, and modern practice.