Training-Free Vector Quantization via Gaussian VAEs

1. Introduction

• GQ outperforms previous VQ-VAEs in reconstruction across UNet and ViT architectures.The reported ImageNet rate-distortion comparison covers both architectures.Source: Figure 1 · p. 1, §1 · p. 2
• GQ converts a constrained Gaussian VAE into a VQ-VAE without additional training by matching posterior means to Gaussian-noise codewords.The paper's theory states that quantization error is small when the codebook bitrate exceeds the Gaussian VAE's bits-back coding rate.Source: §1 · pp. 1–2
• Target divergence constraint (TDC) encourages equal per-dimension KL divergence corresponding to the bits-back coding bitrate.The paper introduces TDC to train Gaussian VAEs for effective conversion.Source: §1 · p. 2
• TDC also improves the performance of TokenBridge (Wang et al., 2025).Source: §1 · p. 2

2. Preliminaries

• The paper denotes the image by X, latent by Z, encoder by f(·), and decoder by g(·); each latent dimension's bits-back rate is its KL divergence from the standard normal prior.Source: §2 · p. 2, §2.2 · p. 3

2.1. Vector Quantized Variational Autoencoder

• A VQ-VAE encodes each latent dimension as its nearest codeword in a shared codebook, with codebook size K giving bitrate log K.Its objective combines rate-distortion with a regularization term.Source: §2.1 · p. 2

2.2. Gaussian VAE and Bits-Back Coding

• A Gaussian VAE uses a standard normal prior and factorized Gaussian posterior; its negative ELBO corresponds to a rate-distortion objective with bits-back bitrate and distortion terms.The per-dimension bits-back rate is defined as KL divergence from the prior.Source: §2.2 · p. 3

3.1. Converting a Gaussian VAE into a VQ-VAE

• GQ converts a Gaussian VAE by quantizing each latent dimension’s posterior mean to its nearest entry in a shared, fixed codebook of one-dimensional Gaussian samples.The codebook is generated once; conversion requires no additional training.Source: §3.1 · p. 3, §3 · p. 3

3.2. Theoretical Relationship between the Codebook Size and Quantization Error

• Under the theorem’s bounds on µ_i and σ_i, when log K exceeds R_i, the probability of |ẑ_i−µ_i| ≥ σ_i decays doubly exponentially; below R_i, it rises exponentially toward 1.The theorem defines R_i as the KL divergence of the latent posterior from N(0, 1).Source: §3.2 · p. 3
• Choosing log K close to R_i is the practical guideline; larger K offers no additional benefit, while smaller K significantly increases error.The passage notes that rounding the rate to choose K typically yields sufficiently small reconstruction error.Source: §3.2 · p. 3

3.3. Training a Conversion-ready Gaussian VAE with Target Divergence Constraint

• TDC targets R_i = log K for each dimension, applying larger penalties above log K + α and smaller penalties below log K − α.It extends a prior heuristic with per-dimension rate control.Source: §3.3 · p. 4

3.4. Supporting Multi-dimensional Codebook

• For m > 1, GQ selects among factorized Gaussian codewords using σ_i-weighted distance; at low bitrates, a regularizer encourages selecting codewords farther from zero to address collapse.Without the regularizer, codebook vectors far from zero may never be selected.Source: §3.4 · p. 4

3.5. Practical Guidelines

• The practical recipe trains a Gaussian VAE with latent dimension d = mN and TDC target log 3K, then converts it using a fixed codebook.The target is stated as log 3K in the supplied passage.Source: §3.5 · p. 4

3.6. Improving TokenBridge with Target Divergence Constraint

• TDC can enhance TokenBridge by matching its PTQ bits to R_i; TokenBridge quantizes a vanilla Gaussian VAE without per-dimension KL constraints.Its quantization centers are equal-probability partition centers of N(0, 1).Source: §3.6 · p. 4
• On ImageNet, GQ outperforms other VQ methods on both UNet and ViT architectures.The table compares methods using the same model architecture.Source: Table 1 · p. 5

4.1. Experimental Setup

• The evaluation uses ImageNet training data and ImageNet and COCO validation data, testing reconstruction at codebook sizes K from 2^14 to 2^18 and reporting reconstruction and generation metrics.Images are resized to 256 × 256 for reconstruction and 128 × 128 for generation; evaluated bitrates span 0.07–1.00 bpp.Source: §4.1 · p. 5

4.2. Main Results

