# SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization

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

## Key facts

- Cited by 68 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 50 references that are in the map.
- First posted on arXiv 16 May 2022.

## Builds on (most cited first, 30 of 34)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [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)
- [Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://mlsanity.com/paper/1710.10196/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Pixel Recurrent Neural Networks](https://mlsanity.com/paper/1601.06759/) (2016)
- [MaskGAN: Towards Diverse and Interactive Facial Image Manipulation](https://mlsanity.com/paper/1907.11922/) (2019)
- [Jukebox: A Generative Model for Music](https://mlsanity.com/paper/2005.00341/) (2020)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [MelGAN: Generative Adversarial Networks for Conditional Waveform Synthesis](https://mlsanity.com/paper/1910.06711/) (2019)
- [Parallel WaveGAN: A fast waveform generation model based on generative adversarial networks with multi-resolution spectrogram](https://mlsanity.com/paper/1910.11480/) (2019)
- [Lossy Image Compression with Compressive Autoencoders](https://mlsanity.com/paper/1703.00395/) (2017)
- [Hyperspherical Variational Auto-Encoders](https://mlsanity.com/paper/1804.00891/) (2018)
- [Variable Rate Image Compression with Recurrent Neural Networks](https://mlsanity.com/paper/1511.06085/) (2015)
- [Soft-to-Hard Vector Quantization for End-to-End Learning Compressible Representations](https://mlsanity.com/paper/1704.00648/) (2017)
- [PixelSNAIL: An Improved Autoregressive Generative Model](https://mlsanity.com/paper/1712.09763/) (2017)
- [From Variational to Deterministic Autoencoders](https://mlsanity.com/paper/1903.12436/) (2019)
- [Fast Decoding in Sequence Models using Discrete Latent Variables](https://mlsanity.com/paper/1803.03382/) (2018)
- [Unsupervised speech representation learning using WaveNet autoencoders](https://mlsanity.com/paper/1901.08810/) (2019)
- [Theory and Experiments on Vector Quantized Autoencoders](https://mlsanity.com/paper/1805.11063/) (2018)
- [Vector-quantized neural networks for acoustic unit discovery in the ZeroSpeech 2020 challenge](https://mlsanity.com/paper/2005.09409/) (2020)

## Built on by (most cited first, 30 of 68)

- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [Image and Video Tokenization with Binary Spherical Quantization](https://mlsanity.com/paper/2406.07548/) (2024)
- [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)
- [Online Clustered Codebook](https://mlsanity.com/paper/2307.15139/) (2023)
- [Restructuring Vector Quantization with the Rotation Trick](https://mlsanity.com/paper/2410.06424/) (2024)
- [HQ-VAE: Hierarchical Discrete Representation Learning with Variational Bayes](https://mlsanity.com/paper/2401.00365/) (2023)
- [Vector Quantized Wasserstein Auto-Encoder](https://mlsanity.com/paper/2302.05917/) (2023)
- [Grounded Object Centric Learning](https://mlsanity.com/paper/2307.09437/) (2023)
- [EdVAE: Mitigating Codebook Collapse with Evidential Discrete Variational Autoencoders](https://mlsanity.com/paper/2310.05718/) (2023)
- [Codebook Transfer with Part-of-Speech for Vector-Quantized Image Modeling](https://mlsanity.com/paper/2403.10071/) (2024)
- [Image Understanding Makes for A Good Tokenizer for Image Generation](https://mlsanity.com/paper/2411.04406/) (2024)
- [Cube: A Roblox View of 3D Intelligence](https://mlsanity.com/paper/2503.15475/) (2025)
- [SAN: Inducing Metrizability of GAN with Discriminative Normalized Linear Layer](https://mlsanity.com/paper/2301.12811/) (2023)
- [Continuous Latent Diffusion Language Model](https://mlsanity.com/paper/2605.06548/) (2026)
- [HyperVQ: MLR-based Vector Quantization in Hyperbolic Space](https://mlsanity.com/paper/2403.13015/) (2024)
- [Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning](https://mlsanity.com/paper/2406.03234/) (2024)
- [StraIT: Non-autoregressive Generation with Stratified Image Transformer](https://mlsanity.com/paper/2303.00750/) (2023)
- [Soft Convex Quantization: Revisiting Vector Quantization with Convex Optimization](https://mlsanity.com/paper/2310.03004/) (2023)
- [MGVQ: Could VQ-VAE Beat VAE? A Generalizable Tokenizer with Multi-group Quantization](https://mlsanity.com/paper/2507.07997/) (2025)
- [Unleash LLMs Potential for Sequential Recommendation by Coordinating Dual Dynamic Index Mechanism](https://mlsanity.com/paper/2409.09253/) (2024)
- [Discrete Tokenization for Multimodal LLMs: A Comprehensive Survey](https://mlsanity.com/paper/2507.22920/) (2025)
- [Learning Hierarchical Features with Joint Latent Space Energy-Based Prior](https://mlsanity.com/paper/2310.09604/) (2023)
- [Music Foundation Model as Generic Booster for Music Downstream Tasks](https://mlsanity.com/paper/2411.01135/) (2024)
- [Improving 3D Gaussian Splatting Compression by Scene-Adaptive Lattice Vector Quantization](https://mlsanity.com/paper/2509.13482/) (2025)
- [Enhancing Vector Quantization with Distributional Matching: A Theoretical and Empirical Study](https://mlsanity.com/paper/2506.15078/) (2025)
- [Instella-T2I: Pushing the Limits of 1D Discrete Latent Space Image Generation](https://mlsanity.com/paper/2506.21022/) (2025)
- [AI Mother Tongue: Self-Emergent Communication in MARL via Endogenous Symbol Systems](https://mlsanity.com/paper/2507.10566/) (2025)
- [DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick](https://mlsanity.com/paper/2509.26469/) (2025)
- [VAEVQ: Enhancing Discrete Visual Tokenization through Variational Modeling](https://mlsanity.com/paper/2511.06863/) (2025)

## Abstract (arXiv)

One noted issue of vector-quantized variational autoencoder (VQ-VAE) is that the learned discrete representation uses only a fraction of the full capacity of the codebook, also known as codebook collapse. We hypothesize that the training scheme of VQ-VAE, which involves some carefully designed heuristics, underlies this issue. In this paper, we propose a new training scheme that extends the standard VAE via novel stochastic dequantization and quantization, called stochastically quantized variational autoencoder (SQ-VAE). In SQ-VAE, we observe a trend that the quantization is stochastic at the initial stage of the training but gradually converges toward a deterministic quantization, which we call self-annealing. Our experiments show that SQ-VAE improves codebook utilization without using common heuristics. Furthermore, we empirically show that SQ-VAE is superior to VAE and VQ-VAE in vision- and speech-related tasks.