MQ-VAE: Training Vector-Quantized Networks via Meta Learning

Authors: Peijia Qin, Jianguo Zhang
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Vector-Quantized Networks, Bi-level Optimization, Image Generation

Abstract:
Deep neural networks with discrete latent variables are particularly well-suited for tasks that naturally involve sequences of discrete symbols.
The vector-quantized variational auto-encoder (VQ-VAE) has made significant progress in this area by leveraging vector quantization.
However, while much effort has been put into maximizing codebook utilization, this does not always result in better performance.
Additional challenges include quantization errors in the VQ layer and the lack of direct integration of task loss into the codebook objective.
To address these issues, we propose Meta-Quantized Variational Auto-Encoder (MQ-VAE), a bi-level optimization-based vector quantization framework inspired by meta-learning.
In MQ-VAE, the codebook and encoder-decoder pair are optimized at different levels, with the codebook treated as hyperparameters optimized via hyper-gradient descent.
This approach effectively tackles these challenges within a unified framework.
The evaluation of MQ-VAE on two computer vision tasks demonstrates its superiority over existing methods and ablation baselines.
Code is available at https://anonymous.4open.science/r/MQVAE-B52C.

Decision: reject
Reviewer scores: 3, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1729229614423

thread_type:
official_review

<thread_object>
creator: reviewer_rpsn

<content>
summary:
**Background**: Gradient-based meta-learning (GBML) aims to updated a models' parameters $\theta$ such that, after only several gradient updates with respect to a class in the support set, the loss on this task is minimized. Intuitively, it tries to find a good set of model parameters that can be quickly adapted (after a couple finetuning steps) to a downstream distribution. 

**This work.** This paper proposes to adapt gradient-based meta-learning (GBML) from the typical few-shot learning setting into the training of a VQ-VAE. 

Specifically, the encoder parameters $\phi$ and decoder parameters $\theta$ are trained with a fixed codebook for several steps, the loss is computed with the updated encoder/decoder parameters, and the codebook is updated with respect to this loss. As in GBML, the encoder/decoder parameter updates are undone; however unlike the typical GBML algorithm, the loss is computed again with the updated codebook vectors, and one step of gradient descent is performed to update the encoder / decoder parameters [keeping the codebook now fixed]. This process is repeated until convergence.

Intuitively, this approach tries to find a good set of codebook vectors by looking-ahead to see how the Encoders'/Decoders' parameters will drift. The authors

strengths:
I really enjoyed reading this paper. I see its strengths as:

1. Very well written and easy to understand. A reader with familiarity of gradient-based meta-learning can read a couple paragraphs in the introduction, look at Figure 1, and immediately understand the proposed approach.
2. Comparison to Huh et al. 2023: I was looking for this comparison after reading the introduction, and I appreciate the thorough discourse comparing to a very close method.
3. Finite difference approximation: this isn't new, but it's quite the nice touch to make differentiating through multiple encoder/decoder updates practical.
4. Section 4.3 is also quite good: I've seen very few mentions of the *indirect* gradients that appear during meta-learning (or other transfer learning techniques) that appear when you differentiate through a prior parameter update.

weaknesses:
This paper also presents several weaknesses:

0. Novelty: this paper applies an existing technique (gradient-based meta-learning) to a new setting (VQ-VAE training). That's not an issue in itself -- the ViT paper was incredible even though it applied an existing technique (attention) to a new setting (computer vision). But ViTs were very practical to train and seemed to work really well on computer vision benchmarks. This leads me to my second point:

1. I don't think this method is practical; it is difficult to imagine a researcher or engineer implementing it as they chase sota performance. Implementing a bi-level optimization loop is quite difficult, and many things can go wrong. Even in the meta-learning community, we try to stay away from this as much as possible. Running a bi-level optimization to train a VQ-VAE would be very difficult for real applications. Even with new frameworks like Betty that make it a bit easier, my understanding is that still very difficult and popular training settings like distributed data parallel would not be supported.

 2. The contribution of Huh et al. 2023 in last years ICLR weakens this paper's contribution. I found many of the ideas in that work highly engaging, and while this work advances them in the direction of meta-learning, my sense is that the results / ideas in this paper are less new in the backdrop of Huh et al. 2023.

