Learning to Quantize for Training Vector-Quantized Networks

We appreciate your constructive feedback very much. We provide our response to your review as follows.

Computational Cost and Memory Overheads

We conducted additional experiments to address your concerns. When evaluated on the CelebA dataset with a batch size of 128, the increase in memory usage is marginal and acceptable in practice. For time comparison, we set VQVAE as the baseline, which needs around 3.6 hours to finish the training of 50k steps (and does not improve after that). We find that, MQVAE only requires 2.4 hours to reach the same LPIPS score as VQVAE (around 9.5k steps), and can keep improving after that. This demonstrates that our method converges much faster than VQVAE and is able to outperform baselines with extended training time.

Difficulty of converging using the meta-learning strategy

Empirically, we did not observe stability issues during training. Theoretically, related convergence analyses for this type of gradient-based bilevel optimization algorithm can be found in [5], [6], and the references therein. Our MQ belongs to this type of optimization and is guaranteed to be stable and converge under certain conditions.

Demonstrate how codes are generated and selected and which part is optimized; The codebook distribution; Implementation Details

Please follow this anonymous link https://anonymous.4open.science/r/MQVAE-B52C for the figure illustration and code implementation, and we will open-source them once the paper is accepted.

Compare with VQVAE+Affine+OPT+replace+l_2, and CVQ-VAE

According to the performance reported in [1], MQVAE performs slightly worse than the combination of VQVAE+Affine+OPT+replace+$l_2$. Fortunately, we can show that when combined with additional techniques such as $l_2$ projection, MQVAE can still outperform [1] on the MNIST datasets, as shown in the table below. We also include a comparison with CVQ-VAE [2] for your reference.

Limited Sensitivity Analysis: Have you tried multiple unroll steps (beyond one step) or different finite-difference approximations? How does that affect training time and results?

Our finite-difference (FD) approximation currently supports only one-step unrolling. We conducted additional experiments with alternative approximations, including the conjugate gradient (CG, [3]) and Neumann series (NMN, [4]) methods. Our ablation studies were done on the CelebA dataset. We found the type of approximation largely does not affect.

Downstream Utility and Domain Generalization: Do you think the approach can directly transfer to other discrete representation tasks, such as speech tokens or molecular modeling, without major changes?

Based on our experiments, we believe that our framework can be directly adapted to other discrete representation tasks. Our approach is flexible in accommodating different types of task loss, as demonstrated in our experiments where both MSE loss and perceptual loss have been evaluated. In both cases, MQ shows superiority over VQ.

[1] Huh, Minyoung, et al. "Straightening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks." International Conference on Machine Learning. PMLR, 2023.

[2] Zheng, Chuanxia, and Andrea Vedaldi. "Online clustered codebook." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[3] Rajeswaran, Aravind, et al. "Meta-learning with implicit gradients." Advances in neural information processing systems 32 (2019).

[4] Lorraine, Jonathan, Paul Vicol, and David Duvenaud. "Optimizing millions of hyperparameters by implicit differentiation." International conference on artificial intelligence and statistics. PMLR, 2020.

[5] Pedregosa, Fabian. "Hyperparameter optimization with approximate gradient." International conference on machine learning. PMLR, 2016.

[6] Rajeswaran, Aravind, et al. "Meta-learning with implicit gradients." Advances in neural information processing systems 32 (2019).

We appreciate your constructive feedback very much. We provide our response to your review as follows.

Is the bi-level optimization approach adopted in this paper necessary?

Yes, it is necessary. The two components of our method address distinct challenges. Specifically, the hypernet resolves the issue of codebook collapse and enhances codebook utilization, while the bilevel optimization approach ensures that the task loss gradient reaches the codebook. As demonstrated in our ablation studies (Section 5.3), both components are essential and must be combined to achieve optimal results.

The description of convergence is inconsistent.

We apologize for the confusion. The figure illustrates the workflow of bilevel optimization for a single gradient step, whereas the algorithm blocks detail our specific implementation. Our approach applies a finite difference-based approximation, wherein one gradient descent step approximates the converged solution of the autoencoder.

Thus, a consistent pseudo-algorithm is to replace

"Update $\psi$ using gradient descent: $\nabla_{\psi}\mathcal{L}(\phi-\xi\nabla_\phi\mathcal{L}(\phi, \theta, \psi), \theta-\xi\nabla_\theta\mathcal{L}(\phi, \theta, \psi), \psi)$" in the while loop with

"Copy $\phi^\prime=\phi$ and $\theta^\prime=\theta$, and update $\phi^\prime$ and $\theta^\prime$ using gradient descent by $\nabla_\phi\mathcal{L}(\phi, \theta, \psi)$ and $\nabla_\theta\mathcal{L}(\phi, \theta, \psi)$ until converging, resulting in $\phi^{\prime*}$ and $\theta^{\prime*}$ (the computation graph is retained during descent). Update $\psi$ using gradient descent: $\nabla_{\psi}\mathcal{L}(\phi^{\prime*}, \theta^{\prime*}, \psi)$". This is consistent with the figure description, in which the lower level is optimized until convergence, followed by one update of the upper level. We use a one-step unroll scheme, approximating $\phi^{\prime*} \approx \phi-\xi\nabla_\phi\mathcal{L}(\phi, \theta, \psi)$ and $\theta^{\prime*} \approx \theta-\xi\nabla_\theta\mathcal{L}(\phi, \theta, \psi)$, which corresponds to the pseudo-algorithm presented in our paper.

