Vector Quantized Wasserstein Auto-Encoder
Authors: Long Tung Vuong
Trung Le
He Zhao
Chuanxia Zheng
Mehrtash Harandi
Jianfei Cai
Dinh Phung
Venue: Submitted to ICLR 2023 (2023)

Abstract: Learning deep discrete latent presentations offers a promise of better symbolic and summarized abstractions that are more useful to subsequent downstream tasks. Recent work on Vector Quantized Variational Auto-Encoder (VQ-VAE) has made substantial progress in this direction. However, this quantizes latent representations using the online k-means algorithm which suffers from poor initialization and non-stationary clusters. To strengthen the clustering quality for the latent representations, we propose Vector Quantized Wasserstein Auto-Encoder (VQ-WAE) intuitively developed based on the clustering viewpoint of Wasserstein (WS) distance. Specifically, we endow a discrete distribution over the codewords and learn a deterministic decoder that transports the codeword distribution to the data distribution via minimizing a WS distance between them. We develop further theories to connect it with the clustering viewpoint of WS distance, allowing us to have a better and more controllable clustering solution. Finally, we empirically evaluate our method on several well-known benchmarks, where it achieves better qualitative and quantitative performances than the baselines in terms of the codebook utilization and image reconstruction/generation.

Under review as a conference paper at ICLR 2023
VECTOR QUANTIZED WASSERSTEIN AUTO-ENCODER
Anonymous authors
Paper under double-bling blind review
ABSTRACT
Learning deep discrete latent presentations offers a promise of better symbolic and summarized abstractions that are more useful to subsequent downstream tasks. Recent work on Vector Quantized Variational Auto-Encoder (VQ-VAE) has made substantial progress in this direction. However, this quantizes latent representa- tions using the online k-means algorithm which suffers from poor initialization and non-stationary clusters. To strengthen the clustering quality for the latent rep- resentations, we propose Vector Quantized Wasserstein Auto-Encode (VQ-WAE) intuitively developed based on the clustering viewpoint of Wasserstein (WS) dis- tance. Specifically, we endow discrete distribution over the codewords and learn a deterministic decoder that transports the codeword distribution to the data distri- bution via minimizing WS distance between them. We develop further theories to connect it with the clustering viewpoint of WS distance, allowing us to have a better and more controllable clustering solution. Finally, we empirically evaluate our method on several well-known benchmarks, where it achieves better quali- tative and quantitative performances than the baselines in terms of the codebook utilization and image reconstructionngenerrtion
INTRODUCTION
Learning compact yet expressive representations from large-scale and high-dimensional unlabeled data is important and long-standing task in machine learning (Kingma & Welling 2013 Chen
let al. 20200 Chen & He 2021: Zoph et al. (020). Among many different kinds of methods, Vari-
ational Encoder (VAE) (Kingma & Welling 2013) and its variants (Tolstikhin al., 2017
Alemi et al. 2016: Higgins et al. 2016; Voloshynovskiy et al. 2019) have shown great success in
unsupervised representation learning. Although these continuous representation learning methods
have been applied successfully to various problems ranging from images (Pathak et al.) 2016 Good-
fellow et al. 2014 Kingma et al. 2016), video and audio (Reed et al. 2017 Oord et al. 2016 Kalchbrenner et al. 2017), in some contexts, input data more naturally modeled and encoded as
discrete symbols rather than continuous ones. For example, discrete representations are natural fit for complex reasoning, planning and predictive learning (Van Den Oord et al. (017)). This motivates the need of learning discrete representations, preserving the insightful characteristics of input data.
Vector Quantization Variational Auto-Encoder (VQ-VAR) (Van Den Oord et al. 2017) pioneer generative model, which successfully combines the VAE framework with discrete latent representa- tions. In particular, the vector quantized models learn a compact discrete representation using de- terministic encoder-decoder architecture in the first stage, and subsequently applied this highly com- pressed representation for various downstream tasks, examples including image generation (Esser et al. 2021), cross-modal translation (Kim et al. 2022), and image recognition (Yu et al. 2021 While VQ-VAE has been widely applied to representation learning many areas (Henter et al.
2018 Baevski et al. 2020, Razavi et al.) [2199] Kumar et al. 2019 Dieleman et al. [2018] Yan et al.
2021. J known to suffer from codebook collapse, which has a low codebook usage, i.e. most of embedded latent vectors are quantized to just few discrete codewords, while the other codewords are rarely used, or dead, due to the poor initialization of the codebook. reducing the information capacity of the bottleneck (Roy et al. 2018 Takida et al. 2022 Yu et al. (021).
To mitigate this issue, additional training heuristics were proposed, such as the exponential moving average (EMA) update (Van Den Oord et al. 2017 Razavi et al. (2019), soft expectation maximiza- tion (EM) update (Roy et al.[[0088) codebook reset (Dhariwal et al.) 2020) Williams et al.) 2020). Notably, soft expectation maximization (EM) update (Roy et al. 2018) connects the EMA update

---

Under review as a conference paper at ICLR 2023
with an EM algorithm and softens the EM algorithm with stochastic posterior. Codebook reset randomly reimitializes unused/low used codewords to one of the encoder outputs (Dhariwal et al, 2020) those near codewords of high usage Williaas et al. 20200. Takida et al. (2022) suspects that deterministic quantization is the cause of codebook collapse and extends the standard VAE with stochastic quantization and trainable posterior categorical distribution, showing that the annealing of the stochasticity of the quantization process significantly improves the codebook utilization.
Additionally, WS distance has been applied successfully to generative models and continuous rep- resentation learning (Arjovsky et al. 2017] Gulrajani et al., 2017; Tolstikhin et al. 2017) owing to its nice properties and rich theory. It is natural to ask: "Can we take advantages of intuitive prop- erties of the WS distance and its mature theory for learning highly compact yet expressive discrete representations?" Toward this question, in this paper, we develop solid theories by connecting the theory bodies and viewpoints of the WS distance, generative models, and deep discrete representa- tion learning. In particular, a) we first endow a discrete distribution over the codebook and propose learning a "deterministic decoder transporting the codeword to data distributions" via minimizing the WS distance between them; b) To devise a trainable algorithm, we develop Theorem 3.1 to equivalently turn the above WS minimization to forwardin the data to codeword distributions via minimizing a WS distance between "the latent representation and codeword distributions"; c) More interestingly, our Corollary 3.1 proves that when minimizing the WS distance between the latent representation and codeword distributions, the codewords tend to flexibly move to the cluster- ing centroids of the latent representations with control on the proportion of latent representations associated to centroid. We argue and empirically demonstrate that using the clustering viewpoint of WS distance learn the codewords, we can obtain more controllable and better centroids than using a simple k-means as in VQ- VQ-VAE (cf. Sections 3.1 and 5.2
applies the WS dis- tance to learn more controllable codebook, hence leading to an improvement in the codebook utilization. We conduct comprehensive experiments to demonstrate our key contributions by com- paring with VAE (Van Den Oord et al. (017) and SQ- VAE (Takida et al.) (022] (i.e., the recent work that can improve the codebook utilization). The experimental results show that our VQ-WAE can achieve better codebook utilization with higher codebook perplexity, hence leading to lower (compared with VQ-VAE) or comparable (compared with SQ- VAE) reconstruction error, with sig- nificantly lower reconstructed Fréchlet Inception Distance (FID) score (Heusel et al., 2017). Gener- ally, a better quantizer in the stage-1 1 can naturally contribute to stage-2 2 downstream tasks Yu et al.] 2 Zheng et al.. 2022). To further demonstrate this, conduct comprehensive experiments on Tour benchmark datasets for both unconditional and class-conditional generation tasks. The exper- imental results indicate that from the codebooks of our VQ-WAE, we can generate better images with lower FID scores.
Our method, called Vector Quantized Wasserstein Encoder
2 VECTOR QUANTIZED VARIATIONAL AUTO-ENCODER
Given training set ID = {II, C RRV VQ-VAE Van Den Oord et al. 2017 aims at learning a codebook which is formed by set of codewords C' [ck]Æ_] R KxD on the latent space Z E R., an encoder fe to map the data examples to the codewords, and a decoder fa (i.c., (x 2)) to reconstruct accurately the data examples from the codewords. Given data example x. the encoder fo (i.e., p (z| æ)) associates æ to the codeword fa (I) defined as
c argmin, d, (fe (z).ck)
where de is a metric on the latent space.
The objective function of VQ- is as follows:
Eyar= [de (fa (x)) x)) d= (sg (Se (z)) + 3d. (Sa (z) (C))],
where P_. 1. E 8z, is the empirical data distribution, sg specifies stop gradient, d, is cost metric on the data space, and B is set between 0.1 and 2.0 (Van Den Oord et al. 2017) and d_ (fa (z), C) Eneo d, (fo (z), c).
The purpose of VQ- training is to form the latent representations in clusters and adjust the codewords to be the centroids of these clusters.
2

