Continuous First, Discrete Later: VQ-VAEs Without Dimensional Collapse

TL;DR

This paper identifies dimensional collapse in VQ-VAE models as a key issue and proposes an autoencoder warm-up phase to improve representation capacity, leading to better reconstruction and perceptual quality.

Contribution

It introduces a theoretical framework explaining dimensional collapse and demonstrates that an autoencoder warm-up phase effectively mitigates this problem in VQ-VAEs.

Findings

Warm-up phase restores representation dimension in VQ-VAEs.

Increases effective codebook dimension from 3-5 to 17-19.

Reduces rFID by 17-35% and improves PESQ by 11-14%.

Abstract

While many approaches to improve VQ-VAE performance focus on codebook size and utilization, the effect of dimensional collapse, where trained VQ-VAE representations live in an extremely low-dimensional subspace (1-2% of full rank), remains unaddressed. We show theoretically and empirically that dimension collapse causes a hard loss lower bound that various codebook improvement techniques fail to surpass. Our analytic framework extends the sequential learning effect of Saxe et al. [2014] by introducing ideas from rate-distortion theory and explains how the latent collapse is caused by the VQ suppressing lower-variance directions. Our theory justifies a simple solution: a "warm-up phase" that trains the model as an (unquantized) autoencoder before introducing VQ. On both synthetic experiments and large-scale image (VQGAN) and audio (WavTokenizer) VQ-VAEs, we show that AE Warm-Up successfully restores representation dimension, leading to lower reconstruction and perceptual loss at the same training budget. Across codebook sizes $K \in$ {$2^{10}, 2^{14}, 2^{16}$}, AE warm-up raises VQGAN codebook effective dimension from 3-5 to 17-19 and reduces rFID by 17-35%; on WavTokenizer at $K \in$ {$2^{13}, 2^{14}$}, it raises codebook dimension from 4 to 17-19 and improves PESQ by 11-14%. We empirically characterize how warm-up duration governs the achievable final loss. In agreement with experiment, our theoretical analysis predicts downstream performance as a function of warm-up length, enabling an adaptive criterion for switching from AE Warm-up to VQ-VAE training.