3. The authors highlight the significance of indirect gradients in red throughout their paper (especially Figure 1 and Section 4.3). When I looked at these for a different application (in meta-learning), it seemed as their gradient norms were **many orders of magnitude less** then the direct gradients. In this context, my prior is that $||\frac{\partial \mathcal{L'}}{\partial \mathcal{C}}|| >>> ||\frac{\partial \phi'}{\partial \mathcal{C}} \frac{\partial \mathcal{L}'}{\partial \phi '}||$ and $||\frac{\partial \mathcal{L'}}{\partial \mathcal{C}}|| >>> ||\frac{\partial \theta'}{\partial \mathcal{C}} \frac{\partial \mathcal{L}'}{\partial \theta '}||$, therefore the effect of the indirect gradients would be almost nothing. I think it is a weakness to not have explored the dynamics of these effects.

4. The experimental results are limited and evaluated on non-realistic datasets. I'm very empathetic to limited computational resources, but evaluating on two small datasets (FashionMNIST and CIFAR100) with a 2017-style VQ-VAE does not convey to me that this would work for a more recent VQ-VAE architecture (such as a VQGAN). 

5. The experimental results section seem inconsistent with Huh et al. 2023 -- specifically Table 3 where VQ-VAE + alt + sync is strictly worse than the default VQ-VAE? Is the training setting unstable, or are you reporting that these two techniques actually result in worse performance than the default model for FashionMNIST? 

6. The claim that "MQ-VAE can learn a superior codebook" seems unsupported by the analysis and should be contextualized in the backdrop of recent methods that find not learning the codebook at all can substantively improve performance [1, 2].

[1] Finite Scalar Quantization: VQ-VAE Made Simple, Mentzer et al., ICLR 2024

[2] Language Model Beats Diffusion - Tokenizer is key to visual generation, Yu et al., ICLR 2024

questions:
Please see weaknesses. In particular, addressing:

1. The practicality of this method when implementing a very strong VQ-VAE, and if you believe the results you present on the datasets and architectures that you've chosen provide sufficient evidence to support this claim: "We demonstrate the superiority of MQ-VAE with two computer vision tasks. MQ-VAE significantly outperforms several comparison baselines and ablation methods, highlighting its effectiveness."

2. Inconsistency between Huh et al. 2023 and the results reported on FashionMNIST.

3. Comparison of the magnitude of indirect gradients to the direct ones. If you 0-out the indirect gradients, will the model still perform well? This study would support your assertions that these indirect gradients are crucial for the performance of your approach.

4. The importance of learning a better codebook when recent work has shown that not learning the codebook is actually quite good.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730389403187

thread_type:
official_review

<thread_object>
creator: reviewer_ercu

<content>
summary:
This submission proposed to conduct quantization for hidden emebeddings / vectors / codebook in VAE (Variational Auto-Encoder). Specially, it designed a bi-level optimization framework to update encoder/decoder (full-precision parameters) and quantized hidden embeddings alternatively: Normal gradient descent is used in encoder/decoder optimization while the codebook is optimized using a hyper-gradient.

strengths:
- The method proposed is straightforward and writing is easy to understand.

weaknesses:
- My main concern of the submission is the motivation and necessity of hyper-gradient used in codebook optimization. Though author claimed that traditional optimization with STE is task-unaware, the reason behind the claim is not explained. STE is widely used in quantization training while I have not seen it is challenged for unable of catching task information. 
- Besides, hyper-gradient is proposed to solve the problem in the bi-level optimization. My understanding towards the problem lie at: the update of lower/upper part is based on an imperfect update of the counterpart. Hyper-gradient is able to perceive the interaction between the two levels. My question is: why hyper-gradient is used in update of quantized codebook instead of encoder/decoder, or why it is not used in both side?

questions:
- To prove the necessity of hyper-gradient, author should provide analytical and empirical experiments. Which baseline in Section 5 is using normal gradient descent in both level?
- Please clarify the motivation of hyper-gradient.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730532720452

thread_type:
official_review

<thread_object>
creator: reviewer_fshd

<content>
summary:
This paper proposes a new training algorithm for learning vector-quantized variational autoencoders (VQ-VAEs). The proposed method interprets the training of VQ-VAE as a bi-level optimization and applies a meta-learning-inspired technique to solve the optimization. The idea seems novel and straightforward, and the experimental results look promising.

strengths:
* The bi-level optimization view of VQ-VAE training is persuasive and seems novel. 
* The proposed algorithm is straightforward and achieves strong performance.
* The related works and the background information are carefully covered.

weaknesses:
* An obvious drawback of the proposed method is additional compute overhead and implementation complexity. However, this point is honestly described in Section 5.5 and Table 4. The improvement in various task performance metrics justifies the additional overhead of the algorithm.
* The scale and the diversity of experiments are relatively small compared to modern standards. Including a larger-scale image dataset (e.g., image of size 64x64 or larger) or experiments on a non-image dataset (e.g., medical, audio, time-series, etc.) would greatly enrich the paper.

questions:
* FID evaluation protocol is known to be highly sensitive to implementation details. It would be beneficial in terms of reproducibility to state the detailed procedure involved in FID computation, such as which software package is used and how images are preprocessed.
* Could you include some examples of generated images in the supplementary material so readers can qualitatively appreciate the improvement in generation quality?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675257492

thread_type:
official_review

<thread_object>
creator: reviewer_vwh6

<content>
summary:
The paper proposes to use meta learning to improve training of VQ-VAEs, which are hard to optimize due to the non-differentiability of the vector quantization operation. More specifically the authors propose to optimize the encoder and decoder parameters in the inner loop, and the codebook in the outer loop. The advantage of this approach is that the codebook gets direct feedback from the reconstruction loss, and not indirect feedback as in the original formulation. Different variants of the proposed meta learning algorithm and baselines are compared on CIFAR100 and FashionMNIST reconstruction and image generation. The authors further present results for applying their method to the feature maps of a ResNet.

strengths:
To my knowledge, the application of meta learning to training VQ-VAEs is new. The proposed algorithm is clearly explained, and the paper also proposes some sound tweaks of the algorithm. At a high level, the experiment design makes sense.

weaknesses:
- The experimental results on image reconstruction/generation are not very convincing. The method is only applied to very low resolution, small data sets, and the tasks are quite simple (CIFAR100, FashionMNIST), so it’s unclear how and whether the method generalizes to more modern VQ-VAE designs involving perceptual losses such as VQGAN [1].
- I appreciate that the authors address computational cost in Section 5.5, showing that MQ-VAE requires roughly 4x the compute of the baseline. However, the fair comparison would be to train VQ-VAE 4x longer, but MQ-VAE and VQ-VAE are only compared when trained for the same number of iterations (e.g. Figure 4). In particular it seems that the cheaper alternating optimization and/or synced codebook updates come close to the proposed method (Table 2), so they might match or outperform it in a compute matched comparison.
- It is somewhat unclear whether the FID and IS metrics provided in Table 1 and 3 are reconstruction or generation metrics. Generally the text seems to make not always a clear distinction between the two.
- Some of the methods are “+group” variants. This is not explained as far as I could see, nor is a reference provided.

\
Minor:
- The authors provided code, which is great. Nevertheless, it would be good to provide some details about the model architectures (both VAE and PixelCNN) and optimizers in the paper.

\
[1] Esser et al. "Taming transformers for high-resolution image synthesis." CVPR 2021.

questions:
- All the models have their codebook initialized with k-means. This is not novel, but is usually not required for VQ-VAEs to work reasonably well. Did the authors observe optimization issues when randomly initializing the codebook?
- Some recent works [2, 3] do not use an explicit codebook, but rather just scalar-quantize feature maps. These methods do not suffer from a “task-unaware codebook” since all parameters get gradients from the task loss. Could the authors comment on this?

\
[2] Yu et al. "Language Model Beats Diffusion--Tokenizer is Key to Visual Generation." ICLR 2024
[3] Mentzer et al. "Finite scalar quantization: Vq-vae made simple." ICLR 2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1732549163630

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear reviewers,

We highly value your concerns, and we have submitted the rebuttal version of the paper with major changes highlighted in purple. Based on the valuable feedback from the reviewers, the main concerns regarding MQ-VAE are related to the lack of experiments, clarity on the motivation for using hyper-gradient, and the effectiveness of the hyper-gradient technique. We have made the following efforts to improve our work.

### Additional Experiments with Scaled-up Experimental Settings

We thank all reviewers for raising concerns about the insufficient experimental evaluation. To address your concern, we have conducted additional experiments with larger-scale experimental settings. These include using modern VQ architectures - VQ-GAN and perceptual losses, as well as larger benchmark datasets like CelebA and Imagenet, all at a resolution of 256x256. The results are consistent with those presented in the main paper, demonstrating that our method achieves superior performance in larger-scale experimental settings. These results are provided in Appendix G.

### Motivation and Intuition Behind Meta-learning and Hyper-gradient

Several reviewers raised concerns about the motivation for applying hyper-gradient and questioned why it helps in learning a better codebook. 

**We want to clarify that a more intuitive understanding of our method should be motivated from a holistic meta-learning perspective.** In the spirit of meta-learning, MQ-VAE treats the codebook as hyperparameters, with the objective of improving the downstream training of the encoder and decoder. The effectiveness of the codebook is validated and updated based on the performance potential of the encoder and decoder, analogous to how meta-learning evaluates the effectiveness of model initialization for fast network adaptation [1]. The concept of hyper-gradient emerges when we consider gradient-based meta-learning [1], where the bilevel optimization problem transforms into an equivalent formulation. In this case, a hyper-gradient is used for the upper level (i.e., the codebook).

**We are then further motivated to analyze why the gradient-based meta-learning and the hyper-gradient are particularly effective in our setting.** By carefully examining the hyper-gradient, we show that it addresses three challenges specific to our setting, as outlined in the introduction. We also provide preliminary results of gradient analysis and its connection to existing work in Section 4.3 and Section 4.4. We also provide empirical analysis in Appendix H. We acknowledge that the effectiveness of hyper-gradient may not be fully intuitive when examined by parts. A more in-depth and solid analysis, such as a generalization bound, of each part of the hyper-gradient could be a valuable future direction.

### 3. Empirical Analysis of Hyper-gradient

In our rebuttal version, we also added empirical studies on the hyper-gradient in Appendix H. Specifically, we examined the dynamics of both direct and indirect gradients. Our main finding is that the indirect gradient is non-negligible in terms of its norm. Additionally, the direct and indirect gradients are not highly correlated in terms of cosine similarity, indicating that the hyper-gradient can enrich the gradient guidance for the codebook. For further details, please refer to Appendix H.

### 4. Importance of Learning an Explicit Codebook

Two reviewers (reviewer rpsn and reviewer vWH6) raised questions about methods using implicit codebooks [2, 3], which do not have trainable parameters in the quantization layer but generally perform well. However, these methods belong to a different research branch and have several limitations compared to explicit codebook-based methods. For example, FSQ [2] has reduced expressiveness since there are only a finite number of possible codewords. Additionally, hyperparameter tuning can be difficult as the range of codewords needs to be determined manually for each dimension. These limitations underscore the need for advances in explicit codebook-based methods, and our work aims to address this.

[1] Finn, Chelsea, Pieter Abbeel, and Sergey Levine. "Model-agnostic meta-learning for fast adaptation of deep networks." *International Conference on Machine Learning*. PMLR, 2017.

[2] Mentzer, Fabian, et al. "Finite Scalar Quantization: VQ-VAE Made Simple." *The Twelfth International Conference on Learning Representations*.

[3] Yu, Lijun, et al. "Language Model Beats Diffusion-Tokenizer is key to visual generation." *The Twelfth International Conference on Learning Representations*.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732703669954

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experiment Results

comment:
Dear reviewers,

Thank you for your valuable insights and suggestions. We put the our additional experiment results here for your convenience. Please also view the other additional results, such as generated samples (Appendix G) and gradient dynamic (Appendix I), in our rebuttal version. 

### Main results in Appendix G

we conducted additional experiments with larger-scale experimental settings. These include using modern VQ architectures - VQ-GAN and perceptual losses, as well as larger benchmark datasets like CelebA and Imagenet, all at a resolution of 256x256. The main results are presented in the following table. We refer to our meta-learned VQ-GAN as MQ-GAN.

| Dataset       | Method | FID ↓    | IS ↑     | LPIPS (10⁻¹) ↓ | Perceptual loss ↓ | Perplexity ↑ |
| ------------- | ------ | -------- | -------- | -------------- | ----------------- | ------------ |
| **CelebA-HQ** | VQ-GAN | 25.1     | 2.90     | 1.34           | 0.242             | 79.5         |
|               | MQ-GAN | **23.2** | **2.97** | **1.28**       | **0.218**         | **79.7**     |
| **Imagenet**  | VQ-GAN | 34.5     | 1.83     | 1.26           | 0.163             | 77.4         |
|               | MQ-GAN | **30.7** | **1.92** | **1.15**       | **0.149**         | **80.8**     |

The results are consistent with those presented in the main paper, demonstrating that our method achieves superior performance in larger-scale experimental settings.

### Main results in Appendix E

We provide a comparison of three design choices for applying hyper-gradient descent in Appendix E. The results are presented below. The methods are evaluated in the CIFAR100 dataset.

| Method                                             | FID ↓    | IS ↑     | LPIPS (10⁻¹) ↓ | MSE (10⁻³) ↓ |
| -------------------------------------------------- | -------- | -------- | -------------- | ------------ |
| VQ-VAE                                             | 81.8     | 4.99     | 1.88           | 7.73         |
| MQ-VAE (Hyper-gradient applied to codebook)        | *53.5*   | **6.40** | *1.25*         | *5.02*       |
| MQ-VAE (Hyper-gradient applied to encoder-decoder) | 56.1     | 5.93     | 1.44           | 5.48         |
| MQ-VAE (Hyper-gradient applied to both sides)      | **52.8** | **6.40** | **1.23**       | **4.58**     |

The results show that applying hyper-gradient to the codebook yields better performance than applying to the encoder-decoder pair. Besides, applying hyper-gradient to both sides is slightly more performant but more computationally costly. This again highlights the effectiveness of our design choice of selecting the codebook as the upper level and applying hyper-gradient to it. 

Thank you again for engaging with us and for your valuable feedback.



Best regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work introduces a meta-learning solution for training a vector quantized VAE, named MQ-VAE. The authors highlight the shortcomings of vector quantized VAEs, such as code collapse and gradient computation, and formulate the learning process as a bi-level optimization problem. This leads to a formulation where the codebook of the VQ-VAE is learned in the outer-level problem, while the network parameters are learned in the inner-level optimization. Experiments on CIFAR and MNIST demonstrate that MQ-VAE outperforms VQ-VAE comfortably.  

The paper was reviewed by four experts in the field. The primary concern raised by the reviewers was the lack of evaluations on large-scale problems, which was only partially addressed during the author-reviewer discussion period. Unfortunately, none of the reviewers championed the paper, as the lack of comprehensive empirical evaluation remains unresolved.  

The AC agrees with the reviewers and, sadly, recommends `rejection.` We hope this review helps the authors identify areas for improvement. On a semi-related note, recent advances in bilevel optimization suggest that first-order methods can expedite the optimization process (e.g., Liu et al., *BOME! Bilevel Optimization Made Easy: A Simple First-Order Approach*, NeurIPS'22). Incorporating such methods might make the learning of MQ-VAE faster.