Justification of advantages over VQGAN-LC.

VQGAN-LC also addresses codebook under-utilization and thus exhibits higher codebook usage compared to the original VQGAN. In contrast, our approach employs a hypernet reparameterization, which enables simultaneous updates of all codebook entries.

We appreciate your constructive feedback very much. We provide our response to your review as follows.

Magnitude Comparison

In our experiment, we have found that the magnitude of both indirect and direct are around $10^{-2}$, so none of them dominate each other. Please follow this anonymous link https://anonymous.4open.science/r/MQVAE-B52C for the figure.

Whether [1] can benefit from bi-level optimization

We want to clarify that [1] addresses an essentially different problem. In our formulation, the core idea is to generate a codebook using (1) a learnable embedding and (2) a learnable transformation (e.g., linear projection). The method presented in [1] does not involve learnable embeddings; moreover, its projections are considered part of the backbone autoencoder architecture, whereas in our case, the projection (a hypernet) is integrated into the codebook. Thus, [1] falls outside the scope of our study, and we cannot comment on applying bilevel optimization to [1] based on our main paper's arguments.

Relation to [2]

This work primarily focuses on improving the STE estimator used in the original VQVAE paper. Although it facilitates smoother gradient propagation through the non-differentiable quantization layer, the codebook update still depends solely on the distribution of codes and features and does not incorporate the task loss. Therefore, that work addresses a different issue and is orthogonal to our approach.

Computation Cost and Implementation details

We conducted additional experiments to address your concerns. When evaluated on the CelebA dataset with a batch size of 128, the increase in memory usage is marginal and acceptable in practice. For time comparison, we set VQVAE as the baseline, which needs around 3.6 hours to finish the training of 50k steps (and does not improve after that). We find that, MQVAE only requires 2.4 hours to reach the same LPIPS score as VQVAE (around 9.5k steps), and can keep improving after that. This demonstrates that our method converges much faster than VQVAE and is able to outperform baselines with extended training time.

[1] Zhao, et al. "Image and video tokenization with binary spherical quantization." arXiv preprint arXiv:2406.07548 (2024).

[2] Fifty, et al. "Restructuring Vector Quantization with the Rotation Trick." arXiv preprint arXiv:2410.06424 (2024).

We appreciate your constructive feedback very much. We provide our response to your review as follows.

Why not use a simpler method such as clustering the embeddings or EMA update

One of the advantages and novel aspects of MQ, compared to simpler methods, is that the codebook update follows a more complete loss path in that it directly interacts with the task loss gradient.

For example, the simplest VQVAE is designed to reconstruct the input image using the mean squared error loss. However, updating the codebook by enclosing the code and feature in latent space does not involve the MSE between the image and the reconstructed output. In contrast, our meta-learning-based method avoids this incomplete gradient issue. Please refer to Figure 2 in the main paper for additional details.

What is the training stability and efficiency of the optimization compared to other methods?

Empirically, we did not observe stability issues during training. Theoretically, related convergence analyses for this type of gradient-based bilevel optimization algorithm can be found in [1], [2], and the references therein. Our MQ belongs to this type of optimization and is guaranteed to be stable and converge under certain conditions.

It is not clear if the training is efficient.

We conducted additional experiments to address your concerns. When evaluated on the CelebA dataset with a batch size of 128, the increase in memory usage is marginal and acceptable in practice. For time comparison, we set VQVAE as the baseline, which needs around 3.6 hours to finish the training of 50k steps (and does not improve after that). We find that, MQVAE only requires 2.4 hours to reach the same LPIPS score as VQVAE (around 9.5k steps), and can keep improving after that. This demonstrates that our method converges much faster than VQVAE and is able to outperform baselines with extended training time.

I am confused about the details of the hyper-net. In lines 287-292, it is described as an MLP. Does it mean it takes a single embedding and generates a single code entry? How does it produce the whole codebook?

In this case, the MLP will separately project each codebook entry with respect to their dimensionality, meaning if the embedding is $e_n, n=1...,N$, the generated codebook is $MLP(e_n), n=1,...,N$. Nevertheless, if any $MLP(e_n)$ is selected and receives a gradient, the MLP will be updated, resulting in a different codebook being generated by this updated MLP.

In lines 69-74, what is the "specific task"? What is "task-aware"? Is it the VQVAE loss?

The specific task depends on the benchmark. For instance, in image reconstruction tasks, the task loss is the MSE loss in the case of VQVAE or a combination of perceptual loss, MSE loss, and adversarial loss in the case of VQGAN. Task awareness means that the gradient used to update the codebook directly incorporates the task loss. One contribution of this work is to establish this connection. In Figure 2, the task loss reaches the codebook through various pathways; previous VQ methods updated the codebook solely based on the gradient of selected codes toward selected features, without incorporating the task loss.

[1] Pedregosa, Fabian. "Hyperparameter optimization with approximate gradient." International conference on machine learning. PMLR, 2016.

[2] Rajeswaran, Aravind, et al. "Meta-learning with implicit gradients." Advances in neural information processing systems 32 (2019).