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

2024-10-18

AI Summary

To address embedding collapse—a common issue in jointly training vector-quantized (VQ) embeddings and latent diffusion models for discrete data generation—this paper proposes the first end-to-end stable co-learning framework. Methodologically, it introduces a joint optimization objective combining a consistency matching (CM) loss with the variational lower bound, and designs two key components: a shifted cosine noise schedule and stochastic embedding dropout, both aimed at enhancing robustness and uniformity in the embedding space. Experimentally, the method achieves state-of-the-art performance on standard benchmarks including FFHQ and LSUN Churches/Beds. Notably, on ImageNet class-conditional generation, it attains an FID score of 6.81 within only 50 sampling steps—demonstrating superior efficiency, stability, and generalization over existing discrete latent diffusion approaches.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

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.

Problem

• Enhancing discrete data modeling via continuous latent diffusion
• Preventing embedding collapse in joint embedding-diffusion training
• Improving image generation quality with consistency-matching loss

Innovation

• Combines joint embedding-diffusion variational lower bound
• Uses consistency-matching loss for stable training
• Employs shifted cosine noise schedule strategy