HQ-VAE: Hierarchical Discrete Representation Learning with Variational Bayes
Authors: Yuhta Takida
Yukara Ikemiya
Takashi Shibuya
Kazuki Shimada
Woosung Choi
Chieh-Hsin Lai
Naoki Murata
Toshimitsu Uesaka
Kengo Uchida
Wei-Hsiang Liao
Yuki Mitsufuji
Venue: Accepted by TMLR ()

Abstract: Vector quantization (VQ) is a technique to deterministically learn features with discrete codebook representations. It is commonly performed with a variational autoencoding model, VQ-VAE, which can be further extended to hierarchical structures for making high-fidelity reconstructions. However, such hierarchical extensions of VQ-VAE often suffer from the codebook/layer collapse issue, where the codebook is not efficiently used to express the data, and hence degrades reconstruction accuracy. To mitigate this problem, we propose a novel unified framework to stochastically learn hierarchical discrete representation on the basis of the variational Bayes framework, called hierarchically quantized variational autoencoder (HQ-VAE). HQ-VAE naturally generalizes the hierarchical variants of VQ-VAE, such as VQ-VAE-2 and residual-quantized VAE (RQ-VAE), and provides them with a Bayesian training scheme. Our comprehensive experiments on image datasets show that HQ-VAE enhances codebook usage and improves reconstruction performance. We also validated HQ-VAE in terms of its applicability to a different modality with an audio dataset.

Published in Transactions on Machine Learning Research (03/2024)
HQ-VAE: Hierarchical Discrete Representation with Variational Bayes
Learning
Yuhta Takida', Yukara Ikemiyal, Takashi Shibuya', Kazuki Shimada!, Woosung Choi', Chich-Hsin Lait, Naoki Muratal, Toshimitsu Uesakal, Kengo Uchida", Liao WeiHsiang', Yuki Mitsufuji
SSoyAA, Tokyo, Japan, 'Sony Croup Corporation, Tokyo, Japan
Reviewed on OpenReview: https: openreview. net/ forum? id= ja
Abstract
Vector quantization (VQ) is a technique deterministically learn features with discrete codebook representations. It is commonly performed with variational autoencoding model, VQ-VAE, which can be further extended to hicrarchical structures for making high-fidelity reconstructions. However, such hierarchical extensions of VQ-VAE often suffer from the codebook/layy /layer collapse issue, where the codebook is not efficiently used to express the data, and hence degrades reconstruction accuracy. To mitigate this problem, we propose a novel unified framework to stochastically learn hierarchical discrete representation on the basis of the variational Bayes framework, called hierarchically quantized variational autoencoder VAE). HQ- LLQ- naturally generalizes the hierarchical variants of VQ- VQ-VAE, such as VQ- VAE-2 and residual-quantized VAE (RQ-VAE), and provides them with Bayesian training scheme. Our comprehensive experiments image datasels show that HQ- HA-VAE enhances codebook usage and improves reconstruction performance. We also validated HQ- VAE in terms of its applicability to different modality with an audio dataset.
1 Introduction
Learning representations with discrete features is one of the core goals in the field of deep learning. Vector quantization (VQ) for approximating continuous features with set of finite trainable code vectors is a common way to make such representations (Toderici et al. 2016 Theis et al., 2017 Agustsson et al. 2017).
It has several applications, including image compression (Williams et al. 2020 Wang et al. 2022) and audio codecs (Zeghidour et al. 22021 Défossez el al. 2022). VQ-based representation methods have been improved with deep generative modeling, especially denoising diffusion probabilistic models (Sohl-Dickstein et all 2015 Ho et al. 2020 [Song al.] 2020; Dhariwal & Nichol 2021
Hoogeboom et al.] 2021 Austin et al. 2021 and autoregressive models vaan den Oord et al.) 2016 Chen et al. 2018 Child et al. 2019). Learning the discrete features of the target data from finitely many representations enables redundant information to be ignored, and such a lossy compression can be of assistance in training deep generative models on large-scale data. After compression, another deep generative model, which called a prior model, can be trained on the compressed representation instead of the raw data. This approach has achieved promising results in various tasks, c.g., unconditional generation tasks (Razavi et al., 2019 Dhariwal et al. 2020 Esser
et al. 2021b a Rombach al.) (2022), lext-to-image generation {Ramesh et al. 2021 Gu al., 2022 Lee et al. 2022a and textually guided audio generation (Yang et al. 2022; Kreuk et al. 2022). Note that the compression performance of VQ limits the overall generation performance regardless of the performance of the prior model.
Vector quantization usually achieved with a vector quantized variational autoencoder (VQ-VAE) (van den Oord al.) 2017). Based on the construction of van den Oord al. (2017), inputs are first encoded and quantized with code vectors, which produces a discrete representation of the encoded feature. The discrete representation is then decoded to the data space to recover the original input. Subsequent, developments

---

Published in Transactions on Machine Learning Research (03/2024)
incorporated hierarchical structure into the discrete latent space to achieve high-fidelity reconstructions. In particular, Razavi et al. (2019) developed VQ-VAE into a hierarchical model, called VQ-VAE-2. In this model, multi-resolution discrete latent representations are used to extract local (e.g, texture in images) and global (c.g., shape and geometry of objects in images) information from the target data. Another type of hierarchical discrete representation, called residual quantization (RQ), was proposed to reduce the gap between the feature maps before and after the quantization process (Zeghidour et al. 2021 Lee al. 2022a
Despite the successes of VQ-VAE in many tasks, training of its variants is still challenging. It is known that VQ- DA- suffers from codebook collapse, problem in which most of the codebook elements are not being used at all for the representation (Kaiser et al., 2018 Roy et al., 2018 Takida et al. 2022b). This inefficiency may degrade reconstruction accuracy, and limit applications to downstream tasks. The variants with hierarchical latent representations suffers from the same issue. For example, Dhariwal et al. (2020)
reported that it is generally difficult to push information to higher levels in VQ-VAF-2; i.c., codebook collapse often occurs there. Therefore, certain heuristic tochniques, such as the exponential moving average (EMA)
update (Polyak & Juditsky 1992) and codebook reset. (Dhariwal et al.; 2020), are usually implemented to mitigate these problems. Takida et al. (202bb) claimed that the issue is triggered because the training scheme of VQ-VAE does not follow the variational Bayes framework and instead relies on carefully designed heuristics. They proposed stochastically quantized VAE (SQ-VAE), with which the components of VQ- VAF, ie., the encoder, decoder, and code vectors, are trained within the variational Bayes framework with an SQ operator. The model was shown to improve reconstruction performance by preventing the collapse issue thanks to the self-annealing effect (Takida et al., 2022b), where the SQ process gradually tends to a deterministic one during training. We expect that this has the potential to stabilize the training even in a hicrarchical model, which may lead to improved reconstruction performance with more efficient codebook usage.
Here, we propose Hierarchically Quantized VAE (HQ-VE)) a general variational Bayesian model for learning hicrarchical discrete latent representations. Figure | illustrates the overall architecture of HQ-VAE. The hierarchical structure consists of a bollom-up and lop-down path pair, which helps to capture the local and global information in the data. We instantiate the generic VAQ by introducing two types of top-down layer. These two layers formulate hicrarchical structures of VQ-VAE-2 and residual-quantized VAE (RQ- VAE) within the variational scheme, which we call SQ- VAE-2 and RSQ- VAE, respectively. HQ-VAE can be viewed as a hierarchical version (extension) of SQ-VAF, and it has the favorable properties of SQ-VAE (e.g.. the self- eefffannegiing effect). In this sense, it unifies the current well-known VQ models in the variational Bayes framework and thus provides a novel training mechanism. We empirically show that HQ-VAE improves upon conventional methods in the vision and audio domains. Moreover, we applied HQ-VAEs Es to generalive lasks on image datasets to show the feasibility of the learnt discrete latent representations. This study is the first attempt at developing variational Bayes on hierarchical discrete representations.
this paper, uppercase letters (P, Q) and lowercase letters (p, q) respectively denote the prob- ability mass functions and probability density functions, calligraphy letters (P, Q) the joint probabilistic distributions of continuous and discrete random variables, and bold lowercase and uppercase letters (e.g., æ and Y) vectors and matrices. Moreover, the ith column vector in Y is written as ys: [N] denotes the set of positive integers no greater than N E N, and J and L denote the objective functions of HQ- (QEV and conventional autoencoders, respectively.
Throughout
2 Background
We first review VQ-VAF and its extensions to hierarchical latent models. Then, we revisit SQ- SA-V which serves as the foundation framework of HQE.
To discretely represent observations & E DR codebook B is used that consists of finite trainable code vectors {b,) (by e Rd).. A discrete latent variable is constructed to be in a dy-tuple of B, i.e., Z E Bd,, which is later decoded to generate data samples. A deterministic encoder and decoder pair is used to connect the observation and latent representation, where the encoder maps æ to Z and the decoder recovers from Z by using decoding function fo Rdsxd; RDP. An encoding function, denoted
VQ-VAE.
K
2

---

Published in Transactions on Machine Learning Research (03/2024)
GΦ : RR. Rdxd, and deterministic quantization operator are used together as the encoder. The encoding function first maps æ to Z Rdxxd; then, the quantization operator finds the nearest neighbor of 2, for i E [dæ], i.e., z; = arg ming, || b& ll The trainable components (encoder, decoder, and codebook) are learned by minimizing the objective,
LvQ-VAE
fo( + + - sg. 2 2
(1)
where sg[-] is the srop-gradient operator and 3 is a hyperparameter balancing the two terms. The codebook is updated by applying the EMA update to ||sg| 2 3
VQ- Razavi et al. 2019 incorporated hierarchical structure into the discrete latent space in VQ- VAE to model local and global information separately. The model consists of multiple levels of latents so that the top levels have global information, while the bottom levels focus on local information, conditioned on the top levels. The training of the model follows the same scheme as the original VQ-VAE (e.g., stop-gradient, EMA update, and deterministic quantization).
RQ-VAE. As an extension of VQ, RQ was proposed to provide a finer approximation of by taking into account information on quantization gaps (residuals) (Zeghidour et al.) 2021 Lee et al. 2022a). With RQ, L code vectors assigned to each vector 2; (i E [d=]), instead of increasing the codebook size K. To make multiple assignments, RQ repeatedly quantizes the target feature and computes quantization residuals, denoted as RR Namely, the following procedure is repeated L times, starting with Ro Z: 21, arg min,, || b and Ra = R - Zí. By repeating RQ, the discrete representation can be refined a coarso-to-fine manner. Finally, RQ discretely approximates the encoded variable as 2 Et= Zr. where the conventional VQ is regarded as special case of RQ with 1
SQ-VAE. Takida et al. 2022b proposed a variational Bayes framework for learning the VQ- VAE com- ponents to mitigate the issue of codebook collapse. The resulting model, SQ-VAE, also has deterministic accoing//ecodding functions and a trainable codebook like the above autoencoders. However, unlike the deterministic quantization schemes of VQ and RQ, SQ-VAE adopts stochastic quantization (SQ) for the encoded features to approximate the categorical posterior distribution P(Z[z). x)) More precisely, it defines stochastic dequantization process p.a (ž, Z) = NN22; z,, IT), which converts a discrete variable z; into a continuous one ă; by adding Gaussian noise with a learnable variance s². By Bayes' rule, this process is associated with the reverse operation, i.e., SQ, which is given by PL = by 2) x exp b to this variational framework, the degree of stochasticity in the quantization scheme becomes adaptive. This allows SQ-VAE to benefit from the ellect of self-armealing, where the SQ process gradually approaches a deterministic one as s² decreases. This generally improves the efficiency of codebook usage. Other related works in the literature of discrete posterior modcling are found in Appendix A
28 Thanks
3 Hierarchically quantized VAE
In this section, we formulate the generic HQ-VAE model, which learns a hierarchical discrete latent rep- resentation within the variational Bayes framework. It serves backbone of the instances of HQ-VAE presented in Section
To achieve a hierarchical discrete representation of depth L, we first introduce L groups of discrete latent variables, which are denoted as Z := (Z. For each / E [L], we further introduce a trainable codebook,
B := (b),)k consisting of Kj dy-dimensional code vectors, i.e., b%. E Rd for k E [Ki]. The variable Z, is represented dy-tuple of the code vectors in B;; namely, Z, B Similarly to conventional VAEs, the latent variable of each group is assumed to follow pre-defined prior mass function. We set the prior as i.i.d. uniform distribution, defined as P(z,, = by)= = 1/K, for i i [d]]. The probabilistic decoder is set to a normal distribution with a trainable isotropic covariance matrix, p= ZI.z.) = N'x:; fo ZI..), o²I) with a decoding function fo xd .... + Xdz + RD. To generate instances, it decodes latent variables sampled from the prior. Here, the exact evaluation of Po Z..z/a:) is required to train the generative model with the maximum likelihood. However, this is intractable in practice. Thus, we introduce an approximated posterior on Z given a and derive the evidence lower bound (ELBO) for maximization instead.
3

