Learning to Quantize for Training Vector-Quantized Networks
Authors: Peijia Qin
Jianguo Zhang
Venue: ICML 2025 poster (2025)
Keywords: Vector Quantization
Meta-Learning
Generative Models
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: TLDR: We propose Meta Quantization (MQ), a bi-level optimization-based VQ training framework inspired by meta-learning.

Abstract: Deep neural networks incorporating discrete latent variables have shown significant potential in sequence modeling.
A notable approach is to leverage vector quantization (VQ) to generate discrete representations within a codebook.
However, its discrete nature prevents the use of standard backpropagation, which has led to challenges in efficient codebook training.
In this work, we introduce **Meta-Quantization (MQ)**, a novel vector quantization training framework inspired by meta-learning.
Our method separates the optimization of the codebook and the auto-encoder into two levels.
Furthermore, we introduce a hyper-net to replace the embedding-parameterized codebook, enabling the codebook to be dynamically generated based on the feedback from the auto-encoder.
Different from previous VQ objectives, our innovation results in a meta-objective that makes the codebook training task-aware.
We validate the effectiveness of MQ with VQVAE and VQGAN architecture on image reconstruction and generation tasks.
Experimental results showcase the superior generative performance of MQ, underscoring its potential as a robust alternative to existing VQ methods.

