Training-Free Vector Quantization via Gaussian VAEs

Vector-quantized variational autoencoders (VQ-VAEs) are discrete autoencoders that compress images into discrete tokens. However, they are difficult to train due to discretization. In this paper, we propose a simple yet effective technique dubbed Gaussian Quant (GQ), which first trains a Gaussian VAE under certain constraints and then converts it into a VQ-VAE without additional training. For conversion, GQ generates random Gaussian noise as a codebook and finds the closest noise vector to the posterior mean. Theoretically, we prove that when the logarithm of the codebook size exceeds the bits-back coding rate of the Gaussian VAE, a small quantization error is guaranteed. Practically, we propose a heuristic to train Gaussian VAEs for effective conversion, named the target divergence constraint (TDC). Empirically, we show that GQ outperforms previous VQ-VAEs, such as VQGAN, FSQ, LFQ, and BSQ, on both UNet and ViT architectures. Furthermore, TDC also improves previous Gaussian VAE discretization methods, such as TokenBridge. The source code is provided in https://github.com/tongdaxu/VQ-VAE-from-Gaussian-VAE.

Many machine learning systems compress images into short sequences of discrete tokens, similar to how words represent text, but training these image tokenizers is often unstable because the model must make hard discrete choices. This paper introduces Gaussian Quant (GQ), a simple way to avoid that difficulty: first train a smoother Gaussian autoencoder, then convert it into a discrete autoencoder without any extra training. The conversion works by creating a random set of Gaussian code vectors and replacing each continuous image representation with the closest code vector. The authors prove that this conversion can keep the error small when the codebook is large enough relative to the information rate of the Gaussian autoencoder. They also propose a training rule called the target divergence constraint, which makes the Gaussian autoencoder easier to convert into a high-quality tokenizer. Experiments show that GQ produces better image reconstructions than several existing discrete autoencoder methods across both convolutional and transformer-based architectures. The same training rule also improves other methods that convert Gaussian autoencoders into discrete tokenizers. Overall, this work offers a simpler and more reliable route to building image tokenizers, which are important components for modern image generation and representation learning systems.