Unveiling And Addressing Dimensional Collapse In Vector Quantization Models Via Codebook Regularization

While recent advancements in Vector Quantization (VQ) models have successfully achieved complete codebook utilization, a critical bottleneck remains largely unexplored: the effective dimensionality of the codebook embedding space. We observe that discrete codebook representations tend to degenerate into low-dimensional subspaces, characterized by significantly lower effective rank than continuous representations during quantization. Through comprehensive spectral analysis, we identify that this dimensional collapse stems from the suppression of low-variance components inherent to the vector quantization process, thereby severely limiting the expressive capacity of VQ models. To mitigate this fundamental issue, we propose a simple yet effective codebook regularization strategy designed to restore low-variance components, effectively bridging the spectral gap between discrete codebook spaces and continuous representations. Extensive experiments demonstrate that this regularization objective is compatible with diverse VQ training paradigms, yielding significant improvements in reconstruction fidelity and downstream performance in autoregressive image generative models.

Many AI systems that generate images, videos, and audio rely on a technique called vector quantization, which compresses complex data into a compact set of discrete codes — similar to how a dictionary reduces language to a fixed vocabulary. However, we discovered that these codebooks suffer from a hidden problem: even when all codes are actively used, the codes themselves collapse into a narrow, repetitive pattern, wasting most of their capacity to capture fine details. Through mathematical analysis, we traced this collapse to the suppression of subtle, low-energy features during the compression process. To fix this, we developed a simple regularization method that encourages each dimension of the codebook to carry unique information, effectively "waking up" the dormant capacity without disrupting the learning process. Our method consistently improves image reconstruction quality and generation fidelity across multiple existing systems, suggesting it could serve as a general-purpose upgrade for AI models that rely on discrete representations.