Gaussian Mixture Vector Quantization with Aggregated Categorical Posterior

## Abstract

Abstract The vector quantization is a widely used method to map continuous representation to discrete space and has important application in tokenization for generative mode, bottlenecking information and many other tasks in machine learning.
Vector Quantized Variational Autoencoder (VQ-VAE) is a type of variational autoencoder using discrete embedding as latent. We generalize the technique further, enriching the probabilistic framework with a Gaussian mixture as the underlying generative model. This framework leverages a codebook of latent means and adaptive variances to capture complex data distributions. This principled framework avoids various heuristics and strong assumptions that are needed with the VQ-VAE to address training instability and to improve codebook utilization. This approach integrates the benefits of both discrete and continuous representations within a variational Bayesian framework. Furthermore, by introducing the Aggregated Categorical Posterior Evidence Lower Bound (ALBO), we offer a principled alternative optimization objective that aligns variational distributions with the generative model. Our experiments demonstrate that GM-VQ improves codebook utilization and reduces information loss without relying on handcrafted heuristics.

## 1 Introduction

Variational autoencoders (VAEs) (Kingma & Welling, ) were originally designed for modeling continuous representations; however, applying them to discrete latent variable models is challenging due to non-differentiability. A common solution is to use gradient estimators tailored for discrete latent variables. The REINFORCE estimator (Williams, ) is an early example, providing an unbiased estimate of the gradient but suffering from high variance. Alternatively, methods such as the Gumbel-Softmax reparameterization trick (Jang et al., ; Maddison et al., ) allow for a continuous relaxation of categorical distributions. While these methods introduce bias into the gradient estimation, they offer the benefit of significantly lower variance, thereby improving training stability.

Vector Quantized Variational Autoencoders (VQ-VAEs) (Van Den Oord et al., ) extend the VAE framework to discrete latent spaces by discretizing continuous representations through a codebook via straight-through estimator (STE) (Bengio et al., ). Beyond the inherent variance-bias tradeoff in gradient estimation, VQ-VAEs are known to suffer from codebook collapse, wherein all encodings converge to a limited set of embedding vectors, resulting in the underutilization of many vectors in the codebook. This phenomenon diminishes the information capacity of the bottleneck. Takida et al. (); Williams et al. () hypothesized that deterministic quantization is the cause of codebook collapse and introduced stochastic sampling, leading to a entropy term in the log-likelihood lower bound. While high entropy is generally beneficial, it is inherently incompatible with Gumbel-Softmax gradient estimation. Various handcrafted heuristics have been proposed to mitigate this issue, including batch data-dependent k-means (Łańcucki et al., ), replacement policies (Zeghidour et al., ; Dhariwal et al., ), affine parameterization with alternate optimization (Huh et al., ), and entropy penalties (Yu et al., ). However, as these heuristics do not derive from the evidence lower Bound (ELBO), they cannot be unified within the variational Bayesian framework, rendering them ad-hoc solutions lacking a coherent foundation in variational inference.

In our work, we propose a Gaussian mixture prior based on VQ-VAE within a variational Bayesian framework, namely Gaussian Mixture Vector Quantization (GM-VQ, see Figure 1), combining the benefits of both discrete and continuous representations while avoiding handcrafted heuristics and strong assumptions. Additionally, to optimize the model and ensure compatibility with the gradient estimation errors inherent to Gumbel-Softmax, we modify the ELBO by replacing the conditional categorical posterior with an aggregated categorical posterior, resulting in an novel lower bound, the Aggregated Categorical Posterior Evidence Lower Bound (ALBO), which minimizes estimation error while preserving codebook utilization. Concretely, our contributions are as follows:

- •
To the best of our knowledge, we are the first to apply the Gaussian mixture prior formulation on VQ-VAE with strict adherence to the variational Bayesian framework.
- •
We introduce Aggregated Categorical Posterior Evidence Lower Bound (ALBO), which is explicitly designed to be compatible with Gumbel-Softmax gradient estimation.
- •
We conduct experiments demonstrating improved codebook utilization and reduced information loss without relying on handcrafted heuristics.

Figure: Figure 1: Overview of GM-VQ. First, the encoder deterministically maps the input to proxy latents, which are then used to retrieve corresponding codewords from the codebook and generate noise. The codewords and noise are then combined to form the continuous latents. Finally, these continuous latents are passed through the decoder to produce the final output.
Refer to caption: x1.png

## 2 Preliminaries

### 2.1 Deterministic VQ-VAE

The VQ-VAE (Van Den Oord et al., ), based on the VAE framework (Kingma & Welling, ), learns discrete latent representations via vector quantization (Gray, ). For simplicity, we represent discrete latents with a single random variable $\hat{{\mathbf{z}}}$ here, though in practice, we extract latent features of various dimensions.

Given an encoder output $\hat{{\mathbf{z}}}$ from a high-dimensional input ${\mathbf{x}}$ and a latent embedding space (codebook) $\mathbf{M}\in\mathbb{R}^{C\times L}$, composed of $C$ row vectors $\bm{\mu}_{i}\in\mathbb{R}^{L}$, selects the discrete latent variable (codeword) as $j=\arg\min_{i}\|\hat{{\mathbf{z}}}-\bm{\mu}_{i}\|$, yielding the approximate posterior distribution $q({\textnormal{c}}\mid{\mathbf{x}})\in\mathbb{R}^{C}$:

$$ $q({\textnormal{c}}\mid{\mathbf{x}})=[\bm{1}_{\mathrm{i=j}}]_{i=1}^{C}=\begin{ cases}1,&\text{if }i=\arg\min_{i}\|\hat{{\mathbf{z}}}-\bm{\mu}_{i}\|,\\ 0,&\text{otherwise}.\end{cases}$ (1) $$

Based on the one-hot distribution $q({\textnormal{c}}\mid{\mathbf{x}})$, we have query vector ${\mathbf{c}}_{q}=q({\textnormal{c}}\mid{\mathbf{x}})$ for codebook, the quantized latent representation is expressed as ${\mathbf{z}}_{c}={\mathbf{c}}^{T}_{q}\mathbf{M}=\bm{\mu}_{j}\in\mathbb{R}^{L}$.

To handle the non-differentiability of the quantization process, the straight-through estimator (STE) (Bengio et al., ) is applied, with the assumption $\frac{\partial{\mathbf{z}}_{c}}{\partial\hat{\mathbf{z}}}={\bm{I}}$ to allow gradient flow. This assumption holds reasonably well when ${\mathbf{z}}_{c}$ does not deviate significantly from $\hat{\mathbf{z}}$; however, greater deviations introduce increased bias. To mitigate this issue and enhance gradient estimation, an additional discretization loss was introduced:

$$ $\mathcal{L}_{\text{discretization}}(\hat{\mathbf{z}},{\mathbf{z}}_{c})=\|\hat{ \mathbf{z}}-\text{sg}[{\mathbf{z}}_{c}]\|^{2}+\alpha\cdot\|\text{sg}[\hat{ \mathbf{z}}]-{\mathbf{z}}_{c}\|^{2}.$ (2) $$

Here, $\text{sg}[\cdot]$ represents the stop-gradient operator, and $\alpha$ adjusts the balance between minimizing the discrepancies between $\hat{\mathbf{z}}$ and ${\mathbf{z}}_{q}$.

### 2.2 Stochastic VQ-VAE

One problem of VQ-VAE is that the learned discrete representation uses only a fraction of
the full capacity of the codebook, a problem known
as codebook collapse. To solve this problem, stochastic sampling was introduced by modifying the approximate posterior , shifting from a one-hot representation to a distribution proportional to the negative squared distance between the encoder output $\hat{{\mathbf{z}}}$ and the codewords $\bm{\mu}_{c}$(Roy et al. (); Sønderby et al. (); Shu & Nakayama ()):

