Improving Vector-Quantized Image Modeling with Latent Consistency-Matching Diffusion

TL;DR

This work tackles end-to-end learning of discrete image representations with diffusion models by training in a continuous latent space attached to a vector-quantized VAE. The proposed VQ-LCMD framework blends a joint embedding-diffusion variational bound with a consistency-matching loss, plus a shifted cosine noise schedule and random embedding dropout to mitigate embedding collapse. Empirical results on FFHQ, LSUN, and ImageNet demonstrate improved generation quality over discrete-state baselines and strong conditional performance, including a competitive ImageNet FID at 50 steps. The approach offers a practical, end-to-end pathway for high-fidelity discrete-data generation using continuous-space diffusion in latent representations.

Abstract

By embedding discrete representations into a continuous latent space, we can leverage continuous-space latent diffusion models to handle generative modeling of discrete data. However, despite their initial success, most latent diffusion methods rely on fixed pretrained embeddings, limiting the benefits of joint training with the diffusion model. While jointly learning the embedding (via reconstruction loss) and the latent diffusion model (via score matching loss) could enhance performance, end-to-end training risks embedding collapse, degrading generation quality. To mitigate this issue, we introduce VQ-LCMD, a continuous-space latent diffusion framework within the embedding space that stabilizes training. VQ-LCMD uses a novel training objective combining the joint embedding-diffusion variational lower bound with a consistency-matching (CM) loss, alongside a shifted cosine noise schedule and random dropping strategy. Experiments on several benchmarks show that the proposed VQ-LCMD yields superior results on FFHQ, LSUN Churches, and LSUN Bedrooms compared to discrete-state latent diffusion models. In particular, VQ-LCMD achieves an FID of 6.81 for class-conditional image generation on ImageNet with 50 steps.

Figures (12)

• Figure 1: Training procedure of VQ-LCMD. An image is compressed into a sequence of discrete tokens ${\mathbf{x}}$ using a pre-trained VQ-VAE. VQ-LCMD learns to generate the discrete latent representations ${\mathbf{x}}$ using the consistency-matching (CM) loss, diffusion loss, and reconstruction loss.
• Figure 1: Results of ablation studies on the FFHQ dataset
• Figure 2: Shifted cosine noise schedule with different shifting factors $s$, $\lambda(t) = \log\mathrm{SNR}(t)$.
• Figure 3: VQ-LCMD samples for unconditional generation
• Figure 4: Generated samples of VQ-LCMD with $\omega$ ranging from 0 to 4 on ImageNet.