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

arXiv: https://arxiv.org/abs/2506.15078 · HTML: https://mlsanity.com/paper/2506.15078/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 43 references that are in the map.
- First posted on arXiv 18 Jun 2025.

## Builds on (most cited first, 29 of 29)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Perceptual Losses for Real-Time Style Transfer and Super-Resolution](https://mlsanity.com/paper/1603.08155/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction](https://mlsanity.com/paper/2404.02905/) (2024)
- [Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation](https://mlsanity.com/paper/2406.06525/) (2024)
- [Autoregressive Image Generation using Residual Quantization](https://mlsanity.com/paper/2203.01941/) (2022)
- [Vector-quantized Image Modeling with Improved VQGAN](https://mlsanity.com/paper/2110.04627/) (2021)
- [Jukebox: A Generative Model for Music](https://mlsanity.com/paper/2005.00341/) (2020)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [Addressing Representation Collapse in Vector Quantized Models with One Linear Layer](https://mlsanity.com/paper/2411.02038/) (2024)
- [SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization](https://mlsanity.com/paper/2205.07547/) (2022)
- [Scaling the Codebook Size of VQGAN to 100,000 with a Utilization Rate of 99%](https://mlsanity.com/paper/2406.11837/) (2024)
- [Hierarchical Quantized Autoencoders](https://mlsanity.com/paper/2002.08111/) (2020)
- [Online Clustered Codebook](https://mlsanity.com/paper/2307.15139/) (2023)
- [Regularized Vector Quantization for Tokenized Image Synthesis](https://mlsanity.com/paper/2303.06424/) (2023)
- [Not All Image Regions Matter: Masked Vector Quantization for Autoregressive Image Generation](https://mlsanity.com/paper/2305.13607/) (2023)
- [Vector Quantized Wasserstein Auto-Encoder](https://mlsanity.com/paper/2302.05917/) (2023)
- [XQ-GAN: An Open-source Image Tokenization Framework for Autoregressive Generation](https://mlsanity.com/paper/2412.01762/) (2024)

## Built on by (most cited first, 5 of 5)

- [WeTok: Powerful Discrete Tokenization for High-Fidelity Visual Reconstruction](https://mlsanity.com/paper/2508.05599/) (2025)
- [VQ-Transplant: Efficient VQ-Module Integration for Pre-trained Visual Tokenizers](https://mlsanity.com/paper/2607.19575/) (2026)
- [Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization](https://mlsanity.com/paper/2602.01140/) (2026)
- [UniWeTok: An Unified Binary Tokenizer with Codebook Size mathit2128 for Unified Multimodal Large Language Model](https://mlsanity.com/paper/2602.14178/) (2026)
- [StableVQ: Practical Guidelines for Stable Vector-Quantized Tokenizer Training](https://mlsanity.com/paper/2609.26774/) (2026)

## Abstract (arXiv)

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.