Structured Variational Inference in VAEs

The mean-field approximation, where $q_\phi(z|x) = \prod_i q_\phi(z_i|x)$, simplifies VAE training considerably by assuming independence among the latent variables $z_i$ given the input $x$. However, this assumption is often too restrictive. The true posterior $p_\theta(z|x)$ might exhibit complex dependencies between latent dimensions, and forcing $q_\phi(z|x)$ to be factorized can prevent it from accurately modeling these correlations. This discrepancy, often part of the "amortization gap," can limit the VAE's ability to learn rich representations and generate high-fidelity data. Structured variational inference offers a way to address this limitation by explicitly modeling dependencies within the approximate posterior.

Modeling Dependencies in the Approximate Posterior

Structured variational inference aims to enrich the family of distributions $q_\phi(z|x)$ by allowing for correlations among the latent variables $z_1, \dots, z_D$. Instead of a fully factorized form, $q_\phi(z|x)$ is designed to capture some statistical structure. This allows $q_\phi(z|x)$ to be a more accurate approximation of the true posterior $p_\theta(z|x)$, potentially leading to a tighter Evidence Lower Bound (ELBO) and improved model performance.

The core idea is to define $q_\phi(z|x)$ using a model that can represent dependencies. Common approaches include autoregressive models and normalizing flows, both of which allow for flexible and expressive posterior distributions.

Comparison of mean-field and structured (autoregressive) approximate posteriors. In the mean-field case, latent variables $z_i$ are conditionally independent given $x$. In the structured autoregressive case, $z_i$ depends on preceding $z_j$ (for $j<i$) and $x$.

Autoregressive Models for $q_\phi(z|x)$

One powerful way to introduce structure is to model $q_\phi(z|x)$ autoregressively. This means the distribution over the latent vector $z = (z_1, \dots, z_D)$ is factorized as a product of conditional distributions:

$$q_\phi(z|x) = q_\phi(z_1|x) \prod_{j=2}^D q_\phi(z_j | z_{<j}, x)$$

Here, $z_{<j}$ denotes $(z_1, \dots, z_{j-1})$. Each conditional distribution $q_\phi(z_j | z_{<j}, x)$ can be parameterized by a neural network that takes $x$ and the previously sampled latent variables $z_1, \dots, z_{j-1}$ as input. For instance, if each $q_\phi(z_j | z_{<j}, x)$ is a Gaussian, its mean $\mu_j$ and standard deviation $\sigma_j$ would be functions of $x$ and $z_{<j}$:

$$\mu_j, \log \sigma_j = f_j(x, z_{<j}; \phi_j)$$

This structure allows $q_\phi(z|x)$ to capture arbitrary dependencies, provided the conditioning neural networks $f_j$ are sufficiently expressive. Sampling from such a model is sequential: first sample $z_1 \sim q_\phi(z_1|x)$, then $z_2 \sim q_\phi(z_2|z_1, x)$, and so on. While calculating the density $q_\phi(z|x)$ is straightforward (a product of $D$ terms), the sequential sampling can be slow if $D$ is large.

Techniques like Inverse Autoregressive Flows (IAFs), which you may recall from Chapter 3, provide a way to implement such expressive autoregressive models where sampling can be parallelized, significantly speeding up the process. In IAFs, $z$ is obtained by transforming a noise vector $\epsilon$ (where $\epsilon_j$ are independent) using an autoregressive transformation: $z_j = g_j(\epsilon_j; h_j(x, \epsilon_{<j}))$.

Utilizing Normalizing Flows for Expressive Posteriors

Normalizing Flows, also discussed in Chapter 3 (Section 3.5 "Normalizing Flows for Flexible Priors and Posteriors"), offer another general and powerful framework for constructing complex posterior distributions. A normalizing flow transforms a simple base distribution $q_0(u)$ (e.g., a standard multivariate Gaussian) through a sequence of invertible transformations $f_1, \dots, f_K$:

$$z = f_K \circ \dots \circ f_1(u), \quad u \sim q_0(u)$$

The density of $z$ can be computed using the change of variables formula:

$$q_\phi(z|x) = q_0(u) \left| \det \left( \frac{\partial (f_K \circ \dots \circ f_1)}{\partial u} \right) \right|^{-1}$$

The parameters of these transformations $f_k$ (and potentially the base distribution $q_0$) are learned as part of $\phi$, and are typically conditioned on $x$. This allows $q_\phi(z|x)$ to learn highly flexible distributions. The important point is that the transformations $f_k$ are designed such that their Jacobians (and thus the determinant) are computationally tractable. Examples include planar flows, radial flows, and more sophisticated flow architectures like RealNVP, MAF, and IAF.

Using normalizing flows for $q_\phi(z|x)$ can significantly increase the expressiveness of the inference network, allowing it to better match the true posterior and thereby tighten the ELBO.

Implications and Trade-offs

Adopting structured variational inference has several important consequences:

1. Improved ELBO and Model Quality: A more flexible $q_\phi(z|x)$ can provide a tighter lower bound on the true log-likelihood $\log p_\theta(x)$. This often translates to better generative performance, such as sharper generated samples and higher likelihood scores on test data. The representations learned in the latent space may also become more meaningful as the inference network better captures the underlying data manifold.

2. Increased Computational Complexity: The primary trade-off is computational cost.

   • Training: Parameterizing and optimizing a structured $q_\phi(z|x)$ involves more parameters and often more complex computations per iteration (e.g., computing Jacobian determinants for flows or sequential conditioning for autoregressive models).
   • Inference: Sampling from some structured posteriors (like standard autoregressive models) can be slower if not parallelizable.

3. Model Design Choices: You now have more choices to make regarding the architecture of $q_\phi(z|x)$. For autoregressive models, this includes the ordering of latent variables and the architecture of the conditional networks. For normalizing flows, it involves selecting the type and number of flow layers. These choices can impact performance and computational load.

4. KL Divergence Term: The KL divergence $D_{KL}(q_\phi(z|x) || p(z))$ in the ELBO might become more challenging. If $p(z)$ is a standard Gaussian and $q_\phi(z|x)$ is a complex distribution (e.g., from a normalizing flow), the KL divergence may no longer have an analytical solution. In such cases, it often needs to be estimated, for example, by sampling $z \sim q_\phi(z|x)$ and calculating $\mathbb{E}_{q_\phi(z|x)}[\log q_\phi(z|x) - \log p(z)]$.

When to Consider Structured Inference

Structured variational inference is particularly beneficial when:

• You suspect strong correlations or dependencies exist among the true latent factors of variation in your data, which a mean-field VAE fails to capture.
• Standard VAEs yield unsatisfactory results, such as overly smooth or blurry generated samples, or low likelihoods, and you hypothesize that the inference network is a bottleneck.
• The application demands highly accurate posterior inference, for instance, in scenarios where the latent variables themselves are objects of direct interest.

While introducing structure adds complexity, the potential gains in model expressiveness and performance often justify the additional overhead, especially for challenging datasets or when aiming for state-of-the-art results. The techniques discussed here, such as autoregressive models and normalizing flows for $q_\phi(z|x)$, are foundational for building more sophisticated and powerful VAEs. As we proceed, you'll see how these improved inference mechanisms can be combined with other advanced VAE components.