Latent-Variable Representations › Variational Inference: Making Latents Probabilistic

This section builds the variational machinery step by step: the latent-variable model and its intractable marginal, the evidence lower bound with its two-term decomposition, and the two techniques, reparameterization and amortization, that make the objective trainable by gradient descent.

This is the modeling direction we want: generate an observation $x$ from a latent $z$, and infer a distribution over $z$ after observing $x$. The factorization is statistical, not automatically causal. Calling $z$ a physical cause requires additional structural assumptions and validation through time or interventions. Even without that causal interpretation, the direction is useful: it supports sampling, conditioning, and later models of action-conditioned latent transitions.

Training means adjusting $\theta$ so that $p_\theta(x)$ is high on real observations. For a generic nonlinear neural-network decoder, the marginal usually has no tractable closed form, even with a low-dimensional continuous $z$; analytic special cases such as the linear-Gaussian model do exist. The ELBO gives us a differentiable surrogate for the generic case.

The classical move, which you met in the filtering chapter, is to approximate the posterior with variational inference. Introduce a family of distributions $q_\phi(z \mid x)$, the approximate posterior or encoder, and optimize it to stay close to the true posterior $p_\theta(z \mid x)$. There is a pleasing symmetry in the setup: the generative model runs in the direction $z \to x$, from state to observation, while the inference model runs in the reverse direction $x \to z$, from observation to belief. In a Bayes filter, the direct analogue of $p_\theta(x \mid z)$ is the observation likelihood, and inference combines that likelihood with a prior belief in the measurement update. A separate transition model supplies the prior belief when temporal dynamics are introduced. Here the observation and inference directions are learned together, and the next section shows how a single objective trains them.

We cannot minimize $\mathrm{KL}(q_\phi(z \mid x) \,\|\, p_\theta(z \mid x))$ directly because the true posterior contains the intractable marginal: by Bayes' rule, $p_\theta(z \mid x) = p_\theta(x, z) / p_\theta(x)$, and it is exactly $p_\theta(x)$ that we cannot compute. The trick is to rewrite the log marginal density in a form that sidesteps it. For the following identity, assume $q_\phi(z \mid x)$ is positive wherever $p_\theta(x,z)$ is positive for the fixed $x$, and that the displayed expectations exist. Start from the definition of the marginal and multiply and divide by $q_\phi(z \mid x)$:

This manipulation puts the chosen approximation $q$ inside the integral for $\log p_\theta(x)$, creating an expectation over $q$ of the ratio $p_\theta(x, z)/q_\phi(z \mid x)$. The logarithm of an expectation is hard, so apply Jensen's inequality, which holds because $\log$ is concave:

The right-hand side is the evidence lower bound, or ELBO. Two facts make it the workhorse of representation learning.

First, the gap is exactly the quantity we could not compute:

For fixed generative parameters $\theta$, maximizing the ELBO over $\phi$ reduces the KL gap to the true posterior within the chosen variational family. Joint optimization over $\theta$ and $\phi$ also raises the bound used to fit the generative model, although neither quantity is guaranteed to improve monotonically at every stochastic update. When the gap is zero, the encoder has recovered the exact posterior for that model.

Second, the ELBO splits into two interpretable terms. Expand $\log p_\theta(x, z) = \log p_\theta(x \mid z) + \log p_\theta(z)$:

The first term rewards reconstruction; the second penalizes each approximate posterior for deviating from the prior. This defines a reference measure for sampling and indirectly pressures the aggregated posterior toward that prior. It does not guarantee hole-free coverage, plausible interpolation, or semantic distance. Those geometric properties require direct evaluation or additional structure. What the KL does provide is an explicit price for information: every posterior deviation from the prior has to earn its reconstruction benefit.

It is worth being precise about what the regularization term is not. It does not require encoded training data to fill the prior, and prior density is not semantic geometry. Interpolation and metric alignment remain empirical questions.

Two engineering problems stand between the ELBO and gradient descent, and both have elegant, now-standard solutions.

Problem 1: a black-box draw blocks ordinary pathwise backpropagation to the distribution parameters. Score-function estimators are possible, but often have high variance. The reparameterization trick instead writes a draw as a differentiable transform of parameter-free noise. For a Gaussian encoder,

where $\mu_\phi$ and $\sigma_\phi$ are network outputs and $\odot$ is elementwise multiplication. All randomness now lives in $\varepsilon$, which does not depend on any parameter, and gradients flow through $\mu_\phi$ and $\sigma_\phi$. The conceptual shift is subtle but important: instead of sampling a value of $z$, we sample a source of randomness and then compute $z$ deterministically from it. The function from parameters to samples is now smooth, so the expectation over $\varepsilon$ can be estimated by a handful of Monte Carlo draws, often just one, and each estimate carries useful gradients. This trick was one key ingredient that made the variational autoencoder of Kingma and Welling (2013) and Rezende, Mohamed, and Wierstra (2014) practical; amortized inference, introduced next, supplies the other ingredient needed to encode new observations efficiently. Reparameterization recurs throughout this book whenever a stochastic layer must be threaded through a differentiable computation.

Problem 2: a separate optimization per data point. Traditional local variational inference often optimizes a dedicated set of variational parameters for each observation. That is poorly suited to a world model that must encode new observations under a tight latency budget, because a long inner optimization loop may not fit between camera frames. Instead we amortize inference: one encoder network $q_\phi(z \mid x)$ learns to output the variational parameters for any $x$. The term emphasizes that the cost of learning an inference procedure is shared across the data distribution rather than paid through a fresh optimization for every observation. Amortization trades per-observation optimization cost for a single upfront training cost, and it is what lets a learned world model ingest a fresh camera frame in one forward pass. The trade has a known downside worth naming: the encoder can be systematically imprecise on inputs it has rarely seen, a form of amortization error that dedicated per-observation optimization can reduce, though not necessarily eliminate. Local optimization may remain incomplete, and the chosen variational family can still be misspecified. For agents with tight latency budgets, amortized inference is often the practical default; whether to add bounded per-observation refinement depends on the available compute, the approximation error, and the cost of an inference mistake.

The closed form of the Gaussian KL shows what the regularizer asks of the encoder. Each dimension is minimized when the encoder's mean $\mu_j$ is zero and its variance $\sigma_j^2$ is one. The KL pressure pulls each encoding distribution toward zero mean and unit variance. The reconstruction term pushes the other way by rewarding informative codes.