---

Under review as a conference paper at ICLR 2023
3
CONTROLLABLE CODEBOOKS WITH WASSERSTEIN QUANTIZATION
Decoler Fa
Per Rahe
WefffttP.PPo
62
W., F.)
Deriving OP
1.
CK
Decoder
n
Eute, (1
Ensoder
We Pex)
Clsstes of emibedded vectoss asoociate %
Z
(a)
(b)
Figure/ 1: (a): Illustration our VQ- derivation. We depart with the minimization of the WS distance the data space in (1) and further turn it to minimizing the reconstruction error in (2) and the WS distance on the latent space in (3); (b): Visualisation of the embedding space with WS regularization. The output of the encoder f=( is distributed and moved to codewords c in which the cardinalities o- (k) (i.e., the number of latent representation which are assigned to t codeword) are proportional ma. At the same time, the codewords tend to flexibly move to the clustering centroids of the latent representations (cf. Corollary 3.1.
We present the theoretical development of our VQ-WAE framework which connects the viewpoints of the WS distance, generative models, and deep discrete representation learning in Section 3.1 Specifically, we propose to learn a "deterministic decoder transporting the codeword to data distri- butions" via minimizing the WS distance between them (Figure D (Top)). We then turn the above WS minimization to push-forwarding the data to codeword distribution via minimizing a WS dis- tance between "the latent representation and codeword distributions" (Figure Ta (Bottom)). We prove that when minimizing the WS distance between the latent representation and codeword distri- butions, the codewords tend flexibly move to the clustering centroids of the latent representations with a control on the proportion of latent representations associated with a centroid (Figure [b). Based on the theoretical development, we devise practical algorithm for VQ-WAE Section 3.2
3.1 THRORETICAL DEVELOPMENT
Given training set D (21, NN) RV, wish learn codebook C' {c+) C R on a latent space and an encoder to map each data example to given codebook, preserving in- sightful characteristies carried in the data. We first endow discrete distribution over the codewords
P Fe,r
K Thdor with the Dirac delta function ó and the weights π AK-1 ( >
|| 1].
We aim learn a decoder function fa : Z ++ X (i.e., mapping from the latent space Z R the data space X)), the codebook C, and the weights to minimize:
min min Wa. (fa#Pe..,PP))
(1)
C f
where P., = ENN Sr,, is the empirical data distribution and d, is cost metric the data space.
We interpret the optimization problem (OP) in Eq. [ as follows. Given a discrete distribution Po.. on the codewords, use the decoder fa to map the codebook C' to the data space and consider Wa, (fa#Pe, IP.) ) as the codebook-de data distortion w.r.t. fa- We subsequently learn fa to minimize the codebook-das distortion given I. Pe,n and finally adjust the codebook C' and to minimize the optimal codebook-dane distortion. To offer more intuition for the OP in Eq. (T) we introduce the following lemma.
3

---

Under review as a conference paper at ICLR 2023
Lemma 3.1. Let C* (g n², and få be the optimal solution of the OP in Eq. 77. Assume T, ; and få are also the optimal solution of the following OP:
K < C,N, then C = (cx)x
min min mii
Ede de (Ins fa (Co(n))),
(2)
n=
where Dx is the set of assignment functions {1,..., NJ + (1, K) such that the cardinalities |o-1 (k) L...., K are proportional to Thi k = ,... K
Lemma 3.1 states that for the optimal solution C = { and fa of the OP in Eq. {fā (ct))
k= become the optimal clustering centroids of the optimal clustering solution which min- imizes the distortion. Inspired by Wasserstein Auto-Encodee (Tolstikhin et al. (017). we establish the following theorem to engage the OP in ① with the latent space.
Theorem 3.1. We can equivalently
the optimization problem in 7 to
min min Car.fa F
EnsR, dd (fa (fe (x)) [x)]
(3)
where fe is deterministic discrete encoder mapping data example x directly to the codebook.
First, we learn both the codebook C and the weights Second, ours seeks deterministic discrete encoder fe mapping data example x directly to a codeword, concurring with vector quantization and serving our further derivations, whereas Theorem 1 in Tolstikhin et al. (2017) involves proba- bilistic/stochusicc encoder mapping to a continuous latent distribution (i.e., larger space to search). More importantly, our proof is totally different from that in Tolstikhin et al. (017)) (all proof details are given in Appendix A).
Additionally, fe a deterministic discrete encoder mapping a data example æ directly to codeword. To make it trainable, we replace fe by a continuous encoder fe X - Z and arrive the OP:
min min (Ez=P, (d, (Sa (Qc (Ja (x))) =x) + XWa, (Je#Pa, Pa.))) C Sa.fe
; ,c) is quantization operator which returns the closest where Qe (Se (x)) argmin.ecdd (fe (x), codeword to Se (x) and the parameter À > 0.
Particularly, we can rigorously prove that the two optimization problems of interest in 33 and are equivalent under some mild conditions in Theorem 3.2 This rationally explains why we could solve the OP in (4) for our final tractable solution.
Theorem 3.2. If we seek fa and fe in family with infinite capacity (c.g., the family of all measur- able functions), the three OPs of interest in and are equivalent.
Moreover, the OP 4 conveys important meaningful interpretations. Specifically, by minimizing Wa. (F x, Pa.z) w.r.t. C, π, we aim to learn the codewords that are clustering centroids of fa #P according to the clustering viewpoint of OT as shown Corollary [3.]] and similar to VQ- (AE, we quantize fe (x) to the closest codeword using Qc (fe (x)) argminegcd, (Se (x),e) and try to reconstruct from this codebook.
Corollary 3.1. Consider minimizing the second term: minfe.c. Wa. (SetP2, in 4 given and assume K < N, its optimal solution f and C' are also the optimal solution of the OP:
min min Σ de (fe (Ix), (o())) fe,CoEs
(5)
n=1
where En is the set of assignment functions {1,..., N) + {1. K) such that the cardinalities 0-1 (1)) 1,..., K are proportional to Tk, k = I,... K.
Corollary 3.1 indicates the aim of minimizing the second term Wa, (fa#Pr.P... c: in (4. By which, we adjust the encoder fa and the codebook C such that the codewords of C become the clustering
'Egg, a the nearest assignment: -1 -1 (k) (fe (x) = ca It argmin,d, (fa (x)) cx) is set of latent representations which quantized to th codeword.
4

---

Under review as a conference paper at ICLR 2023
centroids of the latent representations {fe (Ix)), minimize the badebook-latens distortion (see Figure (ighht). Additionally, at the optimal solution, the optimal assignment function o², which indicates how latent representations (or data examples) associated with the clustering centroids (i.e., the codewords) has a valuable property, i.e., the cardinalities {aa) (A)) k = 1,... K are pro- portional Tky k 1,.... K.
Remark: Recall the codebook collapse issue, i.e. most of embedded latent vectors are quantized to just few discrete codewords while the other codewords are rarely used. Corollary 3.1 give impor- tant properties: (1) we can control the number of latent representations assigned to each codeword by adjust n, guaranteeing all codewords are utilized, (2) codewords become the clustering centroids of the associated latent representations to minimize the codebook-late distortion, to develop our VQ-WAR framework.
3.2 PROPOSED FRAMEWORK
One of crucial aims of learning meaningful and well-distributed codewords is to make use of each individual codeword efficiently by solving the OP in (4). Specifically, we wish the latent represen- tations are more uniformly associated with the codewords. Based on Corollary 3.1 pointing out that the numbers of latent representations associated with the codeword is proportional to Trs we hence fix π as a uniform distribution (i.e., P.,,= = Σx= 1 equally by the model, hence boosting the perplexity the codebook usage.
K 5 to make all the codewords utilized
We now present the practical method based on the OP in with P. iteration, we sample a mini-batch 21, & and then solve the OP in (4) by updating fa, fe and C based on this mini-batch as follows. Let denote P, B Σ j=ر of embedded vectors. over the current batch. Basically, we learn the optimal transportation plan P by solving:
Ex= t At each
6n; as the empirical distribution
Wa, (fetPo, Pc.))
min PeT(Ia:Ic)
(P,De,)),
(6)
where 1B = # B is the vector of atom masses of IP, lc = | is the vector of atom masses of P. r (1B. is the set of feasible transportation plans, and De.r = [dg (x,, (x),,,)), E B8x is the cost matrix.
The pseudcode of our VQ-WAE summarized in Algorithm 1 We use the copy gradient trick (Van Den Oord et al. 2017) to deal with the back-propaganl from decoder to encoder for reconstruction term while Wasserstein regularization term Wa, (fe#PP, IP,, IP,,)) can be optimized directly without further manipulation. Additionally, Wa. (fe# P Pe,x) term is only utilized in the training phase.
Algorithm 1 VQ-WAE
1: Initialize: encoder fa. decoder fa and codebook C.
2: for iter in iterations do
3: Sample mini-batch of samples Ti,-:; forming the empirical batch distribution Ps // i + B for i= 1, 1,..,, B // Nearest neighbor assignment
4: Encode: Zi = fe (z)-n)
5: Quantize: C = arg ming d2 (z)
6: Decode: 2 fa(c,+))
7: Optimize fe, fa and C by minimizing the objective in
B
Σ
dd (Z., + A Wa. Pe.a
minper(Ia (l))) (P.De.))
8: end for
9: Return: The optimal Sex fa and C.
4
RELATED WORK
Variational Auto-Encoder (VAE) was first introduced by Kingma & Welling Kingma & Welling 2013) for learning continuous representations. However, learning discrete latent representations has
5

---

Under review as a conference paper at ICLR 2023
proved much more challenging because it is nearly impossible to accurately evaluate the gradients which are required to train models. To make the gradients tractable, one possible solution is to apply the Gumbel Softmax reparameteriiztioo trick (Jang et al. 2016) to VAE, which allows to estimate stochastic gradients for updating the models. Although this technique has a low variance, it brings up high-bias gradient estimator. Another possible solution is employ the REINFORCE algorithm [Williams] 1992), which is unbiased but has high variance. Additionally, the two techniques can be complementarily combined (Tucker al. 2017).
To enable learning the discrete latent codes, VQ-VAE (Van Den Oord et al. 2017) uses deterministic encoder/decoder and encourages the codebooks become the clustering centroids of latent repre- sentations. Additionally, the copy gradient trick is employed in back paok-propgggting gradients from the decoder to the encoder (Bengio 2013). Some further works were proposed to extend VQ-VAF, notably (Roy et al. 2018 Wu & Flierl| 2020). Particularly, Roy et al. (2018) uses the Expectation Maximization (EM) algorithm in the bottleneck stage to train the VQ- VQ-VAE for improving the qual- ity of the generated images. However, to maintain the stability of this approach, we need to collect a large number of samples on the latent space. Wu & Flierl (022)) imposes noises on the latent codes and uses a Bayesian estimator to optimize the quantizer- based representation. The introduced bottleneck Bayesian estimator outputs the posterior mean of the centroids to the decoder and per- forms soft quantization of the noisy latent codes which have latent representations preserving the similarity relations of the data space. Recently, Takida et al. (2022) extends the standard VAE with stochastic quantization and trainable posterior categorical distribution, showing that the annealing of the stochasticity of the quantization process significantly improves the codebook utilization.
Wasserstein (WS) distance has been widely used in generative models (Arjovsky al., 2017 Gul- rajani et al. 2017 Tolstikhin et al. 2017). Arjovsky et al. Arjovsky et (2017) uses a dual form of WS distance develop Wasserstein generative adversarial network (WGAN). Later, Gulrajani al. (017) employs the gradient penalty trick to improve the stability of WGAN, In terms of theory de- velopment, mostly related to our work is Wasserstein Auto-Encoder (Tolstikhin et al. (017) which aims to learn continuous latent representation preserving the characteristics of input data.
EXPERIMENTS
In this section, we conduct extensive experiments to show the effectiveness of our proposed method compared to other advances.
Datasets: we empirically evaluate the proposed VQ- WAE in comparison with VQ-VAE VAE Den Oord et al. (017) that is the baseline method and recently proposed SQ-VAE Takida al., 2022) which is the state-of-the-art work of improving the codebook usage, on four different bench-
(Van
mark datasets: CIFARIO (Van Den Oord et al. (017). MNIST (Deng) 2012), SVHN (Netzer al.
2011 CelebA (Liu et al. 2015) and the high-resolution images dataset FFHQ Karras et al. (2019).
utilize the same architectures and hyper-parameter for all methods. Additionally, in the primary setting, we use the codeword (discrete latent) dimensional- ity of 64 and codebook size |C| 512 for all datasets except FFHQ with codeword dimensionality of 256 and |C*| = 1024, while the hyper-paun anper-paammeerss [B,T,AL }} are specified as presented in the orig- inal papers, Le., 3 0.25 for VQ- VAE and VQ-GAN (Esser et al. 2021). r 1 for SQ-VAR and À = 1 for our VQ- TOE. The details of the experimental settings are presented in Appendix C
Implementation:
For fair comparison,
5.1 RESULTS ON BENCHMARK DATASETS
In order quantitatively assess the quality of the reconstructed images, we report the results on most common evaluation metrics, including the pixel-level peak signal-to-noise ratio (PSNR), patch- level structure similarity index (SSIM), feature-level LPIPS (Zhang et al., 2018), and dataset-level Fréchlet Inception Distance (FID) (Heusel et al., 2017). ). We report the test-set reconstruction results on four datasets in Table I With regard to the codebook utilization, we employ perplexity score which is defined as e En Pr, log where Pce ΣK,, No. (i.e., No, is the number of latent
Noy
representations associated with the codeword c;) is the probability of the 2th codeword being used. Note that by formula, perplexity max = |C| as P(c) becomes to the uniform distribution, which means that all the codewords utilized equally by the model.

---

Under review as a conference paper at ICLR 2023
Table I: Reconstruction performance (4: the lower the better and 1: the higher the better).
Dataset
Model
Latent Size
SSIM f
PSNR
LPIPS
rFID +
Perplexity
CIFARIO
VQ- VAE
8 x 8
0.70
23.14
0.35
77.3
69.8
SQ-VAE
8 x 8
0.80
26.11
0.23
55,4
434.8
VQ-WAE
8 x 8
0.80
25.93
0.23
54.9
505.0
MNIST
VQ- VAE
8 x 8
0.98
33.37
0.02
4.8
47.2
SQ-VAI VAEE
8 x
0.99
36.25
0.01
3.2
301.8
VQ-WAE
8 x 8
0.99
35.61
0.01
2.4
507.7
SVHN
VQ- VAE
8 x 8
0.88
26.94
0.17
38.5
114.6
SQ VAE
8 x
0.96
35.37
0.06
24.8
389.8
VQ-WAE
8 x 8
0.96
34.67
0.06
22.6
486.0
CELEBA
VQ-VAE
16 X 16
0.82
27.48
0.19
19.4
48,9
SQ-VAE
16 x 16
0.89
31.05
0.12
14.8
427.8
VQ-WAE
16 x 16
0.88
30.08
0.13
13.6
508.0
FFHQ
VQ-GAN
16 16
0.6641
22.24
0.12
4.42
423
VQ-WAE
16 X 16
0.6648
22.45
0.1245
4.20
1022
We compare VQ-WAE with VQ-VAE, SQ-VAE and VQ-GAN for image reconstruction in Table All instantiations of our model significantly outperform the baseline VQ-VAE under the same com- pression ratio, with the same network architecture. While the latest state-of-the of-the-ar art SQ- VAE holds slightly better scores for traditional pixel- and patch-level metrics, rFID scores which evaluate the image quality at the dataset level. Note that our VQ-WAE signif- icantly improves the perplexity of the learned codebook. This suggests that the proposed method significantly improves the codebook usage, resulting better reconstruction quality. Finally, to complete the assessment, the qualitative results are visualized in Figure ( (Appendix B
method achieves much better
5.2 DETAILED ANALYSIS
We number of ablations to analyze the properties of VAE, SQ-VAE and VQ-WAE in order to assess if our VQ-WAE can simultaneously achieve (i) efficient codebook usage, (ii) reasonable latent representation.
5.2.
CODEBOOK USAGE
Table 2: Distortion and Perplexity with different codebook sizes.
Dataset
MNIST
CIFAR10
[C|
64
128
256
512
64
128
256
512
Perplexity rFID
47.8
70.3
52.0
47.2
24.3
44.9
85.1
69.8
VQ-VAE
5.9
6.2
5.2
4.8
86.6
78.9
73.6
69.8
Perplexity rFID
47.4
85,4
184.8
301.8
59.5
113.2
220.0
434.8
SQ-VAE
4.7
4.3
3.5
3.2
71.5
66.9
62.6
55,4
Perplexity
63.8
127.7
255.1
507.7
63.4
126.1
252.0
505.0
VQ-WAE
rFID
5.6
3.8
2.8
2.4
73.5
68.5
60.3
54.9
We observe the codebook utilization three methods with different codebook sizes (64, 128, 256, 512) MNIST and CIFAR10 datasets. Particularly, we present the reconstruction performance for different settings in Table 2 and the histogram latent representations over the codebook in Figure??
As discussed in Section 3.1 and Section 3.2 the number used centroids reflects the capability of the latent representations. In other words, it represents the certain amount of information is pre-

---

Under review as a conference paper at ICLR 2023
C
512
[C] 256
[C]
VQ-VA
Patskeey
Teskoke 243
enfor
[C]=
512
256
SQNAS
SQNARE
128
512
64
256
WO
WO-AL
Nartokk
(a) MNIST.
(b) CIFARIO.
Figure 2: Latent distribution over the codebook on test-set.
served in the latent space. By explicitly defining the numbers of latent representations associated with the codebooks to be uniform (i.e., fixing т in (4) as a uniform distribution) in the Wasserstein regularization term, VQ-WAE is able maximize the information in the codebooks, hence improv- ing the reconstruction capacity. It can be seen from Figure that the latent distribution of 'Q-WAE over the codebook is nearly uniform and the codebook's 's perplexity almost reaches the optimal value (i.e., the value of perplexities reach to corresponding codebook sizes) in different settings. It is also observed that as the size of the codebook increases, the perplexity of codebook of VQ-WAF: also in- creases, leading to the better reconstruction performance (Table [ in line with the analysis in (Wil & Flier!| 2018). SQ-VAE also has good codebook utilization as its perplexity proportional the size of the codebook. However, its codebook utilization becomes less efficient when the codebook size becomes large, especially in low texture dataset (i.e., MNIST).
On the contrary, the codebook usage of VAEV is less efficient, i.e., there are many zero entries in its codebook usage histogram, indicating that some codewords have never been used (Figure 2.. Furthermore, Table also shows the instability of VQ-VAE's reconstruction performance with different codebook sizes.
5.2.2 VISUALIZATION OF LATENT REPRESENTATION
To better understand the codebook's representation power, we employ t-SNE (van der Maaten & Hinton 2008) to visualize the latent representations that have been learned by VQ-VAE, SQ- -Q-E and VQ-WAE on the MNIST dataset with two codebook sizes of 64 and 512. Figure 3l shows the latent distributions of different classes in the latent space, in which the samples are colored accordingly to their class labels. Figure 3 shows that representations from different classes of VQ- WAE are well clustered (i.e., each class focuses on only one cluster) and clearly separated to other classes. In contrast, the representations of some classes in VQ-VAE and SQ-VAE are distributed to several clusters and or mixed each other (Figure 3b.. Moreover, the class-clusters of SQ- VAE are uncondensed and tend to overlap with each other. These results suggest that the representations learned by VQ-WAE better preserve the similarity relations of the data space better than the other models.
5.2.3 IMAGE GENERATION
As discussed in the previous section, VQ-WAE able to optimally utilize its codebook, leading to meaningful and diverse codewords that naturally improve the image generation. To confirm this ability, we perform the image generation on the benchmark datasets. Since the decoder reconstructs images directly from the discrete embeddings, only need model prior distribution over the discrete latent space (i.e., codebook) to generate images.
We employ a conventional autoregressive model, the CaN-based PixeICNN (Van den Oord al. 2016), to estimate a prior distribution over the discrete latent space of VQ- VAE, SQ-VAE and VQ- WAE on CIFAR10, MNIST, SVHN and CelebA. The details of generation settings are presented
8

---

Under review as a conference paper at ICLR 2023
|C| = 64
2
|C| = 512
Mised
(a) VQ-VAF (B = 0.25)
(b) SQ-VAE (r = 1e
(c) VQ-WAR (2 1.0)
Figure 3:
The t-SNE feature visualization the MNIST dataset.
Table 3: FID scores of generated images.
Dataset
VQ-Model
Generation
Latent Size
|C|
unconditional
class-conditional
CIFAR10
VQ-VAE
PixelCNN
8 x 8
512
117.49
117.16
SQ-VAE
PixeICNN
8
512
103.78
90.74
VQ-WAE
PixelCNN
8 x 8
512
87.62
88.93
MNIST
VQ-VAE
PixeICNN
8 x 8
512
27.01
25.56
SQ- VAE
PixelCNN
x 8
512
8.93
4.94
VQ-WAE
PixeICNN
8 8
512
8.17
3.96
SVHN
VQ-VAE
PixeICNN
8 8
512
62.13
64.24
SQ-VAE
PixeICNN
8 x 8
512
31.26
36.41
VQ-WAE
PixeICNN
8
512
30.64
34.24
CELEBA
VQ- VAE
PixeICNN
16 x 16
512
42.0
SQ-VAE PixelCNN
16 16
512
29.5
VQ-WAE
PixelCNN
16 x 16
512
28.8
in Section 3.2 of the supplementary material. The quantitative results in Table indicate that the codebook of VQ-WAE leads to better generation ability than VQ-VAE and SQ-VAE.
CONCLUSION
In this paper, inspired by the nice properties and mature theory of the WS distance allowing be applied successfully to generative models and continous representation learning, we propose Vec- tor Quantized Wasserstein Auto-Encoder (VQ-WAE), which endows discrete distribution over the codewords and learns deterministic decoder that transports the codeword distribution to the data distribution via minimizing a WS distance between them. We then developed theoretical analysis to show the equivalence of this WS minimization to another OP regarding push-forwarding the data distribution to the codeword distribution, which can be realized by minimizing a WS distance be- tween the latent representation and codeword distributions. We conduct comprehensive experiments to show that our VQ-W WAE utilizes the codebooks more efficiently than the baselines, hence leading to better reconstructed and generated image quality.

---

Under review as a conference paper at ICLR 2023
7 REPRODUCIBILITY STATEMENT
We provide the implementation of our framework in the supplementary material.
REFERENCES
Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint 1612.0061,00010, 2016.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Pro- ceedings of Machine Learning Research, pp. 214 223. PMLR, 2017.
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A frame- work for self-supervised learning of speech representations. Advances Neural Information Processing Systems, 33:1 18::12449-12445 12400, 2020.
Yoshua Bengio. Estimating or propagating gradients through stochastic neurons. arXiv preprint arXiv: 2013.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15750-15758, 2021.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 01-22292-2300 2013.
Li Deng. The mnist database of handwritten digit images for machine learning research [best of the web]. IEEE signal processing magazine, 1916)114 2012.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and llya Sutskever. Jukebox: A generative model for music. arXiv preprint arXiv:2005.00341 2020.
Sander Dieleman, Aaron den Oord, and Karen Simonyan. The challenge of realistic music generation: modelling raw audio at scale. Advances in Neural Information Processing Systems, 31, 2018.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pp. 12873-12883, 2021.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanis- las Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy. Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78): 1-8, 2021.
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. Advances in neural information processing systems, 29, 2016.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
10

---

Under review as a conference paper at ICLR 2023
Gustav Eje Henter, Jaime Purenzo-Truebee Xin Wang, and Junichi Yamagishi. Deep encoder- decoder models for unsupervised learning of controllable speech synthesis. arXiv preprint 1rXfv::807711470, 2018.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. 2016.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameteriiztioo with gumbel-softmax, arXiv preprint Xrviv: 1611.01 144, 2016.
Nal Kalchbrenner, Aäron Oord, Karen Simonyan, Ivo Danihelka, Oriol Vinyals, Alex Graves, and Koray Kavukeuoglu. Video pixel networks. International Conference on Machine Learning, pp. 1771-1779. PMLR, 2017.
Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In Proceedings the JEEE/CVF conference on computer vision and pattern recognition, pp. 4401-4410, 2019.
Taehoon Kim, Gwangmo Song, Sihaeng Lee, Sangyun Kim, Yewon Seo, Soonyoung Lee, Se- ung Hwan Kim, Honglak Lee, and Kyunghoon Bae. L-verse: Bidirectional generation between image and text. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16526-16536, 2022.
Diederik Kingma and Max Welling. arXiv: 14, 2013.
Auto-encoding variational
bayes. arXiv preprint
Durk P Kingma, Tim Salimans, Rafal Jozefowiez, Xi Chen, Ilya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 29, 2016.
Kundan Kumar, Rithesh Kumar, Thibault Boissiere, Lucas Gestin, Wei Zhen Teoh, Jose Sotelo, Alexandre de Brébisson, Yoshua Bengio, and Aaron C Courville. Melgan: Generative adversarial networks for conditional waveform synthesis. Advances in neural information processing systems, 32, 2019.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. In Proceedings of the IEEE international conference on computer vision, pp. 7730-3738, 2015.
Deep learning face attributes in the wild.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv: 1609,03999 2016.
Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. Context encoders: Feature learning by inpainting. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2536-2544, 2016.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019.
Scott Reed, Aäron Oord, Nal Kalchbrenner, Sergio Gómez Colmenarejo, Ziyu Wang, Yutian Chen, Dan Belov, and Nando Freitas. Parallel multiscale autoregressive density estimation. In Interna- tional conference on machine learning, pp. 2912-2922 PMLR, 2017.
Aurko Roy, Ashish Vaswani, Arvind Neelakantan, and Niki Parmar. Theory and experiments on vector quantized autoencoders. arXiv preprint arXiv: 1805. 1063, 2018.

---

Under review as a conference paper at ICLR 2023
Filippo Santambrogio. Optimal transport for applied mathematicians, Birkäuser, NY, 55(58-63):94, 2015.
Yuhta Takida, Takashi Shibuya, Weihsiang Liao, Chich-Hsin Lai, Junki Ohmura, Toshimitsu Ue- saka, Naoki Murata, Shusuke Takahashi, Toshiyuki Kumakura, and Yuki Mitsufuji. SQ-VAE: Variational Bayes on discrete representation with self-annealed stochastic quantization. In Ka- malika Chaudhuri, Stefanie Jegelka. Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 20987-21012. PMIR, 17-23 Jul 2022.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. arXiv preprint arXiv:1711.01555 2017.
George Tucker, Andriy Mnih, Chris Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 30, 2017.
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et Con- ditional image generation with pixelenn decoders. Advances in neural information processing systems, 29, 2016.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne, 2008
Slava Voloshynovskiy, Mouad Kondah, Shideh Rezaeifar, Olga Taran, Taras Holotyak, and Danilo Jimenez Rezende. Information bottleneck through variational glasses. arXiv preprint arXiv:1912.008800. 2019.
Ronald J Williams.
Simple statistical gradient-followinn algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-255 1992.
Will Williams, Sam Ringer, Tom Ash, David MacLeod, Jamie Dougherty, and John Hughes. Hier- archical quantized autoencoders. Advances in Neural Information Processing Systems, 33:4524- 4535, 2020.
Hanwei Wu and Markus Flierl. Variational information bottleneck on vector quantized autoencoders. arXiv preprint 1808.0:800.0/000 2018
Hanwei Wu and Markus Flierl. Vector quantization- based regularization for autoencoders. In Pro-
ceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 6380-6387, 2020.
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using vq-vae and transformers. arXiv preprint arXiv:2 104. 10157, 2021.
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved vggan. In International Conference on Learning Representations, 2021.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features perceptual metric. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 586-595, 2018
Chuanxia Zheng, Long Tung Vuong, Jianfei Cai, and Dinh Phung. Movq: Modulating quantized vectors for high-fidelity image generation. Advances in Neural Information Processing Systems, 35, 2022.
Barret Zoph, Golnaz Ghiasi, Tsung- Yi Lin, Yin Cui, Hanxiao Liu, Ekin Dogus Cubuk, and Quoc Le. Rethinking pre-training and self-training. Advances in neural information processing systems, 33: 3833-3845, 2020.
12

---

Under review as a conference paper at ICLR 2023
APPENDIX
This appendix organized as follows:
In Section A
present all proofs for theory developed in the main paper.
In Section we present additional experimental results on the high-quality image dataset FFHQ.
In Section present experimental settings and implementation specification of VQ-
WAE.
A THEORETICAL DEVELOPMENT
Given a training set D (x1, (2,,,..,222) 2ND C. RRV we wish learn set of codebooks {c*] RXxD latent space Z and encoder to map each data example given codebook, preserving insightful characteristics carried in data. We first endow a discrete distribution over the codebooks as Pe,n k= Tedce with the Dirac delta function 6 and the weights Ak. We aim to learn decoder function fa Z + X' (i.e., mapping from the latent space Z C IR to the data space 2),, the codebooks C, and the weights π, to minimize:
min min Wa, (atPPe.a,PP.
C fe
where Fr En 5z,, is the empirical data distribution and dy is cost metric on the data space. Lemma A.1. (Lemma 3.1 in the main paper) Les C = {ckbe , Tt, and få be the optimal solution of the OP in Eq. (7. Assume K N, then C = {ck)k x², and fä are also the optimal solution of the following OP:
Σ de (#ns Sa (Co(m))),
min min min fa C dEE
(8)
where En is the assignment functions {1,..., N) + {1, K) such that the cardinalities o- (k) k I,..,K are proportional to Tai k = 1, ... K.
Proof of Lemma A.I
It is clear that
fa#Po c, falea)
Therefore, we reach the following OP:
N
min min Wa
8am
C fu
By using the Monge definition, we have
Wa.
Σ
8
TkdSa((w)
T:T min
E (d (x, T (x)))
n=1
1=
Σ de (In, (xx))
min
Since T#P, SaHP.T T (In) fa (cx) for some k. Additionally, - (Ja(cx))) =1.,,,, I,.., K proportional to TRi h 1,..., K. Denote {1..... N} {l,..., K). such that T (Ix) fa(Co(m))) Vi vi L, L,,,, we have a E Σn. It follows that
N
Wd.
Σ
5g Σ muδsatew)
min = N dEE,
Ede de (Ins fa (Co(n)))
N
==1
k=1
n=1
13

---

Under review as a conference paper at ICLR 2023
Finally, the the optimal solution of the OP in Eq.
(7) is equivalent to
N
min min min C. ae
dz (Fn, fa (Cr(n))),
n=l
which directly implies the conclusion.
Theorem A.1. (Theorem 3.1 in the main paper) We can equivalently uurn the optimization problem
min C...a #F.
min
EnwR, [de (fa (fe (x)) .x)].
(10)
where Fe is deterministic discrete encoder mapping data example x directly to the codebooks.
Proof of Theorem[A.]
We first prove that the OP of interest is equivalent to
min C CFff#PP-Ree.
min
Pe [dz (Sa (c),
(11)
where Je is stochastic discrete encoder mapping data example directly the codebooks. To this end, we prove that
Wa. (fa# IP2)
min feife
Fx (d. (c) x))]
(12)
where fa is stochastic discrete encoder mapping data example directly to the codebooks.
Let fc be a stochastic discrete encoder such that fe#P, = Pe,, (i.e., x Pa and ~ fe (x) implies P..,)) We consider TId as the joint distribution of (x, e) with x ~ Pa and c ~ fa (x). We also consider Tff as the joint distribution including (x.x²) ~ "fad where x ~ Pa,C ~ fa (x),
and &' (c). This follows that 7fe,d E T (fa#Po,n o,5 Pa) which admits fa# o,a and IP, 2 as its marginal distributions. We also have:
E Ep-P. Idz (Sa (c), x]]] = E(a,oraa [dz (Sa (c) x]] ② Bloaorrro Idz x2))
(T B(are")arad (da (x,
min
Wa, (JatPo, Po,,, Pn).
Note that we have the equality in (1) due to (id, fa) # Tfcdd
Therefore, we reach
min I.J.
B on Rc(s) (d (fa ()) 2 Wa. (fa#Pe
Let Tfed E I (SatPe e, Pa). Let 7/e. E r (fatPP c,T: PE,x)) be a deterministic coupling such that e ~ Pe,x and x fa imply xc ~ Te,fe- Using the gluing lemma (see Lemma 5.5 in Santambrogio (0155). there exists joint distribution E (Pe,,, fa#Po. (x..P..PP)) which admits Tood and fag as the corresponding joint distributions. By denoting Jd, E T (PP,, P.,,)) as the marginal distribution of over P2,, Pe,,, we then have
E ((2,") [dg (x,, x')] = Ego,'' aya (dg (x, x')] = B(o.))
(dg (x,)
= Blo (de (fa x1)] = EppP, (de (fa (c),
min
E EawP,o-ffte fe(x)) Idz (fa (c),
fa =P.
where = xx).
This follows that
Wa. f...P..
min
E B(a,"")~sad [d (a, (x. x')]
(
> min
BawBs. (do (fa (c) x)]
PP.
14

---

Under review as a conference paper at ICLR 2023
This completes the proof for the equality in Eq. 122. which means that the OP of interest in is
equivalent to
min Ja fe: #P Por
min
E P (x) [d (fa (c),
(13)
We now further prove the above OP is equivalent to
min C.a.fa Jere PP
min
ExwP, dd (fa (fe (x)) .x)]
(14)
where is deterministic discrete encoder mapping data example x directly to the codebooks.
It is obvious that the OP in (14) is special case of that in (13) when we limit to search for deter- ministic discrete encoders. Given the optimal solution C*= . fel and fe of the OP in 13
we show how to construct the optimal solution for the OP in 14). Let construct C =
fa² = fd Given æ ~ Pz. let us denote fe² (x) = argmin.d, (fr (c), Thus, fr² is deterministic discrete encoder mapping data example z directly the codebooks. We define
2 = Pr (fa² (z) = c x ~ P2,) k I,..,,K, K, meaning that f2 = P From the construction of fe², we have
ERPP, [d. (J2² (Je² (x)) m)]
ca [d (fă (o),
Furthermore, because C²2, T²2 ,²2, , and are also a feasible solution of the OP 14 we have
ENNP, [dz (fa² (fe² (I)) x>] 2 BewPr, dd (fă (c) xx)]
This means that
EawP, de (fa² (fe² (I)) x)] Bxwds fz (x) dd (fā (c) x0)]
and C*2 n x=2, a are also the optimal solution of the OP 14).
Additionally, fa deterministic discrete encoder mapping data example x directly to the code- books. To make it trainable, we replace fe by continuous encoder fe X + 2 and arrive the following OP:
min min {Eeup, (de (Sa (Qc (Je (x))), ,x)] XWa, (Je#Po, Pa,))} C. fa.fa
(15)
where Qe (fa (x)) argmin.ecd. (fe (z) e) is quantization operator which returns the closest codebook to Je (x) and the parameter > O.
We now propose and prove the following lemma that is necessary for the proof of Theorem A.2
Lemma A.2. Consider C, far and fe as feasible solution of the OP in (15). Let us denote fe(x) argmin,d, (fe())) Qz(x), then fe(x) is Borel measurable function.
Proof of Lemma A.2
We denote the A on the latent space as
A = {a dz cx) <
(=:0((2)) c+)
A is known as Voronoi cell w.r.t. the metric d.. If we consider continuous metric d,, A is measurable set. Given Borel measurable function B, we prove that f-1 (B) is a Borel measurable set the data space.
Let n (ci..., (ck) = (Gi,.... J. we prove that f- (B) = = u (A,,).
Indeed, take f-r (B), then fz( E B, implying that fe(z) Qc(x) ct for some j = I,..., m. This means that fax E Ai, for some 1,..,, Therefore, we reach for (B) C Uf (Ai,).
We now take U (At,). ). Then fe (x) A,, for I,... m, hence fc(x)= Qo (x) Ci for some j = 1, m, Thus, fe(x)) C or equivalently x f-1 (B), implying f-1 (B) C 1E (Asy).
Finally, we reach f (B)= (Ar,). which concludes our proof because f. is measurable function and A; are measurable sets.
15

---

Under review as a conference paper at ICLR 2023
Theorem A.2. (Theorem 3.2 in the main paper) If we seek fa and fe in family with infinite capacity (e.g. the space of all measurable functions), the three OPs of interest in ( 10 and 75 are equivalent.
Proof of Theorem A.2.
Given the optimal solution O* ,"',"",,, ,ā',
and fe of the OP in 15 we conduct the optimal so-
lution for the OP in 10p. Let conduct C=2 C+ .el fa We next define fe² (x)
argmin, (fe (z)
Qc-1 (fz (x) Qc (fa' (z)). We prove that C _T= fá and
f=² are optimal solution of the OP 10 By this definition, we yield # Po and
hence Wa, #P=2#P2. x()) = Therefore, we need to verify the following:
(i) fe is Borel-measurable function.
(ii) Given feasible solution C, n. fa. and fe of (10), have
EzoP, dd (fa² (fc² (x)) x)] Easr, dd (fa (fe (z)) x))
(16)
We first prove (i). a direct conclusion because the application of Lemma A.2 to C+1 x= ,ā!
and fo!
We next prove (ii). We further derive as
ExsF,
[d (Jā² (fee (x:) []] + AWd, Po
= ENns ddz (fá² (fe² [m))))])
BaNZ dr (fā) (Qo-z (f² (z)))
E
[d (fā² (Qe-1 (f² (z))) ,z]
S B [d (ft² (Qo-1 (fr² (z))) .=)] + AWd, #Pa P....
(17)
Moreover, because J. #P = Per which is discrete distribution over the set of codewords C, we obtain Qc(fe(x)) = fz)... Note that C, π, fa. and fe is also a feasible solution of (15) because fo also a specific encoder mapping from the data space the latent space, we achieve
Exwr, [de (fa (Qc (fe (x))) x)] + AWd. (f.HP,)PP..) > Eunz, [d (fāl (Qc-1 (fzz) (x)) =]) + AWa, PP.,
Noting that FeHPa = Pe,, and (c(F((x)) = fz() we arrive at
EawRz [do (fa (fe (x)) x]]
2 Exnz, [dz (fā² (Qo-1 (Fa) (z))) =]] + +WW, (f: Po..
(18)
Combining the inequalities in 17 and (8), we obtain Inequality (16)
EaaN, - [de (Jā² (Je² (x)) x)] S EaoP, [d (Ja (Ja (x)) xx]
(19)
This concludes our proof.
Corollary A.1. (Corollary 3.1 in the main paper) Consider minimizing the second term: minj..c Wa, (fe#Pa, Pe,n) in (15) given and assume K N, its optimal solution f and C* are also the optimal solution of the following OP:
min min
Ede de (fe (In), Co(())
(20)
C.,C DEE
n=1
where En is the set of assignment functions fl..., N) + fl.... K. such that the cardinalities a- (k)) 1,..., K are proportional to Tk. k = 1, .... K.
16

---

Under review as a conference paper at ICLR 2023
Proof of Corollary A.1
By the Monge definition, we have
N
K
Wa, (fe#P Poxx) = Wa.
1
Σ
((m))
Σ Tidde
min U
Ee Ids (z, T (=))
N
n=1
min
de (fe (Zn) , (fe (In))).
N T
Since T# (Je#Pr) Pe.,, T (Je (kn)) c& for some k. Additionally, 7 (cm)l ,k I,... K are proportional to Thi k L,... K. Denote {I,..., N} + {I,..., K} such that T (fe (En)) Co(n), Vi 1,.,, N, we have a E Ex. also follows that
N
K
N de (Je (vn) co(r))
Wa.
Txdez
min N dEE
1
Σ n=1
N
n=1
km
B VISUALIZATION OF RECONSTRUCTION RESULTS
(a) imput
(6) VOGAS
(6 WQ-YAR
Figure 4: Reconstruction results for the FFHQ dataset.
Qualitative assessment: We present the reconstructed samples from FFHQ chigh-resolution im- ages) for qualitative evaluation. It can be clearly seen that the high-level semantic features of the input image and colors are better preserved with VQ-WAE than the baseline. Particularly, we notice that VQGAN often produces repeated artifact patterns in image synthesis (see the hair of man second column in Figure 4 while VQ-WAE does not. This is because VQ-GAN is lack of diversity the codebook, which will be further analyzed in Section 5.2. Consequently, the quantization operator embeds similar patches into the same quantization index and ignores the variance in these patches (e.g., VQ-GAN reconstructs the background in third column of Figure as hair of woman).
C EXPERIMENTAL SETTINGS
C.I VQ-MODEL
Implementation: For fair comparison, utilize the same framework architecture and hyper- parameters for both VQ-VAE VAE and VQ-WAE. Specifically, we construct the VQ-VAE and VQ- WAE models as follows:
17

---

Under review as a conference paper at ICLR 2023
For CIFARIO, MNIST and SVHN datasets, the models have encoder with two convolu- tional layers of stride and filter size x 4 with ReLU activation, followed by residual blocks, which contained x 3, stride 1 convolutional layer with ReLU activation followed by a x convolution. The decoder was similar, with two of these residual blocks followed by deconvolutional layers.
For CclebA dataset, the models have encoder with two convolutional layers of stride 2 and filter size of 4 x with ReLU activation, followed by residual blocks, which contained 3 3, stride 1 convolutional layer with ReLU activation followed by 1 x 1 convolution. The decoder was similar, with two of these residual blocks followed by two deconvolutional layers.
For high-quality image dataset FFHQ, we utilize the well-known VQGAN framework Esser et al. (2021) as the baseline.
We only replace the regularization module of VQ- VQ-VAE i.e. two last terms of objective func- tion: d₂ (sg (Se (x)) C) Bd. (fe (x), sg (C)) by proposed by Wasserstein regularization XWa, [Je# (e,.)) 1) in Eq. 4 for VQ-WAF. Additional, we employ the POT library (Flamary al. (021)) compute WS distance for simplicity. However, our VQ-WAE does not require op- timal transport map the from WS distance in (6) to update the model. Therefore, we can employ wide range of speed-up algorithms to solve optimization problem (OP) in (6) such as Sinkhorn algorithm (Cuturi 2013) or entropic regularized dual for (Genevay et al.) 2016).
Hyper-parameters: following (Takida al. (022), we adopt the adam optimizer for training with: learming-rate is e-1, batch size of 32, embedding dimension of 64 and codebook size |C| = 512 for all datasets except FFHQ with embedding dimension 256 and |C| 1024. Finally, we train model for CIFARIO, MNIST, SVHN, FFHQ in 100 epoches and for CelebA in 70 epoches respectively.
C.2 GENERATION MODEL
Implementation: It is worth to noting that we employ the codebooks learned from reported VQ- models to extract codeword indices and employ PixeICNN (Van den Oord al.) 2016 with the same setting for generation for all VQ-VAE, SQ-VAE and VQ-WAE. In particular, we feed PixeICNN over the "pixel" values of the 8 1-channel latent space for CIARIO, MNIST, SVHN, and 16 x 16 1-channel latent space for CelebA.
sapper-arrameter: we adopt the adam optimizer for training with: learning-ratee is 3e-4, batch size of 32. Finally, we PixeICNN over the "pixel" values of the 8 x 8 1-channel latent space in 100 epoches.
18


Reviews:
[{"review_id": "nlGcvQ_AYFp", "reviewer": "Reviewer_zbTA", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a vector quantized autoencoder that minimizes the Wasserstein distance between the encoded samples and the code word. In short, the loss function is simply the quantized autoencoding loss plus a penalty term that measures the Wasserstein distance between the empirical distribution of the encoded samples and the empirical distribution of the code words. The authors then evaluate their approach on the generative modeling for CIFAR10, MNIST, SVHN, and CelebA datasets showing competitive performance compared to VQ-VAE and SQ-VAEs. ", "summary_of_the_review": "While the paper has exciting aspects, I think its current state lacks sufficient quality for ICLR. In particular, the paper can benefit from reorganization. Also, the authors could provide a finer prior work section that better places their work among the existing approaches. Finally, detailed discussions on various choices throughout the paper could benefit this paper. To summarize, I think this paper has potential, however, it is not ready for publication yet. ", "strength_and_weaknesses": "### Strengths\n\n* The idea of using Wasserstein distance for vector quantization is interesting. \n\n### Weaknesses\n\n* The paper is not well-written, and the flow of the paper can significantly benefit from a significant rewriting. \n\n* The paper's novelty is limited to applying the Wasserstein distance in the latent space of an autoencoder with a moving sparse target distribution (i.e., the codeword distribution). \n\n* Given the deterministic nature of both the encoder and decoder and applying the Wasserstein distance in the embedding space, the paper is more closely related to the Sliced-Wasserstein Auto-Encoders (SWAE), Kolouri, et al. 2018, as opposed to the WAE that uses probabilistic encoder/decoders. \n\n* The choice of uniform $\\pi$s is not well justified in the paper. Also, the authors may want to use a different symbol for $\\pi$, because $\\pi$ is often used to denote the `transport plan' in the optimal transport-related literature and in calculation of the Wasserstein distance. \n\n* The stop gradient used in VQ-VAE could also be used in your formulation. In particular, if you freeze the encoder, the optimal codeword could be directly obtained from the barycentric projection of the transportation plan. Some discussions on these aspects of the training could improve the quality of the paper. \n\n* Eq (3) is missing a $f_d$ for the inner minimization. \n\n* The paragraphs following Eq (3) state $\\bar{f}_e$ is deterministic repeatedly. There are several instances of repeated reworded statements throughout the paper, which flags that the paper could use significant polishing to improve the flow and communicate the concepts more concisely. \n\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "ebpNfodLk0", "reviewer": "Reviewer_9V99", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper extends the Wasserstein autoencoder (WAE) framework due to Tolstikhin et al. (2017) to incorporate discrete representations as considered by the vector quantized variational autoencoder (VQ-VAE) due to van den Oord et al. (2017). It begins with a theoretical development leading to the optimization problem (4) and then propose an approximate solution procedure in Algorithm 1. The performance of the proposed algorithm is tested on the standard benchmark, namely CIFAR10, MNIST, SVHN, and CelebA. Comparisons are made with VQ-VAE and SQ-VAE (Takida et al., 2022), a stochastic variant of the former. While for the traditional pixel- and patch-level metrics SQ-VAE exhibits better performances, VQ-WAE has merits in dataset-level metrics such as rFID, and Shannon entropy of the latent codeword distribution. The latter phenomenon can be understood as that VQ-WAE yields better codebook utilization and thus prevents codebook collapse problem.", "summary_of_the_review": "There are gap between the theory developed and the actual algorithm, both of which have unresolved flaws.", "strength_and_weaknesses": "## Strengths\n\nApplication of the WAE framework to discrete representation learning is reasonably demonstrated with various evaluation metrics. The derivation of objective function (4) from the Wasserstein distance minimization problem (1) is natural at least intuitively.\n\n\n## Weaknesses\n\n1. The major weakness is the gap between the theory and algorithm. The optimization problems (1), (2) and (4), which are claimed to be equivalent to each other in Theorem 3.2, are basically about fitting a $K$-support discrete latent distribution $\\sum_{k=1}^K \\pi_k \\delta_{c_k}$ to the data distribution by minimizing the Wasserstein distance between them. In particular, the mass $\\pi_k$ assigned to cordwood $c_k$ is an optimization variable to be fitted. However, in the actual algorithm (Algorithm 1), it is fixed as $1/K$ and only $c_k$'s are fitted. So Algorithm 1 in fact solves a different problem from (4). This distinction is important since the so-called \"codebook collapse\" occurs precisely because some $\\pi_k$ is zero but is destined to be prevented if the proportion of codewords is forced to be uniform. In fact, that the optimal transport problem (1) is equivalent to the $K$-means clustering has been known at least since 1982 (Pollard, D., 1982. Quantization and the method of k-means. IEEE Transactions on Information theory, 28(2), pp.199-205). In light of this result, the modified problem that Algorithm 1 solves is the $K$-means clustering problem under the constraint that each (latent) cluster contains the same number of sample (latent) points. I do not think that this modified problem is necessarily better or more desirable than the original $K$-means problem (1). Suppose (1) is solved exactly and there is a codebook collapse. Then this means that strictly less than $K$ codewords best represent the data. Is this bad? So the case that codebook collapse becomes problematic is when the algorithm is trapped in a bad local minimum. This is a problem of the algorithm, rather than the sin of codebook underutilization. In this sense, the role of the added constraint is to *regularize* the algorithm to avoid bad local minima. On the other hand, Takida et al. (2022) avoids the same problem by using stochastic quantization but not constraining cluster sizes.\n\n2. The proof of Theorem 3.2 is incomplete. The proof only states that the (3) $\\leq$ (4). The possibility of strict inequality cannot be ruled out because the constructed variables $(C^{*2}, \\pi^{*2}, f_d^{*2}, \\bar{f}_e^{*2})$ are only feasible for (3) but not necessarily optimal. Further, measurability of $\\bar{f}_e^{*2}$ needs to be proved.\n\n3. The objective of Algorithm 1 is not differentiable because of the $\\arg\\min$ operator in Line 5 and the Wasserstein distance penalty (6).  How this is put to the gradient learning pipeline is not stated. Is the  gradient passing trick due to van den Oord et al. (2017) used to deal with the $\\arg\\min$ operator? Or the Gumbel-softmax trick? is (6) smoothed with entropy regularization? If the answer to the first and the third questions are yes (which I presume), then the real contribution of this paper appears to be replacing the last two terms in the VQ-VAE objective with (a smoothed version of) (6); see Sect. B.1. \n\n4. Finally, I wonder what are the authors' opinion about some peaks of codewords in Fig. 2a when $|C|$ is large in the MNIST dataset.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "d9wlLaRBeMj", "reviewer": "Reviewer_7LM4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "A new vector quantization method is proposed, named Vector Quantized Wasserstein Auto-Encoder (VQ-WAE). The presented VQ-WAE employs the Wasserstein (WS) distance in both the observation x space and the latent z space to encourage matching, mimicking the existing Wasserstein Auto-Encoder. Experiments on MNIST, CIFAR10, SVHN, and CelebA datasets are conducted to demonstrate the superiority of the VQ-WAE to the existing VQ-VAE and SQ-VAE methods.", "summary_of_the_review": "The main weakness of the proposed VQ-WAE lies in its mismatch between theory and implementation. Besides, it's not clear if the presented VQ-WAE works indeed better than existing VQ methods like the VQ-GAN.", "strength_and_weaknesses": "Strength:\n\n(1) The paper is well-written in general.\n\n(2) The empirical experimental results seem to be convincing. \n\nWeaknesses:\n\n(1) The theoretical derivations are not sound, because, in both the VQ-VAE and the presented VQ-WAE, many (like 8 $\\times$ 8) latent codes correspond to one input image $x$, whereas the presented theorems are derived based on one latent code per x. In other words, there is a mismatch between theory and implementation.\n\n(2) No explicit demonstrations like the reconstructed images are shown, which makes it difficult to evaluate the quality of the vector quantization. For example, is the reconstruction quality better than that of the VQ-GAN?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject