Variational Inference with Implicit Models

The expressiveness of the approximate posterior $q_\phi(z|x)$ is a significant factor in Variational Autoencoder (VAE) performance. Standard choices, such as diagonal Gaussians, often fall short of capturing the true complexity of $p_\theta(z|x)$. Using implicit models for the variational posterior or even the prior offers a powerful alternative.

What Are Implicit Models?

An implicit model is one where we can easily sample from the distribution but cannot easily (or at all) evaluate its probability density function (PDF) or probability mass function (PMF) for a given point. Think of a generative neural network: you feed it random noise, and it produces a complex output. You can get samples, but what's the $\log q(z)$ for a specific $z$? That's often intractable.

Formally, if we want an implicit posterior $q_\phi(z|x)$, we define it through a deterministic transformation $g_\phi$ of a simple noise variable $\epsilon$ (e.g., $\epsilon \sim \mathcal{N}(0, I)$) and the input $x$:

$$z = g_\phi(\epsilon, x)$$

While we can generate samples $z$ from $q_\phi(z|x)$ by first sampling $\epsilon$ and then applying $g_\phi$, the density $q_\phi(z|x)$ itself is not directly accessible. This is in stark contrast to, say, a Gaussian posterior where $q_\phi(z|x) = \mathcal{N}(z | \mu_\phi(x), \Sigma_\phi(x))$, and we can directly compute $\log q_\phi(z|x)$.

The Challenge: The ELBO and Implicit Densities

The Evidence Lower Bound (ELBO) for a VAE is:

$$L(\theta, \phi; x) = E_{z \sim q_\phi(z|x)}[\log p_\theta(x|z)] - KL(q_\phi(z|x) || p(z))$$

Let's expand the KL divergence term:

$$KL(q_\phi(z|x) || p(z)) = E_{z \sim q_\phi(z|x)}[\log q_\phi(z|x) - \log p(z)]$$

If $q_\phi(z|x)$ is implicit, the term $E_{z \sim q_\phi(z|x)}[\log q_\phi(z|x)]$ (the negative entropy of $q_\phi$) becomes problematic because $\log q_\phi(z|x)$ is unknown. This is the central difficulty when using implicit models for the variational posterior.

If the prior $p(z)$ is also chosen to be an implicit model (perhaps to represent a complex target structure in the latent space), then the $E_{z \sim q_\phi(z|x)}[\log p(z)]$ term also becomes intractable through direct evaluation. However, if $p(z)$ is simple (e.g., $\mathcal{N}(0, I)$), this second term can still be estimated by sampling $z \sim q_\phi(z|x)$ and evaluating $\log p(z)$.

Tackling Intractable Log-Densities

So, how do we optimize the ELBO if we can't evaluate $\log q_\phi(z|x)$? Several strategies have emerged, with adversarial training being particularly prominent.

1. Density Ratio Estimation / Adversarial Approaches: The KL divergence, or parts of it, can often be rewritten or approximated using techniques reminiscent of Generative Adversarial Networks (GANs). The core idea is to train a discriminator (or critic) network $D_\psi(z)$ (or $D_\psi(z, x)$ if conditioned on $x$) to distinguish between samples from $q_\phi(z|x)$ and samples from $p(z)$. The objective for $q_\phi$ (specifically, the parameters of $g_\phi$) is then to produce samples $z$ that "fool" this discriminator.

   For instance, the term $KL(q_\phi(z|x) || p(z))$ can be estimated or optimized indirectly. If $p(z)$ is a simple distribution, the main issue is the entropy term $H(q_\phi(z|x)) = -E_{z \sim q_\phi(z|x)}[\log q_\phi(z|x)]$. Various f-divergences, including the KL divergence, can be expressed using a variational representation that involves a discriminator. For example, the Jensen-Shannon divergence, $JS(q_\phi || p)$, which is minimized by the original GAN objective, is one such case.

   The Adversarial Variational Bayes (AVB) framework, which we will discuss in the next section, provides a specific mechanism. It often involves training a separate network $T(z)$ to approximate $\log q_\phi(z|x)$ or directly estimates the KL term. The point is that $q_\phi(z|x)$ is implicitly defined through $z = g_\phi(\epsilon,x)$, and its parameters are updated based on gradients flowing from both the reconstruction term and the adversarial term approximating the KL divergence.

2. Kernel Density Estimation (KDE): One could, in principle, draw many samples from $q_\phi(z|x)$ and use KDE to estimate the density $q_\phi(z|x)$ at any point $z$. Then, $\log q_\phi(z|x)$ could be approximated. However, KDE suffers severely from the curse of dimensionality and requires a very large number of samples, making it impractical for typical latent space dimensions in VAEs.

3. Likelihood-free Inference Methods: The broader field of likelihood-free inference (also known as Approximate Bayesian Computation, ABC) deals with situations where the likelihood (or in our case, the posterior density) is intractable but sampling is possible. Some techniques from this field can inspire approaches for VAEs with implicit posteriors, often involving comparing summary statistics of data generated from different models.

The following diagram illustrates how an implicit posterior might be integrated into a VAE, using an adversarial approach to handle the KL divergence term.

Information flow in a VAE with an implicit posterior $q_\phi(z|x)$ and an adversarial mechanism for the KL divergence. The sampler $g_\phi$ generates latent codes $z$, which are used for reconstruction. These samples, along with samples from a prior $p(z)$, are fed to a discriminator $D_\psi$ to compute an adversarial loss that helps shape $q_\phi(z|x)$.

Implicit Priors $p(z)$

The same philosophy can be applied to the prior distribution $p(z)$. Instead of a fixed, simple prior like $\mathcal{N}(0, I)$, one might want to learn a more complex prior, perhaps one that itself is an implicit model. Adversarial Autoencoders (AAEs), which we'll touch upon in Chapter 7, often use an adversarial loss to match the aggregated posterior $q(z) = E_{p_{data}(x)}[q_\phi(z|x)]$ to a chosen prior $p(z)$, and this $p(z)$ could be implicitly defined by samples from another generator. If both $q_\phi(z|x)$ and $p(z)$ are implicit, then the $KL(q_\phi(z|x) || p(z))$ term usually relies entirely on adversarial or density ratio estimation techniques.

Advantages and Downsides

Advantages:

• Highly Expressive Posteriors/Priors: Implicit models can represent arbitrarily complex distributions, potentially capturing multi-modal or non-Gaussian characteristics of the true posterior $p_\theta(z|x)$ or a desired prior $p(z)$. This can lead to tighter ELBOs (if they can be estimated) and better generative performance.
• Improved Sample Quality: A more accurate posterior can lead to higher-quality reconstructions and generated samples, as the latent space might be better structured.
• Flexibility: Avoids restrictive parametric assumptions on the form of $q_\phi(z|x)$.

Downsides:

• Training Instability: Adversarial training components can be notoriously difficult to stabilize and tune, requiring careful balancing of generator and discriminator updates, choice of loss functions, and network architectures.
• Evaluation Challenges: Evaluating the true ELBO becomes difficult. While we can optimize a surrogate objective, knowing the actual tightness of the bound is harder. Model comparison often relies on sample quality, held-out likelihood estimation using methods like Importance Weighting (covered earlier), or downstream task performance.
• Increased Complexity: The overall model complexity increases due to the additional networks (e.g., the discriminator or critic) and more intricate training loops.

Moving Forward

Variational inference with implicit models opens up a rich avenue for enhancing VAEs. By defining $q_\phi(z|x)$ (and potentially $p(z)$) via a sampling procedure $z = g_\phi(\epsilon, x)$, we can model distributions far more complex than standard Gaussians. The main hurdle, the intractability of $\log q_\phi(z|x)$, is typically addressed using adversarial training schemes. The next section on Adversarial Variational Bayes (AVB) will provide a more concrete example of how these ideas are put into practice to create more powerful VAEs. These methods represent a significant step towards bridging the gap between the tractable but limited posteriors of basic VAEs and the highly flexible but often unstructured latent spaces of models like GANs.