PCA-VAE: Differentiable Subspace Quantization without Codebook Collapse

2026-02-21

AI Summary

This work addresses key limitations of vector-quantized variational autoencoders—namely, the non-differentiability of quantizers, reliance on straight-through estimators, and codebook collapse—by introducing PCA-VAE. The proposed method replaces conventional vector quantization with a differentiable PCA bottleneck trained online via Oja’s rule, thereby eliminating the need for a discrete codebook, commitment loss, or lookup noise. This approach learns orthogonal latent representations ordered by variance and, to our knowledge, is the first to integrate differentiable PCA into generative modeling as a quantization alternative. On CelebA-HQ, PCA-VAE achieves superior reconstruction quality compared to VQ-GAN and SimVQ while using only 1/10 to 1/100 of the latent bit budget, and naturally yields semantically interpretable latent dimensions corresponding to attributes such as pose, lighting, and gender.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• vector quantization
• codebook collapse
• non-differentiable quantizer
• straight-through estimator

Innovation

• PCA-VAE
• differentiable quantization
• codebook collapse
• online PCA
• interpretable latent space