VQ-Diffusion: Conditional Latent Diffusion in Discrete Space

VQ-Diffusion is a conditional latent diffusion model developed by the University of Science and Technology of China and Microsoft. Unlike most common diffusion models that operate in continuous space, VQ-Diffusion performs its noising and denoising processes on a quantized latent space composed of a discrete set of vectors.

Architecture and Technical Implementation

VQ-Diffusion utilizes a combination of VQ-VAE for dimensionality reduction and an encoder-decoder transformer for the diffusion process.

VQ-VAE Latent Space

Images are first encoded into discrete tokens or embedding vectors using a VQ-VAE encoder. The process involves splitting images into patches and replacing each patch with the closest entry from a fixed-size vocabulary codebook. VQ-Diffusion specifically employs the VQGAN variant from Taming Transformers, using a pre-trained VQ-VAE that remains frozen during the diffusion training process.

The Forward Diffusion Process

In the forward process, latent tokens transition from a clean state to a noised state. Each token can:

  1. Remain the same.
  2. Be resampled to a different latent vector (with equal probability).
  3. Be masked.

Once a token is masked, it remains masked. The process is governed by hyperparameters $\alpha_{t}$, $\beta_{t}$, and $\gamma_{t}$, where $\gamma_{t}$ is the probability of a token becoming masked, and $\alpha_{t} + \beta_{t}$ is the probability it stays the same. The transition to any individual non-masked latent vector occurs with probability $\beta_{t}$.

Approximating the Reverse Process

To reverse the noise and generate images, an encoder-decoder transformer is used to approximate the classes of the un-noised latents ($x_{0}$) conditioned on a prompt ($y$).

  • Encoder: A CLIP text encoder with frozen weights.
  • Decoder: A transformer that provides unmasked global attention to all latent pixels and outputs the log probabilities of the categorical distribution over vector embeddings.

Because the decoder predicts the entire distribution of un-noised latents in a single forward pass, it provides global self-attention over the current state $x_{t}$.

VQ-Diffusion vs. Other Generative Models

VQ-Diffusion provides a technical alternative to both continuous diffusion models and autoregressive (AR) models.

Comparison with Continuous Diffusion

Most contemporary diffusion models are continuous, iteratively adding Gaussian noise and approximating the reverse process via a U-Net to predict that noise. In contrast, VQ-Diffusion operates on discrete values, where predicting the distribution for $x_{0}$ directly is a more clear objective than predicting noise.

Comparison with Autoregressive (AR) Models

VQ-Diffusion addresses three primary pain points associated with AR transformer-based image models: linear decreases in inference speed as resolution increases, error accumulation, and directional bias.

  • Inference Speed: AR models factor image probability such that each pixel is conditioned on previous pixels in raster scan order. VQ-Diffusion's inference complexity is superior as long as the number of diffusion steps is fewer than the number of latent pixels. For the ITHQ dataset (32x32 latent resolution), VQ-Diffusion is trained up to 100 steps, resulting in an approximate 10x big-O improvement. In practice, it can be up to 15 times faster than AR methods while achieving better image quality.
  • Error Correction: Unlike AR models that require teacher-forcing, VQ-Diffusion learns to correct incorrectly predicted tokens because training involves both masking and replacing latent pixels with random tokens.
  • Global Context: VQ-Diffusion provides global context on $x_{t}$ while predicting $x_{t - 1}$, whereas AR models are restricted by masked attention to avoid seeing future pixels.

Implementation and Further Optimizations

VQ-Diffusion is integrated into the Hugging Face diffusers library via the VQDiffusionPipeline. Models have been trained on various datasets including ITHQ, CUB-200, Oxford-102, MSCOCO, Conceptual Captions, LAION-400M, and ImageNet.

Further optimizations include a faster inference strategy using a time stride $\Delta t$ to skip reverse diffusion steps. Additionally, "Improved Vector Quantized Diffusion Models" introduces discrete classifier-free guidance to improve sample quality and an alternative inference strategy to address the joint distribution issue.

Sources