Gaussian Mixture VQ: Probabilistic Vector Quantization

• Gaussian Mixture VQ (GM-VQ) is a probabilistic vector quantization framework that employs a Gaussian mixture prior to map continuous representations to discrete codebooks, improving generative modeling.
• It introduces an aggregated categorical posterior ELBO (ALBO) that replaces per-sample entropy with the entropy of the aggregated posterior, ensuring effective codebook utilization.
• GM-VQ optimizes the network end-to-end without separate commitment losses, demonstrating significant reduction in reconstruction error and robust performance on benchmarks like CIFAR-10 and CelebA.

Gaussian Mixture Vector Quantization (GM-VQ) is a probabilistic framework extending vector quantized variational autoencoders (VQ-VAE) for mapping continuous representations to discrete codebooks, crucial for generative modeling, information bottlenecking, and discrete tokenization in machine learning. GM-VQ introduces a Gaussian mixture as the generative prior, utilizes adaptive variances to capture complex data structure, and replaces heuristic objectives in prior VQ-VAE works with a unified Bayesian optimization approach. The introduction of the Aggregated Categorical Posterior Evidence Lower Bound (ALBO) enables improved codebook utilization and reduced reconstruction error without reliance on handcrafted regularization or codebook management heuristics (Yan et al., 2024).

1. Probabilistic Generative Model and Inference

GM-VQ defines a hierarchical latent variable model comprising a categorical latent code $c \in \{1, \dots, C\}$ with prior $p(c) = \pi_c$, typically uniform ($\pi_c = 1/C$). Conditional on $c$, a continuous latent $z \in \mathbb{R}^L$ is drawn from an isotropic Gaussian: $p(z|c) = \mathcal{N}(z; \mu_c, \sigma_z^2 I)$, where $\{\mu_c\}_{c=1}^C$ serves as the codebook of latent means and $\sigma_z^2$ is a fixed variance hyper-parameter (the limit $\sigma_z^2 \to 0$ recovers deterministic VQ). Observed data $x \in \mathbb{R}^D$ are generated by a decoder $D_\theta$ as $p(x|z) = \mathcal{N}(x; D_\theta(z), \sigma_x^2 I)$. The joint generative model factorizes as $p(x, z, c) = p(x|z) p(z|c) p(c)$.

Inference is performed via a variational posterior $q(c, z|x) = q(c|x) q(z|c, x)$. The discrete component $q(c|x)$ is parameterized with categorical logits $\alpha_c(x)$ derived from a Mahalanobis‐style distance between an encoder “proxy” $\hat{h}(x)$ and codebook entry $\mu_c$, scaled by an adaptive per-sample weight $w(x) = \operatorname{Softplus}(\tilde{w}(x))$:

$$\alpha_c(x) = -\frac{1}{2} (\hat{h}(x) - \mu_c)^{\top} \operatorname{diag}[w(x)] (\hat{h}(x) - \mu_c)$$

$$\pi_c(x) = \frac{\exp[\alpha_c(x)]}{\sum_{c'} \exp[\alpha_{c'}(x)]}$$

The continuous posterior $q(z|c, x) = \mathcal{N}(z; \mu_c, \Sigma_c(x))$ uses a diagonal covariance $\Sigma_c(x) = \sigma_c^2(x) I$, with adaptive variance

$$\sigma_c^2(x) = \left( \frac{ \| \hat{h}(x) - \mu_c \|^2 }{L} \right) / (2 \sigma_z^2)$$

This yields a “soft-assignment” of $x$ to each codeword, alongside adaptive uncertainty per assignment.

2. Aggregated Categorical Posterior ELBO (ALBO)

GM-VQ introduces an alternative to the standard evidence lower bound (ELBO) by aggregating the categorical posterior over data. Standard ELBOs include the term $-H[q(c|x)]$, incentivizing high entropy that conflicts with effective low-temperature Gumbel-Softmax sampling. GM-VQ replaces this per-example entropy with entropy of the aggregated posterior $q(c) = \mathbb{E}_{x \sim p_{\text{data}}}[q(c|x)]$, yielding the ALBO objective:

$$E_{ALBO} = \mathbb{E}_{x \sim p_{\text{data}}} \mathbb{E}_{q(c|x) q(z|c, x)} [ \log p(x, z, c) - \log q(c) ]$$

This objective aligns the variational posterior with the generative model and eliminates the adverse entropy component.

3. Loss Function and Optimization

The practical GM-VQ loss is formulated as:

$$L_{GM-VQ} = \mathbb{E}_x \mathbb{E}_{c|x}[ \| x - D_\theta(z(c, x)) \|^2 ] + \gamma \left( \mathbb{E}_x \mathbb{E}_{c|x}[ \| \bar{z}(c, x) - \mu_c \|^2 ] + \beta \, KL[q(c) \| p(c)] \right)$$

Here, $z(c, x) = \mu_c + \sigma_c(x) \odot \epsilon$, $\epsilon \sim \mathcal{N}(0, I)$, and $\gamma$, $\beta$ are positive hyper-parameters absorbing fixed decoder and latent variances ($\sigma_x^2$, $\sigma_z^2$). This loss is optimized fully end-to-end without handcrafted commitment losses or codebook management steps.

4. Training Procedure

The GM-VQ training pipeline for each mini-batch involves:

• Encoding $x$ via the encoder to produce $(\hat{h}(x), w(x))$, compute Mahalanobis-style logits $\alpha_c(x)$, and evaluate soft assignments $\pi_c(x)$.
• Sampling “soft” one-hot vectors $\tilde{c}$ using Gumbel-Softmax$(\pi(x), \tau)$, with annealing $\tau$ from $2.0$ to $0.1$ throughout training.
• Drawing $z = \tilde{c}^{\top}M + \sigma_c(x) \odot \epsilon$ for straight-through approximation.
• Decoding to $x̂ = D_\theta(z)$.
• Computing $L_{GM-VQ}$ and backpropagating gradients through parameters $\theta$, $\phi$, and codebook $M$ directly.
• Optimization is performed via AdamW with a cosine schedule and linear warm-up.

No separate exponential moving average, commitment loss, or post-hoc clustering is necessary.

5. Empirical Evaluation

Experimental results on CIFAR-10 and CelebA benchmark datasets demonstrate substantial performance improvement by GM-VQ over previous VQ-VAE approaches:

Model	MSE (CIFAR-10 / CelebA)	Perplexity (CIFAR-10 / CelebA)
VQ-VAE	5.65 / 10.02	14.0 / 16.2
VQ-VAE + replace	4.07 / 4.77	109.8 / 676.4
GM-VQ	3.13 / 1.38	731.9 / 338.6
GM-VQ + Entropy	3.11 / 0.97	878.7 / 831.0

GM-VQ achieves approximately 50% reduction in reconstruction error relative to vanilla VQ-VAE and attains high codebook perplexity, indicating usage of nearly the full codebook (perplexity $\gg C$). Entropy regularization parameter $\beta$ produces higher perplexity and slightly lower MSE as $\beta$ increases, indicating stronger codebook utilization.

6. Analysis, Extensions, and Limitations

By adopting a Gaussian mixture prior over latent codes and maintaining learnable, small conditional variances, GM-VQ generalizes deterministic quantization to a fully probabilistic model in which codebook entries, mixing weights, and variances are optimized under a unified objective. The aggregated categorical posterior mitigates conflicts arising in gradient estimation when using Gumbel-Softmax and high-entropy code assignments; the harmful $-H[q(c|x)]$ term is removed from the training objective. GM-VQ eliminates the need for post-hoc codebook replacement heuristics, commitment losses, and cluster management found in previous VQ-VAE frameworks. Code collapse is naturally prevented.

Potential extensions include end-to-end learning of $\sigma_z$ and mixture weights $\pi_c$, as well as architectural hierarchies formed by stacked GM-VQ layers. Key limitations remain: decoder variance $\sigma_x^2$ is fixed, effectiveness depends on the accuracy of Gumbel-Softmax straight-through approximation, and further bias reduction may be achievable via control variates. A plausible implication is that the principled Bayesian framework enables broadening of quantized representations without sacrificing tractable training or differentiation (Yan et al., 2024).