$$ $q({\textnormal{c}}\mid{\mathbf{x}})=\operatorname{Softmax}\left(-\frac{\|\hat{ {\mathbf{z}}}-\bm{\mu}_{c}\|^{2}}{2\sigma^{2}}\right).$ (3) $$

To estimate gradients, the Gumbel-Softmax trick (Jang et al., ; Maddison et al., ) approximates the categorical distribution as $q^{g}({\textnormal{c}}\mid{\mathbf{x}})=\operatorname{Softmax}_{\tau}(\log q_{
i}({\textnormal{c}}\mid{\mathbf{x}})+{\mathbf{g}}_{i})$, where ${\mathbf{g}}_{i}$ are independent and identically distributed (i.i.d.) samples from Gumbel(0,1). The index $j$ is selected as $j=\arg\max_{i}q_{i}^{g}({\textnormal{c}}\mid{\mathbf{x}})$, resulting in the quantized representation ${\mathbf{c}}_{q}=[\bm{1}_{i=j}]_{i=1}^{C}$.

In this framework, the gradient assumption shifts from $\frac{\partial{\mathbf{z}}_{c}}{\partial\hat{{\mathbf{z}}}}=\displaystyle{\bm{
I}}$ to $\frac{\partial{\mathbf{c}}_{q}}{\partial q^{g}({\textnormal{c}}\mid{\mathbf{x}
})}=\displaystyle{\bm{I}}$. This gradient estimation requires a low entropy in $q({\textnormal{c}}\mid{\mathbf{x}})$ to be accurate . When entropy $q({\textnormal{c}}\mid{\mathbf{x}})$ has a higher entropy the output distribution becomes more uniform. In this scenario, the Gumbel-Softmax trick outputs values that are less peaked, leading to noisy estimates for gradient-based optimization because the model becomes less certain about which category is the most likely. This uncertainty increases the inaccuracy in the gradient estimates. Thus, with a fixed temperature, high entropy in $q({\textnormal{c}}\mid{\mathbf{x}})$ leads to gradient estimation errors. However, this leads to a problem.

Assuming a uniform prior, the variational bound is:

$$ $-\log p({\mathbf{x}})\leq\mathbb{E}_{{\textnormal{c}}\sim q({\textnormal{c}} \mid{\mathbf{x}})}\left[-\log p({\mathbf{x}}\mid{\textnormal{c}})\right]-H(q({ \textnormal{c}}\mid{\mathbf{x}}))+\log C,$ (4) $$

which introduces a negative entropy loss that promotes a high-entropy posterior distribution, conflicting with low entropy requirement in Gumbel-Softmax trick and resulting in incompatible high gradient estimation bias.

## 3 GM-VQ: Gaussian Mixture Vector Quantization

To solve the problem mentioned above, we propose a stochastic vector quantization method using using Gaussian mixture distribution, which we refer to as Gaussian Mixture Vector Quantization (GM-VQ), with the following encoding and decoding process:

$$ ${\mathbf{x}}\xrightarrow{E_{\bm{\theta}}}\hat{{\mathbf{z}}}\xrightarrow{ \mathbf{\mathbf{M}}}{\textnormal{c}}\xrightarrow{\bm{\varepsilon}}{\mathbf{z}} \xrightarrow{D_{\bm{\phi}}}\tilde{{\mathbf{x}}}$ (5) $$

Figure: Figure 2: Probabilistic Graphical Model depicting the Gaussian Mixture Vector Quantization (GM-VQ) for the generative model (left) and the inference model (right). The codebook $\mathbf{M}$ plays a dual role, being shared between both the generative and inference models.
Refer to caption: x2.png

In our framework, GM-VQ comprises an encoder $E_{\bm{\theta}}$ and a decoder $D_{\bm{\phi}}$, parameterized by the deep neural network weights $\bm{\theta}$ and $\bm{\phi}$, respectively, and interconnected via a codebook $\mathbf{M}$, which contains the means of Gaussian mixtures. The term $\bm{\varepsilon}$ denotes the Gaussian noise introduced into the process.

Unlike previous works adapted VQ-VAE with Gaussian mixture priors (Takida et al., ; Williams et al., ), which directly transmit categorical representations to the decoder, our method explicitly adds noise to the categorical variables during training, feeding continuous latent variables into the decoder. While Liu et al. () theoretically demonstrate that discrete representations can enhance generalization and robustness during evaluation, the codewords in our codebook are continuously updated throughout training, leading to a non-static categorical representation. By explicitly introducing noise, our method allows the decoder to adapt in parallel with the evolving codebook.

### 3.1 Generative Model

The generative model is defined by the joint distribution:

$$ $p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}})=p({\mathbf{x}}\mid{\mathbf{z}}) \,p({\mathbf{z}}\mid{\textnormal{c}})\,p({\textnormal{c}}),$ (6) $$

where ${\mathbf{x}}\in\mathbb{R}^{D}$ represents the observed data, ${\mathbf{z}}\in\mathbb{R}^{L}$ is a continuous latent variable, and ${\textnormal{c}}\in\{1,2,\ldots,C\}$ is a discrete latent variable indicating the mixture component. The latent variable ${\mathbf{z}}$ follows a Gaussian mixture distribution, with the means stored in a codebook $\mathbf{M}\in\mathbb{R}^{C\times L}$, where each row $\bm{\mu}_{c}$ represents a codeword for component $c$.

First, the discrete variable c is sampled from a categorical distribution, where $\bm{\pi}$ represents the prior probabilities of the mixture components, typically assumed to be uniformly distributed. Then, given c, the continuous latent variable ${\mathbf{z}}$ is drawn from a multivariate Gaussian distribution with mean $\bm{\mu}_{c}$ and isotropic covariance matrix $\sigma_{{\mathbf{z}}}^{2}{\bm{I}}$. As $\sigma_{{\mathbf{z}}}^{2}$ approaches zero, this setup converges to deterministic quantization, similar to VQ-VAE.

Finally, the observed data ${\mathbf{x}}$ is generated conditionally on ${\mathbf{z}}$, where the decoder $D_{\bm{\theta}}$ maps the continuous latents ${\mathbf{z}}$ to the mean of the observation distribution, modeled as a Gaussian with fixed variance $\sigma_{{\textnormal{x}}}^{2}{\bm{I}}$.

### 3.2 Variational Inference and Posterior Estimation

To approximate the intractable posterior $p({\textnormal{c}},{\mathbf{z}}\mid{\mathbf{x}})$, we employ a variational posterior $q({\textnormal{c}},{\mathbf{z}}\mid{\mathbf{x}})$ following standard variational autoencoder methods. Without simplifying assumptions, we employ variational inference with an approximate posterior. We use the chain rule to factorize the posterior:

$$ $q({\textnormal{c}},{\mathbf{z}}\mid{\mathbf{x}})=q({\textnormal{c}}\mid{ \mathbf{x}})q({\mathbf{z}}\mid{\mathbf{x}},{\textnormal{c}}),$ (7) $$

allowing for conditional dependencies between ${\mathbf{z}}$ and c given ${\mathbf{x}}$. To achieve this, we introduce the following variational distributions:

- 1.
Variational Distribution over c:
$q({\textnormal{c}}\mid{\mathbf{x}})=\text{Categorical}\big{(}\bm{\pi}({\mathbf
{x}})\big{)},$
(8)
where $\bm{\pi}({\mathbf{x}})=(\bm{\pi}_{1}({\mathbf{x}}),\bm{\pi}_{2}({\mathbf{x}}),
\ldots,\bm{\pi}_{C}({\mathbf{x}}))$ are the posterior probabilities over the mixture components, parameterized by an encoder $E_{\bm{\theta}}$ and codebook $\bm{M}$.
- 2.
Variational Distribution over ${\mathbf{z}}$:
We model $q({\mathbf{z}}\mid{\mathbf{x}},{\textnormal{c}})$ as a multivariate Gaussian distribution centered at the codeword $\bm{\mu}_{c}$, with a covariance $\mathbf{\Sigma}_{c}({\mathbf{x}})$ that depends on ${\mathbf{x}}$:
$q({\mathbf{z}}\mid{\mathbf{x}},{\textnormal{c}})=\mathcal{N}\big{(}{\mathbf{z}
};\bm{\mu}_{c},\mathbf{\Sigma}_{c}({\mathbf{x}})\big{)},$
(9)
where $\mathbf{\Sigma}_{c}({\mathbf{x}})=\sigma^{2}_{c}({\mathbf{x}}){\bm{I}}$ is an isotropic covariance matrix.

### 3.3 Training Objective and Optimization

Figure: Figure 3: Gradient Bias vs. Entropy Relationship
Refer to caption: extracted/5922865/img/bias_perplexity_relationship_2.png

To train the model, we typically maximize the ELBO using Gumbel-Softmax gradient estimation. However, the presence of entropy in ELBO leads to poor Gumbel-Softmax gradient estimation(see Figure 3). Given that entropy can be particularly detrimental to Gumbel-Softmax gradient estimation, it is critical to mitigate its impact during training.

To resolve this, Yu et al. () introduced two contrasting entropy measures derived from mutual information (Krause et al., ):

$$ $\mathcal{L}_{\text{entropy}}=H(q({\textnormal{c}}\mid{\mathbf{x}}))-H(\mathbb{ E}_{{\mathbf{x}}\sim p({\mathbf{x}})}q({\textnormal{c}}\mid{\mathbf{x}})).$ (10) $$

In the context of vector quantization, the first term of this loss reduces uncertainty when mapping observed data ${\mathbf{x}}$ to c and the second term encourages more discrete latent variables to be used, which decreases the codebook collapse problem.

However, this approach remains an additional heuristic that cannot be fully incorporated into the variational Bayesian framework. To address this limitation, we propose an alternative ELBO formulation.

#### 3.3.1 Aggregated Categorical Posterior Evidence Lower Bound

We introduce the Aggregated Categorical Posterior Evidence Lower Bound (ALBO) as an alternative to the traditional ELBO:

$$ $\mathcal{E}_{\text{ALBO}}(x)=\mathbb{E}_{q({\textnormal{c}})\,q({\mathbf{z}} \mid{\mathbf{x}})}\left[\log\frac{p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}} )}{q({\textnormal{c}})}\right]\leq\log p({\mathbf{x}}).$ (11) $$

The ALBO provides a lower bound on the log-likelihood $\log p({\mathbf{x}})$ (for derivation, see Appendix A.2).

Here, $q({\textnormal{c}})$ represents the marginal distribution over c, induced by the data distribution $p({\mathbf{x}})$ and the approximate posterior $q({\textnormal{c}}\mid{\mathbf{x}})$, defined as $q({\textnormal{c}})=\int q({\textnormal{c}}\mid{\mathbf{x}})p({\mathbf{x}})\,d
{\mathbf{x}}$. However, since the true data distribution $p({\mathbf{x}})$ is not accessible in practice, we have to rely on a finite dataset instead. To balance computational complexity and accuracy, we use a mini-batch approximation. For a mini-batch $\mathcal{B}$ of size $|\mathcal{B}|$, the marginal distribution $q({\textnormal{c}})$ is approximated as:

$$ $q^{(\mathcal{B})}({\textnormal{c}})=\frac{1}{|\mathcal{B}|}\sum_{{\mathbf{x}} \in\mathcal{B}}q({\textnormal{c}}\mid{\mathbf{x}}).$ (12) $$

This mini-batch approximation is computationally efficient and well-suited for stochastic optimization in deep learning.

For simplicity, $\sigma_{{\mathbf{x}}}^{2}$ and $\sigma_{{\mathbf{z}}}^{2}$ are typically fixed and not treated as learnable parameters. Thus, the objective function is constructed by minimizing the negative ALBO, resulting in the GM-VQ loss $\mathcal{L}_{\text{GM-VQ}}({\mathbf{x}})$:

$$ $\displaystyle\mathcal{L}_{\text{GM-VQ}}({\mathbf{x}})=\mathbb{E}_{q({\mathbf{z }}\mid{\mathbf{x}})}\|{\mathbf{x}}-D_{\bm{\theta}}({\mathbf{z}})\|^{2}+\gamma \cdot\mathcal{L}_{\text{reg}}({\mathbf{x}}),$ (13) $$

where the regularization term $\mathcal{L}_{\text{reg}}({\mathbf{x}})$ is defined as:

$$ $\mathcal{L}_{\text{reg}}({\mathbf{x}})=\mathbb{E}_{q({\textnormal{c}})q({ \mathbf{z}}\mid{\mathbf{x}})}\|{\mathbf{z}}-\bm{\mu}_{c}\|^{2}+\beta\cdot D_{ \mathrm{KL}}(q(c)\|p(c)),$ (14) $$

with $\beta$ and $\gamma$ as non-negative hyperparameters controlling the balance between reconstruction and regularization. Detailed derivation can be found in Appendix A.3.

This loss function consists of three main components:

- •
$\mathbb{E}_{q({\mathbf{z}}\mid{\mathbf{x}})}\|{\mathbf{x}}-D_{\bm{\theta}}({
\mathbf{z}})\|^{2}$: Reconstruction loss, standard in autoencoder frameworks.
- •
$\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})}\|{\mathbf{z}}-
\bm{\mu}_{c}\|^{2}$: Latent regularization, ensures alignment between the latent variables and the learned codewords $\bm{\mu}_{c}$.
- •
$D_{\mathrm{KL}}(q({\textnormal{c}})\|p({\textnormal{c}}))$: Entropy term, enhances overall codebook utilization and prevents individual $q({\textnormal{c}}\mid{\mathbf{x}})$ from drifting towards high entropy, thereby reducing gradient estimation errors.

#### 3.3.2 Parameterization of the Variational Distributions

##### Variational Categorical Distribution q ⁢ ( c ∣ 𝐱 ) 𝑞 conditional c 𝐱 q({\textnormal{c}}\mid{\mathbf{x}}) italic_q ( c ∣ bold_x )

To compute $q({\textnormal{c}}\mid{\mathbf{x}})$, we use the encoder $E_{\bm{\phi}}$ to generate a proxy representation $\hat{{\mathbf{z}}}({\mathbf{x}})\in\mathbb{R}^{L}$ and raw weights $\hat{{\mathbf{r}}}({\mathbf{x}})\in\mathbb{R}^{L}$, which are then activated into positive weights $\hat{{\mathbf{w}}}({\mathbf{x}})$ using the Softplus function $\zeta(\cdot)$:

$$ $\hat{{\mathbf{z}}}({\mathbf{x}}),\,\hat{{\mathbf{r}}}({\mathbf{x}})=E_{\bm{ \phi}}({\mathbf{x}}),\quad\hat{{\mathbf{w}}}({\mathbf{x}})=\zeta\big{(}\hat{{ \mathbf{r}}}({\mathbf{x}})\big{)}.$ (15) $$

