# NSVQ: Mitigating Codebook Collapse by Stabilizing Encoder Drift in Vector Quantization

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 29 references that are in the map.
- First posted on arXiv 9 Jun 2026.

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

- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [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)
- [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)
- [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)
- [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)
- [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)
- [Theory and Experiments on Vector Quantized Autoencoders](https://mlsanity.com/paper/1805.11063/) (2018)
- [Online Clustered Codebook](https://mlsanity.com/paper/2307.15139/) (2023)
- [Regularized Vector Quantization for Tokenized Image Synthesis](https://mlsanity.com/paper/2303.06424/) (2023)
- [QLIP: Text-Aligned Visual Tokenization Unifies Auto-Regressive Multimodal Understanding and Generation](https://mlsanity.com/paper/2502.05178/) (2025)
- [Bridging Continuous and Discrete Tokens for Autoregressive Visual Generation](https://mlsanity.com/paper/2503.16430/) (2025)
- [Deep Embedded K-Means Clustering](https://mlsanity.com/paper/2109.15149/) (2021)
- [Scalable Training for Vector-Quantized Networks with 100% Codebook Utilization](https://mlsanity.com/paper/2509.10140/) (2025)
- [CODA: Repurposing Continuous VAEs for Discrete Tokenization](https://mlsanity.com/paper/2503.17760/) (2025)
- [Quantize-then-Rectify: Efficient VQ-VAE Training](https://mlsanity.com/paper/2507.10547/) (2025)

## Abstract (arXiv)

Vector quantization is central to modern generative modeling pipelines, but large-codebook VQ models often suffer from codebook collapse. We identify encoder drift as a key driver of this failure: as the encoder moves the latent distribution, sparsely updated code vectors can lag behind, lose assignments, and increase quantization error, creating a feedback loop through the straight-through estimator. We propose NSVQ, a non-stationary-aware VQ training strategy that combines a dense non-stationary embedding loss, codebook replacement, and stage-wise encoder freezing. NSVQ first helps the codebook track encoder drift during early training, then freezes the encoder to consolidate the codebook under a fixed latent geometry, and finally reintroduces adversarial refinement. Experiments on ImageNet-1k show that NSVQ improves reconstruction quality while maintaining full codebook utilization. On ImageNet-1k at 128×128 with 65,536 codes, NSVQ reduces rFID from 2.39 to 2.10 compared with SimVQ, while both methods maintain 100% utilization. Additional latent diffusion experiments show that NSVQ also improves downstream ImageNet generation FID.