Early Quantization Shrinks Codebook: A Simple Fix for Diversity-Preserving Tokenization

2026-03-17

AI Summary

This work addresses the loss of diversity in discrete representations caused by codebook and embedding collapse in vector-quantized generative models. It presents the first systematic analysis of this issue, identifying improper random initialization and insufficient encoder capacity as the two primary causes of collapse. The authors propose simple yet effective strategies to mitigate these problems and rigorously evaluate their approach through empirical studies and ablation experiments on both synthetic and real-world datasets, clarifying the conditions that trigger different collapse phenomena and their respective impacts. The proposed method substantially enhances the diversity of discrete tokens and improves training stability, offering both theoretical insights and practical guidance for the reliable application of vector quantization in generative modeling.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Knowledge Representation and Reasoning: Qualitative Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Vector quantization is a technique in machine learning that discretizes continuous representations into a set of discrete vectors. It is widely employed in tokenizing data representations for large language models, diffusion models, and other generative models. Despite its prevalence, the characteristics and behaviors of vector quantization in generative models remain largely underexplored. In this study, we systematically investigate the issue of collapses in vector quantization, where collapsed representations are observed across discrete codebook tokens and continuous latent embeddings. By leveraging both synthetic and real datasets, we identify the severity of each type of collapses and triggering conditions. Our analysis reveals that random initialization and limited encoder capacity result in tokens collapse and embeddings collapse. Building on these findings, we propose potential solutions aimed at mitigating each collapse. To the best of our knowledge, this is the first comprehensive study examining representation collapsing problems in vector quantization.

Problem

• vector quantization
• representation collapse
• codebook collapse
• embedding collapse
• diversity preservation

Innovation

• vector quantization
• representation collapse
• codebook
• tokenization
• generative models