Dimensional Collapse in VQVAEs: Evidence and Remedies

Vector-Quantized Variational Autoencoders (VQVAEs) have enabled strong performance in generative modeling by mapping continuous data to learnable codes. In this work, we identify a surprising yet consistent phenomenon that we term dimensional collapse: despite using high-dimensional embeddings, VQVAEs tend to compress their representations into a much smaller subspace, typically only 4 to 10 dimensions. We provide an in-depth analysis of this phenomenon and reveal its relation to model performance and learning dynamics. Interestingly, VQVAEs naturally gravitate toward this low-dimensional regime, and enforcing higher-dimensional usage (e.g., via rank regularization) could lead to degraded performance. To overcome this low-dimensionality limitation, we propose Divide-and-Conquer VQ (DCVQ), which partitions the latent space into multiple low-dimensional subspaces, each quantized independently. By design, each subspace respects the model's preference for low dimensionality, while their combination expands the overall capacity. Our results show that DCVQ overcomes the inherent dimensional bottleneck and achieves improved reconstruction quality across image datasets. Effective Dim 0.002 0.004 0.006 Val Loss CelebA, CNN, f=4 0 5 10 15 Effective Dim 0.0025 0.0050 0.0075 0.0100 Val Loss CelebA, ViT, f=4 0 5 10 15 20 25 Effective Dim 0.03 0.04 0.05 Val Loss CelebA, CNN, f=16 0 5 10 15 20 25 Effective Dim 0.03 0.04 0.05 0.06 Val Loss CelebA, ViT, f=16 0 10 20 30 40 50 Effective Dim 0.05 0.10 0.15 Val Loss CIFAR10, CNN, f=4 0 5 10 15 20 Effective Dim 0.05 0.10 0.15 0.20 Val Loss CIFAR10, ViT, f=4 0 5 10 15 20