The unnormalized log probabilities ${\bm{l}}_{c}({\mathbf{x}})$ are then computed based on the Mahalanobis-like distance between $\hat{{\mathbf{z}}}({\mathbf{x}})$ and each codeword $\bm{\mu}_{c}$:

$$ $\displaystyle{\bm{l}}_{c}({\mathbf{x}})$ $\displaystyle=-\frac{1}{2}\left(\hat{{\mathbf{z}}}({\mathbf{x}})-\bm{\mu}_{c} \right)^{\top}\text{diag}(\hat{\mathbf{w}}({\mathbf{x}}))\left(\hat{{\mathbf{z }}}({\mathbf{x}})-\bm{\mu}_{c}\right)$ (16) $\displaystyle=-\frac{1}{2}\sum_{i=1}^{L}\hat{{\mathbf{w}}}_{i}({\mathbf{x}}) \left(\hat{{\mathbf{z}}}_{i}({\mathbf{x}})-\bm{\mu}_{c,i}\right)^{2}.$ $$

The posterior probabilities $\bm{\pi}_{c}({\mathbf{x}})$ are obtained via softmax:

$$ $\bm{\pi}_{c}({\mathbf{x}})=\operatorname{Softmax}\big{(}{\bm{l}}_{c}({\mathbf{ x}})\big{)}=\frac{\exp\big{(}{\bm{l}}_{c}({\mathbf{x}})\big{)}}{\sum_{c^{ \prime}=1}^{C}\exp\big{(}{\bm{l}}_{c^{\prime}}(x)\big{)}}.$ (17) $$

This formulation ensures that components with codewords closer to $\hat{{\mathbf{z}}}({\mathbf{x}})$ have higher posterior probabilities.

##### Variational Continuous Distribution q ⁢ ( 𝐳 ∣ 𝐱 , c ) 𝑞 conditional 𝐳 𝐱 c q({\mathbf{z}}\mid{\mathbf{x}},{\textnormal{c}}) italic_q ( bold_z ∣ bold_x , c )

The variance $\bm{\sigma}^{2}_{{\textnormal{c}}}({\mathbf{x}})$ in $q({\mathbf{z}}\mid{\mathbf{x}},{\textnormal{c}})$ reflects the uncertainty in assigning ${\mathbf{x}}$ to component c. It is parameterized based on the squared distance between the encoder’s output $\hat{{\mathbf{z}}}({\mathbf{x}})$ and the codeword $\bm{\mu}_{c}$:

$$ $\bm{\sigma}^{2}_{c}({\mathbf{x}})=\frac{\|\hat{{\mathbf{z}}}({\mathbf{x}})-\bm {\mu}_{\textnormal{c}}\|^{2}/L}{2\sigma^{2}},$ (18) $$

where $\sigma^{2}$ is the scalar variance from the generative model and $L$ is the latent dimensionality.

This parameterization allows $q({\mathbf{z}}\mid{\mathbf{x}},{\textnormal{c}})$ to adapt its variance based on the distance between $\hat{{\mathbf{z}}}({\mathbf{x}})$ and $\bm{\mu}_{c}$, ensuring a flexible representation of uncertainty. As proxy latent $\hat{{\mathbf{z}}}({\mathbf{x}})$ approaches $\bm{\mu}_{c}$, the variance $\bm{\sigma}^{2}_{c}({\mathbf{x}})$ decreases, indicating higher confidence in the assignment to component $c$. Conversely, when the distance increases, the variance grows, signaling greater uncertainty.

##### Reparameterization and Codebook Update

To enable backpropagation through the discrete sampling of c, we use the Gumbel-Softmax reparameterization trick. The discrete latent variable c is computed as $j=\arg\max\,\operatorname{Softmax}_{\tau}(\log q({\textnormal{c}}\mid{\mathbf{
x}})+{\mathbf{g}})$, where ${\mathbf{g}}$ are i.i.d. Gumbel(0,1) samples, yielding the quantized representation ${\mathbf{c}}_{q}=[\bm{1}_{i=j}]_{i=1}^{C}$.

For the continuous latent variable ${\mathbf{z}}$, we apply the standard VAE reparameterization: ${\mathbf{z}}=\bm{\mu}_{c}+\bm{\sigma}_{c}({\mathbf{x}})\odot\bm{\epsilon}$, with $\bm{\epsilon}\sim\mathcal{N}(0,{\bm{I}})$. In the ALBO framework, we sample from $q({\mathbf{z}}\mid{\mathbf{x}})$, combining the quantized codeword and noise as ${\mathbf{z}}={\mathbf{c}}_{q}^{T}\mathbf{M}+\bm{\sigma}_{c}({\mathbf{x}})\odot
\bm{\epsilon}$.

This formulation ensures that all codewords $\bm{\mu}_{c}$ are naturally updated during optimization, preventing the codebook collapse problem without the need for additional commitment loss functions or exponential moving averages. It enables the model to manage deviations from the codewords in a controlled and efficient manner.

## 4 Related Work

Variational Autoencoders (VAEs) (Kingma & Welling, ) comprise a generative model and a recognition model, bridged by a latent variable typically modeled with a multivariate Gaussian prior. While the generative component is well-known, the recognition model effectively learns continuous representations from data (Zhang et al., ; Yang et al., ; Zhao et al., ; Higgins et al., ).

To address discrete representation learning, Vector Quantized VAE (VQ-VAE) (Razavi et al., ; Van Den Oord et al., ) employs vector quantization (Gray, ) to discretize latent embeddings under a uniform prior. Since the quantization process is non-differentiable, techniques like the straight-through estimator (Bengio et al., ) are used to approximate gradients, introducing potential bias. VQ-VAE also incorporates a discretization loss to mitigate these issues, but challenges such as codebook underutilization and information loss remain. Sønderby et al. (); Shu & Nakayama () introduced stochastic sampling based on the negative distance and applied Gumbel-Softmax for gradient estimation. Later, Karpathy (); Esser et al. () proposed an encoder that directly outputs the posterior, applying Gumbel-Softmax without conditioning on the codebook.

While many existing VAEs utilize Gaussian mixture priors (Liu et al., ; Bai et al., ; Falck et al., ; Guo et al., ; Jiang et al., ; Dilokthanakul et al., ; Nalisnick et al., ), our approach is distinct in its close connection with vector quantization. Specifically, we reuse the means from the codebook, whereas in other GMM-based models, the posterior means are typically learned transiently, conditioned on different components or networks. Although Takida et al. (); Williams et al. () also explore Gaussian mixture priors in VQ-VAE, they modify the reconstruction loss and feed discrete latents directly into the decoder, deviating from strict adherence to the ELBO. Furthermore, prior works often rely on simplifying assumptions on variational posterior, such as mean-field (Liu et al., ; Falck et al., ; Figueroa, ; Jiang et al., ) or Markovian assumptions (Takida et al., ; Williams et al., ).

Previous works (Tomczak & Welling, ; Hoffman & Johnson, ; Makhzani et al., ) have attempted to modify the ELBO framework by averaging the objective over data distribution, but they retain the original variational conditional distribution, leaving the entropy term of the conditional posterior it intact. In contrast, our approach is specifically motivated by the gradient estimation error in categorical latents by using an aggregated categorical posterior instead of the conditional categorical posterior.

Beyond these theoretical advancements, vector quantization (VQ) and its related concepts have also seen extensive applications across various domains. Here, we highlight some recent works. UniMoT (Zhang et al., ) introduces a VQ-driven tokenizer that converts molecules into molecular token sequences, while VQSynergy (Wu et al., ) integrates VQ for drug synergy prediction. In image and video generation, MAGVIT-v2 (Yu et al., ) applies VQ-VAE to achieve high-fidelity reconstructions, and VAR (Tian et al., ) leverages VQ to advance visual autoregressive learning. These studies demonstrate the pivotal role of VQ in both theoretical advancements and practical applications.

