Variational Autoencoder

Introduction

A central challenge in machine learning systems that interact with the physical world is uncertainty. Examples include robotics, autonomous vehicles, and the embodied agents that came to be called "Physical AI" around 2025. Real-world interaction is inherently stochastic. Sensor noise, unobserved physical properties, and environmental variability all conspire against deterministic models. Even classical robotic systems handle this stochasticity with probabilistic methods, Kalman filtering being the canonical example. Other approaches push further by treating system states as full probability distributions rather than point estimates.

The Variational Autoencoder (VAE) is one such approach. It synthesizes ideas from information theory and Bayesian inference into a model that learns a continuous latent distribution characterized by a mean (\(\mu\)) and a variance (\(\sigma^2\)). The variance of this latent distribution quantifies how uncertain the model is about the latent representation of a given input.

Variational Inference

In Bayesian statistics, Variational Inference (VI) transforms the problem of posterior inference, which typically involves high-dimensional intractable integrals, into a constrained optimization problem.

While deterministic control algorithms rely on point-estimates (for example, "the center of mass is exactly at coordinates \(\mathbf{r}\)"), VI treats the state as a probability distribution. Instead of seeking a single value, we search for the parameters of a distribution \(q_{\boldsymbol{\psi}}(\boldsymbol{z})\) that minimize the Kullback-Leibler (KL) divergence from the true, unknown posterior \(p(\boldsymbol{z} \mid \boldsymbol{x})\).

Since the true posterior is analytically intractable for complex models, we restrict our search to a family of simpler distributions (such as Gaussians). By maximizing the Evidence Lower Bound (ELBO), the system finds the "best fit" distribution that balances data fidelity with prior beliefs.

For a robot, this approximation is one ingredient in uncertainty-aware control. The posterior variance measures the spread of the distribution and provides a quantitative signal that can inform decisions such as whether to proceed with an action or trigger a fallback. The relation between this spread and epistemic uncertainty (uncertainty due to limited model knowledge, as distinct from aleatoric uncertainty inherent in the observation noise) is not automatic. Such a signal complements rather than replaces classical safety mechanisms. We return to the reliability of this reading, and its limits, in the demo section below.

The variational framework introduced here rests on measure-theoretic foundations that are developed across our probability section. The Radon-Nikodym theorem licenses speaking of densities \(q(\boldsymbol{z})\) and \(p(\boldsymbol{z} \mid \boldsymbol{x})\) with respect to a common dominating measure. Conditional expectation, defined rigorously as a Radon-Nikodym derivative, is what gives the posterior \(p(\boldsymbol{z} \mid \boldsymbol{x})\) its meaning, and it is this posterior that the variational distributions below approximate.

In practice we work directly with explicit Gaussian densities, Monte Carlo samples, and reparameterized gradients. The measure-theoretic construction is not a layer that the code instantiates. Rather, it guarantees that these computations approximate a well-defined object, and the VAE's amortized inference is a special case of the broader variational principle.

Variational Autoencoder (VAE)

From Factor Analysis to Nonlinear Generative Models

In classical factor analysis (FA), we model centered observed data \(\boldsymbol{x} \in \mathbb{R}^D\) through a latent variable \(\boldsymbol{z} \in \mathbb{R}^K\) with \(K \ll D\), drawn from the standard Gaussian prior \(p(\boldsymbol{z}) = \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\), and a linear-Gaussian likelihood: \[ p(\boldsymbol{x} \mid \boldsymbol{z}) = \mathcal{N}(\boldsymbol{x} \mid \boldsymbol{Wz}, \sigma^2 \boldsymbol{I}) \] where \(\boldsymbol{W} \in \mathbb{R}^{D \times K}\) is the factor loading matrix. General factor analysis allows an arbitrary diagonal noise covariance. The isotropic choice \(\sigma^2 \boldsymbol{I}\) shown here is the special case known as probabilistic PCA, and it is the noise model that the VAE decoder below inherits. Because the prior is Gaussian and the likelihood is Gaussian with mean linear in \(\boldsymbol{z}\), posterior inference is tractable. Given an observation \(\boldsymbol{x}\), the posterior \(p(\boldsymbol{z} \mid \boldsymbol{x})\) is Gaussian in closed form. However, this linearity also limits the model's expressiveness, since real-world data distributions are rarely well-captured by linear mappings.

