PCA-VAE: Differentiable Subspace Quantization without Codebook Collapse

TL;DR

PCA-VAE introduces a fully differentiable, PCA-based latent quantization method that surpasses traditional VQ in quality, interpretability, and efficiency, avoiding codebook collapse and complex regularizations.

Contribution

It presents a simple, mathematically grounded PCA-based alternative to vector quantization for autoencoders, improving stability and interpretability without codebooks or regularization.

Findings

Outperforms VQ-GAN and SimVQ in reconstruction quality on CelebAHQ.

Uses 10-100x fewer latent bits than traditional VQ methods.

Produces interpretable latent dimensions such as pose and lighting.

Abstract

Vector-quantized autoencoders deliver high-fidelity latents but suffer inherent flaws: the quantizer is non-differentiable, requires straight-through hacks, and is prone to collapse. We address these issues at the root by replacing VQ with a simple, principled, and fully differentiable alternative: an online PCA bottleneck trained via Oja's rule. The resulting model, PCA-VAE, learns an orthogonal, variance-ordered latent basis without codebooks, commitment losses, or lookup noise. Despite its simplicity, PCA-VAE exceeds VQ-GAN and SimVQ in reconstruction quality on CelebAHQ while using 10-100x fewer latent bits. It also produces naturally interpretable dimensions (e.g., pose, lighting, gender cues) without adversarial regularization or disentanglement objectives. These results suggest that PCA is a viable replacement for VQ: mathematically grounded, stable, bit-efficient, and semantically structured, offering a new direction for generative models beyond vector quantization.