## 5 Experiments

In this section, we provide a comprehensive analysis of our experiments using the proposed GM-VQ model, focusing on its performance in image reconstruction tasks across two benchmark datasets CIFAR10 and CelebA.

### 5.1 Experimental Setup

We use Gumbel-Softmax for gradient estimation throughout our experiments, adopting an annealing schedule similar to Takida et al. (); Jang et al. (). The temperature starts at 2.0 and is gradually reduced to 0.1 during training.

Our architecture and hyperparameters closely follow the setup in Huh et al. (). For the CIFAR10 dataset (32x32 image size), we use a compact architecture consisting of convolutional layers followed by vector quantization, similar to standard autoencoders. For CelebA (resized to 128x128), a deeper network is employed to manage the higher resolution. Both architectures incorporate a VQ layer to quantize the latent representations into discrete codes.

Models are trained for 100 epochs on both datasets, using a batch size of 256. We employ AdamW (Loshchilov, ) as the optimizer, with a maximum learning rate of 1e-2 for CIFAR10 and 1e-4 for CelebA. The learning rate follows a cosine decay schedule with linear warmup as in Huh et al. (), with 10 epochs of warmup starting at a factor of 0.2, followed by 90 epochs of cosine decay. All models are initialized using K-means clustering for codebook initialization, following Esser et al. (). We use 1024 codes with a latent dimension of 64 across all experiments, applying the same weight decay to both the encoder-decoder and the codebook. During evaluation, no noise is added to the reconstructed ${\mathbf{x}}$ for maximum likelihood estimation, and latents ${\mathbf{z}}$ are directly sampled from the codebook $\mathbf{M}$ without extra noise.

We report Mean Squared Error (MSE) as the metric for reconstruction quality, which measures the mean pixel-wise difference between the original and reconstructed images. Additionally, we report perplexity, defined as $2^{H(q)}$, to evaluate the diversity of codebook usage. Higher perplexity indicates a more balanced use of the available codes. Notably, this perplexity is not based on the entropy of individual codes $q({\textnormal{c}}\mid{\mathbf{x}})$ but on the average entropy across a batch of categorical distributions $q^{(\mathcal{B})}({\textnormal{c}})$, with perplexity computed per batch and then averaged across all batches.

We compare our GM-VQ model against several baseline methods commonly used in vector quantization-based image reconstruction. The primary baseline is the standard VQ-VAE (Van Den Oord et al., ). Variants include VQ-VAE + $l_{2}$ (Yu et al., ), which stabilizes training through $l_{2}$ normalization, and VQ-VAE + replace (Zeghidour et al., ; Dhariwal et al., ), which replaces unused code vectors with random embeddings to avoid codebook collapse. SQ-VAE (Takida et al., ) introduces stochastic quantization for improved code diversity, while Gumbel-VQVAE (Karpathy, ; Esser et al., ) employs Gumbel-softmax for smoother gradient updates. Lastly, VQ-VAE + affine + OPT (Huh et al., ) addresses codebook covariate shift with affine parameterization and alternating training. Our primary model, GM-VQ, was tuned by fixing $\beta=1$ and selecting the best $\gamma$, while the variant GM-VQ + Entropy was tuned with higher entropy regularization ($\beta>1$) and the fixed $\gamma$, promoting more balanced codebook usage.

### 5.2 Performance Comparison

**Table 1: Comparison of methods on CIFAR10 and CELEBA datasets using MSE and Perplexity metrics.**
| Method | CIFAR10 | CELEBA |  |  |
| --- | --- | --- | --- | --- |
|  | MSE (10^-3) $\downarrow$ | Perplexity $\uparrow$ | MSE (10^-3) $\downarrow$ | Perplexity $\uparrow$ |
| VQVAE | 5.65 | 14.0 | 10.02 | 16.2 |
| VQVAE + $l_{2}$ | 3.21 | 57.0 | 6.49 | 188.7 |
| VQVAE + replace | 4.07 | 109.8 | 4.77 | 676.4 |
| VQVAE + $l_{2}$ + replace | 3.24 | 115.6 | 4.93 | 861.7 |
| VQVAE + Affine | 5.15 | 69.5 | 7.47 | 112.6 |
| VQVAE + OPT | 4.73 | 15.5 | 7.78 | 30.5 |
| VQVAE + Affine + OPT | 4.00 | 79.3 | 6.60 | 186.6 |
| SQVAE | 3.36 | 769.3 | 9.17 | 769.1 |
| Gumbel-VQVAE | 6.16 | 20.3 | 7.34 | 96.7 |
| GM-VQ | 3.13 | 731.9 | 1.38 | 338.6 |
| GM-VQ + Entropy | 3.11 | 878.7 | 0.97 | 831.0 |

We evaluate the performance of GM-VQ on the CIFAR10 and CelebA datasets, comparing it against several baseline methods. Table 1 presents the results, using MSE for reconstruction accuracy and perplexity metrics.

In the CIFAR10 dataset, GM-VQ achieves an MSE of 3.13, a significant improvement over the standard VQVAE (MSE 5.65) and variants like VQVAE + $l_{2}$ (MSE 3.21) and VQVAE + replace (MSE 4.07). In terms of codebook utilization, GM-VQ achieves a perplexity of 731.9, considerably higher than VQ-VAE + replace (perplexity 109.8), indicating more efficient and diverse code usage. This highlights GM-VQ’s ability to mitigate codebook collapse and ensure robust code assignments.

On the CelebA dataset, GM-VQ excels with an MSE of 1.38, significantly outperforming the baseline VQVAE (MSE 10.02) and VQVAE + replace (MSE 4.77). GM-VQ also maintains strong codebook diversity with a perplexity score of 338.6. The GM-VQ + Entropy variant further enhances performance, achieving the lowest MSE of 0.97 and a perplexity of 831.0. This shows that entropy regularization effectively promotes balanced codebook usage without sacrificing reconstruction quality.

In summary, across both datasets, GM-VQ and GM-VQ + Entropy consistently outperform all baseline models in terms of both reconstruction accuracy and codebook utilization. These results demonstrate the model’s robustness and its ability to maintain a balance between reconstruction fidelity and efficient code usage.

### 5.3 Impact of Entropy Regularization

Figure: Figure 4: Box plots showing the impact of entropy regularization on reconstruction quality (MSE) and codebook utilization (Perplexity) for the GM-VQ model. The left panel demonstrates the general trend of decreasing MSE with increasing entropy, the right pane shows the rise in perplexity with higher perplexity.
Refer to caption: x3.png

To further demonstrate the compatibility of the aggregated posterior with increased entropy, we conducted experiments across different entropy regularization levels using GM-VQ on the CIFAR10 dataset.

The box plots in Figure 4 provide a summary of how changes in entropy affect both Mean Squared Error (MSE) and Perplexity. In the left panel, we observe a general trend where MSE decreases as entropy increases, indicating a tendency toward improved reconstruction quality, though not uniformly across all entropy levels. Meanwhile, the right panel shows that higher entropy promotes more effective codebook usage, with perplexity rising as entropy grows, reflecting more balanced code assignments.

This trend suggests that while increased entropy yields better codebook utilization (higher perplexity), it also drives improvements in reconstruction accuracy (lower MSE).

## 6 Conclusion