The Variational Autoencoder (VAE) extends factor analysis by replacing the linear mapping \(\boldsymbol{Wz}\) with an arbitrary nonlinear function parameterized by a neural network: \[ p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z}) = \mathcal{N}(\boldsymbol{x} \mid f_d (\boldsymbol{z}; \boldsymbol{\theta}), \, \sigma^2 \boldsymbol{I}) \] where \(f_d(\cdot\,; \boldsymbol{\theta})\) is a decoder network with parameters \(\boldsymbol{\theta}\).

This nonlinear generative model can represent far richer data distributions, but it comes at a cost. The posterior \(p_{\boldsymbol{\theta}}(\boldsymbol{z} \mid \boldsymbol{x})\) is no longer analytically tractable. By Bayes' rule, \[ \begin{align*} p_{\boldsymbol{\theta}}(\boldsymbol{z} \mid \boldsymbol{x}) &= \frac{p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z})\, p(\boldsymbol{z})}{p_{\boldsymbol{\theta}}(\boldsymbol{x})}, \\\\ p_{\boldsymbol{\theta}}(\boldsymbol{x}) &= \int p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z})\, p(\boldsymbol{z})\, d\boldsymbol{z}, \end{align*} \] and the integral over the latent space that defines the marginal likelihood \(p_{\boldsymbol{\theta}}(\boldsymbol{x})\) is intractable when \(f_d\) is a deep network. This intractability is the central challenge that motivates the variational approach.

Amortized Variational Inference

We therefore introduce an approximate posterior computed by a recognition network (also called the inference network or encoder) that is trained simultaneously with the generative model. We restrict the approximate posterior to a tractable family: a Gaussian with diagonal covariance, yielding \[ q_{\boldsymbol{\phi}} (\boldsymbol{z} \mid \boldsymbol{x}) = \mathcal{N}\!\left(\boldsymbol{z} \,\middle|\, f_{e,\mu} (\boldsymbol{x} ; \boldsymbol{\phi}), \, \operatorname{diag}\!\left(f_{e,\sigma}(\boldsymbol{x};\boldsymbol{\phi})\right)^{2}\right) \] where \(f_{e,\mu}\) and \(f_{e,\sigma}\) are the encoder's output heads that produce the mean vector and the vector of positive standard deviations respectively, and \(\boldsymbol{\phi}\) denotes all encoder parameters.

This approach is called amortized inference. Rather than running a separate optimization procedure to compute the posterior for each individual data point (as in classical variational inference), we train a single neural network that directly maps any input \(\boldsymbol{x}\) to the parameters of its approximate posterior. The cost of inference at test time is thus reduced to a single forward pass through the encoder, with the computational investment "amortized" across the entire training set.

The Evidence Lower Bound (ELBO)

We cannot maximize the intractable marginal likelihood \(p_{\boldsymbol{\theta}}(\boldsymbol{x})\) directly, so we derive a tractable lower bound. For a single observation \(\boldsymbol{x}\), define the evidence lower bound (ELBO) as follows.

To see that this is indeed a lower bound on the log-evidence, assume that \(\log p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z})\) and \(\log q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})\) are integrable under \(q_{\boldsymbol{\phi}}(\cdot \mid \boldsymbol{x})\), so that the ELBO is finite. This forces \(p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z}) \gt 0\) for \(q_{\boldsymbol{\phi}}\)-almost every \(\boldsymbol{z}\), so the ratio \(p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z}) / q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})\) is positive almost surely. Applying Jensen's inequality to the convex function \(-\log\), which amounts to moving the concave \(\log\) outside the expectation, we obtain \[ \begin{align*} L(\boldsymbol{\theta}, \, \boldsymbol{\phi}\mid \boldsymbol{x}) &= \int q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) \log \frac{p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\, d\boldsymbol{z} \\\\ &\leq \log \int q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) \frac{p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\, d\boldsymbol{z} \\\\ &\leq \log \int p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z})\, d\boldsymbol{z} \\\\ &= \log p_{\boldsymbol{\theta}}(\boldsymbol{x}). \end{align*} \] The third step can only increase the integral, since the integrand equals \(p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z})\) where \(q_{\boldsymbol{\phi}} \gt 0\) and vanishes elsewhere. That is, \(L(\boldsymbol{\theta}, \boldsymbol{\phi} \mid \boldsymbol{x}) \leq \log p_{\boldsymbol{\theta}}(\boldsymbol{x})\) for every \(q_{\boldsymbol{\phi}}\) satisfying the integrability assumption. Maximizing the ELBO with respect to \(\boldsymbol{\phi}\) tightens the bound. Whenever \(q_{\boldsymbol{\phi}}(\cdot \mid \boldsymbol{x})\) is absolutely continuous with respect to the posterior, the ELBO-KL decomposition identifies the gap \(\log p_{\boldsymbol{\theta}}(\boldsymbol{x}) - L(\boldsymbol{\theta}, \boldsymbol{\phi} \mid \boldsymbol{x})\) as \(D_{\mathbb{KL}}\!\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) \,\|\, p_{\boldsymbol{\theta}}(\boldsymbol{z} \mid \boldsymbol{x})\right)\). Maximization with respect to \(\boldsymbol{\theta}\), by contrast, raises this lower bound on the log-likelihood rather than the log-likelihood itself.