• GQ achieves state-of-the-art reconstruction in most cases across UNet and ViT on ImageNet and COCO, while outperforming prior Gaussian-VAE discretization methods.Against pre-trained VQ-VAEs on ImageNet, GQ also outperforms while requiring less training.Source: §4.2 · pp. 5–6, Table 2 · p. 6
• For image generation, GQ has higher codebook usage and entropy than VQGAN, FSQ, LFQ, and BSQ, with generation FID and IS comparable to FSQ and better than the others.Under limited computational resources, autoregressive generation is more efficient than diffusion in both FID and IS.Source: §4.2 · p. 6, Table 6 · p. 7

4.3. Ablation Study

• Among conversion alternatives, GQ achieves the best PSNR, SSIM, and rFID; for codebook dimension m = 1, bisection search makes it asymptotically faster.Source: §4.3 · p. 6, Table 4 · p. 7
• GQ works best when log K matches the bits-back rate Ri, and its performance is nearly invariant to the random seed used to sample the fixed Gaussian codebook.Source: §4.3 · p. 6
• With TDC, GQ outperforms GQ applied to a vanilla Gaussian VAE: PSNR is 31.25 vs. 26.43 dB and rFID is 0.372 vs. 0.978.The two Gaussian VAEs have similar reconstruction performance before quantization, while the vanilla model’s per-dimension bits-back rates vary more widely.Source: §4.3 · pp. 6–7
• Training GQ approaches directly with a fixed Gaussian codebook or a Gumbel-Softmax target fails to converge reliably, while fine-tuning after Gaussian-VAE initialization adds only marginal improvement.Source: Table 9 · p. 7, §4.3 · p. 8
• TDC’s performance is minimally affected by α and β when α ≤ 0.5; larger codebook dimension m improves reconstruction, and the chosen ω depends on bitrate.The reported settings are ω = 2.0 for bitrates ≤ 0.50 bpp and ω = 0.0 for 1.00 bpp.Source: §4.3 · p. 8, Table 10 · p. 8

• VQ-VAEs are difficult to train directly because discretization prevents gradient descent and can cause codebook collapse, motivating techniques including commitment loss, EM, STE, and Gumbel-Softmax.Source: §5 · p. 8
• Fixed-codebook methods FSQ, LFQ, and BSQ still rely on STE, while TokenBridge and ReVQ convert pretrained Gaussian VAEs without constraining their divergence.Source: §5 · p. 8
• Reverse channel coding methods simulate distributions stochastically, whereas VQ-VAE requires deterministic quantization; GQ is reported to outperform MRC on quantization error and supports bisection search for m = 1.Source: §5 · p. 8
• TDC is related to MIRACLE / HiFiC and IsoKL: the former uses one λ and controls the range of Ri less effectively, while the latter controls Ri strictly but has numerical instability and suboptimal performance.Source: §5 · p. 8

• GQ trains a constrained Gaussian VAE, then converts it into a VQ-VAE by matching posterior means to a Gaussian-noise codebook, without additional training.Source: §6 · p. 8
• The authors report that GQ outperforms VQGAN, FSQ, LFQ, and BSQ, while TDC improves TokenBridge.Source: §6 · pp. 8–9
• GQ is more complex to implement and adds a hyperparameter; this work uses a standard single-scale architecture and focuses on reconstruction, leaving the relationship between reconstruction and generation for future work.Source: §6 · p. 9

• The approach focuses on reconstructing existing images at limited bitrate.Source: Impact Statement · p. 9
• Because the model is essentially not generative, the authors say its ethical concerns are not obvious.Source: Impact Statement · p. 9
• The authors note that the decoder’s GAN module might have negative effects involving misrepresentation and trustworthiness.Source: Impact Statement · p. 9

A. Proof of Main Results

• With bounded |μ_iσ_i| and |μ_i| + |σ_i|, the probability of |ẑ_i−μ_i| ≥ σ_i decays doubly exponentially as log K exceeds the bits-back rate R_i.The theorem holds for a fixed R_i = D_KL(q(Z_i|X)||N(0, 1)).Source: A. Proof of Main Results · p. 13
• When log K falls below the bits-back rate by t nats, the probability of |ẑ_i−μ_i| ≥ σ_i increases exponentially with t.Source: A. Proof of Main Results · p. 15

B.1. Mean-KL Parametrization

• Mean-KL parametrization uses two m-dimensional outputs to allocate the target log K across dimensions and determine posterior means and variances.Source: §B.1 · p. 16

B.2. Stable Mean-KL Parametrization

