# EdVAE: Mitigating Codebook Collapse with Evidential Discrete Variational Autoencoders

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 53 references that are in the map.
- First posted on arXiv 9 Oct 2023.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [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)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [LSUN: Construction of a Large-scale Image Dataset using Deep Learning with Humans in the Loop](https://mlsanity.com/paper/1506.03365/) (2015)
- [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)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Jukebox: A Generative Model for Music](https://mlsanity.com/paper/2005.00341/) (2020)
- [Deep Evidential Regression](https://mlsanity.com/paper/1910.02600/) (2019)
- [Ladder Variational Autoencoders](https://mlsanity.com/paper/1602.02282/) (2016)
- [Lossy Image Compression with Compressive Autoencoders](https://mlsanity.com/paper/1703.00395/) (2017)
- [Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images](https://mlsanity.com/paper/2011.10650/) (2020)
- [Bailando: 3D Dance Generation by Actor-Critic GPT with Choreographic Memory](https://mlsanity.com/paper/2203.13055/) (2022)
- [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)
- [Discrete Variational Autoencoders](https://mlsanity.com/paper/1609.02200/) (2016)
- [Fast Decoding in Sequence Models using Discrete Latent Variables](https://mlsanity.com/paper/1803.03382/) (2018)
- [Evidential Deep Learning for Open Set Action Recognition](https://mlsanity.com/paper/2107.10161/) (2021)
- [Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks](https://mlsanity.com/paper/2305.08842/) (2023)

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

- [Next Token Prediction Towards Multimodal Intelligence: A Comprehensive Survey](https://mlsanity.com/paper/2412.18619/) (2024)
- [Wavelet Latent Diffusion (Wala): Billion-Parameter 3D Generative Model with Compact Wavelet Encodings](https://mlsanity.com/paper/2411.08017/) (2024)
- [Representation Collapsing Problems in Vector Quantization](https://mlsanity.com/paper/2411.16550/) (2024)
- [CODA: Repurposing Continuous VAEs for Discrete Tokenization](https://mlsanity.com/paper/2503.17760/) (2025)
- [Elucidating the design space of language models for image generation](https://mlsanity.com/paper/2410.16257/) (2024)
- [Balance of Number of Embedding and their Dimensions in Vector Quantization](https://mlsanity.com/paper/2407.04939/) (2024)
- [BioSerenity-E1: a self-supervised EEG model for medical applications](https://mlsanity.com/paper/2503.10362/) (2025)
- [ELVAE: Evidential Learning-Based Variational Autoencoder for Uncertainty-Aware Generation](https://mlsanity.com/paper/2608.10398/) (2026)
- [MARRS: Masked Autoregressive Unit-based Reaction Synthesis](https://mlsanity.com/paper/2505.11334/) (2025)
- [Semantics Meet Signals: Dual Codebook Representationl Learning for Generative Recommendation](https://mlsanity.com/paper/2511.20673/) (2025)
- [UniGRec: Unified Generative Recommendation with Soft Identifiers for End-to-End Optimization](https://mlsanity.com/paper/2601.17438/) (2026)
- [Early Quantization Shrinks Codebook: A Simple Fix for Diversity-Preserving Tokenization](https://mlsanity.com/paper/2603.17052/) (2026)
- [LASERS: LAtent Space Encoding for Representations with Sparsity for Generative Modeling](https://mlsanity.com/paper/2409.11184/) (2024)
- [Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization](https://mlsanity.com/paper/2602.01140/) (2026)
- [Investigating Permutation-Invariant Discrete Representation Learning for Spatially Aligned Images](https://mlsanity.com/paper/2604.01843/) (2026)
- [Collapse-Free Prototype Readout Layer for Transformer Encoders](https://mlsanity.com/paper/2604.03850/) (2026)

## Abstract (arXiv)

Codebook collapse is a common problem in training deep generative models with discrete representation spaces like Vector Quantized Variational Autoencoders (VQ-VAEs). We observe that the same problem arises for the alternatively designed discrete variational autoencoders (dVAEs) whose encoder directly learns a distribution over the codebook embeddings to represent the data. We hypothesize that using the softmax function to obtain a probability distribution causes the codebook collapse by assigning overconfident probabilities to the best matching codebook elements. In this paper, we propose a novel way to incorporate evidential deep learning (EDL) instead of softmax to combat the codebook collapse problem of dVAE. We evidentially monitor the significance of attaining the probability distribution over the codebook embeddings, in contrast to softmax usage. Our experiments using various datasets show that our model, called EdVAE, mitigates codebook collapse while improving the reconstruction performance, and enhances the codebook usage compared to dVAE and VQ-VAE based models. Our code can be found at https://github.com/ituvisionlab/EdVAE .