Global Context with Discrete Diffusion in Vector Quantised Modelling for Image Generation

TL;DR

The paper addresses the limited global context of autoregressive VQ-VAE image generation by introducing VQ-DDM, a two-stage framework that first compresses images into a discrete codebook and then learns a discrete diffusion prior over these codes. By employing a discrete diffusion process with a categorical transition on the codebook and a diffusion-based reverse model, the approach captures global structure while remaining computationally efficient. A key contribution is the ReFiT strategy, which rebuilds and fine-tunes the discrete codebook to dramatically improve code usage and reconstruction quality, enabling competitive performance with far fewer parameters. The proposed method delivers high-fidelity image generation and robust inpainting capabilities, with practical speed advantages over traditional diffusion models, and opens avenues for extending discrete diffusion to other data modalities.

Abstract

The integration of Vector Quantised Variational AutoEncoder (VQ-VAE) with autoregressive models as generation part has yielded high-quality results on image generation. However, the autoregressive models will strictly follow the progressive scanning order during the sampling phase. This leads the existing VQ series models to hardly escape the trap of lacking global information. Denoising Diffusion Probabilistic Models (DDPM) in the continuous domain have shown a capability to capture the global context, while generating high-quality images. In the discrete state space, some works have demonstrated the potential to perform text generation and low resolution image generation. We show that with the help of a content-rich discrete visual codebook from VQ-VAE, the discrete diffusion model can also generate high fidelity images with global context, which compensates for the deficiency of the classical autoregressive model along pixel space. Meanwhile, the integration of the discrete VAE with the diffusion model resolves the drawback of conventional autoregressive models being oversized, and the diffusion model which demands excessive time in the sampling process when generating images. It is found that the quality of the generated images is heavily dependent on the discrete visual codebook. Extensive experiments demonstrate that the proposed Vector Quantised Discrete Diffusion Model (VQ-DDM) is able to achieve comparable performance to top-tier methods with low complexity. It also demonstrates outstanding advantages over other vectors quantised with autoregressive models in terms of image inpainting tasks without additional training.

Figures (12)

• Figure 1: FID v.s. Operations and Parameters. The size of the blobs is proportional to the number of network parameters, the X-axis indicates FLOPs on a log scale and the Y-axis is the FID score.
• Figure 2: The proposed VQ-DDM pipeline contains 2 stages: (1) Compress the image into discrete variables via discrete VAE. (2) Fit a prior distribution over discrete coding by a diffusion model. Black squares in the diffusion diagram illustrate states when the underlying distributions are uninformative, but which become progressively more specific during the reverse process. The bar chart at the bottom of the image represents the probability of a particular discrete variable being sampled.
• Figure 3: Reconstruction images $384\times384$ from ImageNet based VQ-GAN and ReFiT
• Figure 4: Reconstruction images of CelebA HQ $256\times256$ from VQ-GAN and ReFiT.
• Figure 5: Steps and corresponding FID during the sampling. The text annotations are hours to sample 50k latent feature maps on 1 NVIDIA 2080Ti GPU