• Stable Mean-KL parametrization addresses NaNs associated with excessively large |μ_i| or small κ_i, using r1 = 0.1 and r2 = 0.01.r1 ensures κ_i ≥ r1/m, while r2 shrinks μ_i toward 0.Source: §B.2 · p. 16

B.3. Results of Stable Mean-KL Parametrization

• Stable Mean-KL achieves performance comparable to TDC on UNet models.Source: §B.3 · p. 16, Table 12 · p. 16
• On ViT models, Stable Mean-KL performs significantly worse than TDC.Source: §B.3 · p. 16, Table 12 · p. 16
• The authors choose TDC for the final model; Stable Mean-KL remains a simpler, effective alternative when only UNet models are required.It avoids constructing an empirical R(λ) model.Source: §B.3 · p. 16

C.1. Details of Training and Distortion Objective

• UNet and ViT VQ-VAEs were trained on ImageNet validation for about 24 hours on 8 H100 GPUs, with architecture-specific epochs, batch sizes, and learning rates.UNet: 30 epochs, batch size 128, learning rate 1 × 10^-4; ViT: 40 epochs, batch size 256, learning rate 4 × 10^-7.Source: §C.1 · p. 17
• For the distortion objective, following Stable Diffusion, UNet models use w1 = 1.0 and w2 = 0.75; consistent with BSQ, ViT models use w1 = 0.1 and w2 = 0.1.For image generation, VQ-VAEs are first trained on 128 × 128 images, then an autoregressive transformer is trained for 100 epochs.Source: §C.1 · p. 17

C.2. Details of Hyper-parameters

• GQ uses K = 2^16; its m and ω settings are (16, 2.0), (8, 2.0), and (4, 0.0) at 0.25, 0.5, and 1.0 bpp, respectively.Source: §C.2 · p. 17

D.1. Detail in Codebook Usage Hyper-parameter ω

• At 1.00 bpp, ω = 0.0 gives good-enough codebook usage and rFID, while at 0.50 bpp it causes codebook collapse; ω = 2.0 gives the best entropy and rFID there.Source: §D.1 · p. 17, Table 14 · p. 18

D.2. The Quantization Error in Pixel Space

• The pixel-space error bound assumes a Lipschitz decoder and scales with its constant c3; measured quality loss from GQ remains reasonable.Source: §D.2 · p. 18

D.3. Quantized Latent Visualization

• The t-NSE visualization compares GQ-quantized latents with unquantized Gaussian VAE latents and shows them as quite similar.Source: §D.3 · p. 18, Figure 3 · p. 20

D.4. More Generation Results

• GQ's generation advantage is reported across bitrate regimes and datasets.The appendix presents generation results for different bitrates and for FFHQ.Source: §D.4 · p. 19, Table 17 · p. 19, Table 18 · p. 20

D.5. Prior-Posterior Mismatch

• 0.00033 bits-per-pixel is the estimated prior-posterior mismatch, about 0.1% of total bpp; best and actual bpp differ insignificantly at a 0.01 scale.The authors describe bitrate waste from the mismatch as negligible.Source: §D.5 · p. 20, Table 19 · p. 20

D.6. Complexity

• PyTorch GQ encodes at 12 FPS versus 104 FPS for Gaussian VAE, while CUDA GQ reaches about 80 FPS with negligible overhead; GQ adds no decoding overhead.The codebook can be generated on the fly, leaving GQ with the same parameter size as the vanilla Gaussian VAE.Source: §D.6 · pp. 20–21

D.7. Asymptotic Complexity

• For m = 1, sorting the scalar codebook enables binary-search quantization, but autoregressive softmax makes the asymptotic complexity advantage practically limited.Most reverse channel coding methods are described as requiring exponential complexity in the bits-back rate.Source: §D.7 · p. 22

E.1. Additional Qualitative Results and Failure Cases

• GQ achieves superior visual quality in additional ImageNet results at 0.25 bpp.The figure caption also notes that none of the approaches correctly reconstructs the license plate.Source: §E.1 · p. 22, Figure 4 · p. 23
• At low bitrate, text remains challenging for VQ-VAEs, and none of the compared approaches reconstructs the residential vehicle's license plate.The figure caption specifies the evaluation at 0.25 bpp.Source: §E.1 · p. 22, Figure 4 · p. 23
• The appendix gives linear-scan and binary-search procedures for selecting a codebook value and index from μ_i in a sorted codebook.The binary-search procedure narrows the ordered codebook to two candidates and compares their distances to μ_i.Source: §E.1 · p. 22