Gaussian Mixture Vector Quantization with Aggregated Categorical Posterior

• 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.

The soundness of the proposed ALBO is questionable. See Questions below.

• 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)

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

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).