Variational inference

Variational inference approximates a posterior distribution by minimizing a Kullback-Leibler divergence over a chosen family of probability distributions. For an unnormalized posterior density $h$ and a trial probability density function $q$, minimizing ${\mathbb{E}}_q[\log q-\log h]$ is equivalent to minimizing $\mathrm{KL}(q\Vert h/\int h)$.

The variational approximation gap is the part of the evidence lower bound deficit that remains even after optimizing separately for an observation within a specified variational family. With the generative model fixed and finite $\log p_{\theta }(x)$, defineThe equality uses the usual evidence lower bound identity for variational inference. A supremum allows for a best value that is not attained. The amortization gap separately measures failure of the shared inference rule to reach this value.^[ref]

An original example shows what independence can cost. Stipulate a centred two-dimensional multivariate normal distribution as the posterior distribution, with covariance matrixRestrict the variational family to centred independent normal distributions, $q=N(0,\mathrm{diag}(v_1,v_2))$ with $v_1,v_2>0$. This is a simple mean-field variational inference restriction. Writing $s=1-{\rho }^2$, the Kullback-Leibler divergence between normal distributions becomesEach variance contributes a term $v_i/s-\log v_i$, minimized at $v_i=s$. ConsequentlyFor $\rho =0.6$, the best independent variances are both $0.64$, and the unavoidable deficit is about $0.2231$ nat. Matching both marginal variances to one instead would give about $0.3394$ nat. The reverse Kullback-Leibler divergence therefore chooses narrower independent marginals here, while still failing to represent the correlation.

Enlarging the family to include the full covariance matrix permits $q=p$ and reduces this example's gap to zero. More generally, inclusion ${\mathcal{Q}}_1\subseteq {\mathcal{Q}}_2$ cannot increase the optimum approximation gap at a fixed model. Merely training the same shared rule longer cannot remove a strictly positive family restriction. Numerical improvements are not proofs that the infimum has been reached.

A normalizing flow constructs a probability distribution by passing a simple reference random variable through invertible differentiable transformations. The change of variables formula supplies its probability density function: if $Z=f(U)$ and $f$ is a diffeomorphism between equal-dimensional spaces, thenThe Jacobian determinant accounts for local expansion or contraction of volume. Composing several transformations gives a flexible distribution while retaining a computable density, useful for a richer approximate posterior distribution in variational inference.^[ref]

For an original nonlinear example, let $U$ have the standard normal distribution with density $\phi$, and set $f(u)=u+u^3$. Since $f^{\prime }(u)=1+3u^2>0$ and $f(u)$ tends to opposite infinities at the two ends of the real line, every output has exactly one preimage. Its probability density function satisfiesAt $u=1$, the output is $z=2$ and $q_Z(2)=\phi (1)/4\simeq 0.06049$. The factor $1/4$ reflects how strongly the transformation stretches nearby inputs. Integrating after the substitution $z=f(u)$ cancels this factor against $dz=f^{\prime }(u)du$, leaving total probability one.

For $z_k=f_k(z_{k-1})$, the chain rule givesGenerating transformed samples alone does not evaluate their probability density function; evaluating the determinant factors is essential. The single-preimage formula also requires invertibility: the map $u\mapsto u^2$ folds two real inputs together, so it cannot be substituted into this formula unchanged.

Amortized inference learns one shared rule that maps observations to approximate posterior distributions, spreading the work of fitting that rule across many inputs. Instead of separately optimizing variational parameters for every observation $x$, a recognition network returns the parameters of $q_{\varphi }(z\mid x)$. The shared parameters $\varphi$ are learned across data; evaluating the rule for a new input is then a different operation from fitting it.^[ref]

A scalar Gaussian model gives an exact example. Let $z\sim N(0,1)$ and $x\mid z\sim N(z,1)$. Completing the square givesThe shared rule ${\mu }_{\varphi }(x)=ax+b$, with $a=1/2$, $b=0$ and fixed variance $1/2$, represents the exact posterior for every real $x$. Evaluating $x/2$ replaces a separate optimization for each observation in this stipulated model. It does not imply that every nonlinear latent-variable model has an exact rule this simple.

