Masked Vector Quantization
Authors: David D Nguyen
David Liebowitz
Surya Nepal
Salil S. Kanhere
Venue: Submitted to ICLR 2023 (2023)
Keywords: generative models
dropout
vector quantization
autoencoder
discrete representations

Abstract: Generative models with discrete latent representations have recently demonstrated an impressive ability to learn complex high-dimensional data distributions. However, their performance relies on a long sequence of tokens per instance and a large number of codebook entries, resulting in long sampling times and considerable computation to fit the categorical posterior. 
To address these issues, we propose the Masked Vector Quantization (MVQ) framework which increases the representational capacity of each code vector by learning mask configurations via a stochastic winner-takes-all training regime called Multiple Hypotheses Dropout  (MH-Dropout). On ImageNet 64$\times$64, reduces FID in existing vector quantization architectures by up to $68\%$ at 2 tokens per instance and $57\%$ at 5 tokens. These improvements widen as codebook entries is reduced and allows for $7\textup{-}45\times$ speed-up in token sampling during inference. As an additional benefit, we find that smaller latent spaces lead to MVQ identifying transferable visual representations where multiple can be smoothly combined.

Under review as a conference paper at ICLR 2023
MASKED VECTOR QUANTIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Generative models with discrete latent representations have recently demonstrated an impressive ability to learn complex high-dimensional data distributions. How- ever, their performance relies on long sequence of tokens per instance and a large number of codebook entries, resulting in long sampling times and considerable computation to fit the categorical posterior. To address these issues, we propose the Masked Vector Quantization (MVQ) framework which increases the repre- sentational capacity of each code vector by learning mask configurations via a stochastic winner-takes-all training regime called Multiple Hypotheses Dropout (MH-Dropout). On ImageNet 64x64, MVQ reduces FID in existing vector quan- tization architectures by up to 68% at tokens per instance and 57% 5 tokens. These improvements widen as codebook entries is reduced and allows for 7-45x speed-up in token sampling during inference. As an additional benefit, we find that smaller latent spaces lead to MVQ identifying transferable visual representations where multiple can be smoothly combined.
INTRODUCTION
In deep generative modelling, the choice of latent representation is an important consideration due to trade-ofTs in sample stability, quality, model size and compatibility with different modalities. Generative models with continuous latent random variables that assume parametric distributions have led the field in sample quality for the past decade (Vahdat & Kautz 2020 Donahue & Simonyan
2019). Despite their impressive performance, these models can be challengigg to stabilise, resulting
in problems such as posterior collapse (Lucas et al. 2019a He et al.
2019 Lucas et al. 2019b).
Interest in discrete representations address these challenges has seen a revival recently with the
development of several discrete autoencoders (Van Den Oord et al. [201]] Razavi et al, (2019] Ramesh
ct al. 2021 Esser et al. 102 Nichol et al. 2021 Rombach et al. (022)) with improved stability
high-dimensional visual and audio domains. This approach maps each instance to discrete
sequence codebook indices (tokens) using vector quantisation (VQ) (Van Den Oord et al. 201 7)
Gumbel-softmax (Jang al. 2016) Concrete distributions (Maddison et al. (0166). In secondary training stage, an autoregressive probabilistic model, such as a PixeICNN (Van Oord et al. 2016) or
Transformer Vaswani al. 2017), learns the categorical posterior, representing the distribution of observable token sequences.
Despite these improvements, growing the representational capacity of discrete autoencoders remains tied to increasing the number of tokens assigned to each instance and the number of total codebook entries. Increasing both hyper-parameters can scale performance to high quality images as seen in Razavi et al. (019) where the authors used 1280 tokens per image. However this resulted in unrealistically long ancestral sampling times and significantly higher computational resource use to fit the categorical posterior. More recent works, such as VQGAN (Fisser et al. 2021) and Stable Diffusion (Rombach al., 2022). introduced a discriminator (Goodfellow al., 2014 and diffusion models (Sohl- Dicksttin et al. 2015) to reduce token length to 256 per image however total sampling times remain rather long at 7258 seconds per image with a single Nvidia Titan X. Whilst demonstrating impressive quality, these times may continue to hinder the use of discrete autoencoders in challenging domains such as video generation and real-time applications particularly where compute resources are constrained.
To address this problem, we explore whether each instance can be compressed into a shorter sequence of tokens using smaller codebook. We achieve this goal by introducing Masked Vector Quantization

---

Under review as a conference paper at ICLR 2023
Vector Quantization
primary code vector
D
Masked Vector Quantization
MH-Dropout
m
81
2
primary code vector
m
1
0
0
random samples
decoder
91
D
y)
transform & upsample
secondary code vector
; e + f(e' O f(e'Om) m2)
-
base concept
factor of variation
Figure MVQ (bottom) compared VQ (top). unique pair of primary and secondary code vectors can encode up to 2 reconstructions. Multiple Hypotheses Dropout trains sampled mask configurations of the secondary code vector to represent different factors of variation (age, hair, skin tone, background, etc.). For example, by comparing D(e) to Z. the masked secondary code vector represents older age and make-up skin tone.
(MVQ), novel variant of VQ that allows the masked configuration of each codebook vector to be individually mapped to separate instances. More precisely, in our framework each instance is encoded with three components: a primary code e, a secondary code e' and mask m,. as shown in Figure
In comparison to standard VQ. each primary code can represent a further 2 instances, where D' is the embedding dimension of the secondary code.
During training, the best primary and secondary code vector for each instance can be found using standard nearest neighbor lookup from the codebook. It is, however, computationally infeasible to search for the best mask across all DD possibilities in latent spaces of dimension typically used in practice. To overcome this, we introduce Multiple Hypotheses Dropout (MH-Dropout) novel variant of dropout (Hinton al. 2012) that incorporates multiple hypotheses training (Guzman- Rivera et al. 2012 Rupprecht al. 2017). During the forward pass,
2D' masks randomly sampled, yielding latent representations. During the backward pass, we use winner-takes-all reconstruction loss where only the best of the representations affects the gradient.
The MVQ framework improves existing VQ architectures, such as VQ-VAE2 (Razavi et al. 2019) and VQGAN (Esser et al. 2021), particularly when the number tokens per instance and codebook size is reduced. Across multiple medium resolution datasets, we observe FID reductions up to 82% tokens per instance, 57% at tokens and 14% at 17 tokens. These improvements consistently grow as the number of codebook entries decrease and also reduce token sampling times by 58-87% on a consumer grade GPU and 80-97% on CPU.
We also find that the dimensions of masked codes learn characteristics that can be smoothly inter- polated, combined and transferred to other primary codes. These features are sometimes clearly interpretable (such as the presence of eyeglasses or smile) and can correspond to individual dimensions as seen in Figure 2
Our paper is outlined as follows: In Section 2, review related work discrete autoencoders, dropout and multiple hypotheses training. In Section 3, we present background on the popular vector quantization framework which work builds upon. In Section 4, introduce Masked Vector
2

---

Under review as a conference paper at ICLR 2023
base concept
color depth
smile
glasses
m
1
0
add factors of variation by unmasking
Figure 2: Each dimension of secondary code vector e' learns interpretable characteristics completely unsupervised. These characteristics can be controlled with mask m as seen above.
Quantization framework and Multiple Hypotheses Dropout. Section we describe the results of extensive experiments on our MVQ framework, benchmarking performance against both VQVAE-2 and VQGAN at multiple compression rates and datasets.
2 RELATED WORK
Our work is closely related three areas of research: discrete autoencoders, dropout and multiple choice learning.
Discrete Latent Models. Autoencoders with discrete latent representations appeared early the work of Hinton & Zemel (993), where the authors proposed learning codebook with stochastic quctt--quanizzttin (VQ) in accordance with the Minimum Description Length (MDL) principle (Rissanen 1987). Recent works have proposed the use of Gumbel-softmal @Jang et al., 2016) or Concrete distribution (Maddison al. 2016) learn codebooks with continuous reparameterization. In the work of Van Den Oord et al. (2017). the authors propose training codebooks with winner- takes-all VQ and an autoregressive model to learn the categorical posterior. Others have since combined codebooks with hierarchical latents Razavi et al. (2019). transformers (Ramesh al.) 2021), generative adversarial networks (GANs) (Esser et al. 2021 and diffusion models (Nichol al.
2021 Rombach et al. 2021 2022) to produce remarkable high quality data reconstructions.
However, contemporary discrete autoencoders come with a trade-off between compression and computational resources to learn the categorical posterior. Larger latent representations reduce reconstruction error, but require more computational resources and are slower to sample from, as outlined in Section Our MVQ framework addresses this by learning more efficient latent representations that uses MH-Dropout, a variant of binary dropout (Hinton et al. 2012).
Dropout. Binary dropout regularizes deep neural networks by randomly dropping hidden units prevent co-adaption of features. This can be seen as computationally tractable approach to bagging during inference by forming ensembles of sub-networks (Srivastava al., 2014p. In contrast to binary dropout, MH-Dropout actively recovers these representations during inference, leading to accurate multiple predictions by the ensemble. This increases compression as the same distribution represented with significantly fewer parameters.
similar approach using variational dropout (Kingma al. 2015) appears in Molchanov al. (2017), where the parameters of large supervised networks are significantly reduced, at minor cost in accuracy. MH-Dropout samples masks during inference similar to Monte-Carlo dropout (Gall 2016) and Concrete dropout (Gal al.] 2017). These, however, primarily developed to quantily uncertainty for tasks in reinforcement and supervised learning.
Multiple Choice Learning. Multiple Hypotheses Dropout is inspired by work from multiple choice learning (MCL) which was first introduced by GGuzman-Riiii Rivera et al. (2012) to address problems with multiple outputs given the same input. These approaches propose initializing multiple support vector machines Guuzman-Riven Rivera et al. 2012) or multi-layer perceptrons (MLPs) (Rupprecht et al.
2017
Lee et al. 2017 Nguyen et al.
2021 to increase the number of predictors per input. MH-
3

---

Under review as a conference paper at ICLR 2023
Dropout takes different approach, generating predictions by sampling from ensemble of possible sub-networks which is significantly more scalable.
MCL techniques commonly criticized for suffering from modal collapse. This problem arises when a subset predictors are not trained due to poor weight initialization. MH-Dropout addresses this by back- pockprropggtting through shared set of weights which also forces the network to adopt general features for multiple outputs.
MASKED VECTOR QUANTIZATION
Our solution extends the popular VQ autoencoder framework (Hinton & Zemel 1993 Van Den Oord et al. 2017) so we begin this section by reviewing the concept. We then introduce our proposed architecture MVQ and MH-Dropout. For simplicity, throughout this section we use only one primary and one secondary token per instance when describing the algorithm. In practice, sequences of tokens are often used.
3.1 BACKGROUND: VECTOR QUANTIZATION
The winner-takes-all VQ framework utilizes a latent codebook containing a mixture of K centroids eş E RD where 1... ... K. Given a dataset = " of n instances, this framework aims to map each instance discrete latent variable which refers to the codebook indices (tokens).
To achieve this, primary encoder is used to encode each instance to a vector y. This vector is replaced with the indices of the nearest vector from the latent codebook, forming discrete token The reparameterization process involves a simple lookup against the codebook indices which yields equivalent vector representation y. The representation y is then input the decoder 'D which outputs reconstruction the instance f.
The goal is to minimize the reconstruction loss between each instance and reconstructed output:
(x, x) = |xx x
(1)
This is made possible by back fack-propagating through the codebook using the straight-through gradient estimator (Bengio et al. 2013) which passes the gradients directly from the decoder to the encoder
with the following codebook loss:
Ly ||| ||sgly] - -112 + Blly sg[>113
where sg refers the stop-gradient operation. This is commonly referred to "commitment loss" (Van Den Oord et al., 2017). An additional stage is required to learn the categorical posterior distribution over the tokens using autoregressive probabilistic model with cross-entropy loss.
common approach to reducing both reconstruction and codebook loss is to increase the number of codebook entries K or number of tokens per instance S. This scales representational capacity by decreasing the expected distance between y s y, at the cost of training larger autoregressive models and longer sampling times. The MVQ framework scales representational capacity, but incurs these costs at much lower rate as codebook size and number of tokens grow, as we verify in Section 4
3.2 INCREASING REPRESENTATIONAL CAPACITY WITH MASKED VECTORS
Our MVQ framework builds upon the representational capacity of each latent VQ centroid by introducing a secondary code vector c' e RR' and allowing its masked configurations be mapped to each instance.
masked configuration secondary code vector is defined as:
where m E (0. D' a binary mask and @ is an element-wise or Hadamard product.
Let M = {m...., M be the set of all Jmax 2 possible masks. If each mask in M is multiplied element-wise with the same code vector e', this constructs the set of all 2D' possible masked configurations:
C (CI.... C.mas) {(e' m1 ...., (e' m mJaxx)
4

---

Under review as a conference paper at ICLR 2023
3
argmin|le;
yll
argmin |
yاا
E2
ex
Sample masks
representation
Codebooks
m
m
e
3
argmin llej
y'1l2
e/
f
Figure 3: The MVQ architecture depicting the reconstruction of an instance during training. MVQ introduces secondary branch to the VQ framework where secondary code and J/ random masks are sampled. Only the best of representations passed to the decoder.
Thus in this framework, each instance mapped from token representation ý using two latent code vectors from separate codebooks and binary mask from M,, (e, e', m)):
y = e f(e' O m) e + f(c)
(5)
where f is a non-linear multi-layer neural network. The complete architecture is illustrated in Figure where the outputs of secondary encoder 8' is quantized into e' using secondary codebook.
The training objective is to find the best combination of (e, e', m) that minimizes the 1.2-norm between y R y. As the number of codebook entries is generally small, simple codebook look-up is feasible for code vectors and e'. However when D' is high dimensional, it is computationally impractical to enumerate through all possible 2 masks in M.
To resolve this, we use MH- MH-Dropout which approximates the best mask by using multiple hypotheses training. In this approach, only < 2 masks are randomly sampled from M and only the best mask which results in the nearest representation ý = is used during back-propagation. This algorithm is described in the following section.
3.3 MULTIPLE HYPOTHESES DROPOUT
fr} MH-Dropout can be seen as efficiently converting f into an ensemble of sub-networks {f1,... ...
using the J masked code vectors as input. As seen in Figure $ each sub-network results from the elimination of input connections by each mask, just as classic dropout sets connections to zero. Thus, sampling J masks to find the nearest representation (in L norm) to y is equivalent to sampling J sub-network predictions (hypotheses) to find the nearest representation.
Only the best of the sub-network hypotheses is selected as the final representation and only the gradients of the respective best sub-network are back-propagge propppppppa.dd
In the forward pass, J masked code vectors are sampled from C and passed to f to give / sub-network hypotheses {.(cc),..... f(,))) The best sub-network is the with its hypothesis closest vector y in L norm.
Only this aubneetwork is addressed in the backward pass. Gradients for the other sub-networks eliminated by multiplying their respective loss by zero, so the objective function is:
Σ wy ($ j=1
(6)
Wi 8(j arg min | y112)
where wy gates gradient of sub-networks by using 8, which is the Kronecker-delta returning for the best hypothesis and ( otherwise.
5

---

Under review as a conference paper at ICLR 2023
Multiple Hypotheses Dropout
e
y
argmin||3)
f(5)
9,
f
random dropout
eliminate connections
multiple hypotheses
select nearest
representation
Figure 4: A forward pass through MH-DDropout layer where J random masks eliminate hidden connections, creating sub-networks that output J/ hypotheses.
The overall MVQ objective function combines the MH-Dropout loss, reconstruction error and two VQ codebook losses:
LMVQ = (xe 8) Lys + Lasly, e') + Lmhh((11J.y)
This objective function can be flexibly modified include a discriminator or perceptual loss (Johnson al. 2016 as seen in Section
3.4 DISCUSSION
MH-Dropout mitigates two related problems in multiple choice learning and discrete autoencoders. Multiple choice learning algorithms commonly criticized for suffering from modal collapse. This occurs when a subset of predictors are not trained, due to poor weight initialization for example, and then used during inference. The subset of all possible token sequences are not trained. This is usually mitigated by introducing probabilistic model to learn the categorical posterior over the used tokens, however it requires additional computational resources as token sequence length and number of codebook entries increase.
problem can be found in VQ autoencoders, when
MH-Dropout avoids modal collapse because the shared propagation. The sub-networks with hypotheses not closest to the primary representation will still be trained since their weights are shared with the winning aub-network. This forces the model to adopt set of weights for f and the secondary code vectors that generalize to multiple outputs. This is also why searching the ensemble sub-networks by sampling masks is effective and creates beneficial trade-off between computational resource use and finding the best mask (see Section 4.2).
weights are jointly trained in back-
We observe that degenerate outputs do not occur for any sampled mask applied to secondary code vectors during inference. It is therefore not necessary learn a distribution (as is done for VQ-VAE codebooks), significantly reducing the size of categorical posterior models.
masks for generation
MH-Dropout has the binary dropout (Srivastava al. 2014) property of tending to learn sparse secondary code vectors, with smaller numbers of larger values. Specific dimensions of the code vectors often correspond to distinct. interpretable characteristics such as smile and glasses, in Figure - These characteristics can be smoothly combined and transferred to other primary code vectors will demonstrate in the next section.

---

Under review as a conference paper at ICLR 2023
A similar approach to ours is product quantization (PQ) Jegou al. (2010), which also creates
multiple representation for each code vector. The difference is PQ breaks vectors into fixed partitions, whereas ours conducts a stochastic search of the projections of the vector using mask operation.
EXPERIMENTS
Here focus understanding the impact of adopting MVQ and MH-Dropout by integrating them into two popular VQ architectures, VQVAE-2 and VQGAN. We discuss our experimental setup in Section 4 We investigate the effects of increasing sampled masks MVQ quality in Section 4.2
and benchmark quality and reconstruction error across a wider range of datasets in Section 4.3 We then demonstrate that MVQ identifies features that can be transferred and combined in Section 4.4
4.1 SETUP
Model Comparisons. Benchmarking generative models is challenging task due to the wide range of factors that effect performance. To fairly assess MVQ we compare its performance to hierarchical VQ latents proposed in Razavi et al. (211)), and keep as many factors equivalent as possible. Thus, our comparative architecture to VQVAE-2 simply replaces the top-bottom hierarchical VQ codebooks with MVQ, and is referred to MVQVAE. Similarly, the hierarchical codebook version of VQGAN is compared to MVQGAN, but uses the same encoder and decoder from VQVAE-2 due to the GPU memory requirements of diffusion probabilistic models (Ho et al. 2020). The VQVAE2 models use PixeICNN as the categorical posterior model, while the VQGAN models use a Transformer model as proposed in Rombach al. 202Ip. Secondary codes are generated by conditioning on the sequence of primary codes, As previous works use different prramtpsr and datasets, we re-train and evaluate their models based on open-source implementations.
Slyper-parameters For each model, we vary two hyper- parameters to study the impact of masked code vectors: the number of codebook entries and tokens per image. There are a total of K codebook entries, with each primary and secondary codebook having K/2 // entries. The number of tokens per image S + 5" where and S' are the numbers of primary and secondary tokens respectively. Given image x E Rxx=x3
extend previous work of Esser al. (2021): Rombach et al. 2021) by studying wider range of dlonn-samplen factors, F = v(h x w)/S, from 7-64. Every set of comparative models are trained for equal epochs. For further hyper-parameters, see Appendix A.
Datasets. Experiments are conducted on medium resolution image datasets: FashionMNIST 28 x8828 (FMNIST) (Xiao et al., 2017). CIFAR 10 32x32 (Krizhevsky et al. 2009), CelebA 64x64 (Liu et al. 2018) and ImageNet 64 64 Deng et al. (2009). Using the down- soow--sapping factors provided above. the number of primary tokens per image are in the range 1-16. Due to this short length, we use 1 secondary token per image.
Metrics. Reconstruction and sample quality is assessed using Fréchet Inception Distance (FID) (Heusel et al. 2017) which measures similarity between real and generated samples. Reconstruction error assessed using perceptual loss (LPIPS) for VQGAN based models. For all the above metrics, a lower value indicates better performance. F-score measures the coverage of diversity and quality in single number per the work of Sajjadi et al. 2018). Due to lack of space, sample FID (FID*), negative log-likelihode (NLL), F-score (F3=1)8) are reported in Appendix A.2
4.2 IMPACT OF INCREASING SAMPLED MASKS PER FORWARD PASS
To explore the effects of masked vectors on quality, we use the FashionMNIST dataset as it allows extensive testing of hyper-paun peee-pparmeeerss due to its small size. Using VQVAB-2 as the bascline, multiple MVQVAE models are trained using increasing numbers of sampled masks per forward pass at high down-sampling factors of J = 28 in Figure 5 5 and F = 14 in Figure (b).
These results show that replacing hierarchical VQ with MVQ improves FID with as little as 16 sampled masks per pass. Increasing the masks per pass consistently improves reconstruction quality across all codebook and token lengths. As a result, MVQVAĒ with only 2 tokens per image outperforms several VQVAF2 models with 5 tokens which indicates better performance at higher
7

---

Under review as a conference paper at ICLR 2023
VQVAE-2
MVQVAE
VQVAE-2
MVQVAE
Codebook
350
350
Entites
250
250
200
200
16
64
100
100
60
16
32
128
256
512
1024 2040 4098 8192 16304
Sampled Masks
Number of Sampled Masks
(a) Tokens per Image = 1
(b) Tokens per Image 4
Figure 5: FID & on FashionMNIST validation split. Increasing sampled masks per forward pass consistently improve FID for MVQVAE (ours). VQVAE-2 scores are shown left of each chart.
Tokens per Instance
1
4 +
16
Metric
Dataset
Codebook Entries
16
256
16
256
16
256
VQGAN
266.55
153.45
68.95
46.79
41.92
37.05
CIFARIO
MVQGAN
58.78
57.22
46.70
41.19
36.73
34.84
VQGAN
113.54
87.67
36.07
28.78
26.47
24.28
4
CelebA64
MVQGAN
33.88
31.63
31.85
27.86
23.81
21.19
VQGAN
285.94
184.30
161.19
77.95
56.67
47.99
ImageNet64
MVQGAN
91.03
75.09
68.33
58.01
$2.50
45.47
VQGAN
0.463
0.447
0.423
0.379
0.340
0.271
CIFAR1O
MVQGAN
0.406
0.390
0.377
0.342
0.322
0.266
LPIPS
VQGAN
0.345
0.327
0.288
0.251
0.230
0.186
4
CelebA64
MVQGAN
0.293
0.261
0.270
0.232
0.227
0.186
VQGAN
0.558
0.537
0.531
0.488
0.466
0.406
ImageNet64
MVQGAN
0.514
0.497
0.498
0.472
0.456
0.404
VQGAN
0.132
0.744
0.770
0.882
0.874
0.906
CIFAR10
MVQGAN
0.846
0.815
0.863
0.885
0.892
0.909
F8=8
VQGAN
0.418
0.584
0.860
0.891
0.917
0.937
CelebA64
MVQGAN
0.868
0.881
0.861
0.915
0.932
0.952
VQGAN
0.125
0.428
0.206
0.571
0.787
0.849
ImageNet64
MVQGAN
0.557
0.620
0.703
0.752
0.832
0.872
Table 1: FID, perceptual loss and F-score (Fa=s) for MVQGAN (J = 2048) and VQGAN. MVQ outperforms VQ particularly when tokens per instance and number of codebook entries reduced.
levels of compression. Despite sampling only a subset of 2D' possible masks, these results verify that MH-Dropout allows to trade-off computational resource use discussed in Section 3.4
4.3 RECONSTRUCTION QUALITY
To understand scalability of performance, we conduct benchmarks across a wider range of datasets. VQGAN is used the baseline as find that the addition of perceptual loss and adversarial discriminator to VQVAE2 consistently improve metrics. Our proposed model, MVQGAN uses fixed J = 2048 sampled masks per pass during training.
In Table we present results which show that introduction of MVQ consistently improves FID, perceptual loss and F-score at all down-sampling factors F in the range 7-64. As seen in the previous section, improvement is greatest at reduced token lengths and number of codebook entries, which again shows that MVQ is more robust at higher compression rates.
8

---

Under review as a conference paper at ICLR 2023
(a) Representation Transfer
(b) Additive Composition
D(e + (Crw))
D(e + f((op))
(ee)
D(e (Cwu) + +((rrighe))
glasses
D(e f ((rrggn))
blonde
hair
D(e))
beard
sun burn
D(e + ((()))
Figure 6: Samples generated by MVQGAN using only 16 codebook entries and 4 + 1 tokens per instance. (a) Transfer of visual representations between two primary codes with different base styles. (b) Additive composition of multiple visual representations.
MVQGAN trained with only 2 tokens per image (J 64) on CelebA64 is reasonably competitive with VQGAN with 17 tokens (J 16). The same pattern is observed with all other datasets for MVQGAN with 5 tokens and 256 codebook entries in comparison to VQGAN with 17 tokens. These improvements allows us to trade a minor decrease in quality (see Appendix A.3 for images) for a more practical benefit of reducing sampling times during inference.
Increasing the down-sampling factor from 16 to 64 for 64x64 images would yield 7x speed-up in token sampling during inference a consumer-grade Nvidia RTX2080, and 45x speed-up for an Intel 19 CPU as seen in the Appendix, Table 6 Similarly for 256x 256 images, increasing the down-samppi wown-sampling factor from 16 32 would yield theoretical 20-24x speed-up. These improvements would be useful for embedded models in edge devices where low inference times are required and in video generation or reinforcement learning research where sampling efficiency low.
4.4 REPRESENTATION TRANSFER AND ADDITIVE COMPOSITION
We demonstrate that MVQ results transferable visual representations that can be smoothly combined in output space. Figure 64 illustrates masked configurations being transferred from primary code vector e to e2. Despite having different base styles, multiple visual representations (such as hair fringe, styles and colours) are smoothly transferred. Figure op depicts grid of 9 images where the center image the decoded primary code vector. Four sampled masked configurations are added to the centre and shown in the orthogonal grid positions. The corner images the result of summing the two masked configurations their two nearest orthogonal images. These corner images demonstrate the smooth combination of multiple visual characteristics.
As these features are learned completely unsupervised, the choice of which features become transfer- able depends on the dataset. For example, FashionMNIST find that shades, patterns, width and arm lengths are transferable (see Appendix A.5)). We also observe that MVQ models with reduced token lengths and codebook entries are better transferability. This because reducing primary VQ centroids increases the reliance on learning factors of variations with secondary masked codes.
5 CONCLUSION
In this paper we introduced the MVQ framework which allows masked configurations to be mapped to instances using a stochastic winner- tales-l training regime called MH-Dropout. MH-Dropout produces code vectors which avoid modal collapse and learn sparse, transferable features. MVQ improves on both reconstruction error and sampling quality in comparison to VQ, and considerably more robust when tokens per instance and codebook entries are reduced.

---

Under review as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
An anonymized codebase for ICLR reviewers has been uploaded to GitFront. The link is provided here This codebase will be open-sourced after publication. To help others replicate our work, additional information on experiment hyper-parameters and metrics are found in Appendix A. and Appendix A.2 respectively.
REFERENCES
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint /rXiv:1308.3432. 2013.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. leee, 2009.
Jeff Donahue and Karen Simonyan. Large scale adversarial representation learning. Advances in Neural Information Processing Systems, 32, 2019.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution im- age synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 12873-12883, 2021.
Yarin Gal. Uncertainty in deep learning. University of Cambridge, 1(3),00 2016.
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in neural information processing systems, 30, 2017.
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27:2672-2680) 2014.
Abner Guzman-Rivera Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: Learning to produce multiple structured outputs. In NIPS, volume 1, pp. 3. Citeseer, 2012.
Junxian He, Daniel Spokoyny, Graham Neubig, and Taylor Keng-Kinkpatrrcc. Lagging inference networks and posterior collapse in variational autoencoders. arXiv preprint arXiv: UrXiv:1901.05534 2019.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. arXiv preprint orXivv/706.. 08500, 2017.
Geoffrey E Hinton and Richard Zemel. Autoencoders, minimum description length and helmholtz free energy. Advances in neural information processing systems, 6, 1993.
Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhutdinov. Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint arXiv:1207.0580. 2012.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv: 16/ 1.0/ 144, 2016.
Herve Jegou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search.
IEEE Iransactions on pattern analysis and machine intelligence, 17-12.117-128, 2010.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In European conference on computer vision, pp. 1917 Springer, 2016.
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterizzation trick. Advances in neural information processing systems, 28, 2015.
10

---

Under review as a conference paper at ICLR 2023
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
Kimin Lee, Changho Hwang, KyoungSoo Park, and Jinwoo Shin. Confident multiple choice learning. In International Conference on Machine Learning, pp. 2014-2023 PMLR, 2017.
Ziwei Liu, Ping Luo, Xiaogang Wang. and Xiaoou Tang. Large-scale celebfaces attributes (celeba) dataset. Retrieved August, 15(2018):11, 2018.
James Lucas, George Tucker, Roger Grosse, and Mohammad Norouzi. Understanding posterior collapse in generative latent variable models. 2019a.
James Lucas, George Tucker, Roger B Grosse, and Mohammad Norouzi. Don't blame the elbo! linear vae perspective on posterior collapse. Advances in Neural Information Processing Systems, 32, 2019b.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712 2016.
Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variational dropout sparsifies deep neural networks. arXiv preprint arXiv: 170100399,, 2017.
David D Nguyen, Surya Nepal, and Salil S Kanhere. Diverse multimedia layout generation with multi choice learning. Proceedings of the 29th ACM International Conference on Multimedia, pp. 218-226, 2021.
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint
2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference Machine Learning, pp. 8821-8831. PMLR, 2021.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances neural information processing systems, 32, 2019.
Jorma Rissanen. Stochastic complexity and the mdl principle. Econometric Reviews, 6(1):85-102, 1987.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. arXiv preprint arXiv:2112.1077 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE International Conference on Computer Vision, pp. 3591- 3600, 2017.
Mehdi SM Sajjadi, Olivier Bachem, Mario Lucic, Olivier Bousquet, and Sylvain Gelly. Assessing generative models via precision and recall. Advances Neural Information Processing Systems, 31, 2018
Jascha Sohl-Dich Sohh-Diccstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256-2265. PMLR, 2015.
11

---

Under review as a conference paper at ICLR 2023
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1958, 2014.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in Neural Information Processing Systems, 33: 1966--99,,,,,, 2020.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International conference machine learning, pp- 1747-1756. PMLR, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 2017.
12


Reviews:
[{"review_id": "AgTxC5SCRB", "reviewer": "Reviewer_GRXC", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": " - This work proposes a Masked Vector Quantization (MVQ) framework to boost per-code capacity, targeting at representing images with shorter sequence or fewer codebook entries. \n\n - In the framework of MVQ, each instance is mapped into three components, including a primary code, a secondary code and a mask. The computational infeasibility is mitigated by Multiple Hypotheses Dropout. \n\n - The authors demonstrate that with MVQ, improvements on FIDs over VQGAN are observed, especially when the number of tokens per instance and codebook size is reduced. The authors also show some properties of learned codes. ", "summary_of_the_review": " - Although this work is studying this important question of vector quantization, lack of solid experiments, such as the feasibility of training a better generative model with this MVQ system, and improvements on longer sequences and larger dataset, are unable to show the effectiveness of this method. Note that the target of VQ is not to solely do better compression, but for handling complex signals that existing models are unable to do well.", "strength_and_weaknesses": "#### **Strength**\n\n - The topic of designing a better vector quantization system is important, which has recently been adopted in a variety of tasks, including generative modeling. \n\n - Vector quantization regimes face a compression-fidelity trade-off, whereas a longer sequence reduces quantization error and increases modeling difficulty. This work proposes to overcome this issue through a novel VQ system.\n\n\n#### **Weaknesses**\n\n - Despite including novel masking operations on the code, the motivation of this work is closely related to Product Quantization[1]. Discussion and comparison against relevant quantization methods should be provided.\n\n - Some details are not self-contained, including the designs of MVQGAN. For example, the authors refer to VQVAE-2 as the architectures of encoder and decoder, and \"All architectures use a Transformer as the categorical posterior model as proposed in Rombach et al. (2021)\", where **all** is quite confusing. The authors need to provide more details to make the work more self-contained.\n\nThe scope of experiments is unable to demonstrate the effectiveness of such modeling power. \n\n - Since the major target of vector quantization is not to do compression, the reconstruction quality, errors, or sampling speed, cannot reflect the learned token representations as reliable. In fact, there exist many adaptations that reduce quantization error, while resulting in a worse token representation that might be more difficult for the probabilistic model to learn. \n\n - Applying an aggressive downsampling rate such as $f$=16 for small images used in the experiments makes less sense, as the obtained sequence length is already short, e.g. 4x4 . A shorter sequence and extremely small codebook hinders the contribution of this work, of which a non-auto-regressive is able to model well(e.g. MaskGIT[2] for 16x16 tokens on ImageNet). \n\n\n\n[1] Product quantization for nearest neighbor search. T-PAMI 2011\n[2] MaskGIT: Masked Generative Image Transformer. CVPR 2022", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "A02ryez9G7g", "reviewer": "Reviewer_qWji", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents Masked Vector Quantization (MVQ), which is an extension of hierarchical Vector Quantized Variational AutoEncoder (VQ-VAE [a]). MVQ introduces mask configuration on the secondary code vector and demonstrates its effectiveness on image generation with shorter sequences (for example, less than 20) per sample.\n\nThe paper uses a training schema named Multiple Hypotheses Dropout (MH-Dropout) to learn the mask configuration. The key idea is to sample a number of binary masks (with Dropout) and take the best one that minimizes the latent reconstruction. \n\n\n[a] Razavi, Ali, Aaron Van den Oord, and Oriol Vinyals. \"Generating diverse high-fidelity images with vq-vae-2.\" Advances in neural information processing systems 32 (2019).", "summary_of_the_review": "As detailed in Section **Strength and Weakness**, I think improving the generation efficiency is important for autoregressive generation. The problem in the paper is well-motivated. \n\nHowever, the current manuscript could be further improved. Particularly, I think it lacks a part to explain the intuition behind the proposed approach and why the proposed approach works. In addition, the difference to VQ-VAE2 is not clear and the experimental results are weak.  \n\nConsidering the above concerns, my current rating is to reject the paper. ", "strength_and_weaknesses": "### Strength\n* [S1]: The approach is well presented, figures are clear and contain enough detailed information to present the approach \n* [S2]: Generating long sequences is indeed computationally expensive (see supplementary material Table 6), the problem is well motivated\n* [S3]: The related work is completed, and to the best of my knowledge, all the related fields are discussed. \n* [S4]: Representations learned by MVQ with short sequences are interesting (Sec. 4.4) \n\n### Weakness \n* [W1]: Although the proposed method (MH-Dropout) shows improvement in image generation with shorter sequences, the reason why it works remains unexplained. I think the paper lacks a discussion on the intuition and the explainability of the proposed approach. \n\n* [W2]: The proposed pipeline is conceptually similar to VQ-VAE2. In fact, in 4.1 Setup, the paper precises that it replaces the top-bottom hierarchical VQ codebooks in VQ-VAE2 with MVQ, which consists of injecting masks into the top-level codes if I understand correctly. In such a case, it would be better for the paper to discuss clearly the difference to VQ-VAE2 and clarify its own technical contribution. \n\n* [W3]: The generation results are weak (Table 1, Figure 7). I am fully convinced that the community needs faster methods, however, I am not sure the community benefits from approaches that sacrifice too much on the generation quality.  \n\n* [W4]: The transformer part is not self-contained. It is not clear whether the primary and the secondary codes are in different scales and how to conduct generation for hierarchical codes. \n \n[a] Razavi, Ali, Aaron Van den Oord, and Oriol Vinyals. \"Generating diverse high-fidelity images with vq-vae-2.\" Advances in neural information processing systems 32 (2019).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "26T87GrrIK", "reviewer": "Reviewer_6yNV", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 5, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "", "summary_of_the_paper": "As a researcher working in the similar field, I definitely appreciate the author's efforts to solve a long-standing problem faced by many using a method that is reasonable according to the nature of the discretization process. The author seek to solve the problem that long sequences of discrete making it very difficult for transformer or other method to model in an autoregressive manner, which leads to difficulty in generative task. The author use a non-linear combination of two set of tokens , the primary and the secondary , to make the output representation semi-continuous and with higher capacity, using a MLP projector and dropout mask. This strategy can be intuitive understood as estimating the posterior of representation distribution given the latent consisting of two different tokens ,in a mathematically non-tight way.  ", "summary_of_the_review": "This is a great piece of work due to its timing in solving an urgent problem in the research field as well as in the industry. The quality of the work is high , the method is novel and clean. One thing I like to point out is that figure is difficult to understand , not sure it is helping or making it more difficult for the readers. I would suggest improve it \n\n", "strength_and_weaknesses": "The authors solve a real problem and pain point confront by many in discrete token based generative model for vision task, using a semi-continuous represention consisting of a primary and a secondary token. They empirically show using different generative task to show their approach improves performance from different aspects. The problem setting is clear, the approach at a high level is great and the results are solid \n\nHowever, the reason why the dropout mask and multi-hypothesis version is use is not probably reasoned. There are many alternative possibilities including a small Bayesian neural network , learnable dropout such as Contextual Dropout or GFlowOut. The reason to use the current masking method to be properly explained. I understand it does not have to be the most sophisticated method in the literature but at least the author need to explain why it cannot be simpler.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject