Adversarial Variational Bayes (AVB)

As we've discussed, the expressiveness of the approximate posterior $q_\phi(z|x)$ is a critical factor in the performance of Variational Autoencoders. Standard VAEs often employ simple, explicit distributions for $q_\phi(z|x)$, like a diagonal Gaussian, primarily because this choice makes the Kullback-Leibler (KL) divergence term $D_{KL}(q_\phi(z|x) || p(z))$ in the Evidence Lower Bound (ELBO) analytically tractable. However, this simplicity can severely limit how well $q_\phi(z|x)$ can model the true, often complex, posterior $p_\theta(z|x)$. Adversarial Variational Bayes (AVB) offers a sophisticated approach to overcome this limitation by enabling the use of highly flexible, implicit approximate posteriors.

The Challenge: Implicit Posteriors and Intractable Densities

To dramatically increase the flexibility of the approximate posterior, we can define it implicitly. Instead of specifying $q_\phi(z|x)$ with an explicit probability density function, we define a sampling procedure:

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

where $g_\phi$ is a neural network (the encoder) parameterized by $\phi$, $x$ is the input data, and $\epsilon$ is a noise variable sampled from a simple distribution, like a standard Normal distribution $p(\epsilon)$. This construction allows $q_\phi(z|x)$ to represent virtually any complex distribution.

The hurdle with such an implicit $q_\phi(z|x)$ is that its probability density function $q_\phi(z|x)$ (and thus $\log q_\phi(z|x)$) is generally unknown and intractable to compute. This is a problem because the ELBO, which we aim to maximize, is:

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

The KL divergence term expands to:

$$D_{KL}(q_\phi(z|x) || p(z)) = \mathbb{E}_{q_\phi(z|x)}[\log q_\phi(z|x) - \log p(z)]$$

Without access to $\log q_\phi(z|x)$, we cannot directly compute or optimize this KL divergence.

Adversarial Variational Bayes: Learning the Log-Posterior

Adversarial Variational Bayes, as introduced by Mescheder, Nowozin, and Geiger (2017), provides an elegant way to handle this intractability. The core idea is to train an auxiliary neural network, often called a critic or discriminator $T_\psi(x,z)$ (parameterized by $\psi$), to approximate the intractable log-density $\log q_\phi(z|x)$.

The Critic Network $T_\psi(x,z)$

The critic network $T_\psi(x,z)$ is not a typical GAN discriminator trying to distinguish "real" from "fake" samples. Instead, it's trained to estimate the log-density of the samples produced by the implicit encoder $g_\phi(x, \epsilon)$. This is achieved by optimizing $T_\psi(x,z)$ with respect to the following objective function, which should be minimized:

$$\mathcal{L}_T(\psi) = \mathbb{E}_{x \sim p_{data}(x)} \left[ \mathbb{E}_{\epsilon \sim p(\epsilon)} [\exp(T_\psi(x, g_\phi(x,\epsilon))) - T_\psi(x, g_\phi(x,\epsilon))] \right]$$

It can be shown that when this objective $\mathcal{L}_T(\psi)$ is minimized, the optimal critic $T^*_\psi(x,z)$ satisfies:

$$T^*_\psi(x,z) \approx \log q_\phi(z|x) - 1$$

The constant $-1$ is not an issue, as it will either cancel out or result in a constant offset in the ELBO, which doesn't affect the gradients with respect to $\phi$ or $\theta$ (except for the entropy part of the ELBO, where it's properly accounted for).

The AVB Objective for VAE Parameters

Once we have the critic $T_\psi(x,z)$ providing an estimate $T_\psi(x,z) + 1 \approx \log q_\phi(z|x)$, we can substitute this into the ELBO formula. The objective for the VAE's encoder parameters $\phi$ (which define $g_\phi$) and decoder parameters $\theta$ (which define $p_\theta(x|z)$) is to maximize:

$$\mathcal{L}_{AVB}(\theta, \phi) = \mathbb{E}_{x \sim p_{data}(x)} \left[ \mathbb{E}_{\epsilon \sim p(\epsilon)} [\log p_\theta(x|g_\phi(x,\epsilon)) + \log p(g_\phi(x,\epsilon)) - (T_\psi(x,g_\phi(x,\epsilon)) + 1)] \right]$$