In summary, the GM-VQ framework extends the traditional VQ-VAE by incorporating a probabilistic structure grounded in a Gaussian mixture model. By employing the ALBO objective, we ensure optimization is well-suited to Gumbel-Softmax gradient estimation. Our empirical results demonstrate the effectiveness of GM-VQ, highlighting its ability to achieve a strong balance between reconstruction quality and codebook diversity.

## Acknowledgments

We sincerely thank Michael Mozer for his valuable suggestions on the idea, as well as his careful revisions and insightful improvements to the manuscript. We are also grateful to Wenhao Zhao for providing experimental results that served as helpful references.

## Appendix A Appendix

### A.1 Implementation Details of Entropy Bias

To investigate the relationship between entropy and the bias in Gumbel-Softmax gradient estimation, we conducted targeted experiments.

We used a multi-layer perceptron (MLP) with hidden layers of size [50, 5], where the input consists of 10 possible actions and the output is a scalar. To obtain the exact gradients, we pass each categorical one-hot action through the non-linear decoder. For the Gumbel-Softmax gradient estimation, we repeated the experiment 50 times to compute the empirical average of the estimated gradients.

To assess the impact of entropy on bias, we varied the input entropy by applying softmax with different temperature ($\tau$) values to a fixed set of unnormalized logits, which determined the input probabilities for the MLP. This setup allowed us to analyze how changes in entropy influence the bias in Gumbel-Softmax gradient estimation.

### A.2 Derivation of ALBO

The Aggregated Categorical Posterior Evidence Lower Bound (ALBO) provides a lower bound on the log-likelihood $\log p({\mathbf{x}})$, similar in structure to the Evidence Lower Bound (ELBO) commonly used in variational inference. This bound is derived through the application of Jensen’s inequality, as shown below:

$$ $\displaystyle\log p({\mathbf{x}})$ $\displaystyle=\log\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x} })}\left[\frac{p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}})}{q({\textnormal{c }})q({\mathbf{z}}\mid{\mathbf{x}})}\right]$ (19) $\displaystyle\geq\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}} )}\left[\log\frac{p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}})}{q({ \textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})}\right]\quad\text{(by Jensen's inequality)}$ $\displaystyle=\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})} \left[\log\frac{p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}})}{q({\textnormal{ c}})}\right]-\mathbb{E}_{q({\mathbf{z}}\mid{\mathbf{x}})}\left[\log q({\mathbf {z}}\mid{\mathbf{x}})\right]$ $\displaystyle\geq\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}} )}\left[\log\frac{p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}})}{q({ \textnormal{c}})}\right]\quad\text{(since }\mathbb{E}_{q({\mathbf{z}}\mid{ \mathbf{x}})}\left[\log q({\mathbf{z}}\mid{\mathbf{x}})\right]\leq 0\text{)}$ $\displaystyle=\mathcal{E}_{\text{ALBO}}({\mathbf{x}}).$ $$

By applying Jensen’s inequality, the logarithm is moved inside the expectation, yielding a tractable lower bound. The term $\mathbb{E}_{q({\mathbf{z}}\mid{\mathbf{x}})}\left[\log q({\mathbf{z}}\mid{
\mathbf{x}})\right]$ represents the entropy of the posterior over ${\mathbf{z}}$. Since entropy is non-positive, it further tightens the bound, ensuring that $\mathcal{E}_{\text{ALBO}}({\mathbf{x}})$ provides a meaningful approximation.

Thus, $\mathcal{E}_{\text{ALBO}}({\mathbf{x}})$ serves as a valid lower bound on $\log p({\mathbf{x}})$, similar to the ELBO in traditional variational inference. The key difference is in the aggregation of the categorical posterior over c, which offers better compatibility with probability-based gradient estimation.

### A.3 Derivation of GM-VQ Loss

Minimizing $\mathcal{L}_{\text{GM-VQ}}({\mathbf{x}})$ ensures that the model effectively reconstructs the data while regularizing the latent distributions to align with the priors, thereby improving generalization. Below, we provide the detailed derivation of the GM-VQ loss.

We start by maximizing the $\mathcal{E}_{\text{ALBO}}({\mathbf{x}})$:

$$ $\displaystyle\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\max}$ $\displaystyle\mathcal{E}_{\text{ALBO}}({\mathbf{x}})$ (20) $\displaystyle=\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\min}$ $\displaystyle-\mathcal{E}_{\text{ALBO}}({\mathbf{x}})$ $\displaystyle=\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\min}$ $\displaystyle-\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})} \left[\log\frac{p({\mathbf{x}},{\mathbf{z}},{\textnormal{c}})}{q({\textnormal{ c}})}\right]$ $\displaystyle=\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\min}$ $\displaystyle\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})} \left[-\log\frac{p({\mathbf{x}}\mid{\mathbf{z}})\,p({\mathbf{z}}\mid{ \textnormal{c}})\,p({\textnormal{c}})}{q({\textnormal{c}})}\right]$ $\displaystyle=\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\min}$ $\displaystyle\mathbb{E}_{q({\mathbf{z}}\mid{\mathbf{x}})}\left[-\log p({ \mathbf{x}}\mid{\mathbf{z}})\right]+\mathbb{E}_{q({\textnormal{c}})q({\mathbf{ z}}\mid{\mathbf{x}})}\left[-\log p({\mathbf{z}}\mid{\textnormal{c}})\right]+D_ {\mathrm{KL}}(q({\textnormal{c}})\|p({\textnormal{c}}))$ $\displaystyle=\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\min}$ $\displaystyle\mathbb{E}_{q({\mathbf{z}}\mid{\mathbf{x}})}\left[\frac{\|{ \mathbf{x}}-D_{\bm{\theta}}({\mathbf{z}})\|^{2}}{2\sigma^{2}_{{\mathbf{x}}}} \right]+\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})}\left[ \frac{\|{\mathbf{z}}-\bm{\mu}_{c}\|^{2}}{2\sigma_{{\mathbf{z}}}^{2}}\right]+D_ {\mathrm{KL}}(q(c)\|p(c))$ $\displaystyle=\underset{\bm{\theta},\bm{\phi},\mathbf{M}}{\arg\min}$ $\displaystyle\mathbb{E}_{q({\mathbf{z}}\mid{\mathbf{x}})}\|{\mathbf{x}}-D_{\bm {\theta}}({\mathbf{z}})\|^{2}+\frac{\sigma^{2}_{{\mathbf{x}}}}{\sigma_{{ \mathbf{z}}}^{2}}\cdot\left(\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{ \mathbf{x}})}\|{\mathbf{z}}-\bm{\mu}_{c}\|^{2}+2\sigma_{{\mathbf{z}}}^{2}\cdot D _{\mathrm{KL}}(q(c)\|p(c))\right)$ $$

Given variance terms $\sigma_{{\mathbf{x}}}^{2}$ and $\sigma_{{\mathbf{z}}}^{2}$ are fixed, we can replace them with positive hyperparameters $\beta$ and $\gamma$, respectively, to simplify this expression, where:

- 1.
$\gamma=\frac{\sigma_{{\mathbf{x}}}^{2}}{\sigma_{{\mathbf{z}}}^{2}}$ controls the balance between reconstruction and latent regularization.
- 2.
$\beta=2\sigma_{{\mathbf{z}}}^{2}$ modulates the strength of the KL divergence regularization.

Thus, the GM-VQ loss can be rewritten as:

$$ $\displaystyle\mathcal{L}_{\text{GM-VQ}}({\mathbf{x}})=\mathbb{E}_{q({\mathbf{z }}\mid{\mathbf{x}})}\|{\mathbf{x}}-D_{\bm{\theta}}({\mathbf{z}})\|^{2}+\gamma \cdot\left(\mathbb{E}_{q({\textnormal{c}})q({\mathbf{z}}\mid{\mathbf{x}})}\|{ \mathbf{z}}-\bm{\mu}_{c}\|^{2}+\beta\cdot D_{\mathrm{KL}}(q(c)\|p(c))\right)$ (21) $$

This formulation, with hyperparameters $\beta$ and $\gamma$, balances the reconstruction fidelity and the regularization of the latent space.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 1
presentation: 2
questions: •	In line 77, “a novel”.

•	In equation 5, is the only novelty here the use of $\epsilon$ ? It would be convenient to clarify what is the main contribution of this paper in that Markov Chain

•	In line 186-188: In previous paragraph you mentioned that $z$ follows a Gaussian mixture but in this sentence you say that you sample it from a Gaussian with mean $\mu_c$. Why is this happening and what value takes $c$ here?
soundness: 1
strengths: I believe that the paper is well contextualized and that a good general view of similar works is given throughout the paper. 

I find studying the use of more complex priors an interesting topic in representation learning that can leverage the performance of models and obtain richer representations.
summary: The present paper introduces an idea to introduce a mixture of gaussians as a prior in a VQ-VAE. Furthermore, they introduce what is called Aggregated Categorical Posterior Evidence Lower Bound (ALBO), which allows compatibility with Gumbel-Softmax gradient estimation. Finally, they present some experiments with CIFAR-10 and CelebA in which they provide MSE and Perplexity as performance measurements.
weaknesses: It results complicated to see how the works in the literature connect with each part of the present proposal. For example, what is the purpose of equation (10)? How does it connect to the next part of the paper? The different ideas of the paper look like unconnected patches and I miss a storyline that helps to understand the motivation of the paper.

One of the main novelties of the paper is claimed to be the use of mixture of Gaussians as prior. However, I miss in section 3 an explanation on how you introduce this prior. The only part that I can find somehow related is in line 335, but in this case $z$ is not sampled from a mixture of gaussians but from a gaussian which mean is a linear combination of different vectors. I would like the authors to clarify better how they use the mixture of gaussians as a prior, since this is an interesting topic. I would recommend the authors to provide a specific subsection or paragraph in Section 3 that explicitly defines the mixture of Gaussians prior and explains how it is incorporated into the model's generative process. This would help clarify the novelty and implementation of this key aspect of the paper.