Sharing can introduce an additional restriction even when the chosen distribution family is adequate. For observations $x=-2$ and $x=2$, the exact posterior means are $-1$ and $1$. The Gaussian family $N(m,1/2)$ can represent each exactly when $m$ is fitted separately. If the shared rule is constrained to output one constant mean for both inputs, its best equal-weight choice is $m=0$. Each case then hasThis original example isolates an amortization gap: the loss from restricting how observation-specific variational parameters are produced, rather than an inability of the per-observation Gaussian family to express the target. The shared mapping's capacity and its training therefore matter alongside the family of probability distributions.

The amortization gap measures how much evidence lower bound is lost by using the distribution returned by amortized inference instead of the best distribution available for that observation in the same variational family. Hold the generative parameters $\theta$ fixed. Write ${\mathcal{L}}_x(q)$ for the evidence lower bound and suppose the optimum $q_x^{\ast }={\mathrm{argmax}}_{q\in \mathcal{Q}}{\mathcal{L}}_x(q)$ is attained. For the shared rule $q_{\varphi }(\cdot \mid x)\in \mathcal{Q}$,The remaining variational approximation gap is $G_{\mathrm{approx}}(x)=\log p_{\theta }(x)-{\mathcal{L}}_x(q_x^{\ast })$, so the total loss isThus a restrictive recognition network or incomplete training can leave an amortization gap even when the variational family can represent the exact posterior distribution.^[ref]

For an original example, stipulate $p(z\mid x)=N(2x,1)$ for equally weighted observations $x=-1$ and $x=1$. The variational family $\mathcal{Q}=\{N(m,1):m\in \mathbb{R}\}$ represents both posterior distributions exactly when fitted separately, so $G_{\mathrm{approx}}=0$. Restrict the shared rule to $q_b(z\mid x)=N(x+b,1)$: it can adjust the intercept but its slope is fixed at one. The Kullback-Leibler divergence between these unit-variance normal distributions is $(b-x{)}^2/2$. Averaging the amortization gap over the two observations givesThe best shared intercept $b=0$ still loses $1/2$ nat on average. A current rule with $b=1$ loses one nat: half a nat can be removed by training this rule better, while the remaining half requires changing its fixed-slope restriction. This separates the variational family's expressive power from the shared mapping's capacity and current training state.

In a numerical comparison, a locally improved per-observation distribution need not be the global optimum $q_x^{\ast }$. Its improvement supplies a lower bound on the true amortization gap; failing to find an improvement does not establish a zero gap. Comparisons must keep both $\theta$ and $\mathcal{Q}$ fixed.

Under fixed-support regularity permitting differentiation under the integral sign, apply the product rule to $\int s_{\lambda }{W}_{\lambda }$. The first derivative term is $\int W_{\lambda }\nabla s_{\lambda }=\mathbb{E}[W_{\lambda }\nabla \log s_{\lambda }]$ and the second is $\mathbb{E}[\nabla W_{\lambda }]$. This proves the displayed identity. For an evidence lower bound, take $W_{\lambda }=\log p(t_{\lambda }(\epsilon ),x)-\log q_{\lambda }(t_{\lambda }(\epsilon ))$, holding noise $\epsilon$ fixed in the inner derivatives. Parameter-independent continuous noise contributes a reparameterization gradient; untransformed discrete choices contribute a score-function gradient estimator. A mixed base law cannot be called parameter independent if its discrete factors retain variational parameters.^[ref]

A guide program is a probabilistic program whose generated probability distribution approximates a model's posterior distribution. Its adjustable parameters are fitted by variational inference. The guide must assign probability only where the target allows it when minimizing $D_{\mathrm{KL}}(q\Vert p)$, but a tractable guide can still miss important target regions.