Here, $z = g_\phi(x,\epsilon)$. The term $p(g_\phi(x,\epsilon))$ is the prior density evaluated at the generated latent sample $z$.

Training Dynamics: A Minimax Game

The overall training process for AVB becomes a minimax (or saddle-point) optimization problem:

1. Update Critic: Minimize $\mathcal{L}_T(\psi)$ with respect to the critic's parameters $\psi$. This step aims to make $T_\psi(x,z)$ a better approximation of $\log q_\phi(z|x) - 1$.
2. Update VAE: Maximize $\mathcal{L}_{AVB}(\theta, \phi)$ with respect to the encoder parameters $\phi$ and decoder parameters $\theta$. This step updates the VAE, using the current critic's estimate for the $\log q_\phi(z|x)$ term.

These two steps are typically alternated during training.

The following diagram illustrates the AVB architecture and flow:

Data flow and objective functions in Adversarial Variational Bayes. The critic $T_\psi(x,z)$ learns to estimate the log-density of the implicit posterior $q_\phi(z|x)$, which is then used in the VAE's objective function. $C$ represents the constant term.

The "Adversarial" Nature of AVB

It's important to clarify what "adversarial" means in the context of AVB. It's not adversarial in the same sense as Generative Adversarial Networks (GANs), where a discriminator tries to distinguish real data from generated data. In AVB:

• The encoder $g_\phi$ (which defines $q_\phi(z|x)$) generates latent samples $z$.
• The critic $T_\psi(x,z)$ tries to accurately estimate $\log q_\phi(z|x)$ for these samples.
• As the encoder $g_\phi$ updates, the distribution $q_\phi(z|x)$ changes, so the critic $T_\psi(x,z)$ must adapt to these changes.

The "adversarial" aspect arises from this interplay: the VAE parameters $(\theta, \phi)$ are updated to maximize an objective that depends on the critic $T_\psi$, while $T_\psi$ is concurrently updated to better model the density defined by $\phi$. This dynamic resembles a two-player game, forming the basis of the saddle-point optimization.

Advantages of AVB

Employing AVB can lead to significant improvements in VAEs:

1. Highly Flexible Posteriors: The primary benefit is the ability to use an implicit $q_\phi(z|x)$ defined by a powerful neural network $g_\phi(x, \epsilon)$. This allows the approximate posterior to capture complex structures, such as multimodality or non-Gaussian shapes, potentially matching the true posterior $p_\theta(z|x)$ much more closely.
2. Direct ELBO Optimization (Approximation): AVB still optimizes an approximation of the ELBO. By using a more expressive $q_\phi(z|x)$, AVB can achieve tighter ELBO values compared to VAEs with simpler, explicit posteriors, provided the critic's approximation of $\log q_\phi(z|x)$ is accurate.
3. Improved Representations and Samples: A more accurate posterior approximation often translates to higher-quality learned latent representations and, consequently, better generative performance (e.g., sharper, more diverse generated samples $x'$).

Challenges and Considerations

While powerful, AVB also introduces some challenges:

1. Training Stability: Minimax optimization problems can sometimes be challenging to train stably. Careful hyperparameter tuning for learning rates and network architectures for both the VAE components and the critic is often necessary.
2. Accuracy of the Critic: The effectiveness of AVB relies on the critic $T_\psi(x,z)$ providing a good estimate of $\log q_\phi(z|x)$. If the critic is poorly trained or lacks capacity, the approximation can be inaccurate, potentially hindering VAE training.
3. Computational Cost: Training an additional neural network (the critic) and performing the minimax optimization adds to the computational overhead per training iteration compared to standard VAEs.
4. Hyperparameter Sensitivity: The interplay between the VAE and the critic can introduce new hyperparameters that require careful tuning.

Despite these challenges, AVB represents a significant step towards more powerful and flexible variational inference in VAEs. By moving beyond simple, explicit posteriors, AVB allows VAEs to model more complex data distributions and learn richer latent representations, pushing the boundaries of what can be achieved with these generative models. It is a prime example of how sophisticated inference techniques can unlock new capabilities in VAEs.