One of the other novelties of this paper is the introduction of the ALBO, whose derivation is in Appendix A.2. However, I believe that this derivation is incorrect. In the penultimate line of the derivation what is inside the parenthesis is false. This is only true when $z$ is a discrete variable. This term is equal to $-H_q(z|x)$ and, since $z$ is a continuous variable, then it can be larger than 0, since this is a differential conditional entropy and not a Shannon conditional entropy (see https://en.wikipedia.org/wiki/Conditional_entropy#Conditional_differential_entropy:~:text=In%20contrast%20to%20the%20conditional%20entropy%20for%20discrete%20random%20variables%2C%20the%20conditional%20differential%20entropy%20may%20be%20negative). Please, revisit the derivation taking into account the distinction between discrete and continuous entropy. This will help ensure the mathematical rigor of the paper.

The experiments are a bit simple. Due to the big amount of applications that VQ-VAE can have, the paper would benefit from experiments in specific applications other than simply reconstruction performance. In fact, I don’t think that reconstruction performance reflects the performance of a VAE (and any of its variations) in almost any task. Vanilla Autoencoders tend to have a lower MSE than VAE (since they are optimized only for that purpose), but they are less powerful models.

I find some issues with mathematical notation. Some variables are not properly defined or there are some imprecisions and mistakes. I list some next:

•	In lines 93-95, maybe I misunderstood this sentence, but you could simply say that $\hat{z}$ is a single random variable with various dimensions.

•	In equation 1, I don’t think the notation \left[1_{i=j}\right]^C_{i=1} is a widely used notation. Some others are some more widely used notations (see for example https://stats.stackexchange.com/questions/436975/compact-notation-for-one-hot-indicator-vectors).

•	Also in equation 1: if j = arg min…

•	In equation 3, I believe that $\sigma$ is not defined.

•	In line 127, notation is incorrect. You do not calculate the softmax of $q_i$ but of $q$. Maybe, it would be easier to include the definition of Softmax (eq. 2 of CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX)

•	In equation 4, why do you assume an uniform prior?

•	In line 185, y ou have not mentioned $\pi$ yet, so it does not make sense the clause "where $\pi$" represents the prior...". This sentence should be reformulated.

•	In line 270, the fact that $\sigma_x$ and $\sigma_z$ are fixed does not imply that "the objective function is constructed by minimizing the negative ALBO". If anything, it causes that maximizing the likelihood becomes equivalent to minimizing the euclidean distance. Thus, the paragraph should be reformulated.

•	In equation 14, how do you calculate this term in practice? I cannot find how you define p(c).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: In the simulations for Fig3, can you elaborate what is the ‘exact gradient’ from ‘passing each categorical one-hot action through the non-linear decoder’? For this experiment, what is the distribution of the inputs and the MLP weights?

For table 1, do all methods use a latent space with the same dimension and same number of mixtures/codebooks?
 
Minor:
In eq. (3), how is softmax defined and how is it different from $\text{softmax}_\tau$ in line 127?
soundness: 2
strengths: Propagation of gradients for VQ-VAEs is a challenging problem and previous approximate approaches - such as straight-through estimators (STE), Gumbel-Softmax approximations or stochastic quantization - can lead to codebook collapse wherein a large fraction of the codebooks remain unused. In contrast, the paper considers a Gaussian Mixture prior model and claims to optimize an evidence lower bound to the model log-likelihood. The suggested model and variational objective is different from previous work that for example decode the quantised latent. The usage of min-batch approximation for the marginal distribution over the mixture component is also new as far as I am aware. Empirical results indeed confirm that the method leads to better reconstruction and codebook usage compared to VQ-VAE models with STE and different codebook lookup tricks.
summary: The submission considers a VAE with a Gaussian mixture piror whose mean parameters can be seen as the codebook in VQ-VAEs, but with a non-degenerate variance on the generative and inference paths. The paper suggests a variational objective involving a mini-batch approximation of the aggregated posterior over the mixture/codebook index. The variational distributions of the mixture index can then be optimized using Gumbel-Softmax approximations. Experiments for CIFAR10 and CelebA show that the suggested approach yields improved MSE for the reconstruction errors compared to previous VQ-VAE models, as well as codebook usage in terms of perplexity.
weaknesses: I’m not convinced that the derivation in Appendix A2 for the ALBO bound is valid. In particular, why is the entropy of the variational distribution negative? If this is a Gaussian, does there have to be a constraint on the variance thereof? 

Is the ALBO still a lower bound on the LLH even if the marginal (aggregated posterior) over c is approximated by a mini-batch approximation? Does the mini-batch size have to be sufficiently large?

The objective contains a KL penalty between the marginal $q(c)$ and the prior $p(c)$. This appears to be similar to an InfoVAE [1] bound that also regularises the marginal distribution of the latents (although commonly with continuous latents), but resort to the MMD to avoid the intractability of computing the KL involving the marginal/aggregated posterior when the latent space is continuous. Such objectives seem to be sensitive to the choice of weight for the different regularising terms. However, I feel that the effect of such hypeparameters (viz $\beta$ and $\gamma$) has not been studied in detail in this paper, for example experimentally.
 
The experimental validation appears a bit limited. In particular, it is not clear how well the generative performance is qualitatively? Also different metrics such as FID values would be useful. In particular, I would expect that the choice of $\beta$ and $\gamma$ affects the reported MSE values. 

 [1] Zhao, Shengjia, Jiaming Song, and Stefano Ermon. "InfoVAE: Information Maximizing Variational Autoencoders." 2017

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: (Q1) What is the key difference between your work and those of Takida et al. (2022) and Williams et al. (2020), regarding adherence to the Bayesian framework?

(Q2) What is meant by a non-static categorical representation, and what are its benefits?

(Q3) How does the continuous updating of the codewords throughout training lead to this non-static categorical representation?

(Q4) Could you specify the level of entropy regularization shown in Figure 4?

(Q5) Could you provide evaluations using perceptual metrics, such as reconstruction FID?

(Q6) Could you show qualitative comparison with examples of reconstructed images?

(Q7) Could you demonstrate that the proposed method reduces the gradient estimation error caused by the Gumbel-softmax technique with some experimental results or mathematical guarantees?
soundness: 2
strengths: * The proposed algorithm is clearly explained.
* The variant of VQ-VAE is theoretically derived by a graphical model and a variational Bayesian framework.
* In the experiments, the proposed method is compared to a sufficient number of VQ-VAE baselines.
summary: This paper proposes a variant of VQ-VAE that utilizes stochastic vector quantization based on a Gaussian mixture prior. 
The model is trained by optimizing a proposed variational lower bound of log-likelihood, referred to as ALBO, which is derived from a new graphical model within a variational Bayesian framework.
The proposed loss function is designed to be compatible with Gumbel-Softmax gradient estimation.
Experiments on two benchmark datasets demonstrate that the proposed method improves reconstruction accuracy and codebook utilization compared to baseline methods.
weaknesses: * The paper does not clearly demonstrate why the proposed method resolves the existing problems. Some parts could be explained more clearly.
  - In the Introduction, the paper claims, "we are the first to apply the Gaussian mixture prior formulation on VQ-VAE with strict adherence to the variational Bayesian framework." However, I find it challenging to identify the key differences compared to the work of Takida et al. (2022) and Williams et al. (2020) in this regard. In Section 4, the paper states, "they modify the reconstruction loss and feed discrete latents directly into the decoder, deviating from strict adherence to the ELBO." However, these operations are theoretically and naturally derived from their own proposed variational lower bound loss within the Bayesian framework. --> (Q1)
  - The main content does not clearly explain the key differences from previous methods or why the proposed framework resolves the issue of gradient estimation error. After reviewing the supplementary material in Section A.2, I noticed that the key difference lies in the aggregation of the categorical posterior over c, which does not conflict with the low entropy requirement for q(c|x) in the Gumbel-Softmax trick. (I think the statement in Section 3.3.2, "prevents individual q(c|x) from drifting towards high entropy," may be too strong.) However, supplementary materials should not be necessary to grasp the logical flow of the paper. This point could be explained in more detail in the main text.
  - I apologize, but I could not fully grasp the meaning of "the codewords in our codebook are continuously updated throughout training, leading to a non-static categorical representation." at the end of p3. --> (Q2), (Q3)
  - In Figure 4, the level of entropy regularization is not clearly specified. --> (Q4)

* The experiments are not sufficient to support the claims.
  - Reconstruction accuracy is measured solely by MSE. However, for image-related tasks such as compression and generation, perceptual metrics are also important. --> (Q5)
  - Showing examples of reconstructed images would further demonstrate the effectiveness of the proposed method. --> (Q6)
  - Since the proposed methods utilize the Gumbel-Softmax technique similar to existing work, there should be, more or less, some degree of gradient estimation error. Demonstrating that the proposed method reduces this error would provide stronger support for the claims. --> (Q7)

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: - The authors should review the notation throughout the paper, as there are typos that may difficult readability. For instance, the parameters of the Decoder neural network are inconsistently denoted as $\theta$ and $\phi$.
- In the graphical model of the generative model, why does $z$ depend on the parameters of the decoder? In its definition, $p(z|c)$ is not parameterized by the decoder.
- In section 3, the authors state ‘By explicitly introducing noise, our method allows the decoder to adapt in parallel with the evolving codebook.’ Could the authors elaborate on this idea? Does the noise act as additional regularization, helping to prevent the decoder from overfitting to a specific codebook?
- How does the performance of the method vary with respect to $\sigma^{2}_{z}$? As it approaches zero, the method converges to the deterministic quantization. I assume that there is a tradeoff in the amount of noise introduced in the method.
- How does the batch size affect the performance of the method, given its influence on the marginalization of $q(c)$?
- I think it would be beneficial to evaluate the entropy of $q(c∣x)$ (either during training or evaluation) to demonstrate that the proposed objective effectively prevents individual $q(c∣x)$ distributions from drifting toward high entropy, thereby reducing gradient estimation errors.
- How is the generation performance of this method? Does it need a two-stage training as the original VQ-VAE?
- What are the potential limitations of this method? 
- I believe it would be beneficial to include an ablation study on the hyperparameters controlling the balance of the loss to characterize the training dynamics.
soundness: 2
strengths: - This work presents an alternative formulation for describing VQ-VAE within a Bayesian framework, avoiding heuristics.
- The method is clearly explained, and the authors provide context through a discussion of relevant state-of-the-art methods.
- The implementation and experimental details are thoroughly outlined.
summary: This work proposes a new formulation for describing VQ-VAE within a Bayesian framework using a Gaussian Mixture prior. This approach avoids the need for heuristics to ensure effective codebook utilization. The authors introduce an alternative training objective, the Aggregated Categorical Posterior Evidence Lower Bound (ALBO), designed to reduce Gumbel-Softmax gradient estimation errors while enhancing codebook usage.
weaknesses: - I believe the experimental evaluation of the method could be improved. While the authors compare their method with several state-of-the-art approaches, the evaluation focuses only on MSE for reconstruction and Perplexity for codebook usage. Including additional generation results would add value (does this model require two-stage training as the original VQ-VAE, or it is not necessary given the new formulation?). An ablation study on the effect of batch size, given its influence on the marginalization of q(c), could clarify its impact on performance (if there’s any). Further codebook analysis would also be beneficial; for instance, are the learned codebook vectors diverse, or are they quite similar, potentially explaining the higher entropy? 
- The paper lacks a concluding discussion and does not explore potential limitations of the method or directions for future research, which I believe would add significant value.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: The tightness of the ALBO in Eq. (11) is not discussed in detail. For example, ALBO is looser than ELBO. What are its advantages/disadvantages?

The parameterization of the variational inference arm should be elaborated. Intuitively, one would model the inference arm as $q(c|x)q(z|x,c)$ following Fig. 2; why use the aggregated $q(c)q(z|x)$ instead in Eq. (11)? What's the definition of $q(z|x)$? 

The optimal inference arm in the ELBO framework is the posterior. What's the optimal solution for the inference arm in the ALBO framework?

What's the influence of the mini-batch approximation?

In Line 294, "$E_{\phi}$" is a typo.
soundness: 1
strengths: The paper is well written and easy to follow in general.

The proposed techniques are original to my knowledge.
summary: The paper proposes a new vector quantization method based on Gaussian Mixture Models (GMMs). Specifically, within the VAE framework, the authors first employ a GMM as the prior and then they design a specifically tailored variational inference arm to implement vector quantization, using a newly introduced ELBO variant called ALBO. Experiments are conducted on the CIFAR10 and CelebA datasets to demonstrate the effectiveness of the proposed method.
weaknesses: The soundness of the proposed ALBO is questionable. See Questions below.