---

Published in Transactions on Machine Learning Research (03/2024)
Codebook Hi
Prior distribution
For posterior
P(Z. 2,)
Prior distribution
Quantizaioo
Bottom-up path
Top-down puth
2. P2((2112 (2,)
II
Output
Hg(o)
2
(b) First lop-down layer.
Input from top layers (Zua-s)
Upsimpling
Isput from
bottom up (us)
Pass
H3(xx)
Cuncar
Pricr distribution
2, +-
Fnc block
Codebook
PZZ Z,)
C Zn
1
11
P(an)ey (R. Zi)
Top-down layer
Quantiztioo
2i
Pa(Z
(2: (i)
Oulpul (impul to nexl layer)
(x)
(c) Injected top- cop-down layer.
L
Input from top layers (Zua-1)
Resblock
Pre-procossing
Decoding (0)
+
Pass
Prior distribution
Calar, Z
Codebook Bt
Pl Zr Z,)
Quamisation
P(a).)p (R:, (2)
Zi (Zi Zij-
Original
Reconstruction
Outpul (inpul next layer)
(a) Overview of HQ-VAE.
(d) Residual top- down layer.
Figure 1: (a) HQ-VAE consists of bottom-up and top-down paths. Red arrows represent the approximated posterior (Equation 3 This process consists of the following three sleps: the boltom-up path generates the features { (x))) at different resolutions; the top-down path fuses each of them with the information processed in the top- down path into Zi.L; the auxiliary variable at the ith layer, Zí is quantized into Z based on the codebook B, Lastly, to complete the autoencoding process, all the quantized latent variables in the top-down path, ZI.L, are fed to the decoder fe. recovering the original data sample. The objective function takes the Kullback Leibler divergence of posterior and prior (in the blue box). (b) First. layer for top-down path. Any HQ- VAE should include this layer as its top layer. An HQ-VAE that includes only this layer in its top-down path reduces to SQ-VAE. (c)-(d) Two types of top-down layers are proposed: injected lop-down and residual lop-down layers. Injected lop-down layer newly injects the higher-resolutionn features generated by the bottom-up path into the latent variables in the previous top-down layers, resulting in a higher resolution (i.e., larger dimension) of Z, and Zi compared to Z Residual top-down layer does not incorporate any new features from the bottom-up path but aims to refine the processed latent variables in top-down path with additional code vectors for representing the quantization residual from the preceding top-down layer. As a result, it maintains the same dimensionality for Z₁ as in the previous layer (i.c., Zi-1). An HQ-VAE that consists entirely of the injected (residual) lop-down layers is analogous to VQ-VAE-2 (RQ-VAE) which is presented in Section 4.2.1 4.3.1
Inspired by the work on hierarchical Gaussian VAEs (onderby et al. 2016 Vahdat & Kautz 2020 Child 2021), we introduce HQE bottom-up and top-down paths, as shown in Figure la The approximated posterior has a lop-down structure (Z) + Zz + ... + Z,.). The bollom-up path first. generates features from x as (x) at different resolutions (r E [R]). The latent variable in each group on the top-down path processod from Z Z, in that order by taking Hz(x) into account. To do so, two features, one extracted by the bottom-up path (Hz(z) and one processed on the higher layers of the top-down path (Zua-),, can be fed to each layer and used to estimate the Z, corresponding to æ, which we define
4

---

Published in Transactions on Machine Learning Research (03/2024)
= ( Z1a-1). The ith group Zi has unique resolution index r; denote it as r(l). For simplicity, we will ignore H in Zi and write Z = Gos (æ, Zia-1). The design of the encoding function G²', lead us to a different modeling of the approximated posterior. A detailed discussion of it presented in the next section.
It should be noted that the outputs of G lie in Rdxxd, whereas the support of Zi is restricted to B To connect these continuous and discrete spaces, we use pair of stochastic dequantization and quantization processes, in Takida et al. (2022b). We define the stochastic dequantization process for each group Pof (2,, Z = (2 Zi,s, 21,,, T) which is equivalent to adding Gaussian noisc to the discrete variable the covariance of which, si1, depends on the index the group l. We hereafter denote the set Z as ZI.L. i.e., Z.. == 1/ Next, we derive a stochastic quantization process in the form the inverse operator of the above stochastic dequantization:
(-
[|
bullē
P.p (21, by bi) exp
(2)
2s
By using these stochastic dequantization and quantization operators, we can connect Z1.7 and ZI.L via Z in a stochastic manner, which leads to the entire encoding process:
L d
QZZ Z1.(=x)
= IIII
P.ę &1, (
Znn-1))Po(((1) 2t).
(3)
i=
The prior distribution ZI and ZEL is defined using the stochastic dequantization process
P(ZI.L, ZT.L)
P(,,,)),
(4)
[=] i=1
where the latent representations are generated from - 1 to L in that order. The generative process does not use Z; rather it uses Z as x fe(
4 Instances of HQ-VAE
Now that we have established the overall framework of HQ-VAE, we will consider two special cases with different top-down layers: injected top-down or residual top-down. We will derive two instances of HQ-VAE that consist only of the injected top-down layer or the residual top-down layer, which we will call SQ-VAE- and RSQ-VAE by analogy VQ-VAF-2 and RQ-VAF. Specifically, SQ-VAF-2 and RSQ-VAF share the same architectural structure VQ-VAE-2 and RQ-VAE, respectively, but the differences lie in their training strategics, discussed in Sections 4.2.2 and 1.3.2 Furthermore, these two layers can be combinatorially
used to define hybrid model of SQ- FA- and RSQ-VAF, as explained in Appendix D Note that the prior distribution (Equation is identical across all instances.
4.1 First top-down layer
The first top-down layer is the top of layer in HAEV As illustrated in Figure lb it takes (z) as input and processes it with SQ. An HQ- VAE composed only of this layer reduces to SQ-VAE,
4.2 Injected tow- layer
The injected top-down layer for the approximated posterior is shown in Figure Ic This layer infuses the variable processed on the top-down path with higher resolution information from the bottom-up path. The Ith layer takes the feature from the boltom-up path (HT* (a)) and the variables from the higher layers in the top-down path as inputs. The variable from the higher layers is first upsampled align with H (x). These two variables are then concatenated and processed in encoding block. The above overall process
5

---

Published in Transactions on Machine Learning Research (03/2024)
corresponds to = (z, Z1:-1) in Section 3 The encoded variable Zi is then quantized into Z with the codebook B, through the process described Equation (2 Finally, the sum of the variable from the top layers and quantized variable Zi is passed through to the next layer.
4.2.1 SQ-VAE-2
An instance of HQ- VAE that has only the injected top-down layers in addition to the first layer reduces to SQ-VAE-2. Note that since the index of resolutions and layers have a onc-to-one correspondence in this structure, r()) = I and L = R. As in the usual VAEs, we evaluate the evidence lower bound (ELBO) inequality: (og)o(x)) 2 -sQ-vAB-((;; Φ s², B), where
s
(SQ.VAE-2(2 Φ. B) E 2l Z.. 2..1w)) [- log po log
Q 21:1(0))
ZI.L) log P(Z ZI.L.
is the ELBO, x² : (of) and : = (B ... Br). Hereafter, we will omit the arguments of the objective functions for simplicity. By decomposing Q and p and substituting paramcterizations for the probabilistic parts, we have
D JSQ. VAE-2 log logo² o² + Eoi 2.ckk))
> (11)L11) 20²
L
11 2s
H(P.((Z(Z)))
(6)
where H()) indicates the entropy of the probability mass function and constant terms have been omitted. The derivation of Equation ( is in Appendix C The objective function (6) consists of reconstruction
term and regularization terms for 1ILL and ZI.L. The expectation w.r.t. the probability mass function
Pa(Z1: = by Zi) can be approximated with the corresponding Gumbel-softmax distribution (Maddison
al. 2017
Jang al. 2017 in reparameterizable manner.
4.2.2
SQ-VAE vs. VQ-VAE-2
VQ- VAE-2 is composed similar fashion to SQ- but is trained with the following objective function:
LvQ-VAE-2 1 fo(
EIG' ZTa-1) sg. Z1]1.
(7)
where the codebooks updated with the EMA update in the same manner the original VQ- VAE. The objective function ( except for the stop gradient operator and EMA update, can be obtained by setting both of and o² to infinity while keeping the ratio of the variances as s B- for I E [L] in Equation (6).
In contrast, since all the paramcters except D and L in Equation 6 are optimized, the weight of cach
term is automatically adjusted during training. Furthermore, SQ-VAE-2 expected to benefit from the
self-ann annealing effect, as does the original SQ-VAE (see Section 5.3
4.3 Residual top-down layer
In this subsection, we set R 1 for demonstration purposes (the general case R in AppendixDi This means the bottom-up and top-down paths are connected only the top layer. In this setup, the bottom-up path generates only a single-resolution feature; thus, we denote H. Hœ(æ) for notational simplicity in this subsection. We design residual top-down layer for the approximated posterior, as in Figure [I] This layer is to better approximate the target fcature with additional assignments of code vectors per encoded vector. By stacking this procedure L times, the feature is approximated as
L
H (x) =
Σ Z
1=1
We empirically found setting Z, to 2, instead sampling Z, from P&a (2.) leads better performance (as reported in Takida et al. 2222hb) therefore, follow the procedure practice.

---

Published in Transactions on Machine Learning Research (03/2024)
Therefore, in this layer, only the information from the higher layers (not the from the bottom-up path) is fed the layer. It is desired that Σ Zr approximate the feature better than E Zr. On this basis, we let the following residual pass through to the next layer:
Go ZI.-1) HΦ(x)
Σ Zv.
(9)
1=1
4.3.1 RSQ-VAE
RSQ- VAE is an instance that has only residual top-down layers and the first layer. At this point, by following Equation and omitting constant terms, we could derive the ELBO objective in the same form as Equation
naive RSQQVVL VAE
2 log 0² + 2.c
| fo (1
Σ 2m1
[ 2 28
-(P..(212.)))
(10)
20
where the numerator of the third term corresponds to the evaluation of the residuals Hœ(x) - - Zy for all E [L] with the dequantization process. However, we empirically found that training the model with this ELBO objective above was often unstable. We suspect this because the objective regularizes ZI.L to make Σ Zy close the feature for all / E [L]. We hypothesize that the regularization excessively pushes information to the top layers, which may result in layer collapse in the bottom layers. To address the issue, modify the prior distribution P to result in weaker latent regularizations, Specifically; we define the prior distribution as joint distribution of (ZI.I. Z) instead of that of (Z.. ZI..), where 2 = Σ Z, From the reproductive property of the Gaussian distribution, continuous latent variable converted from Z via the stochastic dequantization processes, 2 = Σi= Zi, follows a Gaussian distribution:
Paa (2, 1) W
1
21,50
Σ $
(11)
Im
[m/
and we will instead use the following prior distribution:
d.
L
P P(ZI. 2) II
II P
Pe (2.
(12)
We now derive the ELBO using the newly established prior and posterior starting from
lo) 2 -TRSQ-VAE = -F. Q2, log Z1.1. + log ((Z [1x)
(13)
P(Z 2)
The above objective can be further simplified as
D JASQ-VAE 2 log o² + (....l......l)
a fo(Z1.1))1
| 2 2
Σ
L
H(P: 211)
(14)
20
2
2Σ 87
1=1
where the third term is different from the one in Equation quantization error H Hœ(x) - Ele Z with the doquantization process.
10 and its numerator evaluates only the overall
4.3.2 RSQ-VAE vs. RQ-VAE
RQ- VAE and RSQ-VA VAE both learn discrete representations in coarse-to-fine manner, but RQ-VAE uses deterministic RQ scheme to achieve the approximation in Equation 88 in which it is trained with the following objective function:
CRQ-VAE = 1-- - (Z +
H (x) sg Σ Zμ
(15)

---

Published in Transactions on Machine Learning Research (03/2024)
RMSE
LPIPS
SSIM
SQ-VAE2 - VQ-VAB-2
0.03
0.01
0.7
16
64
128
16
64
125
8
32
64
128
Codebook
Codebook
Codchook size
(a) CIFARIO
0.24
0155
0.52
32
64
256
512
32
64
128
256
512
32
128
256
512
Codebooks size
Codchonk size
Codebook size
(b) CelebA-HQ
Figure 2: Impact of codebook capacity reconstruction of images in (a) CIEAR1O and (b) CelebA-HQ. Two layers are lested CIFARIO, three on CelebA-HQ.
Table 1: Evaluation on ImageNett (256x 256) and FFHQ (1024x1024) RMSE (x10²), LPIPS, and SSIM are evaluated using the test set. Following Razavi et al. 2019)) the codebook capacity for the discrete latent space is set to (di Ki) (32 512), (64 512) and (di, Ki) = (322, 512). (642, 512), (128?, 512) for ImageNet and FFHQ, respectively. Codebook perplexity is also listed for each layer.
Reconstruction LPIPS4
Codebook perplexity exp(H(Q(Z,)))
Dataset
Model
RMSEL
SSIM
exp(H(Q(21))
exp(HIQ(Z.))
ImageNet
VQ-VAE-2
6.071 + 0.006
0.265 + 0.012
0.751 = 0.000
106.8 0.8
258.8 1.4
SQ-VAE-2
0.006
0.9
1.7
FFHQ
VQ-VAE-2
4.866 + 0.291
0.323 + 0.012
0.814
0.003
24.6 + 10.7
41.3 4 14.0
310.1 # 29.6
SQ-VAE-2
2.118 + 0.013
0.166 0.002
0.909 = 0.001
123.8 9.0
398.7 # 14.1
411.3 =
where the codebooks are updated with the EMA update the same manner as VQ-VAF. The second term of Equation 15 resembles the third term of Equation (10), which strongly enforces a certain degree of reconstruction even with only some of the information from the higher layers. RQ- VAE benefits from such a regularization term, which leads stable training. However, in ESQ-VAE, this regularization degrades the roconstruction performance. To deal with this problem, we instead use Equation (14) as the objective, which regularizes the latent representation by taking into account of accumulated information from all layers.
Remark. HQ-VAF has favorable properties similar to those of SQ-VAE. The training scheme significantly reduces the number of hyperparameters required for training VQ-VAI variants. First, the objective func- tions 6 and (14) do not include any hyperparameters such 3 in Equations (7) and (15), except for the introduction of a temperat ure parameter for the Gumbel-softnax to enable the reparameterization of L for the
the discrete latent variables (Jang et al.) 2017 Maddison et al. 2017). Please refer to Appendix approximation of the categorical distributions @(Z Z..r Moreover, our formulations do not require heuristic techniques such as stop-gradient, codebook reset, or EMA update. Additionally, HQ-VAE has several other empirical benefits over the VQ-VAE variants, such as the effect of self-anncaling for better codebook utilization and better rale-distortion (RD) trade-off, which will be verified in Section
5 Experiments
We comprehensively examine SQ-VAE-2 and RSQ-VAE and their applicability to generative modeling. In particular, we compare SQ-VAE-2 with VQ-VAE-2 and RAQ with RQ-VAE to see if our framework improves reconstruction performance relative to the baselines. We basically compare various latent capacities in order to evaluate our methods in a RD sense (Alemi et al. 2017 Williams et al. 2020). The error plots the RD curves in the figures below indicate standard deviations based on four runs with different training seeds. In addition, we test HQ-VAL VAE on an audio dataset see if it is applicable a different modality. Moreover, investigate the characteristics of the injected top-down and residual top-down layers through
8

---

Published in Transactions on Machine Learning Research (03/2024)
RSQ-VAE RSQ-VAE (codebook
RQ-VAF onde
RQ-VAE
RQ-VAE code rowd
Perplexity
shane)
shane)
01
0166
RSQ RSQ.VAR VAE
RMSE
0006
SSIM
0.04
0111
0.8
32
[-1
128
RO.VAE
Codebock $120
Codeboolc $120
Codebook
(a.
CIFARIO (L 4)
0.06
0.5
[ 32
0.04
0.5
w. code resett
RMSE
0.00
1,71PS
SSIM
RQ-VAE
0.5
Number Layers
Number Layers
Layers
(c) Codebook perplexity,
(b) CelebA-HQ (K, 32)
Figure 3: Impact of codebook capacity reconstructions of images from (a) CIFAR10 and (b) CelebA-HQ (c) Codebook perplexity at each layer is plotted, where models with 32 layers are trained on CelebA-HQ and all layers share the same codebook.
visualizations. Lastly, we apply HQ-VAFs to generative tasks to demonstrate their feasibility as feature extractors. Unless otherwise noted in what follows, we use the same network architecture in all models and set the codebook dimension d, = 64. Following the common dalaset splits, we utilize the validation sets to adjust the learning rate during training and compute all the numerical metrics on the test sets. The details of these experiments are in Appendix
5.1 SQ-VAE-2 vs. VQ-VAE-2
We compare our SQ-VAE-2 with VQ-VAE-2 from the aspects of reconstruction accuracy and codebook utilization. First, investigate their performance on CIFAR10 (Krizhevsky et al. 2009) and CelebA-HQ (256x256) under various codebook settings, i.e. different configurations for the hierarchical structure and numbers of code vectors (K)). We evaluate the reconstruction accuracy in terms of a Euclidean metric and two perceptual metrics: the root mean squared error (RMSE), structure similarity index (SSIM) (Wang et al. 2004], and learned perceptual image patch similarity (LPIPS) (Zhang et al., 2018). As shown in Figure 2 SQ-VAs VAE-2 achieves higher reconstruction accuracy in all cases. The difference in performance between the two models is especially noticeable when the codebook size is small.
Comparison on large-scale datasets. Next, we train SQ- SAE-2 and VQ-VAE-2 on ImageNet (256x256) (Deng et al. 2009) and FFHQ (1024x1024) (Karras et al. 2019) with the same latent set- tings as in Razavi et al. (2019). As shown Table SQ-VAE-2 achieves better reconstruction performance in terms RMSE, LPIPS, and SSIM than VQ- VAE-2, which shows similar tendencies in the comparisons on CIFAR10 and Cel Furthermore, we mcasure codebook utilization per layer by using the perplexity of the latent variables. The codebook perplexity is defined exp(H(Q (Z))), where Q(Zi) is marginalized distribution of Equation 3 with æ (=)) The perplexity ranges from to the number of code vectors (Ki) by definition. SQ- VAE-2 has higher codebook perplexities than VQ-VAE-2 at all layers. Here, VQ-VAE-2 did not effectively use the higher layers. In particular, the perplexity values at its top layer are extremoly low, which is a sign of layer collapse.
5.2 RSQ-VAE RQ-VAE
9

---

Published in Transactions on Machine Learning Research (03/2024)
Table 2: Evaluation on UrbanSoundSK. RMSE evaluated on the test set. The network architecture follows
the one described in Liu al. (2021. Codebook size is set K 8.
Model
Number of Layers
RMSE I
RQ- VAE
0.506 + 0.018
0.497 + 0.057
RSQ. VAE
0.427 = 0.014
0.314 + 0.013
fs(2,)
( (()))))))) Zz))
32/15,,0
fo fo(2.)
fo - -aa
(b) H (2 (.)) in SQ-VAE-2
(a) Reconstructed images and magnified differences of SQ- 3Q-VAE-2
1.0
11 (2t. (2)
5000
10000
2500
S10R1 7510
(d) (P. Z+)) in RSQ-VAE
(c) Reconstructed images and magnified differences of RSQ-VAE
Figure 4: Reconstructed samples with partial layers in (a) SQ-VAE-2 and (c) RSQ-VAE. The top row shows the components added at each layer. (di, K2) (162, 256), (322, 16), (642,4) and Kg) (32 4), (32 16), (32 256) for latent capacities, of 1, 2, and 3, respectively. Notice that the numbers of bits in these models are equal at each layer. For a reasonable visualization, we apply progressive coding, which induces progressive compression, to VQE-2 (see Appendix F.6p. (b) and (d) Variance parameter sf normalized by the initial value $ and average entropy of the quantization process (H,,)(())))2 at cach layer.
shows the reconstructed images, while the bottom
We compare VAQ with RQ-VAE the metrics as in Section 5.1 the original study of RQ- VAE [Zeghidour et al.] 2021 Lee al. 2022a to prevent codebook collapse, we add a RQ- VAE incorporating this technique to the baselines. We do not apply it to RSQ- VAE because it is not explainable in the variational Bayes framework. In addition, Lee et al. (2022a) proposed that all the layers share the codebook, i.c., B = B for I E [L], to enhance the utility of the codes. We thus test both RSQ-VAB, and RQ- VAE with and without codebook sharing. First, we investigate their performances on CIFAR1O and Hel (256x256) by varying the number of quantization steps (I) and number of code vectors (K)). As shown in Figures 3a and 3b. RSQ-VAE achieves higher reconstruction accuracy in terms of RMSE, SSIM, and LPIPS compared with the baselines, although the codebook reset overall improves the performance of RQ-VAF overall. The performance difference was remarkable when the codebook is shared by all layers. Moreover, there is noticeable difference in how codes are used; more codes are assigned to the bottom layers in RSQ-VAE than to those in the RQ-VAEs (see Figure 3c RSQ-VAB captures the coarse information with relatively small number of codes and refines the reconstruction by allocating more bits the bottom layers.
As codebook reset is used in
Validation on an audio dataset. We validate RSQ-VAE in the audio domain by comparing it with RQ- VAE in an experiment on reconstructing the normalized log-Mel spectrogram in an environmental sound datasot, UrbanSonn8K (Salamon et al.) 2014). We use the same network architecture as an audio generation paper (Liu et al.) (2021), in which multi-scale convolutional layers of varying kernel sizes are deployed to capture the local and global features of audio signals in the timeeffryyency domain (iian et al.
10

---

Published in Transactions on Machine Learning Research (03/2024)
SQ-VAE22
RSQ-VAF
RMSE
LPIPS
0,02
0.12
0,8
981,9,3)
(16,4,3)
(16,4,2
4,3)
9.))
(255,
(2,41 16)
(3,981)
(3,9,
(2,4,
(3.9,81)
(4. 16,16, 256)
Codebook sizes (4G,, KG]
Codebook sizes (As, AG)
Codebook sizes AG.
Figure 5: Comparison of SQ-VAE-2 and RSQ-VAE in three latent-capacity cases. The x-axis in blue (red) is (K1, K2, Kş) for SQ-VAE-2 (RSQ-VAE) Note that (d1, dg) = (16, 32, 64) for SQ-V and (d1, dz) (32, 32, 32) for RSQ-VAF, and the values at the same point on the x-axis indicate the same latent capacity. SQ-VAE-2 outperforms RSQ-VAE at relatively large latent capacities. In contrast, RSQ-VAB achieves better reconstruction performance at higher compression rates.
2021). The codebook size is set to K₁- 8. The number of layers is set to 4 and 8, and all the layers share the same codebook. We run each trial with five different random seeds and obtain the average and standard deviation of the RMSFs. As shown in Table 2 RSQ-VAE has better average RMSEs than RQ-VAF has across different numbers of layers on the audio datasct. We also evaluate the results in terms of perceptual quality by performing subjective listening tests, which are described in Appendix F.2.3
5.3 Empirical study of top-down layers
In this section, we focus on visualizing the obtained discrete representations. This will provide insights into the characteristics of the top-down layers. We train SQ-VAE-2 and RSQ- each with three layers, on CelebA-HQ (Karras et al. 2018). Figure shows the progressively reconstructed images. For demonstration purposes, we incorporate progressive coding (Shu & Ermon 2022) in SQ-VAE-2 to make the reconstructed images only with the top layers interpretable. Note that progressive coding is not applied to cases other than those illustrated in Figure | SQ-V and IRSQ-VAE share a similarity in that the higher layers generate the coarse part of the image while the lower layers complement them with details. However, upon examining Figure 4 we can see that, in the case of SQ-VAE-2, the additionally generated components (bottom row in Figure 4ah in each layer have different resolutions. We conjecture that the different
resolutions H (1) (x), which are injected into the top-down layers, contain different information. This implics that we may obtain more interpretable discrete representations if we can explicitly manipulate the extracted features in the bottom-up path to provide (²"()(2) (x),, giving them more semantic meaning (e.g., texture or color). In contrast, RAE seems to obtain different discrete representation, which resembles more decomposition. This might be due to its approximated expansion in Equation { Moreover, we can see from Figures 4b and 4d that the top-down layers also benefit from the self- effect.
In Appendix we explore the idea of combining the two layers to form a hybrid model. There it is shown that the individual layers in a hybrid model produce effects similar to using them alone. That is, the outputs from the injected top-down layers have better resolution and residual top-down layers make refinements upon certain decompositions. Since these two layers enjoy distinct refining mechanisms, bybrid model may bring more flexible approximation the posterior distribution.
3
Additionally, we compare the reconstruction performances of SQ-VAE-2 and RSQ-VAE with the same ar- chitecture for the bollom-up and lop-down paths. The experimental conditions are the same as in the visualization of Figure 4 We examine three different compression rates by changing the number of code vec- tors K, Interestingly, Tigure]Jshow that SQ-VAR SQ-VAF-2 achieves better reconstruction performance the case of the lower compression rate, whereas RSQ-VAL VAE reconstructs the original images better than SQ- VAE-2 in the case of the higher compression rate.
11

---

Published in Transactions on Machine Learning Research (03/2024)
Table 3: Image generation on FFHQ. i and denote the use of RQQ-rrasformmr and con- textual RQ-Transsoor Transformer as prior model.
Table 4: Image generation on ImageNet.
Model
FIDJ
VQ- VAE-2
Razavi el
2019
31
Model
FID4
DALL-E
Ramnsh et al.
2021
32.01
VQ-GAN
Esser et al. 2021b)
11.4
VQ-GAN
Esser et
2021b
15.78
RQ-VAE
Lee et al. 2022a)
10.38
RQ-VAE
Lee et al.
2022a
7.55
RSQ- VAE
(ours)
9.74
HQ-TVAE
You et al.
2022
7.15
RSQ-VAE
(ours)
8.46
Contextual
RQ Transtormer
Lee et al.
2022b
3.41
SQ- SQ-VAE-2 (ours)
4.51
5.4 Applications of HQ-VAEs to generative tasks
Lastly, we conduct experiments in the vision domain to demonstrate the applicability of HQ- VAE to realistic generative tasks.
First. we train RSQ-VAE on FFHQ (256x256) by using the same encoder-decoder architecture as in Lee
et al. 2022a). We set the latent capacities to L 4, (di, K₁) = (8*, 256) for I E [4] by following their RQ-
VAF. We train prior models, an RQ-Transformer (Tee et al. 2022a) and a contextual RQ-Transformmr (Lee
et al. 2022b), on the latent feature extracted by RAE. Table lists the Fréchet inception distance (FID)
scores for the generated images (refer to Table for more detailed comparisons). According to the table, our RSQ-VAE leads to comparable generation performance to RQ-VAE, even without adversarial training. Figure 16 shows the generated samples from our model on FFHQ.
Next, we train SQ-V on ImageNet (256x256) by using a simple resblock-based sncoder-decodee ar- chitecture. We define two levels of discrete latent spaces whose capacities are (d1, K2) = (162, 2048) and (d2, K2) (32², 1024). For training SQ-VAE-2 training, we use the adversarial training framework achieve perceptual reconstruction with feasible compression ratio. Subsequently, to train the prior model on the hierarchical discrete latent features. As reported in Table our SQ- VAE- 2-based models achieve FID score that is competitive with current state-of-the-art models (refer to Table g for more detailed comparisons). Figure 17 shows the generated images from our model on ImageNet.
Muse (Chang et al. 2023)
6 Discussion
6.1 Conclusion
We proposed HA-VAE, general VAE approach that learns hierarchical discrete representations. HQ-VAE is formulated within the variational Bayes framework as stochastic quantization technique, which (1) greatly reduces the number of hyperparameters to be tuned, and (2) enhances codebook usage without any heuristics thanks to the aelf-annealing effect. We instantiated the general HQ- LLQ-VÆ with two types of posterior approximators for the discrete latent representations (SQ-VAE-2 and RSQ-VAE), These two variants have infomartion passing designs similar to those of as VQ-VAE-2 and RQ-VAE, respectively, but their latent representations are quantized stochastically. Our experiments show that SQ-VAF-2 and RSQ- VAE outperformed their individual baselines with better reconstruction and more efficient codebook usages in the image domain as well as the audio domain.
6.2 Concluding remarks
Replacing VQ-VAE-2 and RQE with SQ- VAE-2 and RSQ- VAE will yeild comparative improvements to many of the previous methods in terms of reconstruction accuracy and efficient codebook usage. SQ-VAE- and RSQ-VAE have better RD curves than the baselines have, which means that they achieve (i) better reconstruction performance for the same latent capacities, and (ii) comparable reconstruction performance at higher compression rate. Furthermore, our approach climinates the need for repetited tumings of many apper-parameters or the use of ad-hoc techniques. SQ SQ-VAE-2 and RSQ VAE are applicable to generative modeling with additional training of prior models as was done in numerous studies. Generally, compression
12

---

Published in Transactions on Machine Learning Research (03/2024)
models with better RD curves are more feasible for the prior models (Rombach et al. 2022 hence, the replacement of VQ-VAF-2 (RQ-VAF) with SQ-VAR-2 (RSQ-VAF) would be beneficial even for generative modeling. Recently, RQ- VAE has been used more often in generation tasks than VQ-VAE- wihch has a severe instability issue, i.., layer collapse (Dhariwal et al.) 2020). Here, we believe that SQ-VAE-2 has potential as a hierarchical model for generation tasks since mitigates the issue greatly.
SQ- SAE- and RSSQ-VAE have their own unique advantages follows. SQ-VAE-2 can learn multi-resolution discrete representations, thanks to the design of the bottom-up path with the pooling operators (see Fig- ure 4a. The results of this study imply that further semantic disentanglement of the discrete representation might be possible by including specific inductive architectural components in the bottom-up path. Morcover, SQ-VAE- out RSQ-VVE at lower compression rates (see Section 5.3) This hints at that SQ-VAE- might be appropriate especially for high-fidelity generation tasks when the prior model is large enough. contrast, one of the strengths of RSQ-VAF, VAE is that can easily accommodate different. compression rates simply by changing the number of layers during the inference (without changing or retraining the modcl). Furthermore, RSQ-VAE outperforms SQ-VAE-2 at higher compression rates (see Section 5.3) These prop-
erties would make RSQ-VAE suitable for application to neural codecs (Zeghidour et al. 2021 Défossez et al.
2022
Acknowledgement
We would like to thank Marc Ferras for many helpful comments during the discussion phase. Besides, we thank anonymous reviewers for their valuable suggestions and comments.
13

---

Published in Transactions on Machine Learning Research (03/2024)
References
Eirikur Agustsson, Fabian Mentzer, Michael Tschannen, Lukas Cavigelli, Radu Timofte, Luca Benini, and Luc V Gool. Soft-to-hard vector quantization for end-to-end lcarning compressible represcntations. In Proc. Advances in Newral Information Processing Systems (NeurIPS), 2017.
Alexander A Alemi, Ben Poole, Ian Fischer, Joshua V Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken ELBO. arXiv preprint or Xiv: 1et.111100044, 2017.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffussin models in discrete state-spaces. Tn Proc. Advances in Newral Information Processing Systems (NeurIPS), volume 34, pp. 17981-17993, 2021.
Shiyue Cao, Yueqin Yin, Lianghua Huang, Yu Liu, Xin Zhao, Deli Zhao, and Kaigi Huang. Efficient- vqgan: Towards high-resolution image generation with efficient vision transformers. In Proc. IEEE/CVF Conference on Compuler Vision and Pullern Recognilion (CVPR), pp. 7368 7377, 2023.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William Freeman. Maskgit: Masked generative image transformer. In Proc. IREE/CVF Conference on Computer Vision and Pallern Recognition (CVPR), pp- 11315 11325, 2022.
Huiwen Chang, Han Zhang, Jarred Barber, Aaron Maschinot, Jose Lezama, Lu Jiang, Ming-Hsuan Yang, Kevin Patrick Murphy, William T. Freeman, Michael Rubinstein, Yuanzhen Li, and Dilip Krishnan. Muse: Text-to-image generation via masked generative transformers. In Proc. International Conference on Ma- chine Learning (ICML), pp. 4055 4075, 2023.
Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. Pixelsnail: An improved autoregressive generative model. In Proc. International Conference on Machine Learning (ICML), pp. 864 872. PMLR, 2018.
Rewon Child. Very deep VAÉs generalize autoregressive models and can outperform them on images. In Proc. International Conference on Learning Represenlation (ICLR), 2021.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse trans- formers. arXiv preprinl Xix::11904, 10509, 2019.
Alexandre Défossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. High fidelity neural audio compression. arXiv preprinl Xic:2220 13438, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In Proc. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 248 255. leee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Proc. Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 8780 8794, 2021.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Πlya Sutskever. Juke- box: A generative model for music. arXiv preprint arXiv:2005.003/ 00341, 2020.
Patrick Esser, Robin Rombach, Andreas Blattmann, and Bjorn Ommer. Imagebart: Bidirectional context with multinomial diffusion for autoregressive image synthesis. In Proc. Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 3532, 2021a.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proc. IEEE/CVP Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12873-181111 12883, 2021b.
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In Proc. IEEE/CVF Conference on 2022.
Computer Vision and Pattern Recognition (CVPR), pp.
14

---

Published in Transactions on Machine Learning Research (03/2024)
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proc. Advonces in Newral Information Processing Systems (NewrIPS), pp. 0840-6851, 2020.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multi- nomial diffusion: Learning categorical distributions. In Proc. Advances in Newral Information Processing Systems (NewrIPS), volume 34, pp. 12454 12465, 2021.
Phillip Isola, Jun- Yan Zhu, Tinghui Zhou, and Alexci A Efros. Image-to-image translation with conditional adversarial networks. In Proc. IEEE Conference on Compuler Vision and Pallern Recognation (CVPR), pp. 1125 1134, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In Proc. International Conference on Learning Representation (ICLR), 2017.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super- resolution. Proc. European Conference on Computer Vision (ECCV), pp. 694 711, 2016.
Lukasz Kaiser, Samy Bengio, Aurko Roy, Ashish Vaswani, Niki Parmar, Jakob Uszkoreit, and Noam Shazeer. Fast decoding in sequence models using discrete latent variables. In Proc. International Conference Machine Learning (ICML), pp. 2390 2399, 2018.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. In Proc. International Conference on Learning Representation (ICLR), 2018.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proc. IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4401 4410, 2019.
Jungil Kong, Jachyeon Kim, and Jackyoung Bac. Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis. Advances in Neural Information Processing Systems, 33:17022 17033, 2020.
Felix Kreuk, Gabriel Synnaeve, Adam Polyak, Uriel Singer, Alexandre Défossez, Jade Copet, Devi Parikh, Yaniv Taigman, and Yossi Adi. Audiogen: Textually guided audio generation. orXiv preprinl arXiv:2209.1558 2022.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Doyup Lee, Chiheon Kim, Saehoon Kim, Minsu Cho, and Wook-Shin Han. Autoregressive image generation using residual quantization. In IEEE/CVF Conference on Computer Vision and Pattern Recognation (CVPR), pp. 11523 11532, 2022a.
Doyup Lee, Chihcon Kim, Sachoon Kim, Minsu Cho, and Wook-Shin Han. Draft-and-revise: Effective image generation with contextual ry-transformer. In Proc. Advances in Neural Information Processing Systems (NeurIPS), 2022b.
Xubo Liu, Turab Iqbal, Jinzheng Zhao, Qiushi Huang, Mark D Plumbley, and Wenwu Wang. Conditional sound generation using neural discrete time-frequency representation learning, In IEEE Int. Workshop Machine Learning for Signal Processing (MLSP), pp. 1 1-6, 2021.
Chris Maddison, Andriy Mnih, and Yee Why Teh. The concrete distribution: A continuous relaxation of discrete random variables. In Proc. International Conference on Learning Representation (ICLR), 2017.
Boris T Polyak and Anatoli B Juditsky. Acceleration of stochastic approximation by averaging. STAM Journal on Control and Optimization, 30(4):838-855 1992.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsca Voss, Alec Radford, Mark Chen, and Tlya Sutskever. Zero-shot text-to-ine to---ge- generation. In Proc. International Conference on Machine Learning (ICML), pp. 8821-8831, 2021.
15

---

Published in Transactions on Machine Learning Research (03/2024)
Ali Razavi, Aaron van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VQ- VQ-VAE-2. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pp. 14866-14877 2019.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proc. IEEE/CVF Conference on Computer Vision and Pallern Recognition (CVPR), pp. 10684 10695, 2022.
Aurko Roy, Ashish Vaswani, Arvind Neelakantan, and Niki Parmar. Theory and experiments on vector quantized autoencoders. orXiu preprint arXiu::805. 1063, 2018.
Justin Salamon, Christopher Jacoby, and Juan Pablo Bello. A dataset and taxonomy for urban sound rescarch. In ACM Int. Conf. on Multimedia (ACM MM). pp. 1041-1044, 2014.
Michael Schoeffler, Sarah Bartoschek, Rabian-Robert Stöter, Marlene Roess, Susanne Westphal, Bernd Edler, and Jürgen Herre. webmushra- a comprehensive framework for web- based listening tests. Journal of Open Research Software, 6(1), 2018
B Series. Method for the subjective assessment of intermediate quality level of audio systems. International Telecomemunication Union Radioocommemicaation Assembly, 2014.
Rui Shu and Stefano Ermon. Bit prioritization in variational autoencoders via progressive coding. In International Conference on Machine Learning (ICML), pp. 20141-20155 2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using noncquilibrium thermodynamics. In Proc. International Conference on Machine Learning (ICML), pp. 2256 2265. PMLR, 2015.
Casper Kaae Sonderby, Tapani Raiko, Lars Maaloe, Soren Kaae Senderby, and Ole Winther. Ladder varia- tional autoencoders. In Proc. Advances in Neural Information Processing Systems (NeurIPS), pp. 3738- 3746, 2016.
Casper Kaac Sønderby, Ben Poole, and Andriy Mnih. Continuous relaxation training of discrete latent variable image models. In Beysian DeepLearning workshop, NIPS, 2017.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Proc. International Conference on Learning Representation (ICLR), 2020.
Yuhta Takida, Wei-Hsiiang Liao, Chieh-Hsin Lai, Toshimitsu Uesaka, Shusuke Takahashi, and Yuki Mitsufuji. Preventing oversmoothing in VAE via generalized variance parameterivation. Neurocompuling, 509:137 156, 2022a.
Yuhta Takida, Takashi Shibuya, WeiHsiang Liao, Chieh-Hsin Lai, Junki Ohmura, Toshimitsu Uesaka, Naoki Murata, Takahashi Shusuke, Toshiyuki Kumakura, and Yuki Mitsufuji. SQ-VAE: Variational bayes on discrete representation with self-annealed stochastic quantization. In Proc. International Conference on Machine Learning (TCML), 2022b.
Lucas Theis, Wenzhe Shi, Andrew Cunningham, and Ferenc Huszár. Lossy image compression with com- prossive autocncoders. In Proc. International Conference Learning Representation (ICLR), 2017.
George Toderici, Sean M O'Malley, Sung Jin Hwang, Damien Vincent, David Minnen, Shumeet Baluja, Michele Covell, and Rahul Sukthankar. Variable rate image compression with recurrent neural networks. In Proc. International Conference on Learning Representation (ICLR), 2016.
Arash Vahdat and Jan Kautz. Nvac: deep hicrarchical variational autoencoder. In Proc. Advances in Neural Information Processing Systems (NewrIPS), volume 33, pp. 19667 19679, 2020.
Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukeuoglu. Pixel recurrent neural networks. In Proc. International Conference on Machine Learning (ICMT), pp. 1747-1756, 2016.
16

---

Published in Transactions on Machine Learning Research (03/2024)
Aäron van den Oord, Oriol Vinyals, and Koray Kavukeuoglu Neural discrete representation learning. In Proc. Advances in Newral Information Processing Systems (NeurIPS), pp. 6306 (316-,31, 2017.
Dezhao Wang, Wenhan Yang, Yueyu Hu, and Jiaying Liu. Neural data-dependent transform for learned image compression. In Proc. IEEE Conference Computer Vision and Pattern Recognition. (CVPR), pp-
17379-17388, 2022.
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility structural similarity. IEEE transactions on image processing. 13(4):600-612, 2004.
Will Williams, Sam Ringer, Tom Ash, John Hughes, David MacLeod, and Jamie Dougherty. Hierarchical quantized autoencoders. Xi preprint arXiv:2002.0811 2020.
Yang Xian, Yang Sun, Wenwu Wang, and Syed Mohsen Naqvi. Multi-scale residual convolutional encoder decoder with bidirectional long short-term memory for single channel speech enhancement. In Proc. 2021.
European Signal Process. Conf. (EUSIPCO),
Dongchao Yang, Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound:
Discrete diffusion model for text-to-sound generation. arXiv preprint arX:::207.00988, 2022.
Tackgeun You, Saehoon Kim, Chiheon Kim, Doyup Lee, and Bohyung Han. Locally hierarchical auto- regressive modeling for image generation. Proc. Advances in Neural Information Processing Systems (NeurIPS), volume 35, pp. 16360-16372, 2022.
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector pector-uantized image modeling with improved VQGAN. In Proc. International Conference on Leurning Representation (ICLR), 2022.
Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. SoundStream: An end-to-end neural audio codec. IEEE Trans. Audio, Speech, Lang. Process. 30:495 507, 2021.
Richard Zhang, Phillip Isola, Alexci A Efros, Eli Shechtman, and Oliver Wang. The unrcasonable effectivenoss of deep features as perceptual metric. In Proc. IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 586-595, 2018.
17

---

Published in Transactions on Machine Learning Research (03/2024)
Contents
A Literature of stochastic posterior modeling for discrete latent,
18
B From SQ-VAE to HQ-VAE
18
Derivations
19
C.1 SQ-VAE-2
C.2 RSQ- VAE
20
D Hybrid model
20
E Gumbel-softmax approximation
22
F Experimental details
22
SQ-VAE-2 vs VQ-VAF-2
22
RSQ-VAE vs RQ-VAE
24
3
Empirical study of top-down layers
28
.4
Applications of HA-VAEs to generative tasks
28
.
Perceptual loss for images
29
6
Progressive coding
30
A
Literature of stochastic posterior modeling for discrete latent
We briefly compare SQ-V.L QAE,, the foundation of HQ- VAE, with two other stochastic discrete models.
Sonderby ct al. 2017) proposed stochastic posterior modeling in the discrete latent sctting as Q(zi = by (x) x 2 by 1. followed by succeeding works (Roy et al. 2018 Williams et al. 2020 Although the approach resembles SQ- VAE in terms of stochastic quantization, they have non-trivial differences as follows. In SQ-
VAE, Takida et al. (202bb) introduced a trainable parameter, denoted as E.p, by considering a dequantization process the inverse operation of quantization. This leads to Q(z; = bş ()) x ] bx) Tx; -bx).
which imposes the self-annealing effect well allows the posterior modeling with more general distances.
Please refer to Takida et al. 2022bb for more details about the difference and some numerical examples.
The logits modeling (LM) in dVAE RRamesh et al. 2021 is another approach to model the discrete posterior distribution (Q(Zlx)) instead vector quantization. LM assumes an encoder that directly maps from x to the logits of Q(z, x). Therefore, it does not have raw explicit features Ż. On the other hand, encoded features Z are directly quantized. When the learnable variance parameter 2² is set (fixed) close to 0, conventional VQ is obtained. We used SQ instead of LM for the posterior modeling in HQ-VAE because SQ serve drop-in replacement of VQ and can be extended to hierarchical and residual quantizations similar to VQ.
SQ-VAE, the
B From SQ-VAE to HQ-VAE
Let B E [bx] (bx E (R))) be a trainable codebook and Z Bd be a discrete latent variable to represent the target data æ. An auxiliary variable, denoted Ž Rd used to connect æ and Z in the SQ framework. In a general VAE with discrete and auxiliary latent variables, the generative process is modeled (2 (Z) with prior distribution Z, 2) P(Z, Z). For variational inference, variational distribution,
18

---

Published in Transactions on Machine Learning Research (03/2024)
denoted Q(Z, Zlz), is used to approximate the posterior distribution P(Z, Z(x). The ELBO becomes
log log 2 pog)o())
DKI @(2 2 P(Z, .))
po(x)P(Z,2(z) [x)
log
Q Z|z)
log [) log (Z 2|x) P(Z, 2)
=Bo(zz2(a))
In the vanilla SQ-VAE (Takida et al. 2022b), the vectors composing Z are assumed to be independent of each other and the quantization process is applied to each vector such that P(Z, 2) I P(z))((ž(1z) and QZ 2|z)) - ((z E.). In contrast, HQ-VAB allows dependencies in Z to be modelled by {2 and Zi Bd (d El= di), and modcling dependencies
decomposing it into L groups, i.c., between {Z, by designing the approximated posterior 2 2 (x) decomposed ways. This paper provides instances in the form of top-down blocks to model the dependencies, which enables graphical models to constructed for (ær, (Z ,) in handy way (see Figure
This framework naturally covers the latent
structures in VQ-VAE-2 and -Q-VAE (see Sections 4.2.2 and 4.3.2
C Derivations
C.1
SQ-VAE-2
The ELBO of VQ--2E-2 is formulated by using Bayes' theorem:
log pa()) 2 lo (a: DKI(2(Z,:)) 11. x3) || P(Z Z1.1 (x))
po(x)P(Z2 ZZ
L.L. Z1.L(x))
= -Egiznd.2uu.le (1)
log Q(Z1.1, Z1.1 (2)
log 1(Z1:1. Z1.1 (0))
-Egizi,, Zmc
log IZ.I.)
P(ZI:L. Zz.L)
1. di
Pop 1sp(21.1) (1)
Pop(21,s12) (i)
P(z1.s) ]
= 2n.
(z Z1:L)
log
+ llog
Pop 21, Zi)
Zi)
1=1 i=1
L d
Pap 21, Z1)
HP (z 2i) log Ki ]]
= =E Zna
lo Z1:L)
Σ
log
(16)
Pap Pop [1,, 21)
Since the probabilistic parts are modeled as Gaussian distributions, the first and second terms can be calculated as
log pe (x) ZI.L) log (x; fo (o(
log log(2TT02)
x
-ff(x))
and
(17)
2
202
Pop (2,,11
IF Q 2..
Pop E2,,1 2x)
2.21
25
||
21, 12
25
21,4
de
-.Eo 2
21,1
(18)
2s
By substituting Equations 17 and instead of sampling it in practical implementation.
18 into Equation
we arrive at Equation
66
where we use 2,
2,
19

---

Published in Transactions on Machine Learning Research (03/2024)
(x)
Injected top-down layer
1y- + 1
1 Lp-1 +
Residual top-down layers
1= Lr
Figure 6: Top-down layers corresponding to the rth resolution of the hybrid model in Appendix D
C.2 RSQ-VAE
The ELBO of RRSQ-VAF is formulated by using Bayes' theorem:
log (2) lo) D IIL 2|x) ) || 1 2)x))
log (a Z I.I., Z x))
2,, (2))
QZZ:: x))
2 I 2 xx)
=Eo(ziL,
log (2: (1.L)
log
P(Z 2)
Pao (2, Pad (2.
Paç(z.. Zi)
log (x) (1.L)
log
log
P(z),))
1=1 i=1
log Pao Pao (2,
L
E,
log (arl (1:L) +
H(P,3 (21. Z1) log Ki
(19)
(2,
==1
Since the probabilistic parts modeled as Gaussian dist.ributions, the second term can be calculated
P= ZZ)
E 2(1)))
E 2(m
2g - -2:ll2+ +
|| - 2:11
2
87
2 87
1
db
-E E 2(a)
12
2il13
+
(20)
8?
2
14] where we use 2 ==H instead of sampling it in a practical implementation. The above the derivation of the ELBO objective the case of A = Appendix D extends this model to the general case with variable R, which is equivalent the hybrid model.
By substituting Equations 17 and 20 into 19
we find that Equation
D Hybrid model
Here, we describe the ELBO of hybrid model where the two lypes of lop-down layers are combinatorially used to build top-down path as in Figure 6 Some extra notation will be needed as follows: L, indicates the number of layers corresponding to the resolutions from the first to rth order; 4, == [L,-1+11 1, L, is
the set layers corresponding to the resolution r; and the output the encoding block the (L,-1 + 1)th Φ
layer denoted as Gr ( Φ (a:), Z.: In Figure 6 the quantized variables Zi, aim approximating the variable encoded at Ly-1
(H) ZT.l.)) 2 Σ 2i = Y,.
(21)
lef.
20

---

Published in Transactions on Machine Learning Research (03/2024)
On this basis, the Ith top-down layer quantizes the following information:
H (x) Φ
(1
2: G (r, Z1:1-1) = (x); -:bman--) Ei= Zp (1> 1).
G
11
0 (H* +
(22)
To derive the ELBO objective, we consider conditional distributions on (ZI.1 YIRR,, where Y, = Eier. Z. From the reproductive property of the Gaussian distribution, the contimuous latent variable converted from Yr via the stochastic dequantization processes, Zz,, follows Gaussian distribution:
(rri Pa. (ūri Ze,) = N N(10 Tr. Σ Z1,i.
Σ
(23)
lEf,
VEer
where &2. = -07))Ee, We will use the following prior distribution to derive the ELBO objective:
R dr IIII [II r=1 t=1 MEL.
P(E..))
P(Z Y.n)
Psę Psę (Br,, Ze.),
(24)
where d, = di for EE With the prior and posterior distributions, the ELBO of the hybrid model can be formulated by invoking Bayes' theorem:
(ogpo(ar) 2 log)o(()) DRI(Q(ZZ.... YL.))) || P(Z {)))))))
log Po(x)P(Z (Z L. 11
Q(Z 1())
Q Z YIRR x)
log (x) Z1:L) log
P 7
Y1.))
R d,
Psp (fr.) (2,.)
di
Poç(z,,1 Zi) P(z).)
= =BolZn,, )(lla)
Po Po x ZI:L)
Σ
log
log
P&7 Vr.i Zt,))
r=1
1=1
1=1
d,
P&p Ur.a Zt.))
=Eo(Zz,,
log log pe
(1.L)
log
Σ i=
H ((P.(<z1.)) Z()) log Ki
Pop Ur.l
(25)
where Y, = (H) 1I. Since we have modelled the dequantization process and the probabilistic decoder as Gaussians, by substituting their closed forms into the above equation, we find that
UHQ-VAE log 2
D
| ZI(Z1.1)11
R G H (
Z.:....) Eiee, Z
1 F
L H H(P (Z) 21))
Σ
+ Q2 1.
20
2 2iEr, 87
1=1
1=1
(26)
where we have used
Pep Tr.i IZe.)
1
E e(Z (i
=Ear Konla)
2 |r.s - yr.il1
+
Dr.i
Dr
Pap (Ø.. ((c.)
2 2 s7
2
2
s
+ d,
= - (onnal)
2 2 of 1l0r. - yr.sl12
(27)
Here, used Y, =Y, instead of sampling it in practical implementation.
21

---

Published in Transactions on Machine Learning Research (03/2024)
Table 5: Impact of decay factor r in the exponential schedule defined in Appendix D on reconstruction of images in CIFAR10. RMSE (x10²) evaluated using the test set. Too large docay factor deteriorates the reconstruction performance.
Model
r= x 10
r= 3 10
r= x 10
r=
x 10
r= 1 10
r= 3 10
SQ-VAE-2 5.223 1 0.016
5.219 1 0.017
5.264 1 0.016
5.356
0.016
7.234 1 0.021
9.083 1 0.019
RSQ-VAL VAE
6.544 + 0.018
6.149 + 0.023
5.697 + 0.022
5.579 + 0.015
5.802 + 0.020
6.439 + 0.019
E
Gumbel-softmax approximation
Since HQ- VAE has discrete latent space, our objective functions (6), 14), and (66 involve expectations
with respective to the categorical variables ZI (in Q(Z Z1)1 As a common method to make the
sampling process of the calegorical variables reparameterizable, we replace the categorical distribution (2) with Gumbel-softmax alternatives. In Gumbel-softmax distribution, temperature parameter (0, Do)
is introduced. The sampling process from Gumbel-softmax distribution is equivalent to that from
categorical distribution as T + 0.
There is a tradc-off with the temperature т. Lower temperatures lead to sampling that closely resembles the original categorical distribution but with larger gradient variances. Conversely, higher temperatures result in greater discrepancies between the categorical and Gumbel-softmax distributions but with smaller gradient variances. It has been shown that annealing the temperature from high to low values during model training
performs well in practice (Jang et al. 2017 Sonderby et al. 2017 Williams et al., 2020). Specifically, [aang exponential scheduling,
et al. (2017) proposed
max Cr, exp(eeppppp) (-))), where denotes the global training
step. Following the approach of Takida et al.] (2022b). we adopt the same annealing schedule with c, = 0 and r r= 10-6 for all the experiments (excopt for Table 5).
Here, we investigate the sensitivity of the choice of the decay factor r. We train two-layer SQ-VAF-2 and RSQ-VAL VAE with codebook capacities of (di, (42, 512), (82, 512) and (di, Ki) (8² 32), (8² 32), respectively. We sweep r while keeping the other settings in Table 5 According the table, a decay factor that is too large, leading to fast convergence to the categorical distribution, deteriorates the reconstruction performance.
F Experimental details
This appendix describes the details of the experiments discussed in Section 5 For all the experiments except for the ones of RSQ-VAE and RQ-VAE on FFHQ and UrbanSound8K in Section architectures for the bottom-up and top-down paths as is described in Figures and To build these paths,
5.2 we construct
we use two common blocks, the Resblock and Convblock by following Child 2021 in Figure 7a these blocks are shown in Figures 7b and 7c Here, we denote the width and height of (z) as w, and h,, respectively, i.e., H (zx) E R xw, X We set Cmid = 0.5 in Figure For all the experiments, we use the Adam optimizer with 31 = 0.9 and 83 0.9. Unless otherwise noted, we reduce the learning rate in half if the validation loss does not improve in the last three epochs.
In HQ-VAE, we deal with the decoder variance o² by using the update scheme with the maximum likelihood estimation (Takida et al.] (022a). We gradually reduce the temperature parameter of the Gumbel-softmax trick with standard schedule exp(10-5. (Jang al. 2017), where t is the iteration step.
We set the hyperparameters of VQ-VAE to the standard values: the balancing parameter & in Equations and 15) is set to 0.25, and the weight decay in EMA for the codebook update is set to 0.99.
Below, we sunmarize the datasets used in the experiments described in Section 5 below.
2he source code is attached in the supplementary material.
22

---

Published in Transactions on Machine Learning Research (03/2024)
Table Notation for the convolutional layers in Figure 7
Notation
Description
Conv (
(Ixl)
Convolutional layer (channel= n, kernel= 1, stride= padding= 0)
Conv (3x3)
Convolational layer (channel= na, kernel= x 3, stride= padding=
Conv (4x4)
2D Convolutional layer (channel= n, kernel= x 4, stride= 2, padding= 1)
ConvT ((
2D Transpose convolutional layer (channel= kemel= x 3, stride= 1, padding= 1)
ConyTy (4x4)
2D Transpose convolutional layer (channel= kernel= x 4, stride= 2, padding- 1)
CIFARIO. CIFAR1O (Krizhessyy et al. 2009) contains ten classes of 32x32 color images, which are separated into 50,000 and 10,000 samples for the training and test sets, respectively. We use the default split and further randomly select 10,000 samples from the training set to prepare the validation set.
CelebA-HQ. He (Karras et al., 2018) contains 30,000 high-resolution face images that are selected from the CelebA dataset by following Karras et al. (2018). We use the default training/valicee wwagg//aiiiiiionnoommmmm s/ooontttttt /test split (24,183/2,993/2,2m 22,2 samples). We preprocess the images by cropping and resizing them to x56x
FFHQ. FFHQ (Karras et al. 2019) contains 70,000 high-rcsolution face images. In Section 5. we split the images into three sets: training (60,000 samples), validation (5,000 samples), and test (5,000 samples) sets. We crop and resize them to 1024x 1024. In Section 5.2 we follow the same preprocessing as Lee let al. 2022a), wherein the images are split training (60,000 samples) and validation (10,000 samples) sets and they are cropped and resized them 256x256.
TetgeNet. ImageNet Deng et al. 2009 contains 1000 classes natural images in RGB scales. We use the default coiniing/vaaiaattinn /test split 1L228,,66/550,00//100,,, samples). We crop and resize the images to 256x256.
UrbanSound8K. UrbanSoundSK (Salamon et al. 2014) contains 8,732 labeled audio clips of urban sound in ten classes, such as dogs barking and drilling sounds. UrbanSound8K is divided into ten folds, and we use the folds 1-8/9/10 as the training/validationn validation/tes split. The duration of cach audio clip is less than 4 seconds. In our experiments, also convert the clips 16 bit and down-sampled them to 22,050 kHz. The 4-second waveform audio clip is converted to Mel spectrogram with shape 80 x 344. We preprocess cach audio clip by using the method described in the paper (Liu et al. 2021
align the lengths input audio, we pad all the audio clips to 4 seconds. We
We extract an 80-dimensional Mel spectrogram by using a short-time Fourier transform (STFT) with frame size of 1024, hop size of 256, and Hann window.
2. We apply dynamic range compression to the Mel spectrogram by first clipping it to minimum value of 1 x 10 and then applying a logarithmic transformation.
F.1 SQ-VAE-2 vs VQ-VAE-2
F.1.1
Comparison on CIFAR10 and CelebA-HQ
We construct the architecture depicted in Figures 1 and To build the top-down paths, we use two injected lop-down layers (i.e., R = 2) with wi = h, and We = lr = 16 for CIFARIO, and three layers (i.e., R = 3) with wi = h1 = 8, W2 = h2 = 16 and w3 = fg = 32 for CelebA-HQ. For the bottom-up paths, we repcatedly stack two Resblocks and an average pooling layer once and four times, respectively, for CIFAR10 and CelebA-HQ. We set the learning rate to 0.001 and train all the models for maximum of 100 epochs with mini-batch size of 32. The sensitivity of SQ-VAF-2 in terms of batch size is investigated in Figure Sa
23

---

Published in Transactions on Machine Learning Research (03/2024)
Cony 2
GeLU
Coor
GeLU
Cony GeIJU
Seal
GeLU
Cony Feaedh
GelUU
Co
Sealk
GeLU
Cony Feaedh
GelJU
Conv
Sealk
GeLU
Resblock
Convblock
(a) Basic blocks.
ConeT 1xll
RoLU
W logs
Reshlock
Reeblock
x2
Comvblock
Char Rell
Conv (600) ReLU
CoarTy Signooid
Resblock
Resblock
Pre-processing
Decodiug
Upsampling
Encoding block
(b) Pre- Pre-processing and decoding
(c) Blocks in top-down layer
Figure 7: Architecture details in Figure C (kxk) and ConvT((xk)
1 Table summarizes the notation for the convolutional layers,
F.1.2 Comparison on large-scale datasets
We construct the architecture depicted in Figures 1 and To build the top-down paths, we two injected top-down layers (i.e., R = 2), with wh = h1 = 32 and h2 64 for ImageNet, and three layers (i.e., R = 3) with W1 = h 32, W2 h2 = 64 and Wy h3 128 for FFHQ, respectively. For the bollom-up paths, we repeatedly stack two Resblocks and an average pooling layer three times and five times respectively for ImageNet and FFHQ. We set the learning rate to 0.0005. We train ImageNet and FFHQ for maximum of 50 and 200 epochs with a mini-batch size of 512 and 128, respectively. Figure 9 and Figure show reconstructed samples of SQ-VA SAEV2 on ImageNet and FFHQ.
F.2
RSQ-VAE VAE vs RQ-VAE
F.2.1
Comparison on CIFAR10 and CelebA-HQ
We construct the architecture as depicted Figures and without injected top-down layers, ie., R = We the resolution of H (z) to w average pooling layer once for CIFAR10 and four times for CelebA-HQ. We set the learning rate 0.001 and train all the models for maximum of 100 epochs with mini-batch size of 32. The sensitivity of ASQQ-VEE in terms of batch size is investigated in Figure 8b
For the bottom-up paths, we stack two Resblocks on the
F.2.2 Improvement in perceptual quality
This experiment use the same network architecture as in Lee et al. (2022a). We set the learning rate to 0.001 and train an RSQ-VAE model for a maximum 300 epochs with a mini-batch size of 128 (4 GPUs, samples for each GPU) on FFHQ. We use our modified LPIPS loss (see Appendix [.5) in the training. For the evaluation, we compute rFID scores with the code provided their repository on the validation set (10,000 samples). Moreover, use the pre-trained RQ-VAE model offered in the same repository for evaluating RQ-VAN VAE.
We will show examples of reconstructed images in Appendix F F.5 after we explain our modified LPIPS loss.
http:/ghhhh.cccm//aaaaaaann/g/// traneformer
24

---

Published in Transactions on Machine Learning Research (03/2024)
SQ-VAB-2
VQ-VAE2
0.00
6037
RMSE
LPISS
0.84
604
322
123
Batch
Butch
(a) SQ-VAE-2 and VQ-VAE-2 on CIFAR10
RSQ-VAE
RQ-VAE
0159
RMSE
LPPS
6:83
1687
Batch size
Batch size
Butch size
(b) RSQ-VAE and RQ-VAE CIFAR1O
Figure 8: Impact of batch size reconstruction of images in CIFAR10. (a) We set the codebook capacity for the discrete space to (di, K2) = (42, 512), (8², 512) with -= 2. (b) We set the codebook capacity for the discrete space to (d,, Ki) = (8² 32), (8², 32) with L = 2. For (a) and (b), we use the same architectures as in Sections 5.] and 5.2 respectively. The plots suggest that batch size should be tuned to achieve better reconstruction performance for all the models.
PIPES
aa) Source images
PIPES
(b) Reconstructed images
Figure 9: Reconstructed samples of SQ- VAE-2 trained on ImageNet
25

---

Published in Transactions on Machine Learning Research (03/2024)
(a) Source images
TOU
(b) Reconstructed images
Figure 10: Reconstructed samples of SQ- VAE-2 trained on FFHQ
26

---

Published in Transactions on Machine Learning Research (03/2024)
100
100
80
80
60
60
Raiing
40
40
20
Rerference
RSQ-VAE
RQ- VAE
Anchor moise)
Rerference (original)
RSQ-VAE
RQ-VAE
Anchor noisc)
(original)
(a) L =
(b) -
Figure 11: Violin plots of MUSHRA listeming test results on UrbanSound8K test set in cases of (a) four layers and (b) eight layers. The white dots indicate the median scores, and the tops and bottoms of the thick vertical lines indicate the first and third quartiles, respectively.
F.2.3 Validation audio dataset
We construct the architecture in accordance with the previous paper on audio generation (Liu et al. 2021 For the top-down paths, the architecture consists of several strided convolutional layers parallel (Xian) et al. 2021). We use four strided convolutional layers consisting of two sub-layers with stride 2, followed by two ResBlocks with ReLU activations. The kernel sizes of these four strided convolutional layers 2 4 x 4, x 6 and 8 x 8 respectively. We add the outputs of the four strided convolutional layers together and pass the result to convolutional layer with a kernel size of 3 x 3. Then, we get the resolution of HΦ(x) to w = 20, h = 86. For the bollom-up paths, we stack a convolutional layer with a kernel size 3 x 3, two Resblocks with ReLU activations, and two transposed convolutional layers with stride 2 and kernel size x 4. We set the learning rate to 0.001 and train the models for maximum of 100 epochs with a mini-batch size of 32.
To evaluate the perceptual quality of our results, we perform subjective listening test using the multiple stimulus hidden reference anchor (MUSHRA) protocol [Series] 2014] on an audio web evaluation tool (Schoef- fler et al. 2018). We randomly select an audio signal from the UrbanSound8K test set for the practicing part. In the test part, we extract ten samples by randomly selecting an audio signal per class from the test set. We prepare four samples for each signal: reconstructed samples from RQ-VAE and RSQ- a white noise signal as a hidden anchor, and an original sample as hidden reference. Because RQ- VAE and RSQ VAE are applied the normalized log- log-Mel spectrograms,
use the same HiFi-GAN vocoder (Kong et al. 2020)) used in the audio generation paper (Liu et al., 2021 to convert the reconstructed spectrograms to the waveform samples. The HiFi-GAN vocoder is trained the training set of UrbanSoundKK from scratch (Liu et al. 2021). As the upper bound of quality of the reconstructed waveform is limited to the vocoder result of the original spectrogram, we use the vocoder result for the reference. After listening to the reference, assessors are asked to rate the four different samples according to their similarity to the reference on scale of to 100. The post-procsssing of the assessors followed the paper (Series) 2014): assessors were excluded from the aggregated scores if they rated the hidden reference for more than 15% of the test signals with a score lower than 90. After the post-screening of assessors (Series) 2014),, total of ten assessors participated in the test. Figure [I shows the violin plots of the listening test. RSQ- VAE achieves better median listening scores comapared with RQ-VAE. The scores of RSQ-VAE, are better by large margin especially in the case of eight layers.
As a demonstration, we randomly select audio clips from our test split of UrbanSound8K, and show their reconstructed Mel spectrogram samples from RQ-VAE and RSQ-VAE in Figure [3] While the samples from RQ- VAE have difficulty in reconstructing the sources with shared codebooks, the samples from RSQ-VAL VAE reconstruct the detailed features of the sources.
27

---

Published in Transactions on Machine Learning Research (03/2024)
Botton-up path
Top-down path
K1 =
Bottom-ap path
Top-down path
Hg(m)
K2 =
K; 16
K3
K4 = 4
K2 16
K₃ = 4
K₃ = 1
Hg(z)
K4=4
Hgla)
Kg
K; m
K7
(x)
K, =
H)(a)
(a) Casel: 61 = (1,2), 4 13, 4} and la = (5,6)
(b) Case2: 4 = 21,,,,,4), 15, 6} and la (7)
Figure 12: Architecture of the hybrid model in Appendix F.3
Table Image generation on FFHQ (a detailed version Table 3
Model
Lasent size (di)
Codebook (Ki)
GAN
rFID]
FTDJ
VQ-GAN
Eaver
2021b
16"
1024
11.4
RO-VAE
et nl. 2022m
g² x
2048 x
7.29
10.38
RSQQ VAE + RQ Transformer (ours)
RSQ-VAR VAE
Contextual NQ-Transformer (ours)
8 x
2048 x 4
8.47
8.46
9.74
F.3 Empirical study of top-down layers
As demonstration, we build two HQ-VAEs by combinatorially using both the injected top-down and the
residual lop-down layers with three resolutions, W1 = h1 16, W2 lt2 = 32 and W3 hy 64. We
construct the architectures as described in Figure 12 and train them on CelebA-HQ. Figure 15 shows the
progressively reconstructed images for each case. We can see the same tendencies as in Figures 4a and 4c
F.4 Applications of HQ-VAEs to generative tasks
In training RSQ-VAE on FFHQ, we borrow the encoder decoder architecture from Lee al. 2022a) which is constructed by adding encoder and decoder block to VQ-GAN (Esser et al. 2021b) to decrease the resolution of the feature map by half. Table compares models, where it can be seem that our model does not rely on the adversarial training framework but is still competitive with RQ-VAE, which is trained with PatchGAN (Isola et al. (017). If we evaluate whole pipelines including the VAE and prior model parts as generative models, our models outperform the baseline model. Figure 16 shows the samples generated from RSQ-VAF, with the contextual Transformer.
model with other VQ-- generative
The encoder decoder architecture for our SQ-VAE-2 on ImageNet is built in the same manner as is shown in Figures I and 7 Our architecture structure most resembles that of VQ- -Q--2-22 because the latent structure is the same. Because ImageNet with resolution of 256 is a challenging dataset to compress to discrete representations with feasible latent sizes, we use adversarial training for reconstructed images. Subscquently, we train two Muse (Chang et al. 2023) prior models to approximate Q(Zi|c) and Q(Z2) where c E [1000] indicates the class label. Table g compares our model with other VQ-based generative models without the technique of rejection sampling terms of various metrics. As summarized in the table, our model competitive with the current state-of-theean models, Figure 17 shows the generated samples from SQ- SQ-VAF-2 with the Muse.
28

---

Published in Transactions on Machine Learning Research (03/2024)
Table 8: Image generation on ImageNet (a detailed version of Table
Latent
endles (Ki)
Model
Params
(d.)
GAN
FFIDJ
FID.
1S1
VQ-VAE-2
Razavi et al.
2019
13.5B
512, 512
31
45
DALL-R Ramesh
32a
8192
32.01
VQ-GAN Esser al.
bb
1.4B
16a
16384
4.98
16.78
78.3
VQ-Diffusion (Gu
al.
2022
370M
16a
16384
4.98
11.89
MaskGIT Chang of mi.
2022
228M
1024
2.28
6.18
182.1
VIT-VQGAT ot
2022
322
8192
1.99
4.17
RQ-VAE Lee
2022a
1.1B
8
3.20
8.71
119.0
RQ-VAE Lee
20222
3.8B
7.55
134.0
Contextual
Translormer
Lee
2022b
8
3.41
HQ-TVAE You all
2022
1.4B
8 16
8192, 8192
2.61
7.15
Efficient V VQGAN
CCao et a.
2023
162
1024
2.34
9.92
82.2
Efficient- VQGAN
Cão et a.
2023
N/A
322
1024
0.95
SQ-VAE-2 (ours)
-0556B
16".32*
20488, 1021
1.73
4.51
276.81
F.5
Perceptual loss for images
We found that the LPIPS loss (Zhang et al. 2018). which perceptual loss for images (Johnson et al. 2016), works well with our VQ-. However, also noticed that simply replacing || fo (Z. ( the objective function of HQ- VAE (Equations (6 and (14) with an LPIPS loss LLPIPS(x, fe (1. leads to artifacts appearing in the generated images. We hypothesize that these artifacts are caused by the max- pooling layers in the VGGNet used in LPIPS. Signals from VGGNet might not reach all of the pixels in backpropagation due to the max-pooling layers. To mitigate this issue, we applied padding-and-trimmiin operation to both the generated image fo(Z1.1) and the corresponding reference image æ before the LPIPS loss function. That LLPIPS (pt [x: pt pt (Z.2))), where pt [ denotes our padding-ane and-trimming operator. The Py Torch implementation of this operation is described below.
import random import torch
import torch.nn nn. functiiional
def padding and trimming(
x_rec, # decoder output reference image
)
H, x.size()
x_rec Fxpad(xx_ree rec, (15, 15, 15, 15), mode='replicate')
x = F. pad(x, (15, 15, 15, 15), mode=' rodee'rrplicate")
- - H_pad, W_pad = sizee
top random. randrangr (0, 16)
bottom H_pad random. ran 16)
left random. ran.rammeeddnngeee 16)
right N_pad random. rradraangeddd (0, 16)
r_rec F. r.interpollate(__reell::: :, top:bottom, left:right],
size=(H, W), mode=/bicubic', aligg_corneee=Faaae)) x F. P.intorpolate (x[;, i, top: left:right], size (H, W), mode=/bicubie bicubic', aliga_cornees==allee
return x_rec,
Note that our padding-and- toimmiiniminggggg operation includes downsampling with a random ratio. We assume that this random downsampling provides a generative model with diversified signals in backpropagation across training iterations, which makes the model more generalizable.
29

---

Published in Transactions on Machine Learning Research (03/2024)
shows images reconstructed by RSQ-VAE model trained with normal LPIPS loss, CLPIPS < fo 21:)), and ones reconstructed by RSQ-VAE model trained with our modified LPIPS loss, CLPIPS (pt [<)] pt (Z.L)ll))))) As shown, our padding-ande and-trimming technique alleviates the artifacts issue. For example, vertical line noise can be seen in the hairs in the images generated by the former model, but those lines are removed from or softened the images generated by the latter model. Indeed, our technique improves rFID from 10.07 to 8.47.
Figure 14
F.6 Progressive coding
For demonstration purposes Figure 4a we incorporate the concept of progressive coding (Ho et al. 2020 Shu & Ermon 2022) into our framework, which helps hierarchical models to be more sophisticated progressive lossy compression and may lead to high-fidelity samples being generated. Here, one can train SQ-VAE-2 to achieve progressive lossy compression (as in Figure 4a by introducing additional generative processes Zt N@r; fo (Z),), for I E [4]. We here derive the corresponding ELBO objective with this concept. A more reasonable reconstruction can be obtained from only the higher layers (i.e., using only low-resolution information (x))
We consider corrupted data क for E [L], which can be obtained by adding noise, for example, ã, = x + El. We here adopt the Gaussian distribution éLd ~ NOO ey) for the noises. Note that {oi is set to be a non-increasing sequence. We model the generative process by using only the top I groups, as pa( fo I) The ELBO is obtained as
D log oî + (o(
| fo( DVi 20
1 Zi
prog
Σ
H(P:((Z)
VAE-2
2x
(28)
In Section 5.3 we simply set 1i = 0 in the above objective when this technique is activated.
This concept can be also applied the hybrid model derived in Appendix D by considering additional generative processes (b(#,,) (ā fowfoooZoooo o. The ELBO objective is as follows:
R Σ
D 2
JHQ-AE
Jprog
log o²
R Y, Eier, Zi 2
R
x
(Z.LL..) + Du, 12
L
+ 2o.... Paake)
H(P:
(11)
(29)
20
2 ZiER, of
In Section F33
we simply set 2) = in the above objective when this technique activated.
30

---

Published in Transactions on Machine Learning Research (03/2024)
Mell fiter bank mndin
ER]
(a) Source.
Mail litter bank inder
index
2
(b) RO- VAF with 4 layers.
Mail litter baikk moder
iodex
RQ-VAE with layers.
Mei fitter baink mndie
iodox
(d) RSQ-RSVAA VAE with layers.
Matl fiterr baank indiex
indox
(e) RAQ with layers.
Figure 13: Mel spectrogram of (a) sources and (b)-(c) reconstructed samples of the UrbanSoundSK dataset. The left and right panels are audio clips of dog barking and drilling, respectively. We can that the RQ- VAEs struggle to reconstruct the sources with shared codebooks. In contrast, the reconstruction of RSQ-VAE reflects the details of the source samples.
31

---

Published in Transactions on Machine Learning Research (03/2024)
(a) Source
(b) RSQ-VAE trained with normal LPIPS loss (rFII== 10.07)
(c) RSQ-VAL VAE trained with our improved LPIPS loss (rFID= 8.47)
Figure 14: Reconstructed samples of FFHQ.
32

---

Published in Transactions on Machine Learning Research (03/2024)
Ie(3.)
Bal)
To(2.a) So(2)) lo( So(2.a) Jo(2.,)) fo(bss)
1o( Io( JolZ.a) fo(2)s)
Io(Ze..) -JolZs,
Iol (2(
1
2
(a) Casel: & = (1,2), l2 = 13, 43 and Es (5.6)
fo(2,)
fo(2.n)
Sa(2
1a( fo(2,)
fo(Zus)
fo(Zza)
Zo(Zu.)) fo(Zaa)
fo(
fo(Zaa)
folZis) folZus)
fe(2.n)
fo(Z.a)
fo( fz(
fo Te. -oa
r=1
r == 2
r = 3
(b) Case2: 4 (1,2,3,4), 62 (5,6) and la (7)
Figure 15: Reconstructed images and magnisied differences of HQ- HQ-VAE CelebA-HQ
Figure 16: Samples of FFHQ from RSQ-VAE with contextual RQ- Transforrrrmer LLc et al. 2022b
33

---

Published in Transactions on Machine Learning Research (03/2024)
Zo
1112)
(a) The used class labels are "0: lench, Tinca Linca", "1: goldfish", and "7: cock".
(Z))
(1(2)
(b) The used class labels are "90: lorikeet", "130: flamingo", and "282 liger cat".
(z))
fo(21.2)
(c) The used class labels are "604: hourglass", "724: pirate, pirate ship", and "751: racer, race racing car".
(1)
f0(21.2)
(d) The used class labels are "850: teddy, teddy bear", "933: cheeseburger", and "980: volcano".
Figure 17: Samples of ImageNet from SQ-VAE-2 with Muse Chang et al. 2023 (Top) Samples generated only with the top layer, fe Zi). (Bottom) Samples generated with the top and bottom layers, fo (Z))).
34


Reviews:
[{"review_id": "5IwYaYmN2O", "reviewer": "Reviewer_bvQe", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is generally well-written\n\n- The experiments are sufficiently rich to demonstrate the efficiency of the proposed approach.\n\nWeaknesses\n- The novelty is somehow limited. This paper seems to be a combination of VQ-VAE-2 for learning hierarchical representations and SQ-VAE for a stochastic matching between the continuous latent code $\\tilde{Z}_i$ and discrete codeword vectors.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes Hierarchically Quantized VAE (HQ-VAE) for learning hierarchical latent representations. The main idea is to use the corresponding bottom-up representation and the top-down representations up to the layer $l-1$ to compute the top-down representation at the layer $l$. Conceptually, this paper is a combination of VQ-VAE-2 for learning hierarchical representations and SQ-VAE for a stochastic matching between the continuous latent code $\\tilde{Z}_i$ and discrete codeword vectors.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have some questions for the authors:\n\n1) What are differences between your approach to learn latent representations and VQ-VAE-2?\n\n2) Section 4.3 needs to be improved for more clarification. What is $H_\\phi(x)$? How can you do $\\sum_{i=1}^l Z_i$ because they have different dimensions?\n\n3) How to compute  the bottom-up path $H^{r(l)}_{\\phi}(x))$? Is it the same as Stage 1: Learning Hierarchical Latent Codes in VQ-VAE?\n\n4) The authors should show the generated images for Table 3.", "broader_impact_concerns": "There are no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vT6N2lQJIs", "reviewer": "Reviewer_MV41", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n\nThe authors have tackled an important and challenging problem. This problem has various applications in computer vision and audio processing and is of interest to the research community. The proposed method is novel overall and leads to several improvements compared to existing schemes. The top-down and bottom-up paths seem to help the model hierarchically learn good codebooks. The method is evaluated on several datasets and compared to leading baselines. \n\n\nWeaknesses\n\nSome aspects of the paper's presentation need improvement, and many parts are hard to follow. Specifically, the English should be improved, and some sections should be reorganized to enhance the paper's readability. I have provided examples for such required changes in the changes bullet. \n\nSome experimental details are missing; for example, is the RMSE evaluated on a left-out test set? If so, what size? \nAn ablation study is missing, as well as an evaluation of how sensitive the approach is with respect to different changes in parameters.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This study presents a new vector quantization method based on a variational autoencoder (VAE). The authors use a variational Bayes framework to develop the hierarchically quantized variational autoencoder (HQ-VAE). HQ-VAE can learn a hierarchical discrete representation is good for efficiently representing data. The authors demonstrate the merits of their approach using image and audio data.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Phrasing: P1 “In the original VQ-VAE…” could be rephrased to \n\nBased on the construction of \\cite{}, inputs are first encoded and…\n\nP2 “local and global information” these terms are not really defined at this point, so it’s hard to understand what they refer to \n\nThe organization of Section 2 could be improved if the subsections were fused and by adding sentences that describe the relation between the methods. Why is the title of dVAE phrased as SQ-VAE vs. dVAE?\n\nP3 “divide up” not sure this is proper phrasing\n\nPlease expand the explanation of the method in the caption of figure 1.\n\nP4 “we give” -> we introduce\n \nIn eq 9 you use l’ but the summation is over l\n\nSome commas are missing after equations, for example after Eq. 14.\n\nP7 Bad phrasing: “the regularization is too strong to regularize”\nAlso, the intuition here isn’t clear neither is the solution to this issue.\n\nP8 In the remark you mention that the method does not require any hyperparameters besides the gumbel softmax and point out eq 6 and 14. However, these hyperparameters do not appear there.\nAlso, isn’t \\beta a hyperparameter? How about the learning rate and batch size?\n\nThe approximation using the Gumbel softmax is not explained.", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Hyfopmv9Z5", "reviewer": "Reviewer_dLpy", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. This work showed how to extend SQ-VAE into a hierarchical architecture and detailed the design choices that made it work well. This is likely useful for practitioners who work with VQ-VAE type of models and seek an easier-to-use alternative. \n\n2. The paper did a pretty comprehensive evaluation comparing the effect of the SQ-VAE-style categorical posterior & training scheme v.s. the usual VQ-VAE learned quantization scheme, in several popular architecture variants. Although perhaps not surprising, the findings in the SQ-VAE paper (less codebook collapse, better code utilization, etc.) are found to continue to hold in the hierarchical architectures. \n\nWeakness:\nIt's unclear to me how the stochastic quantization in SQ-VAE differs from earlier methods like that of Sønderby et al. (2017) and dVAE, and what's the advantage of the proposed stochastic quantization if there is indeed a difference. In referring to the earlier methods of Sønderby et al. (2017) and dVAE in Sec 2.2, the paper claimed that \"such an index-domain modeling cannot incorporate the codebook geometry explicitly into the posterior modeling. In contrast, VQ/SQ-based methods allow us to model the posterior distribution with vector operations\". I'm not sure what it means to \"incorporate the codebook geometry explicitly\" and don't think the assessment is quite correct, because these earlier methods used something very similar to the stochastic quantization of SQ-VAE; specifically, \"For the probabilistic models we let $q(z|x) = Cat(f(x, M))$, where $f(x, M)$ calculates the logits for the categorical distribution using the distance between $v$ and the rows of $M$\" (Sønderby et al. (2017)); here $v = f_{enc}(x)$ is some vector representation of the input, and the distances between $v$ and codebook vectors go through a softmax to form the categorical posterior, just like in SQ-VAE. \n\n------- Sønderby, Casper Kaae, Ben Poole, and Andriy Mnih. \"Continuous relaxation training of discrete latent variable image models.\" Beysian DeepLearning workshop, NIPS. Vol. 201. 2017.\n\nLastly, I know the focus of this work is on VQ, but I'm curious if the authors have any comment on the potential of the scalar quantization approach, especially in light of this paper: [Finite Scalar Quantization: VQ-VAE Made Simple](https://arxiv.org/pdf/2309.15505.pdf); this is a fairly recent paper, so I don't expect any empirical comparisons, and the authors' comment on it won't affect my evaluation of the current submission.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an extension of the recently proposed SQ-VAE (Takida et al., 2022), a VAE with discrete latents learned through probablistic quantization, into a hierarchical architecture. Overall, the hierarchical VAE architecture is similar to that of ladder VAE with bidrectional inference, but uses the technique of SQ-VAE to learn categorically distributed latents with learned codebooks; two main variants are considered, based on VQ-VAE-2 and residual-quantized VAE. The experiments demonstrate enhanced codebook utilization and higher reconstruction quality on image and audio, and additionally better FID on image generation when pairing the learned latent tokens with a transformer prior.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "A clearer explanation of how the stochastic quantization in SQ-VAE differs from earlier methods like that of Sønderby et al. (2017) and dVAE, and why this makes a difference, ideally through a small numerical example rather than hand-wavy arguments.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper extends the vector quantization VAE to a hierarchical setting. 

The reviewers note that the novelty of the paper is limited but the all agree that the paper fills an important gap in the literature and the paper after revisions is of very good quality.
Decision: Accept (poster)