Learning to Quantize for Training Vector-Quantized Networks
Peijia Qin Jianguo Zhang
Abstract
Train quantizer min ((().)
Deep neural networks incorporating discrete la- tent variables have shown significant potential in sequence modeling. notable approach to leverage vector quantization (VQ) generate discrete representations within codebook. How- ever, its discrete nature prevents the use of stan- dard backpropagation, which has led to challenges in efficient codebook training. In this work, we in- troduce Meta- cetta-Quatizztionn (MQ), a novel vector quantization training framework inspired by meta- learning. Our method separates the optimization of the codehook and the auto-encoder into two levels. Furthermore, we introduce hyper-net replace the embedding parameterized embaddinn-parammtterieeee codebook. enabling the codebook to be dynamically gener- ated based on the feedback from the auto-encoder. Different from previous VQ objectives, our in- novation results in a meta-objective that makes the codebook training task-aware. We validate the effectiveness of MQ with VQVAE and VQ- GAN architecture on image reconstruction and generation tasks. Experimental results showcase the superior generative performance of MQ, un- derscoring its potential as a robust alternative to existing VQ methods.
Hyp
Φ ((((
0 (7)))
Hyp
E
Train encoder-decoder pair min [(Φ) 8, 1)) (..
Hyper-net: Hig. Learnable embedding: Fip
Encoder: Fg Decoder: Gg
Figure 1. Our framework is based on bi-level optimization. Meta-quantization learns the codebook and encoder-decoder pair using bi-level optimization framework. At the lower level, the encoder-decoder pair is trained to converge while keeping the codebook fixed. the upper level, the codebook is optimized via hyper-gradient descent using the optimal encoder-decoder pair.
learning latent variables through vector quantization (VQ, Gray (1984)), a process that quantizes features into clusters referred to as codes. The concept of VQNs was first in- troduced with the vector-quantized variational autoencoder (VQVAE, den Oord al. (2017)) in the context of gen- erative models. Subsequent research has demonstrated that training autoregressive priors on discrete representations ob- tained via vector quantization yields highly effective models for image generation (Razavi et al., 2019; Roy et al., 201 8; Ramesh et al., 2021; Esser et al., 2021; Chang et al., 2023). Beyond image generation, VQNs have achieved notable success in speech generation (Dhariwal et al., 2020) and have extended their utility to other domains, such as im- age representation learning (Caron et al., 2020) and speech representation learning (Chung et al., 2020).
1. Introduction
Learning discrete latent variables is favorable for tasks natu- rally modeled sequences of discrete symbols, including language and speech (Vinyals et al., 2015). Vector-quantized networks (VQNs) provide an effective approach for this by
Research Institute of Trustworthy Autonomous Systems and Department of Computer Science and Engineering, Southern Uni- versity of Science and Technology, Shenzhen 518055, China. "Ganggoon Provincial Key Laboratory of Brain-inspired Intelli- gent Computation. Department of Computer Science and Engineer- ing, Southern University of Science and Technology, Shenzhen 518055, China, and also with the Pengeheng Laboratory. Corre- spondence to: Jianguo Zhang <zhanggg@ austech.edu.en...
The training approach for vector-quantized networks (VQNs), as implemented in VQVAE, involves learning a codebook to represent compressed semantic data. The encoder F' maps input data an embedding, which sub- sequently quantized by selecting the nearest neighbor in C. The selected code replaces the embedding and is then passed to the decoder Gg to produce the output. Since the quantization operation introduces a non-differentiable bot- tleneck, a traight-through estimator (STE, Bengio et al. (2013)) is employed to enable gradient flow through the
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Meta-Quantization
VQ layer to the encoder during backpropagation. However, because the backpropagated gradient bypasses the codebook in the STE framework, the codebook is instead optimized using a vector quantization objective. This objective aligns the distribution of the encoder embeddings with the selected codes to achieve distribution matching.
functions of the hyper-net. At the lower level, the tentative optimization steps of the encoder-decoder pair from the first stage are undone, as a better hyper-net is identified in the previous step. Instead, the encoder and decoder are opti- mized for one step using gradient descent with the updated hyper-net. These two levels of optimization are performed iteratively until convergence.
Despite its effectivvness, the current training framework for vector quantization faces several challenges. First, vector quantization often suffers from phenomenon known index collapse, where only small subset of codes remains active during training (Kaiser et al., 2018). This issue arises because frequently optimized token embeddings gradually align more closely with the feature map distributions, while less frequently or never-optimized token embeddings (inac- tive tokens) are excluded from the training process. These inactive tokens remain unused during inference, leading to suboptimal codebook utilization. Second, the codebook is optimized solely using the vector quantization objective, which focuses on distribution matching. However, due to estimator (STE), gradients from the task loss bypass the codebook during backpropagation. As a result, codebook updates are agnostic to the specific task, potentially undermining its effectiveness for the overall ob- jective of image modeling.
Figure 2 provides a detailed illustration of the training procedure and gradient flow in our proposed framework. Specifically, the meta-objective is implemented by unrolling the lower-level training for several steps. with its gradient computed through hypergradient descent (i.e., optimizing through the inner gradient descent path). The hypergradient paths demonstrate how the gradient from the task loss can now reach the codebook through multiple routes. This con- trasts with previous codebook training strategies, where the task loss has no direct gradient impact on the codebook. A practical advantage of our framework is that the hyper-net can be discarded after training, retaining only the generated weights of the codebook. This means the downstream does not suffer from additional inference burden. Furthermore, the bi-level training procedure does not affect downstream tasks, the parameters trained in the first stage are fixed, This makes meta-learned codebooks as efficient as their conventional counterparts.
the
Drawing inspiration from meta-learning (Finn et al., 2017), we propose a novel vector quantization training framework termed Meta Quantization (MQ). Our method builds di- rectly on the vector quantization mechanism and reparam- eterizes the codebook by introducing separate hyper-net Hy. The hyper-net predicts the parameter matrices of the codebook from a learned embedding. While maintaining the same expressivity at the quantization bottleneck, the hyper- net aggregates gradient feedback from all matched codes and generates the codebook holistically. This design acceler- ates convergence without introducing additional complexity to the quantization operation and enables seamless integra- tion with existing approaches. The theoretical advantages (for example, accelerating convergence) of such overpa- rameterized architecture have been studied in Arora et al. (2018). Moreover, we formulate an asymmetric bi-level optimization problem to allow the hyper-net to learn quan- tization in a Mheaa-leannnng fashion. In this framework, the hyper-net functions as a meta-learner (analogous the hy- perparameters in meta-learning), while the encoder-decoder pair acts as a task-specific learner. As illustrated in Figure I, the hyper-net and the encoder-decoder pair are optimized at the upper and lower levels, respectively. At the upper level, the hyper-ne net anticipates the encoder-decode pair's 's future performance by tentatively training them until conver- gence (unrolling for one step as a practical surrogate) while keeping the codebook fixed. Subsequently, the hyper-net is updated via hyper- gradient descent to minimize the loss for one step, leveraging the optimal encoder and decoder as
We evaluate the effectiveness of our framework across the VQVAE and VQGAN architectures on standard image re- construction and generation tasks. The results show that the MQ approach consistently outperforms multiple baselines and ablation methods, underscoring its superiority. Impor- tantly, we modify only the codebook training procedure without altering the quantization mechanism or downstream autoregressive model training. Thus, our work is comple- mentary to ongoing advancements in these areas.
The remainder of this article is organized as follows. Section presents related work. Section 4 details our proposed method. Experimental setup and results are discussed in Section 5. The paper is concluded in Section 6.
2. Related Work
2.1. Vector-Quantized Networks
Vector quantization networks (VQNs), initially proposed for image generation as VQVAE, map continuous embeddings to discrete codebook representations using vector quantiza- tion (VQ). The non-differentiable nature of VQ is addressed via straight-through estimators (STE, Bengio et al. (2013)). However, this presents the index collapse issue during code- book training, where only a small number of codes are ac- tive, To prevent the issue, Lańcucki et al. (2020); Zeghidour et al. (2021); Dhariwal et al. (2020) propose to reset unused codes periodically, In contrast, Kaiser al. (2018) proposes
2

---

Meta- Veta-Quantization
Hyper-gradient descent
00
0Φ d+p
V/ + ( .))
aask ow
Step 2
C (Fp)
Gy
Lias
Ō Mcodebook inp
codebook
0. commit
ä iLiss 0
commit
0Φ
Gradient descent Φ (v.. 0)
Tentatively train and 0 until converging Result in and g*
Gradient descent
Forward path
(xb)
((P) 0 <V {V.L(.. 0,
Gradient path
Step
Ze
Zg
C= = (Fg) (Temporarily fixed)
FΦ
Go
Frask
Indirect hyper- gradient path
Leakk aB
commit
commit
Direct hyper- gradient path
Laak d
Figure Gradient paths meta-quantization. The lower level is optimized by gradient descent. The upper level is optimized by hyper-gradient descent with direct hyper-gradient and indirect hyper-gradient. Here, we use Linsk, Lcommit and condebook to denote the corresponding loss computed with o (1b) and (x²).
performing quantization in lower dimensionality to improve codebook utilization. Gumbel- VQQ (Karpathy, 2021) enables differentiable quantization through continuous approxima- tion of the arg min operator. Other approaches include affine reparameterization (Huh et al., 2023), la normaliza- tion (Yu et al., 2022), probabilistic formulations (Roy et al., 2018; Takida et al., 2022). More recently, VQGAN-LC (Zhu et al., 2024a) scales up the codebook size by initial- izing codebooks with features extracted by a pre-trained vision encoder, and focuses on training a projector that aligns the entire codebook with the feature distributions of the encoder. Sim VQ (Zhu al., 2024b) achieves even supe- rior results by reparameterizing the code vectors through linear transformation layer based on learnable latent basis.
hand, searches for the optimal hyperparameter following the well-studied bi-level optimization based hyperparameter optimization literature (Pedregosa, 2016 6; Finn et al., 2017).
2.2. Bi-level Optimization
Bi-level optimization (BLO) has found wide applicability in various machine learning tasks, with meta-learning (Finn et al., 2017: Rajeswaran et al., 2019) being one of its most notable applications. Other areas where BLO has been suc- cessfully employed include neural architecture search (NAS, Liu et al. (2019); Zhang et al. (2021)) and hyperparameter optimization (HPO, Lorraine et al. (2020); Franceschi et al. (2017)). Despite its broad usage, solving BLO problems remains challenging due to the inherently nested structure of the two optimization tasks. Gradient- BLO (Choe et al., 2023) has garnered significant attention due its scal- ability nigh-dimensional problems with large number of trainable parameters.
Other works, including FSQ (Mentzer et al., 2024) and LFQ (Yu et al., 2024), have explored using non-learnble, implicit codebooks. They replace the vector quantizer with scalar quantization schemes through rounding, making the code- book predefined and non-learnable. However, such methods suffer from reduced expressiveness due the finite set of possible code values (specifically, grid of integer values). Furthermore, FSQ and LFQ require representation to be projected into a reduced dimensionality, leading to signifi- cant information loss. These constraints motivate our focus on explicit codebook improvements. Our method shares certain similarities with them when viewing the codebook as hyperparameter. FSQ and LFQ treat the codebook as predefined hyperparameter; our work, on the other
In this work, we extend the application of gradient-based BLO to propose a novel approach for codebook training within the vector quantization framework. Inspired by meta- learning, our MQ framework treats the codebook as hyper- parameters, which are parameterized by hyper-net. The meta-objective is enhance the training of the encoder and decoder, replacing the previous VQ loss. The effectiveness of the codebook generated by the hyper-net is validated by assessing the performance of the encoder and decoder. This
3

---

Meta-Quantization
process mirrors the practice of validating model initializa- tion in meta-learning A similar strategy involving hyper- gradient descent with inner loop unrolling is employed in our work. a technique that is also commonly found in the meta-learning literature.
Overall, differentiable objective is minimized:
min (xx))~ (trak (Go(hc((a(())))))) 0,8,C
ta Coommit (hcc (hc ((x))
(6)
+B
(7)
3. Preliminary
where a and ] are scalars that balance the loss combination. In this training framework, the decoder optimizes the first term, the encoder optimizes both the first and the middle terms, and the codebook is optimized only by the last term.
A vector-quantized network (VQN) is a neural network that incorporates vector- quatto-quantzzziionn (VQ) layer he()), as described by the following equation:
y = (F)((((()))))) = (i(hhc(z))) (z(z))
4. Methodology
Here, &, denotes the embedding obtained by applying the en- coder Fs (parameterized by Φ) to the input x. 2y represents the quantized embedding, which is obtained by applying the vQ layer he to ze. The output y is generated by the decoder Go (parameterized by 0), which takes bn as input.
4.1. Hyper-Net Reparameterization
We avoid directly optimizing the codebook by introducing a hyper-net H., which takes as input a learnable embedding Eg and predicts the codes in the codebook, which are then used for vector quantization. More formally, this can be expressed as C Hy (E,)). As shown on the right side of Figure 1, the trainable parameters (denoted by e)) in the quantization layer are the hyper-net and the learnable embedding rather than the codebook itself. For notation simplicity, we use Ha to denote H3s (Ea) in the remaining parts of this paper.
The VQ layer he(-) performs quantization on Ze by select- ing a vector from the codebook C = {e1, e2, ................................ ex } based on a distance measure d(.,)). )),
ex
Zg ek:
where == arg min dizes e,)
Here, a learned vector e, is referred to as the code, and the index i denotes the corresponding code. The Fuclidean distance commonly used as the distance measure for -(---) (van den Oord et al., 2017). The quantized embedding 2n subset of C, and updating Zy corresponds to partially updating C.
Once the first stage of image modeling is completed, the hyper-net can be discarded, retaining only the generated codebook for downstream tasks. This codebook acts as direct replacement for the codebook used in previous train- ing procedures. Furthermore, since the generated codebook is a complicated non-linear combination of the learnable embedding through the mapping of the hyper-net, even only a few codes are selected, substantial part of the hyper- net still receives gradients. This allows the hyper-net to generate codebook that better aligns with the embedding distribution.
The task loss CGass(Go (Fa())))) y) is not continuously differentiable due to the arg min operator in lec To address this issue, a (2013)) is applied, where the non-differentiable part öz. ignored:
estimator (STE, Bengio et al. dx, is
ŌLask Oip
Lrnk δy Oze
4.2. A Bi-level Optimization Framework
2
(3)
θy Ong Jeh
In the spirit of meta-learning, we propose to optimize the hyper-net and encoder-decoder pair by solving a bi-level optimization problem. Specifically, we optimize the hyper- net H. at the upper level and the encoder F& and decoder Go at the lower level. In both levels, we consider a loss defined in the same form as in Equation (7), i.e., the sum of (5) Lesski Lcommits and Loodebook.
To ensure an accurate STE, ze and z, are aligned using two additional losses:
Lcommit (Zg. Ze) = d (ze: sg (zy])
(4)
Loodebook (Zy. Ze) = d (sg [ze]. 2g)
Here, sg denotes the stop-gradient operator, which treats the entire term constant with zero partial derivatives. The commitment loss Leommit encourages the embedding to move toward the selected codes, whereas the codebook loss Loodebook. also known as the vector quantization objec- tive, pushes the selected codes toward the centroids of the embedding.
Lower Level At the lower level, we train the encoder F& and decoder Gg by minimizing L(ó, 8, v)) Specifically, we aim find the optimal values of á and @ with şi temporarily fixed, resulting in the following optimization problem:
Φ (v), 0" (15)
arg min L(o, v))
(8)
0.0

---