A reparameterization gradient differentiates an expected value by representing the simulated random variable as a differentiable transformation of parameter-independent noise. For $X_{\eta }=g_{\eta }(\epsilon )$, the identity ${\nabla }_{\eta }\mathbb{E}[F_{\eta }(X_{\eta })]=\mathbb{E}[{\nabla }_{\eta }{F}_{\eta }(g_{\eta }(\epsilon ))]$ requires justified differentiation under the integral sign. For a normal distribution, use $X=m+e^{\rho }\epsilon$ with $\epsilon$ drawn from the standard normal distribution.

Gaussian reparameterization represents a random variable with a normal distribution using a location, a scale and parameter-independent noise:Here $\epsilon$ has the standard normal distribution. The resulting $Z$ has expected value $\mu$ and variance ${\sigma }^2$. The probability distribution is unchanged; the representation separates the random draw from the adjustable parameters. This makes reparameterization gradients available for differentiable expectations, including the Gaussian approximate posterior in a variational autoencoder.^[ref]

For a complete example, fix $\sigma$ and considerHolding the same noise value fixed while differentiating givesAt $\mu =1$, $\sigma =2$, the particular noise draw $\epsilon =1/2$ gives $Z=2$ and derivative estimate $4$, while the true derivative is $2$. An individual sample need not equal the derivative; the equality is in expected value.

Parameterizing a positive scale as $\sigma =e^{\rho }$ also gives $\partial Z/\partial \rho =e^{\rho }\epsilon$. One can differentiate through the deterministic transformation while leaving the noise law fixed. For a general integrand, exchanging derivative and expectation still requires justified differentiation under the integral sign; the representation alone is not a proof of that exchange. Reusing noise when parameters change also connects the construction to common random numbers.

Consider the explicitly specified objectiveand estimate its derivative by a symmetric finite difference with step $h>0$. The choice of shared or independent noise changes the estimator's variance even though both use correct normal distributions at the two parameter values.

With common random numbers, the same $\epsilon$ appears on both sides:Thus $\mathbb{E}[G_{\mathrm{shared}}]=2\mu$ and $\mathrm{Var}(G_{\mathrm{shared}})=4{\sigma }^2$, independent of $h$. Cancellation removes the noisy square term. For this quadratic, the symmetric difference has no truncation bias in exact arithmetic.

For independent standard normals ${\epsilon }_{+},{\epsilon }_{-}$, instead useIf $Y=m+\sigma \epsilon$, expand $Y^2=m^2+2m\sigma \epsilon +{\sigma }^2{\epsilon }^2$. The standard-normal moments $\mathbb{E}{\epsilon }^2=1$, $\mathbb{E}{\epsilon }^3=0$ and $\mathbb{E}{\epsilon }^4=3$ giveAdding the two independent square variances and dividing by $4h^2$ yieldsBoth estimators have the correct expected value $2\mu$, but the independent-noise variance diverges as $h$ tends to zero. At $\mu =\sigma =1$ and $h=0.1$, the variances are $4$ and $302$. An average of $n$ independent replicate estimators divides each variance by $n$.

This is a complete calculation for the stated quadratic, not a universal dominance theorem. At $\mu =0$, $\sigma =1$ and $h=1$, the independent variance is $3$, below the shared value $4$. The covariance principle in common random numbers explains why coupling needs a specified objective and comparison. Very small steps in floating-point evaluation also introduce cancellation concerns that this exact-arithmetic model does not include.

Markov chain variational inference optimizes an initial parametric probability distribution ${\nu }_{\eta }$ by minimizing $\mathrm{KL}(K^t{\nu }_{\eta }\Vert \pi )$ after a fixed number of Markov kernel updates. The best initial probability distribution can depend on $t$; Kullback-Leibler divergence contraction does not imply that optimization commutes with applying the Markov kernel.

Mean-field variational inference restricts the trial probability density function to $q(x)={\prod }_i{q}_i(x_i)$. With the other factors fixed, $q_j^{\ast }(x_j)\propto \exp ({\mathbb{E}}_{q_{-j}}\log \pi (x_j,X_{-j}))$, provided this expression has a finite positive normalization constant. Subtracting the objective at $q_j^{\ast }$ leaves $\mathrm{KL}(q_j\Vert q_j^{\ast })\ge 0$, proving the update is optimal when the terms are well defined.