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.
Figures & tables
Figure 1: Overall architecture of the proposed PCA-VAE. The encoder extracts latent features h from the input image x . The PCA layer performs an online orthogonal projection h^=CC⊤(h−μ)+μ , where C and μ are updated via Oja’s rule and r -fade averaging but treated as stop-gradient variables during VAE backpropagation. The quantized latent h^ is then decoded to reconstruct x^ . The PCA layer supports both global (single-vector) and spatial (multi-patch) latent configurations, each with its own PCA basis.
Figure 2: Normalized reconstruction performance. We compare PCA-VAE (16 × 16, 100% bases) with VQGAN [ 3 ] , SimVQ [ 27 ] , VQ-VAE [ 12 ] , and a VAE [ 4 ] baseline. All VQ models use 16 × 16 latents and 8,912 codebook tokens. Metrics are normalized (PSNR/SSIM min–max, LPIPS/rFID reverse) so that higher is better.
Figure 3: Scaling behavior of PCA-VAE with respect to the fraction of principal bases used (1%–100%) under different latent grid resolutions (1 × 1, 4 × 4, 8 × 8, 16 × 16). The red horizontal line marks the best VQ baseline (SimVQ [ 27 ] , 16 × 16).
Figure 4: Latent bit–budget curves comparing PCA–VAE to VQGAN [ 3 ] , SimVQ [ 27 ] , VQ-VAE [ 12 ] , and AutoencoderKL [ 4 ] . PCA–VAE achieves higher reconstruction quality per bit across PSNR, SSIM, LPIPS, and rFID, often matching VQ performance with 10 × -100 × fewer latent bits.
Figure 5: Latent semantics in PCA-VAE. We modify a single latent coefficient within the range [−2,2] while fixing the others. Each principal direction produces a coherent semantic transition (illumination, head pose, facial structure, shading, hair density), demonstrating interpretable continuous latent axes.
Figure 6: Latent interpolation along the 8th and 10th PCA dimensions. Although many dimensions capture well-defined semantics, some attributes co-vary (e.g., emotion and gender appearance, sunglasses and perceived masculinity). This reveals interpretable yet not strictly disentangled latent factors, likely reflecting data co-occurrence structure rather than model artifacts.
Figure 7: Qualitative reconstruction comparison on CelebA-HQ (256 × 256). Each row displays reconstructions from a different model using the same set of input images (top row). PCA-VAE produces sharp and identity-preserving reconstructions, noticeably improving facial structure (e.g., eye and mouth shapes) and global consistency compared to AutoencoderKL and VQ-VAE v2. SimVQ and VQGAN-FC perform competitively but occasionally exhibit blocky textures or over-smoothed regions, whereas PCA-VAE achieves similarly strong fidelity without vector quantization, codebook lookup, or STE.
Vector Quantized Variational Autoencoders (VQ-VAEs) are fundamental to modern generative modeling, yet they often suffer from training instability and "codebook collapse" due to the inherent coupling of representation learning and discrete codebook optimization. In this paper, we propose VP-VAE (Vector Perturbation VAE), a novel paradigm that decouples representation learning from discretization by eliminating the need for an explicit codebook during training. Our key insight is that, from the neural network's viewpoint, performing quantization primarily manifests as injecting a structured perturbation in latent space. Accordingly, VP-VAE replaces the non-differentiable quantizer with distribution-consistent and scale-adaptive latent perturbations generated via Metropolis--Hastings sampling. This design enables stable training without a codebook while making the model robust to inference-time quantization error. Moreover, under the assumption of approximately uniform latent variables, we derive FSP (Finite Scalar Perturbation), a lightweight variant of VP-VAE that provides a unified theoretical explanation and a practical improvement for FSQ-style fixed quantizers. Extensive experiments on image and audio benchmarks demonstrate that VP-VAE and FSP improve reconstruction fidelity and achieve substantially more balanced token usage, while avoiding the instability inherent to coupled codebook training.
Vector Quantization (VQ) underpins many modern generative frameworks such as VQ-VAE, VQ-GAN, and latent diffusion models. Yet, it suffers from the persistent problem of codebook collapse, where a large fraction of code vectors remains unused during training. This work provides a new theoretical explanation by identifying the nonstationary nature of encoder updates as the fundamental cause of this phenomenon. We show that as the encoder drifts, unselected code vectors fail to receive updates and gradually become inactive. To address this, we propose two new methods: Non-Stationary Vector Quantization (NSVQ), which propagates encoder drift to non-selected codes through a kernel-based rule, and Transformer-based Vector Quantization (TransVQ), which employs a lightweight mapping to adaptively transform the entire codebook while preserving convergence to the k-means solution. Experiments on the CelebA-HQ dataset demonstrate that both methods achieve near-complete codebook utilization and superior reconstruction quality compared to baseline VQ variants, providing a principled and scalable foundation for future VQ-based generative models. The code is available at: https://github.com/CAIR- LAB- WFUSM/NSVQ-TransVQ.git
Hao Lu, Onur C. Koyun, Yongxin Guo +3
Wake Forest University School of Medicine, Winston-Salem, NC, USA
Vector quantization is central to modern generative modeling pipelines, but large-codebook VQ models often suffer from codebook collapse. We identify encoder drift as a key driver of this failure: as the encoder moves the latent distribution, sparsely updated code vectors can lag behind, lose assignments, and increase quantization error, creating a feedback loop through the straight-through estimator. We propose NSVQ, a non-stationary-aware VQ training strategy that combines a dense non-stationary embedding loss, codebook replacement, and stage-wise encoder freezing. NSVQ first helps the codebook track encoder drift during early training, then freezes the encoder to consolidate the codebook under a fixed latent geometry, and finally reintroduces adversarial refinement. Experiments on ImageNet-1k show that NSVQ improves reconstruction quality while maintaining full codebook utilization. On ImageNet-1k at 128×128 with 65,536 codes, NSVQ reduces rFID from 2.39 to 2.10 compared with SimVQ, while both methods maintain 100% utilization. Additional latent diffusion experiments show that NSVQ also improves downstream ImageNet generation FID.
Hao Lu, Yongxin Guo, Onur Koyun +3
Wake Forest University School of Medicine / Advocate Health · Wake Forest University School of Medicine · Advocate Health