Meta- ceta-Quantization
Here, x (16) and 0 (1) denote the optimal solutions for and 0, which are functions of v/ since the lower-level problem does not take & as an argument.
for & (vb and 0 (v). To address this, we adopt a one-step
approximation (Finn et al., 2017):
(V), 2 -
Upper Level At the upper level, the hyper-net Hy is trained by minimizing the loss of the same functional form but using 6 (x) and 8" (v) that were optimally learned at the lower level as arguments. The loss then only depends on ,, and the upper-level optimization problem is formulated as:
(11)
where & is the learning rate for the lower-level problem. One-step unrolled approximated solutions, (((+)= Φ [VoC(o, 8, eb) and ('()) = - -V <VoL( 8, +)), are used as surrogates for (사) and (v)). This is equivalent to introducing a surrogate objective L( Lf 0, v²), 0 {Vo[(o, 0, e)), 16) that closely resembles the upper-level
min L(0" (V), (0)66(())) (v)).
Bi-level Optimization Framework By integrating the two levels of optimization problems, we present the overall bi-level optimization problem as:
objective in Equation (9).
Existing meta-learning methods compute Equation (11) by either backpropagating through the optimization process at the lower level (Finn et al., 2017) or applying implicit dif- ferentiation with Hessian matrix of the inner optimization problem (Rajeswaran et al., 2019). However, the problem size scales, the memory and computational burden grow significantly. Therefore, we employ further approximation by noticing that Equation (11) can be computed using the chain rule, followed by a finite difference approximation (Liu et al., 2019) as:
min L(6 ((b"(w), 0 (v)) 1))
s.t. 6- (b), 0 (v) = arg min C(o, 0, 0.0
(10)
The two levels of optimization problems are mutually depen- dent on each other. The solution to the optimization problem at the lower level, o² (+)) and 0* (1² serves as a parameter for the upper-level problem, while the non-optimal variable ei at the upper level acts a parameter for the lower-level prob- lem. By solving the two interconnected problems jointly, we can learn o² 0", and rf* in an end-to-end manner.
VoLlo - (VVC),00))) (12) = =
(13)
Algorithm 1 Meta- Quantization
Require: Dataset D
(14)
1: Initialize Φ. and vp for Fạ. Go. Hy and Exp
2 v L( 0',
(15)
2: while not converged do
VeL(Ot O 1b) V 1w.
3: Update vb by
gradient descent
as
2
(Vold 0, 6) (V.l(o, 0, v)), 4: (E (E=00 0 using alternated optimization)
V Ve +)
v. ,b,
(16)
5: Update @ and 0 by gradient descent as V ((0,00)))
2c
and VoL(Ø, bb)
where 0² = Φ + eV e v)) 0
0 +
6: end while
L( 0r 1)), and e is a small scalar. The finite dif- ference is applied to approximate the matrix-vector multi- plication term in Equation (14) for efficient computation.
Ensure: 0",0" 0 and 2b*
Optimization Algorithm We employ an efficient gradient-based optimization algorithm to solve the bi-level optimization problem presented in Equation (10), with the two levels optimized iteratively until convergence. Related convergence analyses of this type of gradient- based bi-level optimization algorithm can be found in Pedregosa (2016), Rajeswaran al. (2019), and references therein.
4.3. Gradient Analysis
Essentially, meta- neta-quantization introduces a meta-objective
L(d (v), 0* (4)) b)) in place of the VQ objective for the quantizer training. While shares the same functional form as the previous framework (Equation (7)), the arguments and 0 are set to their optimal values (²(()) and 0 (v) We demonstrate that this meta-objective improves the gradient guidance for db by performing a gradient analysis on the one-step-unrone unrolled surrogate loss using the chain rule.
Gradient descent can be applied directly to the lower-level problem; however, a significant challenge arises at the upper level: computing the hyper- gradient,,,,,, i.e., the gradient of the upper-level objective with respect to vb, can be computa- tionally prohibitive due to the lack of an analytical solution
Define L' 0) = L(d (v), = C(/
5

---

Meta- ceta-Quantization
EV 0, 0),0 EVoL(O, 0,,)),)))) 1).
We then have
pared methods follows Takida al. (2022) with 64 channels. A codebook size of 1024 is used across all methods. For the hyper-net configuration, an MLP is used. It first lifts the 32-dimensional learned embedding to 256 dimensions, then projects it back to 32 dimensions to form the code- book entries for quantization, using Tanh as the activation function.
dL'
0c'
0
Oc
00
Oc
+
x
+
(17)
dup
Oup
dmb
00
The last two terms on the right-hand side, especially 00'
0 and Be
Ogg referred to as the best-response Jacobian in the literature (Choe et al., 2023), capture how the encoder-decoder pair reacts to changes of the quantizer. Therefore, the update of 2b must consider not only the direct gradient from the loss for minimizing quantization error but also additional information about indirect gradients- how the encoder and decoder would respond to changes in the quantizer
Models are trained using Mean Squared Error (MSE) as the reconstruction loss; no perceptual or discriminative losses are used, We employ the Adam optimizer (Kingma & Ba, 2015) with momentum set to (0.9, 0.95) and an initial learn- ing rate of le-4. The learning rate follows a linear warmup for the first 10% of epochs, then a half-cycle cosine de- cay. No weight decay is applied to the quantizer. Training runs for a maximum of 90 epochs, with early stopping if performance saturates.
Od Th3
and
ag ), and their performance potential ( BC and B.C ). The en- Dey
30"
coder and decoder select their best response by conducting gradient descent, which the quantizer accounts for. This fa- cilitates the finding of a globally optimal quantizer, thereby improving its stability and robustness. See also Figure 2 for the gradient path and hyper-gradient path used in the lower and upper levels, respectively. Additionally, we observe that by employing this strategy, v/r can now receive gradient from frskk- For example, the first terms of C', i.e., Classie depend on ø', which in turn depends on vb. This joint effort enables the Ltask to influence yi during backpropagation.
Evaluation metrics include MSE (the same as training loss), LPIPS (Zhang et al., 2018), and model perplexity. Perplexity is defined as 2*((r), where (p) is the entropy over the codebook likelihood; a higher value is preferred as it implies more uniform code usage.
Our method is compared against VQVAE (van den Oord et al., 2017), SQVAE (Takida et al., 2022), and Gumbel- VQVAE (Karpathy, 2021; Esser et al., 2021). Comparisons also include VQ training techniques such as l2 normaliza- tion (Yu et al., 2022), least-recently-used (LRU) replace- ment (Lañcuuci et al., 2020), affine reparameterization, and synchronized training (OPT) (Huh et al., 2023). For other unspecified configurations, experimental settings from Huh et al. (2023) are adopted, and their results are cited for comparison.
5. Experiment
In this section, we first evaluate MQ on image generative modeling, including image reconstruction and image gener- ation, using the VQVAE architecture (van den Oord al., 2017). The experiments are done on small-scale datasets: CIFARIO (Krizhevsky et al., 2009) and CelebA (Liu al., 2015). We then scale up to a larger experimental setting on FFHQ (Karras et al., 2019) and ImageNet (Deng et al., 2009) with VQGAN (Esser et al., 2021), which involves perceptual loss and adversarial loss as task losses. We re- fer to the resulting methods combined with MQ "MQVAE" and "MQGAN", respectively. For the bi-level optimization algorithm implementation, Betty library (Choe et al., 2023). We release our code at GitHub for future research.
Image Reconstruction Table I shows that almost all base- lines significantly improve training stability and the recon- struction performance of generative models. However, as analyzed in the previous section, existing methods suffer from the limitations that lack a direct gradient flow from task loss to codebook, resulting in degraded reconstruction performance. Notably, our method provides the best im- provements, particularly with respect to the MSE, which is the task loss. This verifies the effectiveness of our method for achieving superior image quality and codebook utiliza- tion.
code is mainly based on the
5.1. Evaluation with VQVAE
Setup We first show preliminary experimental results with the VQVAE architecture on CIFAR10 (Krizhevsky et al., 2009) at 32 x 32 resolution and CelebA (Liu et al., 2015) at 128 x 128 resolution. For CelebA, images undergo ran- dom cropping to 140 x 140 pixels, followed by resizing the smaller dimension to 128 while maintaining aspect ratio. No additional augmentations are used for CIFAR10.
Image Generation We extend our results to the image generation task using MaskGIT (Chang et al., 2022) on CelebA, utilizing the codebook trained in the first stage directly. A slimmed-down version of MaskGIT (Huh et al., 2023) is used: VQGAN using 32 channels instead of 128 and transformer using 8 blocks instead of 24. In Table 2, the generation results improve from the baseline by 90.4 FID and 74.8 FID from the hest-performing variation. In
The backbone architecture of the autoencoder for all com-
hhaps//giinnh.cco///aee///aaaa
6

---

Meta-Quantization
Table Comparison between various methods image recon-
Table 3. Reconstruction performance on FFHQ
struction task.
Method
Utilization (%) rFID LPIPS
Method
MSE (10
Perplexity
LPIPS
VQGAN+ (Esser ct al., 2021)
2.3
5.25
0.12
VQVAE (van den Oord et al., 2017)
5.65
14.0
0.54
VQGAN-FCF (Yu et al., 2022)
10.9
4.86
0.11
12 (Yu el al., 2022)
3.21
57.0
0.36
VQGAN-EMAT (Razavi et al., 2019)
68.2
4.79
0.10
LRU (Lancuuci ct al., 2020)
4.07
109.8
0.43
VQGAN-LC (Zhu et al., 2024a)
99.9
4.65
0.10
12+ LRU
3.24
115.6
0.35
MQGAN (Ours)
100.0
4.25
0.08
CIERRIO
SQVAE (Takida et al., 2022)
3.36
769.3
0.39
Gumhel-VQVAA (Karpathy, 2021)
6.16
20.3
0.57
Affine (Huh al., 2023)
5.15
69.5
0.51
Table 4. Reconstruction performance on ImageNet-1k with a reso-
OPT (Huh et al., 2023)
4.73
15.5
0.48
lution of 128 128.
Affine OPT (Huh al., 2023)
4.00
79.3
0.43
MQVAE (ouns)
3.05
783.4
0.29
Method
Utilization (%) rFID LPIPS
VQVAE(yan den Oord al., 2017)
10.02
16.2
0.27
VQGAN (Esser et al., 2021)
1.4
3.74
0.17
12 (Yu al., 2022)
6.49
188.7
0.18
VQGAN-EMA (Razavi et al., 2019)
4.5
3.23
0.15
LRU (Lañcucki al., 2020)
4.77
676.4
0.16
VQGAN-FC (Yu al., 2022)
100.0
2.63
0.13
LRULRU
4.93
861.7
0.15
FSQ (Mentzer et al., 2024)
100.0
2.80
0.13
CELEBA
SQVAE (Takida et al., 2022)
9.17
769.1
LFQ (Yu ct al., 2024)
100.0
2.88
0.13
Gumbel-VQVAE (Karpathy, 2021)
7.34
96.7
VQGAN-LC (Zhu et al., 2024a)
100.0
2.40
0.13
Affine(Huh et al., 2023)
7.47
112.6
0.22
SimVQ (Zhu al., 2024b)
100.0
2.24
0.13
OPT(Huh et al., 2023)
7.78
30.5
0.23
Affine OPT (Huh et al., 2023)
6.60
186.6
0.18
MQGAN (Ours)
100.0
2.13
0.12
MQVAE (Ours)
3.10
877.8
0.14
Table Image generation on FFHQ
contrast, our method improves the FID to 70.5, achieving the best performance among all baselines.
Method
Utilization (%)
FID
VQGAN-FC* (Yu et al., 2022)
10.9
3.23
VQGAN-EMAF (Razavi et al., 2019)
68.2
4.87
VQGAN-LC (Zhu et al., 2024a)
99.9
3.05
Table 2. Image generation on CelebA using MaskGIT.
MQGAN (Ours)
99.9
2.87
Method
FID
MaskGIT (Chang et al., 2022)
90.4
map of the same dimensions (256 tokens), which the U- Net-based decoder uses for image reconstruction. For the hyper-net configuration, an MLP is used. It first lifts the 32-dimensional learned embedding 256 dimensions, then projects it back to 32 dimensions form the codebook entries for quantization, using Tanh as the activation func- tion. The number of codes is set to 65, 536 for ImageNet and 16, 384 for FFHQ, consistent with comparison meth- ods. For initialization, codebook embeddings are initialized using simple Gaussian distribution, avoiding the need for pretrained model in Zhu al. (2024a). The hyper-net transformation uses PoToh'''s default initialization.
l2 (Yu et al., 2022)
81.5
LRU (Lancucki et al., 2020)
79.7
Affine + OPT (Huh et al., 2023)
74.8
MQVAE (Ours)
70.5
5.2. Evaluation with VQGAN
Setup In this subsection, we evaluate the scalability of method using the VQGAN (Esser et al., 2021) archi- tecture, with all compared methods utilizing its original encoder and decoder. These backbone architecture choices directly follow Esser et al. (2021) to ensure comparable results. Experiments were conducted on the ImageNet- 1 (Deng et al., 2009) and FFHQ (Karras et al., 2019) datasets. For ImageNet- input images were processed at 128 x 128 pixels. This involved resizing the image's smaller dimen- sion to 128 while maintaining aspect ratio, followed by 128 x 128 random crop and a 50% probability horizontal flip. FFHQ images were processed at 256 x 256 pixels di- rectly. The U-Net-based encoder (Ronneberger et al., 2015) downsamples input images to a 16 x 16 feature map for each case. The quantizer then converts this into a token
We train for 20 epochs on ImageNet- 1K and 800 epochs on FFHQ, employing early stopping if performance saturates. The Adam optimizer (Kingma & Ba, 2015) with optimizer momentum (0.5, 0.9) used with an initial learning rate of 1e-4. The learning rate undergoes a linear warmup for the first 10% of epochs, followed by half-cycle cosine decay schedule. Evaluation involves several metries. Reconstruc- tion FID is measured using the GAN-trained autoencoder with validation images passed through the quantized autoen- coder; this reflects the FID a Stage transformer could achieve with perfect data modeling. Codebook usage is determined by the fraction of codewords utilized at least

---

Meta-Quantization
once when encoding the validation set, following (Mentzer al., 2024) and (Zhu et al., 2024b). Additionally, genera- tion FID is reported for the second stage involving trained transformer, obtained by decoding representations sampled (potentially class-condiionnally) with the transformer.
Table 7. Ablation study hyper-net.
Hyper-net type
Perplexity
MSE
PSNR
SSIM
Identity
9.3
5.54
25.7
70,5
Linear-1
334.28
3.50
28.4
78.6
Linear-2
52.8
4.25
27.2
75.2
Our method compared against several baselines: standard VQGAN (Esser al., 2021). VQGAN-FC (Yu et al., 2022), VQGAN-EMA (Razavi et al., 2019), VQGAN-LC (Zhu et al., 2024a), Sim VQ (Zhu et al., 2024b), and implicit codebook methods FSQ (Mentzer et al., 2024) and LFQ (Yu et al., 2024).
MLP
877.8
3.10
29.4
82.1
MSE (task loss) on CelebA across all studies.
Effectiveness of Bi-level Optimization In Equation (14), that & controls the influence of task-specific objec- tives. We set the value of & to zero to investigate the effec- tiveness of indirect gradient, denoted by "w/ BLO (5 = 0)". We also provide the results of "w/o BLO"", which denotes using hyper-net reparameterization without bi-level opti- mization in MQ.
the image reconstruction task, we evaluate performance using rFID, LPIPS, PSNR, and SSIM metries on the validation sets of ImageNet and FFHQ. Table 3 and Table present the reconstruction performance for FFHQ and ImageNet, respectively. For FFHQ, results cited from Zhu et al. (2024a) are denoted by t. For Im- ageNet, all results of baselines are cited from Zhu et al. (2024b).
Image Reconstruction
From Table 6, we draw the observation that using a meta- learning-inspine objective significantly improves image modeling ability. This improvement validates the core idea of our approach to directly incorporate task loss into the codebook training via the meta-objective.
We make the key observation that our method consistently achieves the best performance with a codebook utilization rate of over 99% both datasets. When compared with implicit codebooks such as FSQ (Mentzer et al., 2024) and LFQ (Yu et al., 2024) in Table 4, MQGAN also achieves superior performance. This can be attributed to the reduction of expressiveness and model capacity in FSQ and LFQ, while our method follows a learnable, explicit codebook approach, leading to better performance when the codebook utilization rate is high.
Type of Hyper-Net We evaluate the effectiveness of the hyper-net by varying its specific parameterization. We com- pare the following types:
1. Identity: A simple baseline without hyper-net.
2. Linear-1: A linear layer that projects embedding E E K to C E kxd i.e., each code is projected
Image Generation We follow VQGAN, using causal Transformer decoder (Radford et al., 2019) with 24 layers, 16 heads per attention layer, and a latent dimension of 1024. For FFHQ, the FID score is determined using 50K gener- ated images in comparison with the combined training and validation sets of FFHQ. Table 5 displays the unconditional generation results the FFHQ dataset. Compared with baselines, the generative model GPT demonstrates improved performance and codebook utilization with the integration of our MQGAN.
to different dimensionality.
3. Linear-2: A linear layer that projects embedding E E IR xd to C E IR kxd i.e. each code is learned linear combination of the basis E.
4. MLP: An MLP with one hidden layer which first
projects embedding E kk hidden neu- rons of ig xd" and then to C e kxd constituting universal approximator to almost any continuous func- tions.
5.3. Ablation Studies
For all ablative methods, we set d = d' = 32 and K K" = 1024. Table 7 shows that employing a more complex hyper-net consistently leads to better performance. This is attributed to the enhanced meta-learner, resulting in better generative modeling performance. Therefore, we employ an MLP reparameterization in all our experiments.
Table 6. Ablation study on bi-level optimization (BLO)..
Perplexity
MSE
PSNR
SSIM
w/o BLO
836.4
3,65
28.3
77.8
BLO (g-
851.4
3.61
28.8
80.3
w/ BLO
877.8
3.10
29.4
82.1
5.4. Visualization of Embedding and Code Distributions
Unless otherwise specified, we evaluate reconstruction per- formance terms of PSNR and SSIM (image quality), and
Figure shows the distribution of the extracted features and codes for two models: VQGAN and our MQGAN.
8

---

Meta- Meta-Quantization
The visualization is created using t-SNE (Van der Maaten & Hinton, 2008). The codebook entries are highlighted in blue, while the extracted features are shown in red. A larger amount of overlaying denotes higher utilization rate. Results show that MQGAN tends to have a larger overlap between two distributions than naive VQGAN, and is able to enhance code utilization.
(a) Samples from dataset
(b) Samples generated
Figure Qualitative evaluation of reconstruction
(a) Code distribution VQ (b) Code distribution of MQ
of time and memory cost. When evaluated on the CelebA dataset with batch size of 128, the increase in memory usage is marginal and acceptable in practice. For time com- parison, set VQVAE as the baseline, which needs around 3.6 hours to finish the training of 50k steps (and does not improve after that). In Table 8, we observe 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 able outperform baselines with extended training time.
Figure 3. Comparison between code distribution. Notably, MQ has a larger overlap area between code and extracted embedding, resulting a higher code utilization rate.
5.5. Visualization between Direct and Indirect Gradients
The hyper- eperr-gradientee consists of two parts- the direct and indirect gradients- and it is worth seeing whether both of them play an important role. We visualize the norm of direct and indirect gradients of the hyper-gradient during a single run. A similar pattern can be observed in all the experiments. As shown in Figure 4, both of the norms are in a magnitude of 1e-2. Additionally, the cosine similarity remains close to zero during training, therefore, neither of them should be excluded from training.
Table 8. Comparison between various methods on time and mem- ory cost.
Wall time to reach VQVAE (h)
Total wall
Method
Memory (GB)
time (h)
VQVAE MQVAE
7.29
3.6
3.6
7.35
2.4
12.2
600
69
688
6. Conclusion
Figure 4. Visualization between direct and indirect gradients. (a) Comparison in terms of gradient norm. (b) Cosine similarity between the two gradients.
We propose neta-quantizzttin (MQ), a novel bi-level optimization-based vector-quantization framework inspired by meta-learning. Building upon prior VQ mechanisms, method trains the quantizer and the autoencoder at different levels within a bi-level optimization problem. Furthermore, we reparameterize the codebook using a hyper-net, which generates the codebook from a learnable embedding and updates it holistically. In essence, our approach replaces the VQ objective of the codebook with a mela-objective, which aims to optimize the learning potential of the autoencoder. This facilitates gradient flow from the task loss to the quantizer, thereby improving overall performance. Empiri- cal studies across various experimental settings demonstrate that MQ outperforms the prior VQ method and its variants, underscoring its effectiveness.
5.6. Qualitative Evaluation
In Figure 5, we present the reconstruction results at reso- lution of 256 x 256 for our MQGAN. The introduction of diversified gradient paths and a large-scale codebook facili- tates the generation of images, as revealed by the qualitative results.
5.7. Training Cost Comparison
We conducted additional experiments to show the efficiency of MQ training. We We compare VQVAE and MQVAE in terms

---

Meta-Quantization
Impact Statement
Chang, H., Zhang, H., Barber, J., Maschinot, A., Lezama, J., Jiang, I.., Yang, M., Murphy, K. P., Freeman, W. T. Rubinstein, M., Li, Y., and Krishnan, D. Muse: Text- to-image generation via masked generative transformers. In Krause, A., Brunskill, E.,, Cho, K., Engelhardt, B., Sabato, S., and Scarlett, (eds.), International Confer- ence on Machine Leurning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings PMLR, of Machine Learning Research, pp. mlr .press/ 2023. URL https: hwwpt::tpt///ptt v202 / cwwnggaab.b...oo html.
This paper presents work whose goal is to advance the field of Machine Learning. Specifically, our work aims to im- prove existing discrete representation learning methods with potential applications for wide range of tasks, such image generation. There are many potential societal con- sequences of our work, none of which we feel must be specifically highlighted here.
Acknowledgment
This work was supported in part by the National Natural Science Foundation of China (Grant No. 62276121), the Tian funds for Mathematics of the National Science Foundation of China (Grant No. 12326604), and Shen- zhen International Research Cooperation Project (Grant No.
Choe, S. K., Neiswanger, W., Xie, P., and Xing, E. P. Betty: An automatic differentiation library for multilevel opti- mization. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5. 2023. OpenReview.net, 2023. URL https: / /poeerrvv ew wwt.:ettt/fffte BarpiidddtVV___M3MS
GJH2220220913142611022
Chung, Y., Tang, H., and Glass, J. R. Vector-quantized autoregressive predictive coding. In Meng, H., Xu, B., and Zheng, T. F. (eds.), Interspeech 2020, 21st An- nual Conference of the International Speech Commu- nication Association, Virtual Event, Shanghai, China, 25-29 October 2020, pp. 3760-3764. ISCA, 2020, doi: 10.2 10.21493/hnterspecchh.0000 1228. URL htttp::/// /doi. org /tt.:....111.. http://nneesseeeeeee o.g//pcc...////tteeeeeccc...
References
Arora, S., Cohen, N., and Hazan, E. On the optimization of deep networks: Implicit acceleration by overparame- terization. In Dy, J. G. and Krause, A. (eds.), Proceed- ings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässon Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 244-253. PMLR, 2018. URL http: //proceedings mlr. p.r.ss/es/// v80/arora1 8a html.
Deng. J., Dong. W., Socher, R., Li, L., Li, K., and Li, F. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami, Florida, USA, pp. 248-255, IEEE Com-
Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for con- ditional computation. arXiv preprint arXiv: 130..3330 2013.
puter Society, 2009. doi: 10.1109/C 2009. 55066448. 2009.
URL https: /tpp::/ddo... http://oi..ooo oroii.org//00 100 5206848.
Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. Unsupervised learning of visual features by contrasting cluster assignments. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual. 2020. URL hitps:/ /proceedings. neurips /a/ pttpp/papr/222//as//ee
Dhariwal, P., Jun, H., Payne, C., Kim, J. W., Radford, A.,
and Sutskever, I. Jukebox: A generative model for music. ArXiv preprint, abs/2005.00341, 2020. URL https:
/ //arxiv. http://wwwg.abb//0......0 00341.
Esser, P., Rombach, R., and Ommer, B. Taming trans- formers for high-resolution image synthesis. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, pp. 12873-12883. Computer Vision Foundation IEEE, 2021. doi: 10.1109/CVPP4443322200...22222
html.
URL https:://ppeacceest tww.:cceeeeseecccccccoocoooo.eo. Cottent//CpRR//2/hhan///sss Taming______i__________________i
Chang, H., Zhang, H., Jiang, L., Liu, C., and Freeman, W. T. Maskgit: Masked generative image transformer. In IEEF/CVF Conference on Computer Vision and Puttern Recognition, CVPR 2022, New Orleans, LA, USA, June IEEE, 2022. doi: 10.
Transfoormerr__foo__gih----e-e
Imsge_Syntheeis__cvR_22111111
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th Inter- national Conference on Machine Learning, 2017,
18-24, 2022, pp. 0BB9CCVPRR22688.22220100000 URL https: //doi. org/10.1109 1t1.10..1110//////22221 2022. 001003.
10

---

Meta- Meta-Quantization
Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1126- 1135. PMLR, 2017. URL http:/ proceedings.. mlr. m.r..pres/vee wwtpss//vo///inn////
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In ICLR (Poster), 2015. URL http:// Dexiv.or//abie 1412..1 6980.
html.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Franceschi, I.., Donini, M., Frasconi, P., and Pontil, M. Forward and reverse gradient-based hyperparame- ter optimization. In Precup. D. and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Aus- tralia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1165-1173. PMLR, 2017. URL ht.t:: http:///pooeo/igeeeeeee ngs mlm prrss/ v70 Eww.ceschi17..
Lańcucki, A., Chorowski, J., Sanchez, G., Marxer, R., Chen, N., Dolfing, H. J., Khurana, S., Alumãe, T., and Laurent, A. Robust training of vector quantized bottleneck mod- els. In 2020 International Joint Conference on Neural Networks (ICCN), pp- 1-7. IEEE, 2020.
Liu, H., Simonyan, K., and Yang, Y. DARTS: differen- tiable architecture search. In 7th International Con- ference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https:/ /oppnreveeeeeeeeeeeeeeeeeeeeeeee net/ forum? id=S1eYHoC5FX.
Gray, R. Vector quantization. IEEE Assp Magazine, 1()): 4-29, 1984.
Huh, M., Cheung, B., Agrawal, P., and Isola, P. Straight- ening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), International Confer- ence on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 14096-14113. PMLR, 2023. URL https::tp.........p. http///ccediiiccccocc igg mlr. press// v202/huh22 html.
Liu, Z., Luo, P., Wang. X., and Tang, X. Deep learning face attributes in the wild. In 2015 IEEE International Conference Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pp. 3730-3738. IEEE Com- puter Society, 2015. doi: 10.1100//CC222015..442 URL https: //wpi. org/ 10. ICCVV hwt.//11............. 2015..22 425.
Lorraine, J., Vicol, P., and Duvenaud, D. Optimizing millions of hyperparameters by implicit differentiation. In Chiappa, S. and Calandra, R. (eds.), The 23rd In- ternational Conference on Artificial Intelligence and Statistics, AISTATS 2020, 26-28 August 2020, Online [Palermo, Sicily, Italy], volume 108 of Proceedings of Machine Learning Research, pp. 1540-1552. PMLR, 2020. URL. /ttopee/proceeiiii ngs. wwigggmmm........... ml r.press///// v108 / www.oooorrrnaaaaoo html.
Kaiser, L., Bengio,
S., Roy, A., Vaswani, A., Par- mar, N., Uszkoreit, J., and Shazeer, N. Fast decod- ing in sequence models using discrete latent variables. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference Machine Learn- ing, ICML 2018, Stockholmsmässsan, Stockholm, Swe- den, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 2395-2404. PMLR, 2018. URL http: http::t/prccceddnnieee /ptp///eeocecccssse/eee//pcs..se mlr .rres// v80/kaiserl8a. html.
Mentzer, F., Minnen, D., Agustsson, E., and Tschannen, M. Finite scalar quantization: VQ-VAE made simple. In The Twelfth International Conference on Learning Repre- sentations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https:://peeeeeiier net
Karpathy, A.
deep-vector-quaanizatio.. https:/ //tt:://gttttbbbt cow.karpathy/ quantizatannnnonnnnn deep-vector-quie deep-rretoor-uuattiaattooo 2021.
Pedregosa, F. Hyperparameter optimization with approx- imate gradient. In Balcan, M. and Weinberger, K. Q. (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY. USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 737-746. JMLR.org. 2016. URL htt.:httpp///cccec ing...nniig.... wng.ing..mm.........m...mmmm.... ml .t.pesss..sss//////// v48/ r48/pedregosal6. html.
Karras, T., Laine, S., and Aila, T. A bayed generator architecture for generative adversarial networks. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, pp. 4401-4410. Computer Vision Foun- dation / IEEE, 2019. doi: 10.1109/CVRR.2019.0022 URL http:://operaaccessstheeee com/ content_ Dentent_cvvR_2000 cVPR_2009//tt///aarra_
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 18))9, 2019.
Ser__Gnnerative_Avversaate Metad CVPR_201 9__appr..ptl...

---

Meta- Lita-Quantization
Rajeswaran, A., Finn, C., Kakade. S. M., and Levine, S. Meta-learning with implicit gradients. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.). Advances in Neural Infor- mation Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS
2022. wwt.ttts::/prooeeeiiiiiii mlr. press/ v1 v162/takida22e html.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 6306-6315, 2017. URL htttts::/ppooeeeeeee sttrips.cc//appee 201 2 9ce2e 96003999922222229299292999999999 fbc- Abstt
2019, December 8- 14, 2019, Vancouver, BC, Canada, pp.
113-124, 2019. URL hitps: https://tpprseeeeee /ptpceeeeeeens... neurips cc/ httcc/papr//2011////// 0 7720030aa22222222222733332322222 html.
-bsttract..
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C.,, Radford, A., Chen, M., and Sutskever, I. Zero-shot text-to-image generation. In Meila, M. and Zhang. T. (eds.), Proceedings of the 38th International Confer- ence on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Ma- chine Learning Research, pp. 8821-8831, PMLR, 2021. htthttp: http: proceedings mlr www.ppresssss/// ramesh21a. html.
html.
Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11). 2008.
Vinyals, O., Toshev, A., Bengio, S., and Erhan, D. Show and tell: A neural image caption generator. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2015, Boston, MA, USA, June 7-12, 2015, pp. 3156-3164. IEEE Computer Society, 2015. doi: 10.1109/ CVPR CVPR.2015.7298995. URL https: /doi. org/10. 1109/CVPR.200 7298935.
Razavi, A., van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with VQ-VAE-2. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.). Advances in Neural In- formation Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8- 14, 2019, Vancouver, BC, Canada, pp- 14837-14847, 2019. URL wwtps::tts/pppp.t httpppss/ppccceedggg.....
Yu, J., Li, X., Koh, J. Y., Zhang, H.,, Pang. R., Qin, J., Ku, A., Xu, Y., Baldridge, J., and Wu, Y. Vector-quantized image modeling with improved VQGAN. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenRe- view.net, 2022. URL https: / openreview nnt./ forum? id d- Nfxxxx122..
3Fbe2aal118dlbcccddcccc113333333 Abstttttaattta--aaaattaaaa html.
Yu, L., Lezama, J., Gundavarapu, N. B., Versari, L.,, Sohn, K., Minnen, D., Cheng, Y., Gupta, A., Gu, X., Haupt- mann, A. G., Gong, B., Yang, M., Essa, I., Ross, D. A., and Jiang, L. Language model beats diffusion - tok- enizer is key to visual generation. In The Twelfth Inter- national Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net. 2024. URL https: http:://ppppeviveeee net forum? id A2=ggrAANCE4c 4g.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted intervention-MIICAI 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part 18, pp. 234-241. Springer, 2015.
Roy, A., Vaswani, A., Neelakantan, A., and Parmar, N. Theory and experiments on vector quantized autoen- coders. ArXiv preprint, abs/1805. 1063,, 2018. URL https:/ //ar:///aavv.////a org/ abs 1805188110 11063.
Zeghidour, N., Luebs, A., Omran, A., Skoglund, J., and Tagliasacchi, M. Soundstream: An end-to-end neural audio codec. JEEE/ACM Transactions on Audio, Speech, and Language Processing, 30:495-500 2021.
Takida, Y. Shibuya, T., Liao, W., Lai, C., Ohmura, J., Ue- saka, T., Murata, N., Takahashi, S., Kumakura, T., and Mitsufuji, Y. SQ-VAE: variational bayes on discrete rep- resentation with self-annealed stochastic quantization. In Chaudhuri, K., Jegelka, S., Song, L.,, Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Bal- timore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 20987-21012, PMLR,
Zhang, M., Su, S. W., Pan, S., Chang, X., Abbasnejad, M. E., and Haffari, R. idarts: Differentiable architecture search with stochastic implicit gradients. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 12557-12566. PMLR,
12

---

Meta-Quantization
2021. URL http: /ttttp:/pprooe http://pooooddieee ngs www. press/ v139/zhang211 html.
ngs
Zhang, R., Isola, P.. Efros, A. A., Shechtman, E., and Wang. O. The unreasonable effectiveness of deep features perceptual metric. In 2018 IEEE Conference Computer Vision and Putlern Recog- nition, CVPR 2018, Salt Lake City, UT. USA, June 18-22, 2018, pp. 586-595. IEEE Computer Society, 2018. doi: 10.1109/CVPR.2018..00
URL http:
2011///thl//haanng..nneeeeereeee html. Bffectiveness__cPR_..22
Zhu, L., Wei, F, Lu, Y., and Chen, D. Scaling the codebook size of VQ-GAN to 100, 000 with a uti- lization rate of 99%.
Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomezak, M., and Zhang, C. (eds.), Advances in Neural Informa- tion Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December
10 15, 2024, 2024a. URL. http: /ttt::/ppppers
html.
Zhu, Y., Li, B., Xin, Y., and Xu, I.. Addressing representa- tion collapse in vector quantized models with one linear layer. arXiv preprint arXiv:2411.02053 2024b.
13


Reviews:
[{"review_id": "6b1TiKhjz7", "reviewer": "Reviewer_e9kM", "summary": "This paper introduces Meta-Quantization (MQ), by using a hyper-net and bi-level optimization to alternatively train the codebook with the autoencoder in Vector Quantization Networks (VQN). Experiments show MQ has better codebook ultilization, image reconstrcution and generation performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your constructive feedback very much. We provide our response to your review as follows.\n\n> Why not use a simpler method such as clustering the embeddings or EMA update\n\nOne 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.\n\nFor 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.\n\n> What is the training stability and efficiency of the optimization compared to other methods?\n>\n\nEmpirically, 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.\n\n> It is not clear if the training is efficient.\n\nWe 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.\n| Method           | Memory (GB) | Wall time to reach baseline (h) | Total wall time (h) |\n| ---------------- | ----------- | ------------------------------- | ------------------- |\n| VQVAE (baseline) | 7.29        | 3.6                             | 3.6                 |\n| MQVAE            | 7.35        | 2.4                             | 12.2                |\n\n> 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?\n\nIn 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.\n\n> In lines 69-74, what is the \"specific task\"? What is \"task-aware\"? Is it the VQVAE loss?\n\nThe 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.\nTask 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.\n\n[1] Pedregosa, Fabian. \"Hyperparameter optimization with approximate gradient.\" *International conference on machine learning*. PMLR, 2016.\n\n[2] Rajeswaran, Aravind, et al. \"Meta-learning with implicit gradients.\" *Advances in neural information processing systems* 32 (2019).", "claims_and_evidence": "Yes. VQ has optimization issues and this paper tackles this.", "methods_and_evaluation_criteria": "- The optimization seems to be very complicated. Why not use simpler method such as clustering the embeddings or EMA update? What are the advantages against existing methods for better codebook ultilization?\n- What is the training stability and efficiency of the optimization compared to other methods?\n- 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 generate a single code entry? How does it produce the whole codebook?", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Lack of comparison on the training stability and efficiency of the optimization with other methods.", "relation_to_broader_scientific_literature": "This paper is related. It proposes a meta-learning approach for codebook optimization.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n- The MQ method and the optimization algorithm makes sense.\n- Results are strong.\n\nWeekness:\n- It is not clear if the training is efficient.", "other_comments_or_suggestions": "Please provide more details of the hyper-net.", "questions_for_authors": "In lines 69-74, what is the \"specific task\"? What is \"task-aware\"? Is it the VQVAE loss?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kgYOhCofv9", "reviewer": "Reviewer_bZLB", "summary": "The paper proposes to train VQ-VAE under a meta-learning framework. To be more specific, the paper introduces a hyper-network to replace the embedding-parameterized and trains the model with bi-level optimization. Experiments are conducted on image reconstruction and generation tasks. The proposed MQ-VAE improves over the VQ baseline.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your constructive feedback very much. We provide our response to your review as follows.\n\n> Magnitude Comparison\n\nIn 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.\n\n> Whether [1] can benefit from bi-level optimization\n\nWe 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.\n\n>  Relation to [2]\n\nThis 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.\n\n> Computation Cost and Implementation details\n\nWe 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.\n| Method           | Memory (GB) | Wall time to reach baseline (h) | Total wall time (h) |\n| ---------------- | ----------- | ------------------------------- | ------------------- |\n| VQVAE (baseline) | 7.29        | 3.6                             | 3.6                 |\n| MQVAE            | 7.35        | 2.4                             | 12.2                |\n\n[1] Zhao, et al. \"Image and video tokenization with binary spherical quantization.\" arXiv preprint arXiv:2406.07548 (2024).\n\n[2] Fifty, et al. \"Restructuring Vector Quantization with the Rotation Trick.\" arXiv preprint arXiv:2410.06424 (2024).", "claims_and_evidence": "The Gradient Analysis section is quite interesting. However, empirical evidence is lacking. I can only find an ablation in Table 6, which disables the indirect gradient by zeroing out $\\xi$.  However, the improvement is not much. I am curious about the true magnitude of the indirect gradient compared to the direct one. A plot wrt the training iterations may better support this.", "methods_and_evaluation_criteria": "++ Applying the bi-level optimization to training VQ-VAE/VQGAN is reasonable. \n\n++ The evaluation criteria include image reconstruction and image generation.", "theoretical_claims": "The formulation and derivations look correct. There are no particular proofs for theoretical claims.", "experimental_designs_or_analyses": "The experimental designs are sound and comprehensive. Specifically, the baseline methods include VQVAE and VQ-GAN.", "relation_to_broader_scientific_literature": "The difficulty of training VQ-VAE lies in the non-differentiability of the *argmin* operator in the bottleneck. This paper proposes a bi-level optimization approach to mitigate this issue. A recent paper [1] tackles this by bounding the quantization error with a spherical vector quantization-like bottleneck. The vector quantization is implicitly modeled by a linear projection, which resembles the hypernetwork design. However, the model can be trained without bi-level optimization, achieving very similar results on ImageNet 128x128 in Table 4. It looks like these two papers try to tackle the same problem from different perspectives, and it would be good to see if [1] can also benefit from the proposed bi-level optimization technique. [2] also proposes to propagate gradients more smoothly via a rotation. \n\n[1] Zhao, et al. \"Image and video tokenization with binary spherical quantization.\" arXiv preprint arXiv:2406.07548 (2024).\n\n[2] Fifty, et al. \"Restructuring Vector Quantization with the Rotation Trick.\" arXiv preprint arXiv:2410.06424 (2024).", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "==== Strength ====\n\n1. The paper is well-written and easy to understand. The overview figure nicely illustrates the gradient flow.\n\n2. Evaluations are comprehensive. \n\n==== Weaknesses ====\n\nSome weaknesses have been covered above. There is one more weakness that I would love to point out.\n\n1. Bi-level optimization introduces the computational cost overhead. The paper didn't show any result on the computational cost. There are not many implementation details about the paper. Nevertheless, an interesting experiment is to see the comparison between MQ-VAE and regular single-level optimized VQVAE with the same training budget. In other words, the vanilla VQVAE would enjoy a longer training schedule, which will generally lead to constant improvement from my experience.", "other_comments_or_suggestions": "No.", "questions_for_authors": "1. Training cost compared to vanilla VQ-VAE\n\n2. A more fair comparison with VQ-VAE under the same training budget. In other words, the authors should allow the baseline to train for longer iterations.\n\n3. More illustrations on the gradient analysis (see the claims and evidence section).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VD6JzTn6fD", "reviewer": "Reviewer_yGdn", "summary": "This paper proposes a novel vector quantization training framework Meta-Quantization inspired by meta-learning, which decouples the optimization of codebook and autoencoder into two stages, enabling dynamic codebook generation and task-specific training. The proposed method outperforms existing vector quantization approaches on image construction and generation tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your constructive feedback very much. We provide our response to your review as follows.\n\n> Is the bi-level optimization approach adopted in this paper necessary?\n\nYes, 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.\n\n> The description of convergence is inconsistent.\n\nWe 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.\n\nThus, a consistent pseudo-algorithm is to replace\n\n\"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\n\n\"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.\n\n> Justification of advantages over VQGAN-LC.\n\nVQGAN-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.", "claims_and_evidence": "This paper achieves direct backpropagation on the codebook instead of using STE, enabling codebook training task-specific.\nThe experiment in both image generation and reconstruction tasks shows that Meta-Quantization consistently outperforms multiple baselines and ablation methods, validating its superiority.\nThe experiment also shows that in codebook utilization,  the methods achieve the best performance.", "methods_and_evaluation_criteria": "**Methods**\n\n- The paper innovatively introduces a hypernet to take the place of the embedding-parameterized codebook. This substitution circumvents the need for direct optimization of the codebook itself. Moreover, after the first-stage training, only the generated codebook needs to be stored and can be directly applied in the subsequent training process.\n\n- To tackle the optimization problems of the hypernet and the encoder-decoder, the paper employs a two-stage optimization framework. In this framework, the parameters of the two structures are optimized in a hierarchical manner, and an efficient gradient-based optimization algorithm with finite difference approximation is utilized. \n\n  **However**, is the bi-level optimization approach adopted in this paper necessary? Since the Hyper-Network is also composed of linear layers or MLPs, if we perform backpropagation on the Hyper-network parameters $\\psi$ together with the encoder-decoder parameters $\\phi$ and $\\theta$ using a single-step training method, would the performance differ significantly from bi-level optimization? Is there any practical or theoretical reason why bi-level optimization might yield superior results?\n\n**Evaluation Criteria**\n\nEvaluation with VQVAE\n\n- In image reconstruction task, in addition to MSE and LPIPS, this paper uses the perplexity of the model as an evaluation metric to measure the similarity of the codebook. A higher perplexity value indicates a more uniform assignment of codes. The evaluation is carried out on the CIFAR10 and CelebA datasets.\n- In image generation task, the FID evaluation metric is adopted. MaskGIT is applied to the CelebA dataset, and the results are extended to the image generation task, enabling the direct utilization of the codebook trained in the first stage.\n\nEvaluation with VQGAN\n\n- The model is trained on the ImageNet-1K and FFHQ datasets.\n- In image reconstruction task, the evaluation metrics included rFID, LPIPS, PSNR, and SSIM. The assessment is conducted on the validation sets of ImageNet and FFHQ.\n- In image generation task, the FID evaluation metric is used, and the evaluation is performed on the FFHQ dataset.", "theoretical_claims": "This paper references the theoretical analysis of gradient analysis for DART's two-level optimization and generalizes the optimization algorithm, containing no errors in proofs.", "experimental_designs_or_analyses": "The experiments in this paper are comprehensive, conducting evaluations on state-of-the-art models VQVAE and VQGAN. The proposed method is compared against other SOTA models, incorporating codebook utilization comparisons alongside original evaluation metrics. Additionally, implicit codebook methods such as FSQ and LFQ are also analyzed. Experimental results demonstrate that the proposed approach not only efficiently utilizes the codebook but also achieves superior performance across all tasks.\n\nIn ablation studies, experiments are conducted to validate the effectiveness of bi-level optimization and compare different Hyper-Net architectures. Results showed that the MLP-based Hyper-Net outperformed other variants, confirming that more complex Hyper-Net designs consistently yield better performance.", "relation_to_broader_scientific_literature": "This paper adopts the DART method in the field of vector quantization, which has implications for the image tokenization domain but has no impact on other areas.", "essential_references_not_discussed": "Essential references have beed discussed.", "other_strengths_and_weaknesses": "The description of convergence is inconsistent. While Figure 2 states that $\\phi$ and $\\theta$ are trained to convergence before training $\\psi$, Algorithm 1 shows that they are updated together. This discrepancy creates ambiguity regarding the actual optimization procedure implemented in the paper.", "other_comments_or_suggestions": "There are no other comments or suggestions.", "questions_for_authors": "In the introduction, it is mentioned that the codebook utilization in previous methods is low. However, in the experiments (Table 3, 4, 5), the codebook utilization of VQGAN-LC is also quite high. Please provide a justification for this conclusion and explain the advantages of this method over VQGAN-LC in terms of codebook utilization.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tkIXnnnD74", "reviewer": "Reviewer_VZz4", "summary": "This paper proposes an improvement to the STE method for training VQ networks. While the backpropagated gradient bypasses the codebook in the STE framework, this paper proposes Meta Quantization (MQ), which adopts a bi-level optimization strategy and learn quantization with a hyper-net in a meta-learning fashion. This enables the task loss to reach the codebook through multiple routes. It can also discard the hyper-net and retain only codebook weights, without affecting downstream tasks. Empirical evidence suggests better reconstruction and generation quality across various datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your constructive feedback very much. We provide our response to your review as follows.\n\n> Computational Cost and Memory Overheads\n\n\nWe 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.\n| Method           | Memory (GB) | Wall time to reach baseline (h) | Total wall time (h) |\n| ---------------- | ----------- | ------------------------------- | ------------------- |\n| VQVAE (baseline) | 7.29        | 3.6                             | 3.6                 |\n| MQVAE            | 7.35        | 2.4                             | 12.2                |\n\n\n> Difficulty of converging using the meta-learning strategy\n\nEmpirically, 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.\n\n> Demonstrate how codes are generated and selected and which part is optimized; The codebook distribution; Implementation Details\n\nPlease 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.\n\n> Compare with VQVAE+Affine+OPT+replace+l_2, and CVQ-VAE\n\nAccording 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.\n\n| Method                         | MSE ($\\times 10^{-3}$) | LPIPS ($\\times 10^{-1}$) |\n| ------------------------------ | ---------------------- | ------------------------ |\n| CVQ-VAE                        | 2.87                   | 3.73                     |\n| VQVAE+Affine+OPT+replace       | 1.81                   | 2.56                     |\n| VQVAE+Affine+OPT+replace+$l_2$ | 1.74                   | 2.27                     |\n| Ours+$l_2$                     | 1.64                   | 2.18                     |\n\n> 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?\n\nOur 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.\n\n| Method | MSE  | LPIPS |\n| ------ | ---- | ----- |\n| FD     | 3.10 | 0.14  |\n| CG     | 3.24 | 0.14  |\n| NMN    | 2.98 | 0.14  |\n\n> 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?\n\nBased 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.\n\n[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.\n\n[2] Zheng, Chuanxia, and Andrea Vedaldi. \"Online clustered codebook.\" *Proceedings of the IEEE/CVF International Conference on Computer Vision*. 2023.\n\n[3] Rajeswaran, Aravind, et al. \"Meta-learning with implicit gradients.\" *Advances in neural information processing systems* 32 (2019).\n\n[4] Lorraine, Jonathan, Paul Vicol, and David Duvenaud. \"Optimizing millions of hyperparameters by implicit differentiation.\" *International conference on artificial intelligence and statistics*. PMLR, 2020.\n\n[5] Pedregosa, Fabian. \"Hyperparameter optimization with approximate gradient.\" *International conference on machine learning*. PMLR, 2016.\n\n[6] Rajeswaran, Aravind, et al. \"Meta-learning with implicit gradients.\" *Advances in neural information processing systems* 32 (2019).", "claims_and_evidence": "•\tClaim: The paper argues that introducing MSE loss for codebook optimizing by hyper-net could enhance the training.\no\tQuestion: Although the experiments have shown the effectiveness of the generated codebook by hyper-net on image reconstruction and generation. There is little direct discussion of computational cost or difficulty of converging using the meta-learning strategy. \no\tPotential Improvement: It will be helpful to show the performance between a conventional VQN and MQ trained with the same epochs, or figure out the steps needed until they have converged, respectively.\n\n•\tClaim: Hyper-net reparameterization and meta-learning help avoid code collapse.\no\tQuestion: The authors show near-100% code usage on image reconstruction and generation, however, it is not clear about the process of codebook optimization itself. \no\tPotential Improvement: It will be better to demonstrate how codes are generated, selected, and which part is optimized. As mentioned in Sec. 4.1 (only a few codes are selected, a substantial part of the hyper-net still receives gradients).", "methods_and_evaluation_criteria": "•\tMethods: The authors optimize the codebook with a hyper-net and introduce task loss for explicit codebook improvements. This makes sense for index collapse problem and potential better image reconstruction or generation performance.\n•\tEvaluation Criteria: The paper relies on standard generative modeling metrics such as MSE, LPIPS, SSIM, and FID across established datasets, following [Straightening’ 23] by Huh et al. published in PMLR 2023. \n•\tPotential Weakness: Since the method focus on explicit codebook improvements, it will be better to illustrate the codebook distribution, using tools such as tSNE, to directly validate the effectiveness of the proposed method.", "theoretical_claims": "The paper uses the logic of meta-learning, referencing the established notion that bi-level optimization can optimize both model and “hyperparameters” jointly.", "experimental_designs_or_analyses": "Design Strength: The authors compare their method to a variety of widely used VQN baselines on multiple tasks. They measure codebook usage, reconstruction quality, and generation quality, providing a holistic overview.\n\nPotential Weakness: \n1.\tFor reconstruction evaluation in Tab. 1, it follows [Straightening’ 23] by Huh et al. published in PMLR 2023, but some results are missed. E.g. VQVAE+Affine+OPT+replace+l_2 shows an MSE of 1.74, and LPIPS of 0.227, better than MQVAE proposed in this paper with MSE of 3.05 and LPIPS of 0.29. It’s better to involve these missing results and analyze about the gap.\n2.\tRecent research such as CVQ-VAE [online’ 23] by zheng et al. published in ICCV 2023 should also be considered for more comprehensive comparison.", "relation_to_broader_scientific_literature": "The authors reference VQVAE, VQGAN, and relevant codebook-improvement methods (Gumbel-VQ, VQGAN-LC, FSQ, etc.). They also draw comparisons to meta-learning methods and hyperparameter optimization.", "essential_references_not_discussed": "The paper does cite and discuss relevant references, including meta-learning strategies and VQNs that focus on codebook collapse. However, it is better to expand the discussion of recent VQN researches dealing with the collapse and compare with them, such as CVQ-VAE [online’ 23] by zheng et al. published in ICCV 2023.", "other_strengths_and_weaknesses": "Strengths:\n1.\tThis paper proposes a novel vector quantization network, incorporating meta-learning methods and VQNs. \n2.\tThe approach is adaptable for multiple VQN models without affecting downstream tasks, so it could be straightforward to integrate into new pipelines.\n\nAdditional Weaknesses (in more detail):\n1.\tComputational Overhead: The new approach requires unrolled gradient steps (or finite difference approximations). The paper mentions an approximation but does not provide an in-depth breakdown of the training cost or memory usage on large codebooks.\n2.\tLimited Sensitivity Analysis: The ablation studies are somewhat narrow. The authors partially test turning off bi-level optimization or using various hyper-net types, yet do not systematically explore the effect of different unroll lengths or layer widths of the hyper-net.\n3.\tGeneralization Beyond Vision: The paper focuses on visual tasks. It remains unclear for the effectiveness of other task losses and how they affect the codebook optimization by the hyper-net.", "other_comments_or_suggestions": "•\tImplementation Details: Additional clarifications about hyperparameter tuning and computational cost might help replicate the strong results.\n•\tDownstream Utility: While the paper shows standard metrics, some real-world usage scenarios or domain-specific tasks (e.g., image segmentation, speech coding) might illustrate practical advantages.", "questions_for_authors": "1.\tChoice of Unroll Steps: Have you tried multiple unroll steps (beyond one-step) or different finite-difference approximations? How does that affect training time and results?\n2.\tPotential Memory Overheads: Could you provide more precise measurements of how memory/time usage increases for the partial unrolling or the hyper-grad approximations?\n3.\tDomain Generalization: Do you think the approach can directly transfer to other discrete representation tasks such as speech tokens or molecular modeling, without major changes?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All of reviewers lean toward acceptance post-rebuttal. The AC checked all the materials and concurs that the paper has proposed a reasonable and principled approach to tackle the difficulty of training online vector-quantized tokenizers for vision. While there are limitations (e.g., the overall complexity of the bi-level framework), the work presents a valuable attempt to mitigate the issue of training vector-quantized networks, and therefore should be accepted. Please incorporate necessary changes in the final version, and open-source code/models as promised.
Decision: Accept (poster)