VP-VAE: Rethinking Vector Quantization via Adaptive Vector Perturbation

TL;DR

VP-VAE introduces a novel approach to vector quantization in generative models by replacing explicit codebooks with adaptive latent perturbations, leading to more stable training and better performance.

Contribution

The paper proposes VP-VAE, a new paradigm that decouples representation learning from discretization, eliminating the need for a codebook and improving training stability.

Findings

Improved reconstruction fidelity on image and audio benchmarks.

More balanced token usage compared to traditional VQ-VAEs.

Enhanced robustness to inference-time quantization errors.

Abstract

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.