Powerful NVAE: No More Saying VAE-Generated Images are Blurry

Basic Introduction

VAE, or Variational Auto-Encoder, has been introduced in many articles on this blog. You can find many related posts by searching for "Variational Autoencoder" in the search bar. Here is a brief review and analysis.

In my derivation of VAE, we start with a set of samples representing a real (but unknown) distribution \tilde{p}(x). We then construct a parameterized posterior distribution p(z|x), and the two form a joint distribution p(x,z)=\tilde{p}(x)p(z|x). Next, we define a prior distribution q(z) and a generative distribution q(x|z), forming another joint distribution q(x,z)=q(z)q(x|z). Finally, our goal is to make p(x,z) and q(x,z) close to each other, so we optimize the KL divergence between them:

\begin{aligned} KL\big(p(x,z)\big\Vert q(x,z)\big) &= \iint p(x,z)\log \frac{p(x,z)}{q(x,z)} dzdx \\ &= \mathbb{E}_{x\sim \tilde{p}(x)} \Big[\mathbb{E}_{z\sim p(z|x)}\big[-\log q(x|z)\big]+KL\big(p(z|x)\big\Vert q(z)\big)\Big] + \text{const} \end{aligned}

This is the optimization objective of VAE.

Difficulty Analysis

The requirements for p(z|x), q(z), q(x|z) are: 1. They must have analytical expressions; 2. They must be easy to sample from. However, in the world of continuous distributions, there aren’t many such distributions. The most commonly used is the Gaussian distribution, and among those, the "Gaussian distribution with independent components" is the simplest. Therefore, in a typical VAE, p(z|x), q(z), q(x|z) are all set as Gaussians with independent components: p(z|x)=\mathcal{N}(z;\mu_1(x),\sigma_1^2(x)), q(z)=\mathcal{N}(z;0,1), and q(x|z)=\mathcal{N}(x;\mu_2(z),\sigma_2^2(z)).

The problem is that a "Gaussian distribution with independent components" cannot fit arbitrarily complex distributions. Once we choose the form of p(z|x), it is possible that no matter how we tune its parameters, \int \tilde{p}(x)p(z|x)dx and \frac{\tilde{p}(x)p(z|x)}{\int \tilde{p}(x)p(z|x)dx} will never become Gaussian. This means KL\big(p(x,z)\big\Vert q(x,z)\big) can never theoretically be zero. Thus, forcing p(x,z) and q(x,z) to approach each other only yields a rough, averaged result, which is why conventional VAE images are blurry.

Related Improvements

A classic direction for improving VAE is combining it with GANs, such as CVAE-GAN and AGE. Currently, the most advanced result in this direction is likely IntroVAE. Theoretically, this type of work implicitly abandons the assumption that q(x|z) is Gaussian, replacing it with a more general distribution, thereby improving generation. However, I feel that introducing GANs into VAE is a bit like "dealing with the devil"—you gain performance but also inherit GAN’s drawbacks (training instability, etc.), and even then, the performance of these VAEs is still inferior to pure GANs. Another direction is combining VAE with flow models, such as IAF-VAE and my previous work f-VAE. These works use flow models to enhance the expressive power of p(z|x) or q(x|z).

There is also the direction of introducing discrete latent variables, with the typical representative being VQ-VAE (see my "A Brief Introduction to VQ-VAE: Quantized Autoencoders"). VQ-VAE encodes images into discrete sequences using specific encoding tricks, and then uses PixelCNN to model the corresponding prior q(z). As mentioned, when z is continuous, the choices for p(z|x) and q(z) are limited, leading to limited approximation accuracy. But if z is a discrete sequence, p(z|x) and q(z) correspond to discrete distributions, and by using autoregressive models (called Language Models in NLP, PixelRNN/PixelCNN in CV), we can approximate any discrete distribution. Thus, the overall approximation can be much more precise, improving generation. The subsequent VQ-VAE-2 further confirmed the effectiveness of this path. However, overall, the workflow of VQ-VAE differs significantly from conventional VAEs, making it sometimes difficult to view as a VAE variant.

Autoregressive Distribution

We have already analyzed that the difficulty of VAE stems from p(z|x), q(z), q(x|z) not being strong enough. Therefore, the improvement strategy is to enhance them. First, NVAE does not change q(x|z), mainly to maintain parallel generation. Instead, it enhances the prior q(z) and posterior p(z|x) through autoregressive models. Specifically, it groups the latent variables as z=\{z_1,z_2,\dots,z_L\}, where each z_l is still a vector (not a single number), and sets:

q(z)=\prod_{l=1}^L q(z_l|z_{< l}),\quad p(z|x)=\prod_{l=1}^L p(z_l|z_{< l},x) \label{eq:arpq}

The q(z_l|z_{< l}) and p(z_l|z_{< l},x) for each group are still modeled as Gaussian distributions. Thus, overall, q(z) and p(z|x) are modeled as autoregressive Gaussian models. The KL divergence term for the posterior then becomes:

KL\big(p(z|x)\big\Vert q(z)\big)=KL\big(p(z_1|x)\big\Vert q(z_1)\big)+\sum_{l=2}^L \mathbb{E}_{p(z_{< l}|x)}\Big[KL\big(p(z_l|z_{< l}, x)\big\Vert q(z_l|z_{< l})\big)\Big]

Of course, this approach is a straightforward generalization and not an NVAE original; it can be traced back to models like DRAW and HVM from 2015. NVAE’s contribution is proposing a "relative" design for Equation [eq:arpq]:

\begin{aligned} &q(z_l|z_{< l})=\mathcal{N}\left(z_l;\mu(z_{< l}),\sigma^2(z_{< l})\right)\\ &p(z_l|z_{< l},x)=\mathcal{N}\left(z_l;\mu(z_{< l})+\Delta\mu(z_{< l},x),\sigma^2(z_{< l})\otimes \Delta\sigma^2(z_{< l}, x)\right) \end{aligned}

In other words, instead of directly predicting the mean and variance of the posterior p(z_l|z_{< l},x), it models the relative values compared to the prior’s mean and variance. In this case, we have (omitting arguments for simplicity):

KL\big(p(z_l|z_{< l}, x)\big\Vert q(z_l|z_{< l})\big)=\frac{1}{2} \sum_{i=1}^{|z_l|} \left(\frac{\Delta\mu_{(i)}^2}{\sigma_{(i)}^2} + \Delta\sigma_{(i)}^2 - \log \Delta\sigma_{(i)}^2 - 1\right)

The original paper points out that this approach makes training more stable.

Multi-scale Design

Now that the latent variables are divided into L groups z=\{z_1,z_2,\dots,z_L\}, the questions are: 1. How does the encoder generate z_1,z_2,\dots,z_L one by one? 2. How does the decoder utilize z_1,z_2,\dots,z_L one by one? That is, how are the encoder and decoder designed?

NVAE cleverly designs a multi-scale encoder and decoder, as shown above. First, the encoder goes through layers of encoding to obtain the top-level encoding vector z_1, and then gradually moves from the top down to obtain lower-level features z_2,\dots,z_L. As for the decoder, it is naturally a top-down process of utilizing z_1,z_2,\dots,z_L. Since this part shares similarities with the encoder’s process of generating z_1,z_2,\dots,z_L, NVAE directly makes the corresponding parts share parameters. This saves parameters and improves generalization through mutual constraints.

This multi-scale design is reflected in almost all current state-of-the-art generative models, such as StyleGAN, BigGAN, and VQ-VAE-2, proving its effectiveness. Furthermore, to ensure performance, NVAE carefully screened residual module designs, finally settling on the following residual block:

Other Enhancement Techniques

Besides the two obvious features above, NVAE includes many other tricks that improve performance. Here are a few.

Improvements to BN layers. Many current generative models have abandoned BN (Batch Normalization), mostly switching to IN (Instance Normalization) or WN (Weight Normalization). NVAE found through experiments that BN is actually helpful for training but harmful for prediction because the moving average mean and variance used during prediction aren’t good enough. Therefore, after training, NVAE re-estimates the mean and variance by sampling batches of the same size multiple times. Additionally, to ensure training stability, NVAE adds a regularization term to the magnitude of BN’s \gamma.

Application of Spectral Regularization. To stabilize training, NVAE adds spectral regularization to every convolutional layer (refer to my "Lipschitz Constraints in Deep Learning: Generalization and Generative Models"). Adding spectral normalization reduces the model’s Lipschitz constant, making the model’s landscape smoother and more conducive to stable training.

Flow models to enhance distributions. NVAE provides an alternative: modeling the intra-group distribution as a flow model to enhance expressiveness while maintaining parallel intra-group sampling.

Memory-saving tricks. NVAE uses mixed-precision training and promotes their own APEX library. It also enables gradient checkpointing (recomputation) in BN, which reportedly saves 18% of memory with almost no impact on speed.

More Results