# Beyond Stationarity: Rethinking Codebook Collapse in Vector Quantization

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 16 of its 19 references that are in the map.
- First posted on arXiv 21 Feb 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

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

- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Deep Clustering for Unsupervised Learning of Visual Features](https://mlsanity.com/paper/1807.05520/) (2018)
- [SoundStream: An End-to-End Neural Audio Codec](https://mlsanity.com/paper/2107.03312/) (2021)
- [Vector-quantized Image Modeling with Improved VQGAN](https://mlsanity.com/paper/2110.04627/) (2021)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [MAGVIT: Masked Generative Video Transformer](https://mlsanity.com/paper/2212.05199/) (2022)
- [vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations](https://mlsanity.com/paper/1910.05453/) (2019)
- [Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks](https://mlsanity.com/paper/2305.08842/) (2023)
- [Addressing Representation Collapse in Vector Quantized Models with One Linear Layer](https://mlsanity.com/paper/2411.02038/) (2024)
- [Restructuring Vector Quantization with the Rotation Trick](https://mlsanity.com/paper/2410.06424/) (2024)
- [Exploration into Translation-Equivariant Image Quantization](https://mlsanity.com/paper/2112.00384/) (2021)

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

- [A Unified Rate-Distortion Perspective on Vector, Product, and Scalar Quantization](https://mlsanity.com/paper/2609.02107/) (2026)

## Abstract (arXiv)

Vector Quantization (VQ) underpins many modern generative frameworks such as VQ-VAE, VQ-GAN, and latent diffusion models. Yet, it suffers from the persistent problem of codebook collapse, where a large fraction of code vectors remains unused during training. This work provides a new theoretical explanation by identifying the nonstationary nature of encoder updates as the fundamental cause of this phenomenon. We show that as the encoder drifts, unselected code vectors fail to receive updates and gradually become inactive. To address this, we propose two new methods: Non-Stationary Vector Quantization (NSVQ), which propagates encoder drift to non-selected codes through a kernel-based rule, and Transformer-based Vector Quantization (TransVQ), which employs a lightweight mapping to adaptively transform the entire codebook while preserving convergence to the k-means solution. Experiments on the CelebA-HQ dataset demonstrate that both methods achieve near-complete codebook utilization and superior reconstruction quality compared to baseline VQ variants, providing a principled and scalable foundation for future VQ-based generative models. The code is available at: https://github.com/CAIR- LAB- WFUSM/NSVQ-TransVQ.git