Variational Autoencoders (VAEs) as Probabilistic Models

Variational Autoencoders (VAEs) represent an intersection of deep learning and Bayesian inference, offering an approach to unsupervised generative modeling. Unlike standard autoencoders which learn a deterministic mapping to and from a lower-dimensional space, VAEs use a probabilistic perspective, rooted in the principles of latent variable modeling and variational inference.

Think of a VAE as consisting of two main components connected through a probabilistic latent space:

1. Encoder (Recognition Model): Maps an input data point $x$ to a distribution in the latent space. Instead of outputting a single point $z$ in the latent space, it outputs parameters (like mean $\mu$ and variance $\sigma^2$) for a probability distribution over possible latent representations for $x$. This distribution, denoted as $q_{\phi}(z | x)$, is our approximate posterior. It's parameterized by $\phi$, typically the weights of a neural network.
2. Decoder (Generative Model): Maps a point $z$ sampled from the latent distribution back to the original data space. It defines a conditional distribution $p_{\theta}(x | z)$, representing the likelihood of observing $x$ given the latent representation $z$. This is parameterized by $\theta$, usually the weights of another neural network.

The underlying idea is that our observed data $x$ is generated from some unobserved, lower-dimensional latent variables $z$. We assume a prior distribution $p(z)$ over these latent variables, often chosen to be a simple distribution like a standard multivariate Gaussian $\mathcal{N}(0, I)$. The decoder $p_{\theta}(x | z)$ defines how these latent variables generate the complex data we see.

The central challenge is computing the true posterior distribution $p_{\theta}(z | x) = p_{\theta}(x | z) p(z) / p(x)$. The denominator, $p(x) = \int p_{\theta}(x | z) p(z) dz$, is the marginal likelihood or evidence, and it's generally intractable to compute because it involves integrating over all possible latent variables $z$.

This is precisely where Variational Inference comes in. We introduce the encoder network $q_{\phi}(z | x)$ as an approximation to the true, intractable posterior $p_{\theta}(z | x)$. As we learned in Chapter 3, VI reframes inference as an optimization problem. We aim to make $q_{\phi}(z | x)$ as close as possible to $p_{\theta}(z | x)$ by minimizing the KL divergence $D_{KL}(q_{\phi}(z | x) || p_{\theta}(z | x))$. Minimizing this divergence is equivalent to maximizing the Evidence Lower Bound (ELBO), $\mathcal{L}(\theta, \phi; x)$:

$$\mathcal{L}(\theta, \phi; x) = \mathbb{E}_{q_{\phi}(z|x)}[\log p_{\theta}(x | z)] - D_{KL}(q_{\phi}(z | x) || p(z))$$

Maximizing the ELBO $\mathcal{L}(\theta, \phi; x)$ simultaneously trains both the encoder (parameters $\phi$) and the decoder (parameters $\theta$). Let's examine the two terms in the ELBO:

• Reconstruction Term: $\mathbb{E}_{q_{\phi}(z|x)}[\log p_{\theta}(x | z)]$. This term measures how well the decoder can reconstruct the original input $x$ after encoding it into the latent distribution $q_{\phi}(z | x)$ and then sampling a $z$ from that distribution. For Gaussian or Bernoulli likelihoods $p_{\theta}(x | z)$, this term often simplifies to a mean squared error or a binary cross-entropy loss between the input $x$ and the decoder's output $\hat{x}$. It encourages the VAE to learn meaningful latent representations that capture the essential information needed for reconstruction.

• Regularization Term (KL Divergence): $D_{KL}(q_{\phi}(z | x) || p(z))$. This term acts as a regularizer. It measures the divergence between the approximate posterior distribution $q_{\phi}(z | x)$ produced by the encoder and the prior distribution $p(z)$ over the latent variables. By encouraging the encoded distributions to stay close to the prior (e.g., a standard Gaussian), it ensures that the latent space is well-structured and prevents the encoder from collapsing different inputs to distinct, isolated regions (posterior collapse). This regularity is important for the VAE's generative capabilities.

The Reparameterization Trick

A significant challenge arises when trying to optimize the ELBO using gradient ascent: how do we backpropagate gradients through the sampling step involved in the expectation $\mathbb{E}_{q_{\phi}(z|x)}[\cdot]$? The sampling operation itself is stochastic and non-differentiable.

The VAE solves this using the reparameterization trick. Instead of sampling $z$ directly from $q_{\phi}(z | x)$, we express $z$ as a deterministic function of the parameters $\phi$, the input $x$, and an independent random variable $\epsilon$ with a fixed distribution (e.g., $\epsilon \sim \mathcal{N}(0, I)$).

For example, if our encoder $q_{\phi}(z | x)$ outputs the mean $\mu_{\phi}(x)$ and standard deviation $\sigma_{\phi}(x)$ of a Gaussian distribution, we can sample $z$ as:

$$z = \mu_{\phi}(x) + \sigma_{\phi}(x) \odot \epsilon, \quad \text{where } \epsilon \sim \mathcal{N}(0, I)$$

Here, $\odot$ denotes element-wise multiplication. Now, the stochasticity comes from $\epsilon$, which does not depend on the model parameters $\phi$. The path from $\phi$ and $\theta$ to the loss $\mathcal{L}$ is deterministic, allowing gradients to flow back through $\mu_{\phi}(x)$ and $\sigma_{\phi}(x)$ to update the encoder network $\phi$, and through the decoder for $\theta$.

Flow diagram illustrating the VAE architecture. Input x is processed by the encoder network to produce parameters mu(x) and sigma(x) defining the approximate posterior q_phi(z|x). A latent sample z is generated using the reparameterization trick with noise eps. The KL divergence term compares q_phi(z|x) with the prior p(z). The sample z is passed through the decoder network to produce the reconstruction x_hat. The reconstruction loss measures the difference between x and x_hat. Both loss terms contribute to the ELBO objective.

VAEs in the Context of Bayesian Deep Learning

VAEs are a foundation of probabilistic deep learning. While they don't typically place priors directly on the weights of the neural networks like BNNs do, they perform Bayesian inference over the latent variables $z$. They learn a deep generative model $p_{\theta}(x|z)p(z)$ using variational inference for efficient training.

Important connections to BDL include:

• Generative Modeling: VAEs learn the underlying distribution of the data $p(x)$ through the latent variable framework $p(x) = \int p_{\theta}(x | z) p(z) dz$. Once trained, you can generate new data samples by drawing $z \sim p(z)$ and then passing it through the decoder $p_{\theta}(x | z)$.
• Uncertainty Representation: The latent space $z$ inherently captures uncertainty about the factors of variation in the data. The variance $\sigma^2_{\phi}(x)$ output by the encoder for a given $x$ reflects the uncertainty in the latent representation for that specific input.
• Inference Technique: VAEs rely entirely on variational inference, one of the primary techniques we explore for approximate inference in complex Bayesian models, including BNNs. Understanding VAEs solidifies the practical application of VI and the ELBO.

In summary, VAEs provide a scalable and effective method for learning deep generative models with a solid probabilistic foundation based on variational inference. They bridge the gap between complex deep learning architectures and principled Bayesian reasoning about latent structure in data.