ELBO Decomposition: Reconstruction + Regularization

The ELBO can be decomposed into two interpretable terms by expanding the joint \(\log p_{\boldsymbol{\theta}}(\boldsymbol{x}, \boldsymbol{z}) = \log p_{\boldsymbol{\theta}}(\boldsymbol{x} \mid \boldsymbol{z}) + \log p(\boldsymbol{z})\) and rearranging.

The first term is the expected reconstruction log-likelihood. It encourages the decoder to reconstruct the input \(\boldsymbol{x}\) from latent samples \(\boldsymbol{z} \sim q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})\). The second term is the KL divergence between the approximate posterior and the prior \(p(\boldsymbol{z}) = \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\), which regularizes the latent space by penalizing approximate posteriors that deviate from the prior. This regularization is what gives the VAE its structured, continuous latent space. Without it, the encoder could collapse each data point to an isolated delta function, and the model would lose the ability to generate new data by sampling from \(p(\boldsymbol{z})\).

The Reparameterization Trick

To train the VAE end-to-end via gradient-based optimization, we need to differentiate the ELBO with respect to both \(\boldsymbol{\theta}\) and \(\boldsymbol{\phi}\). The gradient with respect to the decoder parameters \(\boldsymbol{\theta}\) requires no reparameterization, since the sampling distribution \(q_{\boldsymbol{\phi}}\) does not depend on \(\boldsymbol{\theta}\). However, differentiating with respect to the encoder parameters \(\boldsymbol{\phi}\) is problematic. The expectation \(\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}[\cdot]\) is taken with respect to a distribution that itself depends on \(\boldsymbol{\phi}\), so we cannot simply interchange the gradient and the expectation.

The reparameterization trick resolves this by expressing the stochastic latent variable \(\boldsymbol{z}\) as a deterministic, differentiable transformation of a parameter-free noise variable. Since the encoder outputs a diagonal Gaussian, the trick takes the following form.

Under this reparameterization, the ELBO becomes \[ L(\boldsymbol{\theta}, \, \boldsymbol{\phi}\mid \boldsymbol{x}) = \mathbb{E}_{\boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{0}, \boldsymbol{I})} \!\left[\log p_{\boldsymbol{\theta}}\!\left(\boldsymbol{x} \mid \boldsymbol{z} = \mu_{\boldsymbol{\phi}}(\boldsymbol{x}) + \sigma_{\boldsymbol{\phi}}(\boldsymbol{x})\odot \boldsymbol{\epsilon}\right)\right] - D_{\mathbb{KL}}\!\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x}) \,\|\, p(\boldsymbol{z})\right). \]

The key observation is that the expectation is now taken with respect to the fixed distribution \(\mathcal{N}(\boldsymbol{0}, \boldsymbol{I})\), which does not depend on \(\boldsymbol{\phi}\). Under the domination hypothesis of the reparameterization gradient theorem, we can therefore interchange the gradient and the expectation: \(\nabla_{\boldsymbol{\phi}} \mathbb{E}_{\boldsymbol{\epsilon}}[\cdot] = \mathbb{E}_{\boldsymbol{\epsilon}}[\nabla_{\boldsymbol{\phi}}(\cdot)]\), and estimate the gradient via Monte Carlo sampling of \(\boldsymbol{\epsilon}\). In practice, a single sample per data point is commonly used, and with minibatch averaging it typically yields a gradient estimate adequate for stochastic gradient descent.

The full VAE training objective over a dataset \(\mathcal{D}\) is therefore \[ \min_{\boldsymbol{\theta}, \, \boldsymbol{\phi}} -\, \mathbb{E}_{\boldsymbol{x} \sim \mathcal{D}} \left[ L(\boldsymbol{\theta}, \, \boldsymbol{\phi}\mid \boldsymbol{x}) \right] \] which is optimized end-to-end by stochastic gradient descent.

Putting It All Together

Together, the nonlinear generative model, the amortized inference network, the ELBO objective, and the reparameterization trick define the Variational Autoencoder.