Enhancing Vector Quantization with Distributional Matching: A Theoretical and Empirical Study

2025/06/18 by Xianghong Fang, Fang, Xianghong, L. B. Guo +18 · 2 citations
Computer Science · #Advanced Data Compression Techniques #Code (set theory) #Codebook #Computer Vision and Pattern Recognition (cs.CV) #Empirical research #FOS: Computer and information sciences #Linde–Buzo–Gray algorithm #Machine Learning (cs.LG) #Neural Networks and Applications #Quantization (signal processing) #Source code #Vector quantization

paper · pdf · doi:10.48550/arxiv.2506.15078

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2025/06/18 · openalex created_date 2025/10/19 · openalex updated_date 2026/08/05

Abstract

The success of autoregressive models largely depends on the effectiveness of vector quantization, a technique that discretizes continuous features by mapping them to the nearest code vectors within a learnable codebook. Two critical issues in existing vector quantization methods are training instability and codebook collapse. Training instability arises from the gradient discrepancy introduced by the straight-through estimator, especially in the presence of significant quantization errors, while codebook collapse occurs when only a small subset of code vectors are utilized during training. A closer examination of these issues reveals that they are primarily driven by a mismatch between the distributions of the features and code vectors, leading to unrepresentative code vectors and significant data information loss during compression. To address this, we employ the Wasserstein distance to align these two distributions, achieving near 100% codebook utilization and significantly reducing the quantization error. Both empirical and theoretical analyses validate the effectiveness of the proposed approach.

Citations

• XQ-GAN: An Open-source Image Tokenization Framework for Autoregressive Generation
• Addressing Representation Collapse in Vector Quantized Models with One Linear Layer
• Scaling the Codebook Size of VQGAN to 100,000 with a Utilization Rate of 99%
• Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation
• Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction
• Finite Scalar Quantization: VQ-VAE Made Simple
• Online Clustered Codebook
• Not All Image Regions Matter: Masked Vector Quantization for Autoregressive Image Generation
• Regularized Vector Quantization for Tokenized Image Synthesis
• Vector Quantized Wasserstein Auto-Encoder
• SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization
• Autoregressive Image Generation using Residual Quantization
• MaskGIT: Masked Generative Image Transformer
• High-Resolution Image Synthesis with Latent Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• Vector-quantized Image Modeling with Improved VQGAN
• Zero-Shot Text-to-Image Generation
• Taming Transformers for High-Resolution Image Synthesis
• Denoising Diffusion Probabilistic Models
• Jukebox: A Generative Model for Music
• Hierarchical Quantized Autoencoders
• Generating Diverse High-Fidelity Images with VQ-VAE-2
• A Style-Based Generator Architecture for Generative Adversarial Networks
• The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
• Neural Discrete Representation Learning
• Image-to-Image Translation with Conditional Adversarial Networks
• Perceptual Losses for Real-Time Style Transfer and Super-Resolution
• U-Net: Convolutional Networks for Biomedical Image Segmentation
• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation

Cited by

• WeTok: Powerful Discrete Tokenization for High-Fidelity Visual Reconstruction
• Distributional Matching for Vector Quantization: A Unified Theoretical and Empirical Framework

Related

• Distributional Matching for Vector Quantization: A Unified Theoretical and Empirical Framework
• Pushing Toward the Simplex Vertices: A Simple Remedy for Code Collapse in Smoothed Vector Quantization
• Representation Collapsing Problems in Vector Quantization
• DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick
• Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks