Learning Discrete Representation with Optimal Transport Quantized Autoencoders
Authors: Xiaoyu Bie
Dexiong Chen
Xiaodong Cun
Xi SHEN
Venue: Submitted to ICLR 2023 (2023)
Keywords: VQ-VAE
Optimal Transport
3D Motion Generation

Abstract: Vector quantized variational autoencoder (VQ-VAE) has recently emerged as a powerful generative model for learning discrete representations. Like other vector quantization methods, one key challenge of training VQ-VAE comes from the codebook collapse, i.e. only a fraction of codes are used, limiting its reconstruction qualities. To this end, VQ-VAE often leverages some carefully designed heuristics during the training to use more codes. In this paper, we propose a simple yet effective approach to overcome this issue through optimal transport, which regularizes the quantization by explicitly assigning equal number of samples to each code. The proposed approach, named OT-VAE, enforces the full utilization of the codebook while not requiring any heuristics. We empirically validate our approach on three different data modalities: images, speech, and 3D human motions. For all the modalities, OT-VAE shows better reconstruction with higher perplexity than other VQ-VAE variants on several datasets. In particular, OT-VAE achieves state-of-the-art results on the AIST++ dataset for 3D dance generation. Our code will be released upon publication.

Under review as a conference paper at ICLR 2023
LEARNING DISCRETE REPRESENTATION WITH OPTIMAL TRANSPORT QUANTIZED AUTOENCODERS
Anonymous authors Paper under double-blind review
ABSTRACT
Vector quantized variational autoencoder (VQ-VAE) has recently emerged as a powerful generative model for learning discrete representations. Like other vec- tor quantization methods, one key challenge of training VQ-VAE comes from the codebook collapse, i.e. only a fraction of codes are used, limiting its reconstruction qualities. To this end, VQ-VAE often leverages some carefully designed heuris- tics during the training to use more codes. In this paper, we propose a simple yet effective approach to overcome this issue through optimal transport, which regu- larizes the quantization by explicitly assigning equal number of samples to each code. The proposed approach, named OT-VAE, enforces the full utilization of the codebook while not requiring any heuristics such as stop-gradient, exponential moving average, and codebook reset. We empirically validate our approach on three different data modalities: images, speech and 3D human motions. For all the modalities, OT-VAE shows better reconstruction with higher perplexity than other VQ-VAE variants several datasets. In particular, OT-VAE achieves state- of-the-art results on the +IST++ dataset for 3D dance generation. Our code will be released upon publication.
INTRODUCTION
Unsupervised generative modeling aims at generating samples following the same distribution as the observed data. Recent deep generative models have shown impressive performance in generat- ing various data modalities such as image, text and audio, owing to the use of a huge number of parameters in their models. The well known examples include VQ-GAN (Esser et al., 2021) for high-resolution image synthesis, DALLE (Ramesh et al., 2021) for realistic image generation from description in natural language, and Jukebox (Dhariwal et al., 2020) for music generation. Sur- prisingly, all these models are based, at least partly, on Vector Quantized Variational Autoencoders (VQ-AAE) (Van Den Oord et al., 2017). The success of VQ-VAE should be mostly attributed to its ability of learning discrete, rather than continuous, latent representations and its decoupling of learning the discrete representation and the prior. The quality of the discrete representation is essen- tial the quality of the generation and our work improves upon the discrete representation learning for arbitrary data modality.
VQ-VAE is a variant of VAEs (Kingma& & Welling, 2014) that first encodes the input data to a dis- crete variable in a latent space, and then decodes the latent variable to a sample of the input space. The discrete representation of the latent variable is enabled by vector quantization, generally through a nearest neighbor look up in a learnable codebook. A new sample is then generated by decoding a discrete latent variable sampled from an approximate prior, which is learned on the space of the en- coded discrete latent variables in a decoupled fashion using any autoregressive model (Van Den Oord et al., 2017). Despite its promising results in many tasks of generating complex data modalities, the naive training scheme of VQ-VAE used in (Van Den Oord et al., 2017) often suffers from code- book collapse (Takida et al., 2022), i.e. only a fraction of codes are effectively used, which largely limits the quality of the discrete latent representations. To this end, many techniques and variants have been proposed, such as stop-gradion along with the commitment and embedding loss ( (Van Den Oord et al., 2017), exponential moving average (EMA) for codebook update (Van Den Oord et al., 2017), codebook reset (Williams et al., 2020) and stochastic variant (SQ- VAE) (Takida et al., 2022).

---

Under review as a conference paper at ICLR 2023
Interestingly, the idea of vector quantization has also been explored in the related field of self- supervised learning, though it generally relies on unsupervised discriminative modeling by only ob- taining data features that can be easily generalized to downstream tasks. The seminal work by Caron et al. (2018) used an encoder and a clustering algorithm to learn discriminative representations of the data. The clusering algorithm used in Caron al. (2018), namely K-means, could be interpreted as an offline version of the vector quantization used in VQ-VAE. Similar to the codebook collapse, some clusters were observed to have a single element, known as cluster collapse. To address this problem, Asano et al. (2020) have proposed an optimal transport (OT) based clustering method to explicitly enforce the equipartition of the clusters. Caron al. (2020) have later proposed an online version of their algorithm for dealing with large-scale datasets.
In this work, we reformulate VQ-VAE under the framework of Wasserstein Autoencoders (WAE) (Tolstikhin et al., 2018), providing a natural connection between distribution match- ing in the latent space VQ- VAV and the clus- tering used in self-supervised learning. Based on this reformulation, we propose to use an on- line clustering method address the codebook collapse issue of training VQ-VAF, by adding the equipartition constraint from Asano et al. (2020) and Caron et al. (2020). The online clus- tering method, inspired by the OT techniques used Caron et al. (2020), assigns each cluster (represented by a code in the codebook) the same number of samples (see Figure 1). Then, we use a Gumbel-softmax trick Jang et al. (2017) to sample from the discrete categorical distribution while easily assessing its gradient. The resulting approach, named OT-VAE, enforces the full utilization of the codebook while not using any heuristics, namely stop-gradien gradgaaddint, EMA, and codebook reset. Unlike SQ- VAE that uses stochastic quantization and dequantization process. our approach explic- itly enforces the equipartition constraint for quantization in a deterministic way while only using a stochastic decoding. To the best of our knowledge, our approach shows the first time that such an equipartition condition, arised in the field of self-supervised learning, is also useful for generative approach on three different data modalities: images, speech and
Optimal
Transport
Figure 1:
Optimal transport (OT) explicitly en- force the equipartition of the clusters.
tasks. We empirically validate 3D human motions. For all the modalities, OT-VAE shows better reconstruction with higher perplex- ity than other VQ-VAE variants on several datasets. In particular, OT-VAE achieves state-of-the-as results on the AIST+ (Li et al., 2021) dataset for 3D dance generation. Overall, our contribution can be summarized below:
We reformulate VQ-VAE instance of WAEs, which provides a connection between distribution matching in the latent space of VQ-VAE and the clustering methods used in self- supervisdd learning.
We propose OT- VAE, novel unsupervised generative model explicitly using the equipar- tition constraint with OT to address the codebook collapse issue in VQ-VAE.
In experiments, without using classic heuristics (such as stop-gradee poogggrddent, EMA, codebook reset etc.), we show that OT-VAE achieves better reconstruction and perplexity than other variants of VQ-VAE for three data modalities: image, speech and 3D human motion.
Using OT-VAE instead of VQ- IQ-VARin in the Bailando model (Li et al., 2022), we obtain state- of-the-art results for 3D Dance generation. Precisely, we improve FIDx from 28.75 to 26.74 and FIDg from 11 1.82 to 9.81 11o on the AIST+ dataset.
2 RELATED WORK
VQ-VAE VQ-VAE framework was first introduced in Van Den Oord et al. (2017), variant of VAE (Kingma Welling, 2014) with a discrete prior. VQ-VAE shows good performance on various generation tasks, which includes: image synthesis (Williams et al., 2020; Esser et al., 2021), text to image generation (Ramesh et al., 2021). motion generation (Li et al., 2022), music genera- tion (Dieleman et al., 2018; Dhariwal et al., 2020) etc. However, a naive training of VAE suffers from the codebook collapse. To alleviate the problem, a number of techniques are commonly used during the training, including stop-gradient along with some losses (Van Den Oord et al., 2017),
2

---

Under review as a conference paper at ICLR 2023
EMA for codebook update (Van Den Oord et al., 2017), codebook reset (Williams et al., 2020), etc. Our work is highly related to VAE (Takida et al., 2022), which also aimed at improving codebook utilization without using those heuristics. SQ-VAE proposed to perform stochastic quan- tization and dequantization at the early stage of the training and gradually anneal the process toward deterministic one. However, a proper prior distribution needs be carefully chosen for different tasks. In contrast, our work uses the same prior throughout different tasks and data modalities.
Wasserstein autoencoders WAES (Tolstikhin et al., 2018) consist of a class of generative models based on the reconstruction from an autoencoder and a regularizer in the latent space that encourages the training distribution to match the prior. The regularizer is defined through divergence between probability distributions. In the previous research, various divergence functions have been proposed, such MMD (Tolstikhin et al., 2018), GAN (Makhzani et al., 2016). Sliced Wasserstein (Kolouri et al., 2018), Sinkhorn (Patrini et al., 2020), and so on. While all of the WAE variants use a con- tinuous prior, work shows that VQ- VAE could be interpreted as an instance of WAEs with discrete prior. By using the Sinkhorn divergence proposed in Patrini et al. (2020), our work suggests natural connection between VQ-VAE and clustering-basee self-supervised learning methods.
Clustering-based self-supervised learning The matching problem between the posterior distri- bution and the discrete prior distribution is related the clustering-based methods for self self-supervised learning (Li et al., 2020b; Alwassel et al., 2020). The seminal work by Caron et al. (2018) shows that the assignments obtained by K-means can be used as pseudo-labels to learn discriminative rep- resentations for images. However, their method could suffer from the cluster collapse issue caused by K-means, where some clusters are collapse to single entity during training. To address this issue, Asano et al. (2020) propose to incorporate an equipartition constraint and cast it as an OT problem. Then, Caron al. (2020) propose Sinkhorn-Knopp (Cuturi, 2013) based online algo- rithm, making it scalable to very large datasets. Our work shows the connection between VQ-VAE and these clustering-basse self-supervised learning methods and demonstrate the first time that these clustering techniques are also useful for generative tasks.
BACKGROUND ON VQ-VAE
VQ-VAE (Van Den Oord et al., 2017) consists of deterministic autoencoders that allow learning a discrete representation of the data. The main building blocks of VQ-VAE are similar to any autoencoders with an additional codebook. The codebook C' defined as set of K trainable vectors CI...., CK in Rh We denote by Cdd C dd a dz-tuple related to the codebook C with d, describing the number of components in the latent space. For a latent variable E Cid., we denote by E C' its i-th entry. The deterministic encoder first maps observation X E X to the same space as Z using a neural network hs followed by a quantization step projecting he () onto Cd2, The quantization process modelled as a deterministic categorical posterior distribution such that 2x(X) = arg min (X)](X); cx Then the objective of VQ- IQ-VAE is defined as
- Rd, xds with trainable parameters Φ.
LvQ-VAE == log pe (x (x)) + [|sg(h, - 2(x)1 (x) sg sg(2()))11 (I)
where po is decoder with parameters 6, sg denotes the stop gradient operator, and 3 is hyperpa- rameter. The hard assignment in the quantization process makes the training of VQ-VAE sensitive to the initialization and prone to local minima. To improve the utilization of the codebook, various strategies have been adopted in previous work, often involving huge amount of work for hyperpa- rameter tuning.
4 METHOD
In this section, we reformulate VQ- VAE instance of WAEs. The new formulation offers natural choice of the prior distribution and the objective function. This yields a simple yet effective autoencoder framework that can learn discrete representations of the data.
3

---

Under review as a conference paper at ICLR 2023
Codebook
L, Normalization
X
2XX
G(Z(X))
Optimal Transport
Gumbell Softmax
Assigmment
Figure 2: Overview of OT-VAE: the encoder maps a sample X to the soft cluster assignment vari- ables (X) given by an OT based clustering algorithm. Then, ((X) fed stochastic decoder, consisting of a Gumbel-softmax sampler and a deterministic decoder, to reconstruct the data sample.
4.1 VQ-VAE AS AN INSTANCE OF WASSERSTEIN AUTOENCODERS
Consider an observation with the true data distribution Pxx on the input space X. Let denote by Pz a fixed prior distribution on latent space Z. A sample from the prior distribution Z is mapped the imput space X' with a transformation parametrized by a neural network denote by Pa the induced marginal distribution, then the OT cost between Px and Pa has been shown by Tolstikhin al. (2018) to take a simpler form:
X.
W We(Px, = inf (ox.n~r(o(X,X)) ren(Px,PG)
inf Ex~rx Ez- G(Z))
Q:Q2=Pz
(2) where e is a cost function such as the Lg cost, Pc) denotes the admissible couplings be- tween Px and Pz. and consider conditional distribution (the encoder) Q(Z) such that its Z marginal Qz == Ex-rx (Q(2) X)] is identical to the prior distribution Pz. Tolstikhin et al. (2018)
showed that learning the generative model G amounts to solving the following objective:
min min Ex-Px R2 x0o fe(x G(2))) + XD(Qz, Pz).
(3)
G Q(Z X)
where À > 0 is Lagrangian multiplier and 'D is an arbitrary divergence between probability distri- butions on Z that controls the discrepancy between the prior and the learned conditional distribution. The resulting objective and models are called Wasserstein autoencoders (WAE).
VQ-VAE could also be interpreted as an instance of WAEs with discrete prior distribution. Let us denote by Cat(K, 1/K) the uniform categorical distribution such that P(r 5) 1/K for Cat(K, 1/x) and j= 1....... K. As noticed by Van Den Oord al. (2017), the prior distribution Pz is assumed to be defined on the latent discrete space {1..... K}d, with d, i.i.d. random variables following the uniform categorical distribution such that IP(Z. = 5) = 1/K any Z ~ Pz. Now let us consider a deterministic encoder hs that encodes each input X E X to h (X) E Rdd We define the posterior distribution (x( arg mink ... [X)x((); ckl the nearest code index. When Q(ZIX) takes the form of (XX and D is the KL divergence, the second term in Equation is constant and we recover the objective VQ-VAE after adding the hard assignment constraint in the posterior. Despite its simplicity and deterministic formulation, the main limitation of VA-VAE that it does not explicitly enforce the equal partition constraint introduced in the prior. In order fully match encoded training distribution the prior, we propose relax the hard assignment and use an optimal transport based divergence.
4.2 OPTIMAL TRANSPORT QUANTIZED AUTOENCODERS
Without loss of generality, we assume d, = 1 in the following and one could generalize our analysis for d, > by assuming (X); to be independent, which also used in Van Den Oord et al.

---

Under review as a conference paper at ICLR 2023
(2017). Instead of using the hard assignment for the posterior distribution, we propose to a soft assignment
Q(Z(X) - (xx) == 1ootmax)_....,x(IIMMAAa cx E A
(4)
where A denotes the probability simplex of dimension K - 1. If we sample 21
i.d.
from Px. Qz could be considered as discrete measure evenly distributed on _,,..., 2n where
2j = (sx)) A On the other hand, the prior distribution also a discrete measure such that = Σx 1/K6u, where us E A denotes the binary vector with only its k-th entry equal to using the one-hot representations. Then, the discrete optimal transport cost between the corresponding Qz and Pz is given by
7
K
min OT( TET((/)))))))
ΣΣ j=1k=1
min Cac(Es, rEM(V=,)))
Σ Jr log Reرn J.k
(5)
where is the KL divergence on A such that c(p,)) q) = >i pi log(m//.)) and II is the transportation polytope of the admissible couplings between Qz and Pz, commonly known in the literature on optimal transport (Peyré et al., 2019)
I(Ya, 1/k) = { T =1= 1 K
Thus, we recover the self-labeling scheme with the equal partition constraint developed in the field of self-supervised learning (Asano et al., 2020). particular when we project both (x( and recover the formulation of Caron al. (2020), which
prototypes e onto the unit L--sphere, amounts to solving an online OT problem for samples within a minibatch as detailed in Section 4.2.1. The only difference is that we perform clustering on the set of small components of the data, such as patches for images, while Caron et al. (2020) clustered the features of the entire images. After the clustering process, latent variable can be sampled from the categorical distribution given by Pj using the Gumbel-softmax relaxation, which will be detailed in Section 4.2.2.
Another way to derive the OT regularization in Equation is to consider a parametrized prior distri- bution instead of the fixed categorical distribution. Specifically, we consider mixture of K Dirac measures with equal mixing masses Zc Ek=1 1/Kde Now, consider (x == ha(X) as the posterior distribution Q(Z Then, the Wasserstein distance between Qz and Pz is equal to
min (2, reI((/,,//x)) J.k
Exxxllho(;;))- .
which admits the same solution as Equation since the denominator of εz* does not depend on k. This formulation suggests that C could be regarded as the Wasserstein Barycenter of (Z( (Cu- turi & Doucet, 2014). The resulting autoencoder could thus be simply interpreted as a Sinkhorn Autoencoder (Patrini al., 2020) when using the entropic regularization to the OT problem.
Based on the above observations, the final objective of the Optimal Transport Quantized Autoen- coders (OT-VAE) can be written as
LoT-VE min min Ex~PxE2 SxxxxxE2- (10 G Q(2Ix)
+X +XOT(Q2,,))
(8)
fut
where the first term corresponds to the reconstruction loss Ere and the second term Cot amounts to clustering the components of the data in the latent space under the equipartition constraint. An overview of OT-VAE illustrated in Figure 2.
4.2. ONLINE CLUSTERING WITH EQUIPARTITION CONSTRAINT
Here, we provide details about how compute the OT loss between Qz and Pz an online fashion following the techniques in Caron et al. (2020). In order to make the clustering algorithm online, we compute the the cluster assignments I' using the features of data components (e.g. image patch features) within batch. We thus assume to be the number of component features within batch and reuse the previous notations. We also assume hg = ho(0)) and G& to have unit T Torms
5

---

Under review as a conference paper at ICLR 2023
as in Yu et al. (2021) for generative modeling and Caron al. (2020) for self-supervised learning. Then, by adding an entropic regularization, the OT loss in Equation is equivalent to solving Caron et al. (2020)
min 5k =H(T),
j.k
where Dzt: -hj C and HT) - - (x - 1) is the entropic regularization with parameter 8 > 0. This problem can be efficiently solved by an iterative matrix scaling algorithm
known as Sinkhorn- -nopp algorithm. More details can be found in Section A.l of the Appendix.
Once a solution 1* to the above problem is found, we can inject it into the OT loss in Equation and obtain:
exp(1/hh3 cx)
OT P2)
ΣΣ ja log 2gk,r, where 2,k,7 Ek==1 (xp(1///))) cx)
(10)
j=1k=1
where we added a temperature parameter 2j. similar Caron al. (2020). This loss function is jointly minimized with the reconstruction loss in Equation 8 with respect to the autoencoder parameters and the codebook C. In our experiments, learnable parameter with its initial value as hyperparameter.
4.2.2 GUMBEL-SOFTMAN RELAXATION
Once the input xy mapped to the soft assignment 2; E A in Equation 4, we need to find way to sample from this categorical distribution with gradient estimator of the parameters. A simple approximation is the Gumbel-softmax relaxation (Jang et al., 2017; Maddison et al., 2017), which has also been used in the training of VQ-VAE (Esser et al., 2021) its variants (Takida al., 2022). Specifically, the Gumbel-softmax function is defined as a vector E A such that
(x(2,) exp(V//(log ZJ사 + gh))
exp(\/((1] Ct + 9k)) for k
K, (11)
exp(1/r((lg Z - gw)) k== exp(t/r((h) cv + gw))
providing continuous differentiable approximation to sample drawn from the categorical distribu- tion with class probabilities žzk when the temperature parameter r' approaching 0. Here, . are i.i.d. samples drawn from Gumbel(0, 1). In practice, we adopt the annealing strategy for r' as in Jang et al. (2017), by starting with a high temperature and annealing to small but non-zero value temperature. In such a way, the gradients at h, can be easily back proppagapp to the encoder using any deep learning framework allowing automatic differentiation.
4.3 PRIOR LEARNING WITH AUTOREGRESSIVE MODELS
Once the discrete representations are learned with VAE, one could use a deep autoregressive model for generation tasks, which has shown state-of-the-a -het---rttttttt performances on different tasks (Diele- man et al., 2018; Esser et al., 2021; Ramesh al., 2021; Dhariwal al., 2020). The principal idea to learn the probability of the encoded discrete representations through the chain rule of proba- bility: Z) p(Zz) [[= p(Z,]Zz--i)) Z1:111). The conditional probability can be parameterized by Generative Pre-trained Transformer (GPT) (Vaswani et al., 2017; Radford et al., 2018).
EXPERIMENTS
In this section, empirically validate OT-VAE on 3 different data modalities: images. speech and 3D human motions. For images, we use CclebA (Liu et al., 2015) and CelebAHQ 2015), which represent continuous and discrete data distributions respectively. For speech, ZeroSpeech 2019 (Dunbar et al., 2019). Finally, we use AIST++ (Li et al., 2021) for 3D dance motion generation.
(Liu et al.,
5.1 IMAGES
Datasets CelebA is face dataset proposed in Liu et al. (2015). Following SQ VAE (Takida et al., 2022), we apply OT-VAE the aligned and cropped version of the dataset, which consists of

---

Under review as a conference paper at ICLR 2023
Table Evaluation on the test set of CelebA (Liu al., 2015). We report MSE (x10-3). Follow- ing SQ-VAE (Takida et al., 2022), the codebook contains 512 codes and each code is with dimension 64. Experiments of reconstruction are repeated three times.
Method
MSE (x10 #){
Perplexity t
Generation FID +
VAE VQ- VAE + EMA
4.79 + 0.01
1,33 + 0,41
VQ- UQ-VAE- + EMA Code Reset
1.62 + 0.36
Gaussian SQ-VAE
0.96 + 0.00
413.2 + 4.9
20.8
OT-VAH (Ours)
0.94 0.00
433.3 +
20.0
Table 2:
Evaluation on the test set of CelebAHQ-Mask (Liu et al., 2015). We report pixel error (%), mloU, and perplexity. Following SQ- SQ-VAE (Takida et al., 2022), the codebook contains 64 codes and each code is with dimension 64. Experiments are repeated three times.
Method
Pixel ervor (%) [
mloU 1
Perplexity
VAE
8.79 + 0.01
55.8 + 0.3
VQ-VAR + EMA
6.95 + 0.14
59.7 + 0.7
46.2 + 2.0
NC SQQVAE
6.63 + 1.38
64.1 5.4
12.6 +.2
VMF SQ-VAE
3,51 + 0.17
74.6
+..6 0.0
52.4 + 0.8
VA- (Ours)
3,40 0.08
75.5 1 0.8
61.1 +: 1.0
202,599 images split into 162,770, 19,867 and 19,962 for train, validation and test. CelebAHQ- Mask contains annotations of face attributes of CelebA, which is thus a categorical image dataset with 19 categories. There are 24, 24,183, 2,993 and 2,82 2,824 images respectively in the train, validation and test set. For both datasets, we adopt the same pre pre-processing as SQ-VAE and report performances on the test set with the model performing the best on the validation dataset.
Evaluation metric For continuous data representation, i.e. CelebA, we report Mean Squared Error (MSE) to quantify the reconstruction quality. For discrete data representation as CelebAHQ-Masse the reconstruction quality is measured using the percentage of pixels that are incorrectly predicted (Pixel error), and the mean of the class-wise intersection over union (mloU). We also report the perplexity on the test set, which can be considered as a measure of the codebook utilization. The perplexity of distribution p is given by the exponential of its entropy: PPI.(p) = II -P(x)
Implementation details We follow the evaluation protocol provided by SQ-VAE (Takida et al., 2022) and keep the same backbone architecture of OT-VAE as SQ-VAE. For both datasets, the hyper-paramters of OT-VAE are set as A=le-3 (regularization weight in Equation 8) and log 1/r (initial temperature in Equation 10). The ablation study of these hyper-parameters on CelebA (Liu et al., 2015) are provided in Appendix C. We also provide more implementation details in Ap- pendix D, which includes training schema, architectures,
ayper-parameters etc.
The results are provided in Table 1 and 2 for CelebA (Liu et al., 2015) and CelebAHQ- Mask (Liu et al., 2015) respectively. We repeat each experiment of reconstruction three times and report the average and standard deviation for each metric. For generation on CelebA, use standard GPT as the generative model. More details can be found in the Appendix D. For both continuous and discrete data distributions, our OT-VAE achieves clearly better reconstruction quality than other variants of VQ- VAE. In particular, the proposed OT- VAE obtains higher perplexity which suggests a more efficient usage of the codebook capacity than SQ- -Q-VAE (Takida et al., 2022). Note that SQ-VAE requires selecting proper priors for different tasks. Using either Naive categorical (NC SQ- VAE) or von Mises-Fisher (vMF SQ- SAF) as the prior, the performance vary significantly on CelebAHH-MMask. In contrast, the proposed OT VAE does not need any selection of priors and its demonstrating its superior robustness.
Results
performances are obtained with the same Our approach also results in better performance for generation, compared to SQ-VAE Note that computing the coupling matrix takes extra computation for training. Precisely, training OT-VAE takes about 102s for epoch CelebA, while training standard VQ- VAE only takes about 89s. For inference, both methods take the same time as the OT coupling matrix is not needed.
'SQ-VEE: https: thub. com/ hoom//so////a// sqvae /taeee/rrreee tttt//eeeee//////
7

---

Under review as a conference paper at ICLR 2023
Table 3: Evaluation on the test set of ZeroSpeech2019 (Dunbar al., 2019). We report MSE (dB²), perplexity, FDSD and KDSD. Following SQ-VAE (Takida et al., 2022), the codebook con- tains 512 codes of dimension 64. Experiments of reconstructions are repeated three times. 'r' and "g' indicate resynthesis and generation respectively.
ZeroSpeech2019
Method
MSE (dB*).
Perplexity f
rFDS
rKDSD &
gFDSD
gKDSD &
Natural speech
5.82
3.03
5.82
3.03
VQ-VAE+ + EMA
34.33 = 1.57
Gaussian SQ-VAE
32.13 + 1.78
440.0 + 2.7
8.95 = 1.11
14.55 + 1.98
10.90
15.09
VA- (Ours)
28.73 = 0.29
446.2 + 6.5
7.29 = 0.44
11.61 + 0.78
10.60
14.55
5.2
SPEECH
Datasets ZeroSpeech2019 (Dunbar et al., 2019) is multi-speaker corpus sampled at 16000 Hz. Following SQ-VAE (Takida et al., 2022). we use the subset train. voice and train_umit as the train set (~20.33 hours), and train train-ppaaaa p-rpaaalllviiicccccoo for the test set (~10 mins). We randomly sample 200 samples from the train set as a validation set to select the best model. The number of total speakers for both train and test is 102. We apply the same pre-processing as in SQ-VAE (Takida et al., 2022) to extract 80-dimensional log-mel spectrogram features for all the speech data. We provide more details in Appendix E.
Evaluation metric Following SQ- VAE (Takida et al., 2022), we report MSE in dB² to measure the reconstruction quality and the perplexity for codebook utilization.
Implementation details We follow the evaluation protocol provided by SQ-VAE (Takida et al., 2022) and keep the backbone architecture of OT-VAE same as VAE. The hyper-paramters of OT-VAE are set as X=le-3 (regularization weight in Equation 8) and log 1/- = 5 (initial temperature in Equation 10). The ablation study of these srper-parameters on ZeroSpeech2019 (Dunbar et al., 2019) is provided in Appendix C. More implementation details such as training schema, architec- tures, and hyper- paper-parameters are given in Appendix E.
Same SQ-VAE (Takida et al., 2022), we repeat each ex- periment on reconstruction three times and report the mean and standard deviation for each metric. For generation, we use standard GPT, more details can be found in the Appendix E. Our proposed OT-VAE significantly outperforms SQ-VAE (Takida et al., 2022) both in terms of the reconstruc- tion quality (lower value both MSE, FDSD and KDSD), the codebook utilisation (6.2 higher in perplexity), and generation quality (lower FDSD and lower KDSD).
Results We show the results in Table 3.
5.3 3D HUMAN MOTIONS
Dataset In order to fully explore the capacity of OT VAE, we also conduct experiments for chal- lenging task, namely 3D dance generation. We use AIST+ (Li et al., 2021), which is the largest publicly available dataset for this task. The dataset contains 992 pieces paired music-motion quences, where 952 are kept for training and 40 are used for evaluation. Motion sequences are provided in the SMPL (Loper et al., 2015) format with 60 FPS. Following Li et al. (2022), we use the 3D coordinates of 24 joints to represent motions. The task aims to generate motions conditioned seeeenceeuee
on music sequences. We compare our method to the state-of-the-are method (Li et al., 2022), based on VQ-VAE. We provide more details about the baseline Li et al. (2022) in Appendix F.
Evaluation metric We measure the dance quality with Frechet Inception Distances (FID) (Heusel et al., 2017) between the predicted motion and all motion sequences (including training and test data) of the AIST++ dataset. Following Li et al. (2021; 2022), we consider two features: kinetic features (Onuma et al., 2008) (denoted 'x') and geometric features (Müller et al., 2005) (denoted as "g'). We extract both features using fairmotion 2, Following Li et al. (2021; 2022), we also report diversity, which is the average feature euclidean distance of generated motions. We additionally anjlddaanheeunnnlaaaia the waihhe wwtilizztiPP. architectures as in Li et al. (2022) for discrete representation learning and generation respectively. We train the proposed OT-VAE using
ffarmotion: https://githubs facebook reseaa rch/ rmotion
8

---

Under review as a conference paper at ICLR 2023
Two samples of dance motion are generated with the proposed OT VAE, which are condi- tioned on different pieces of music from the test set of +IST++ (Li et al., 2022).
Figure 3:
Table 4: Evaluation on the test set of AIST+ (Li et al., 2021) for 3D dance generation. FIDx and FIDg are used to evaluate generation quality. We also report DIV, and DIV, quantify motion diversity for completeness. Following Li et al. (2022), we train two codebooks for upper and lower body respectively. Each codebook contains 512 codes with dimension 512.
Motion Quality FIDg
Motion Diversity
Method
Perplexity
FID₂ 4
DIV, f
DIV, f
Ground Truth
17.10
10.60
8.19
7.45
Motion Reconstruction
VQ-VAE + FMA + Code Reset (I.i et al., 2022)
28.23
12.63
6.80
6.57
138. 14
VT-VAB (Ours)
23.50
9.42
6,98
6.47
243.21
Motion Generation Approaches
Li et al. (Li et al., 2020a) DanceNet (Zhuang et al., 2022) DanceRevolution (Huang et al., 2021)
86.43
43.46
6.85
3.32
69.18
25.49
2,86
2.85
73.42
25.92
3,52
4.87
FACT (Li et al., 2021)
35.35
22.11
5.94
6.18
Motion Generation with GPT
VQ-VAF + FMA + Code Reset (I.i et al., 2022)
28.75
11.82
6.41
6.13
166.84
OT VAE (Ours)
26.74
9.81
6.52
5.89
334.53
the same reconstruction loss as in Li et al. (2022), measuring the reconstruction between the input and reconstructed motion as well as their velocity and acceleration. The hyper-parameter are set as x=le-3 (regularization weight in Equation 8) and log 1/r (initial temperature in Equation 10). Note that the VQ- used Li al. (2022) trained using the stop-gradient operator, EMA and codebook reset. We provide more implementation details in Appendix F, which includes training schema, architectures, byper-parameters.
Results Metrics for both reconstruction and generation are shown in Table 4. For reconstruction, the proposed OT- VAE achieves better performance than Li et al. (2022) which uses a standard VQ-VAE. VAE. Note that numerous heuristics used in Li et al. (2022), while non of these heuristics are employed in OT- and we achieve better reconstruction quality and higher perplexity.
The downstream generation task also benefits from the better reconstruction. Using the same GPT architecture, our OT-VAE achieves state-of- of-the-rrttttttte on this task and outperforms VQ-VAE Li et al. (2022) in terms of motion quality and achieves comparable diversity. We showease and illustrate two samples of generated motion in Figure 3, which demonstrates the high-quality motion generated by our algorithm. In the supplementary material, additionally provide videos containing dance motion along with the input music.
6 CONCLUSION
We proposed OT- a WAE-based generative model with discrete prior and a Sinkhorn di- vergence to match the encoded training distribution and the prior. Our approach first shows the clustering techniques developed in the field of self-supervised learning could be beneficial for gen- erative tasks. With the rapid development of self-supervised learning in recent years, we believe that more algorithms arising from self-supervised learning, such contrastive strategies (Jaiswal al., 2020), knowledge distillation (Grill et al., 2020; Caron et al., 2021) or a direct use of large-scale pre-trained models could be adaptable for generative modeling.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Humam Alwassel, Dhruv Mahajan, Bruno Korbar, Lorenzo Torresani, Bernard Ghanem, and Du Tran. Supe-sup learning by cross-modal audio-video clustering. In Advances Neural Information Processing Systems (Neur[PS), 2020.
Dario Amodei, Sundaram Ananthanarayanan, Rishita Anubhai, Jingliang Bai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Qiang Cheng, Guoliang Chen, et al. Deep speech 2: End- to-end speech recognition in english and mandarin. 2016.
Yuki Markus Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultancous clustering and representation learning. International Conference on Learning Representations (ICLR), 2020.
Mikolaj Binkowski, Danica Sutherland, Michael Arbel, and Arthur Gretton. Demystifying mmd gans. International Conference Learning Representations (ICLR), 2018.
Mikotaj Binkkwski,, Jeff Donahue, Sander Dieleman, Aidan Clark, Erich Elsen, Norman Casagrande, Luis C Cobo, and Karen Simonyan. High fidelity speech synthesis with adversarial networks. International Conference on Learning Representations (ICLR), 2019.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsu- pervised learning of visual features. In Proceedings of the European Conference on Computer Vision (ECCV), 2018.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Un- supervised learning of visual features by contrasting cluster assignments. In Advances Neural Information Processing Systems (NeurIPS), 2020.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the International Conference on Computer Vision (ICCV), 2021.
Jan Chorowski, Ron Weiss, Samy Bengio, and Aäron Van Den Oord. Unsupervised speech rep- resentation learning using wavenet autocncoders. IEEE/ACM transactions on audio, speech, and language processing, 2019.
Marco Cuturi. Sinkhom distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems (NeurIPS). 2013.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International Conference on Machine Learning (ICML), 2014.
Prafulla Dhariwal, Hcewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and llya Sutskever. Jukebox: A generative model for music. arXiv, 2020.
Sander Dieleman, Aaron van den Oord, and Karen Simonyan. The challenge of realistic music generation: modelling raw audio scale. In Advances Neural Information Processing Systems (NeurIPS), 2018.
Ewan Dunbar, Robin Algayres, Julien Karadayi, Mathieu Bernard, Juan Benjumea, Xuan-Nga Cao, Lucie Miskic, Charlotte Dugrain, Lucas Ondel, Alan W. Black, Laurent Besacier, Sakriani Sakti, and Emmanuel Dupoux. The Zero Resource Speech Challenge 2019: TTS Without T. In Proc. Interspeech 2019, pp. 1088-1092, 2019.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the Conference Computer Vision and Pattern Recognition (CVPR), 2021.
Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Berhaard Schölkopf. From variational to deterministic autoencoders. In International Conference on Learning Representa- fions (ICLR), 2020.
10

---

Under review as a conference paper at ICLR 2023
Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a approach to superviue learning. Advances in Neural Information Processing Systems (NeurIPS), 2020.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge local nash equilibrium. In Advances in Neural Information Processing Systems (NeurlPS), 2017.
Ruozi Huang, Huang Hu, Wei Wu, Kei Sawada, Mi Zhang, and Daxin Jiang. Dance revolution: Long-term dance generation with music via curriculum learning. In International Conference on Learning Representations (ICLR), 2021.
Ashish Jaiswal, Ashwin Ramesh Babu, Mohammad Zaki Zadeh, Debapriya Banerjee, and Fillia Makedon. A survey on contrastive self-supervised learning. Techmologies, 2020.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparametrization with gumble-softmax. In International Conference on Learning Representations (ICLR), 2017.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference Learning Representations (ICLR), 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014.
Soheil Kolouri, Phillip E Pope, Charles E Martin, and Gustavo K Rohde. Sliced wasserstein auto-
encoders. International Conference on Learning Representations (ICLR), 2018.
Jungil Kong, Jaehyeon Kim, and Jaekyoung Bae. Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis. Advances in Neural Information Processing Systems (NeurIPS), 2020.
Oleksii Kuchaiev, Boris Ginsburg, Igor Gitman, Vitaly Lavrukhin, Carl Case, and Paulius Micike- vicius. Openseq2seq: extensible toolkit for distributed and mixed precision training sequence- to-sequence models. In Proceedings of Workshop for NLP Open Source Software (NLP-OSS), pp. 41-46, 2018
Jiaman Li, Yihang Yin, Hang Chu, Yi Zhou, Tingwu Wang, Sanja Fidler, and Hao Li. Learning to generate diverse dance motions with transformer. arXiv, 2020a.
Junnan Li, Pan Zhou, Caiming Xiong, and Steven CH Hoi. Prototypical contrastive learning of unsupervised representations. In iclr, 2020b.
Ruilong Li, Shan Yang. David A Ross, and Angjoo Kanazawa. Ai choreographer: Music ditioned 3d dance generation with aist++. In Proceedings of the International Conference on Computer Vision (ICCV), 2021.
Siyao Li, Weijiang Yu, Tianpei Gu, Chunze Lin, Quan Wang, Chen Qian, Chen Change Loy, and Ziwei Liu. Bailando: 3d dance generation by actor-critic gpt with choreographic memory. In Proceedings of the Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the International Conference on Computer Vision (ICCV), 2015.
Matthew Loper, Naureen Mahmood, Javier Romero, Gerard Pons-Moll, and Michael J Black. Smpl: A skinned multi-person linear model. ACM transactions on graphics (TOG), 2015.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv, 2017.
C Maddison, A Mnih, and Y Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations (ICLR), 2017.
Alireza Makhzani. Jonathon Shlens, Navdeep Jaitly, and Ian Goodfellow. Adversarial autoencoders. In International Conference on Learning Representations (ICLR), 2016.
11

---

Under review as a conference paper at ICLR 2023
Meinard Müller, Tido Röder, and Michael Clausen. Efficient content-basee based retrieval of motion capture data. In ACM SIGGRAPH, 2005.
Kensuke Onuma, Christos Faloutsos, and Jessica K Hodgins. Fmdistance: A fast and effective distance function for motion capture data. In Eurographics (Short Papers), 2008.
Giorgio Patrini, Rianne van den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncertainty in Artifi- cial Intelligence, 2020.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning. 2019.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ily Sutskever, al. Improving language un- derstanding by generative pre-training. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alee Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference Machine Learning (ICMI), 2021.
Yuhta Takida, Takashi Shibuya, WeiHsiang Liao, Chieh-Hsii Miiin Lai, Junki Ohmura, Toshimitsu Ue- saka, Naoki Murata, Shusuke Takahashi, Toshiyuki Kumakura, and Yuki Mitsufuji. Sq-vae: Varia- tional bayes on discrete representation with self-annealed stochastic quantization. In International Conference Machine Learning (ICML.), 2022.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. In International Conference on Learning Representations (ICLR), 2018.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. In Advances in
Neural Information Processing Systems (NeurIPS), 2017.
Benjamin van Niekerk, Leanne Nortje, and Herman Kamper. Vector-Quanteen Neural Networks for Acoustic Unit Discovery in the ZeroSpeech 2020 Challenge. In Interspeech, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems (NeurIPS), 2017.
Will Williams, Sam Ringer, Tom Ash, David MacLeod, Jamie Dougherty, and John Hughes. Hierar- chical quantized autoencoders. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
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 (ICLR), 2021.
Wenlin Zhuang, Congyi Wang, Jinxiang Chai, Yangang Wang, Ming Shao, and Siyu Xia. Mu- sic2dance: Dancenet for music-driven dance generation. Transactions on Multimedia Computing, Communications, and Applications (TOMM), 2022.
12

---

Under review as a conference paper at ICLR 2023
Appendix
In the appendix, we provide the following content:
1. Additional background in Appendix A, including Sinkhorn-Knope Algorithm.
2. Pytorch implementation of OT-VAE Appendix
3. Analysis on and т (regularization weight and initial temperature of Lor) for OT-VAE the test set of CelebA (Liu et al., 2015) and ZeroSpeech2019 (Dunbar et al., 2019) in Appendix C.
4. Details of experiments on images in Appendix D.
5. Details of experiments speech data in Appendix E.
6. Details of experiments on 3D dance generation in Appendix F.
Additionally, the supplementary material (the uploaded zip file, suppMat LTVAEEtt_OOV...... zi,), we also provide the source code to reproduce experiments on images in the folder /suppMat_ OTVAE/ code z zip and videos of generated dance in /sspppMa___T/////issuuaalll zip.
A ADDITIONAL BACKGROUND
This section provides additional background including the details about the Sinkhorn-Knopp algo- rithm.
A.1 SINKHORN-KNOW ALGORITHM
The online clustering problem with the equipartition constraint is equivalent to the entropy- regularized OT problem in Equation 9, given by
min (T, D)r - =H((), rem(/=a./k)
(12)
where =IT) (log i 1) and DトE- - c is the cost matrix. This problem can be efficiently solved by the Kankhorn-Knopp algorithm, which is iterative matrix scaling method to approach the double stochastic matrix. Specifically, the f-th iteration of the algorithm performs the following updates:
1/n Sultj
1/K
11 (E)
(e+1)
S (e)
where = eD/e DD= Then, the matrix diag(u(") (Sdiig(())) converges to the solution of Equation 12 This algorithm converges faster with larger strong regularization leads to
when l tends a more convex objective (Peyré et al., 2019). However when & becomes too small, some of the elements of the denominators become null and result in division O. To overcome this stability issue, it is preferable to do the computations in the log-scale (Peyré et al., 2019). This algorithm can be easily adapted to a batch of encoded data features h, leading to scalable and GPU-friendly computation. In our experiments, we also observed that only few iterations or 10) are sufficient to obtain good performance. A Pytorch implementation of this algorithm is provided in Section B.
B PYTORCH IMPLEMENTATION OF VA-VAE
C cocebook
Y features from the encoder)
2. features after ot quant quantization)
tau:
learnable temperature parameter in
13

---

Under review as a conference paper at ICLR 2023
import torch
import torch. nn, functional as import math
Sinkhorn-Knopp algorithm
11 def Log_Sinkhorr (K, eps 0.5):
12
m, = K.shape
13
= k.new_zeross (m,))
a, b, 0, math.log (m / n), eps
15
for
16
u -torch. logsumexp (v. v..view(e (m, 1) K, dim-0) +
17
v -torch. logsumexp (u. (1, n) + K, dim=l) +
return torch. exp (K u.view n) v.view (m, 1))
20
21
def Gumbel-Softmax Sample algorithm <--
22
(logit, gumbel_temporaturr)
#Sampling from uniform distribution on [0, rarch.rand (loggit))
25
sample logit torch. log (-torch, log (U))
26
return m.aoftmax (sample / caanbelltemerratuur dim--1)
27
---> L2 Normalization
29
norm E.normalize (Y, p=2.0)
30
C_norm E.normalize (C, p-2.0)
31
----> Compute probability <---
logit Y_normmmm T) tau.exp()
34
probs torch.softma (logit, dim=-1) log_probs =
35
(logit, dim=-1)
36
---> optimal Transport <---
38
with torch.no_grade
39
ot Leg__inkhorn
40
---> Training <---
isstraining:
43
sampling Gumbel Softmax Gumbel_Softmax Sample (Logit, g temperature)
44
---> Evaluation <---
else:
sampling torch. argmax (probs, dim=-1)
47
48 torch. mm (sampling, C)
49
oT Loss <---
loss -terch.. mean oq_ log_prcbs)
C ABLATION STUDY
In this section, we provide analysis À and (regularization weight and initial temperature of Lot) for OT-VAE on the test set of CelebA (Liu et al., 2015) and ZeroSpeech2019 (Dunbar et al., 2019). Results illustrated in Table 5. We report MSE (x10- for CelebA, MSE (dB2) for ZeroSpeech2019 and perplexity (PPL) for both datasets. Two insights can be found from the table: first, with larger X, PPL importantly boosts which validates the fact that optimal transport explicitly brings equipartition of the clusters; second, le 3 works decently across different modalities while the temperature needs to tuned according to different tasks.
D DETAILS OF EXPERIMENTS ON IMAGES
Pre-processing Following SQ- VAE (Takida et al., 2022), all the images in CelebA (Liu et al., 2015) are center-cropped to 140 x 140 and resized to 64 64, which is setting used in previous
14

---

Under review as a conference paper at ICLR 2023
Table 5: Analysis of À and 7 (regularization weight and initial temperature of Lor) for OT- VAE on the test set of CelebA (Liu et al., 2015) and ZeroSpeech2019 (Dunbar et al., 2019). We report MSE (x10-3) for CelebA, MSE (dB²) for ZeroSpeech2019 and perplexity (PPL) for both datasets. The codebook contains 512 codes and each code is with dimension 64. All the experiments are repeated three times.
CelebA
ZeroSpeech2019
log 1
MSE (x10 1)
PPL 1
MSE (dB?) )
PPL f
0.1
0.96 + 0.0
391.7 + 4.9
32.62 + 0.97
433.9 + 8.8
0.5
0.98 + 0.0
373.3 +: 2.7
31.22 + 0.49
435.8 + 4.6
le-4
0.98 0.0
365.6 16.4
32.65 + 1.06
5
1.0
32..1
28.80 + 0.36
438.8 3.6
10
0.9
1
31.37 0.79
433.8 15.5
0.1
0.95 0.0
4.9
32.00 0.28
448.2 2.1
0.5
0.96 0.0
2.6
30.76 0.41
450.2 1.7
1e-3
1
0.94 0.0
43.
31.96 0.35
452.9 2.0
5
0.95 = 0.0
426.6 + 4.4
28.73 0.29
446.2 + 6.5
10
0.99 0.0
424.3 2.8
32.41 0.66
443.5 + 3.0
0.1
0.96 +996+00 0.0
446.7 16.0
+2.15+0.09
453.4
0.5
0.96 0.0
4.0.2
32.32 0.98
445.2
le-2
1
0.96
448.6 0.6
30.79 0.43
454.9 5.5
1.00 0.0
450.9 3.2
29.07 + 0.52
449.5 1.7
10
1.00 - 0.0
413.7 + 2.2
30.61 + 0.83
440.8 + 2.0
work (Tolstikhin et al., 2018; Ghosh et al., 2020). For CelebAHQ-Mask (Liu et al., 2015), we center- crop the the segmentation maps 128 128 and resized to 64 x 64 with the nearest neighbor interpolation.
Architectures for reconstruction We use the same architectures as SQ-VAE (Takida et al., 2022), which is adopted from the GitHub repository of DeepMind Sonnet and composed of ConvResNet-type encoder and decoder. These networks include convolutional layers, transpose convolutional layers, and ResBlocks, the downsampling rate from the image resolution to quantized layer is Number of residual blocks are 6 for CelebA (Liu al., 2015) and 2 for CelebAHQ- Mask (Liu et al., 2015).
Model architectures for generation As SQ-VAE did not release their implementation for gen- eration, we reimplemented SQ-VAE using the same model architecture as for OT-VAE for a fair comparison. For both SQ-VAE (Takida et al., 2022) and OT-VAE, we train standard GPT model (Vaswani et al., 2017: Radford et al., 2018) to generate images in an auto-regressive fash- ion. The GPT model has 12 layers, 12 heads, and feedforward dimension 768.
Training details Following SQ- SQ-VAE (Takida et al., 2022), the reconstruction loss Ere is the MSE loss for both datasets. We adopt the same optimization schema as SQ-VAE (Takida et al., 2022). we leverage Adam (Kingma Ba, 2015) optimizer with inital learning rate 0.001, batch size 64, 31 = 0.9 and 32 0.99. The learning rate will be halved every 3 epochs if the validation loss is not improving. We train 200 epochs for CelebA and CelebAHQ-Masse We validate on the same capacity of codebooks
SA-V (Takida et al., 2022): the codebook contains 512 codes for CelebA and 64 codes for CelebAHQ-Mast. each code is with dimension 64. The training takes ~ 2 hours for CelebA using single GPU NVIDIA V100-32G. The gumbel softmax decay is 1e-5 in all the datasets. The validation MSE and perplexity (PPL) on CelebA (Liu et al., 2015) during the training illustrated in Figure 4a and 4b respectively, which shows the superiority of the proposed OT- VT- compared to SQ-VAL VAE (Takida et al., 2022).
Sonnet: https wwnnet/eonenteeeeeeeeeeeeeeeeeee
ipybb
15

---

Under review as a conference paper at ICLR 2023
dr SQ-WAC
50 VAE
Aparks
Spark
Validation MSE
Validation PPL
Figure 4: The validation MSE and perplexity (PPL) on CelebA (Liu et al., 2015) during the training. The propose OT VAE obtains better reconstruction with higher perplexity, Experiments are repeated three times.
Figure 5: Generated images CelebA (Liu et al., 2015) with the proposed OT- VAE.
For generation, an image is first encoded into a sequences of 256 tokens through the encoder of OT- VAE and VAE. As standard strategy, SOS token is padded to the beginning of each sequence. During training, we perform next token prediction with the teacher forcing. For generation. For generation, start with the SOS token and use simple multinomial sampling with the predicted
16

---

Under review as a conference paper at ICLR 2023
probability to generate the next token and repeat this step until the full length is reached. We train 300K iterations with AdamW (Loshchilov & Hutter, 2017) optimizer, and set the batch size to 16. We generate 10K images and use torch-fidelity to compute the FID. We provide visual results in Figure 6.
E DETAILS OF EXPERIMENTS ON SPEECH DATA
Pre-processing Following SQ- VAE (Takida et al., 2022), we apply the same pre-processing as van Nickerk et al. (2020) to extract 80-dimensional log-mel spectrogram features for all the speech data. We firstly scale the maximum amplitude of each audio signal to 0.999 and pre-emphasize the scaled audio signals with first-order autoregressive filter (y/ Jt 0.97x,-1- where Ty and y/s indicate the input and output of the filter at time t), then apply 10 ms frame shift, and frequency cutoffs between 50 Hz and 8000 Hz. After that, we clip those bins that are 80.0 dB lower than the maximum and re-scale the log-mel spectrogram by dividing it by 80.
-ointppintttoo FFT with 25 ms Hann window,
The architecture for speech data is the same as SQ- SQ-VAE (Takida et al., 2022), which is with ConyResNet-type encoder/decodeer, whereas the Conv2d for the image data is replaced by Convld accommodate the structure of speech data. The downsampling rate from the mel spectro- gram to the quantized layer is 2.
Architecture
Model architectures for generation Similar to image generation, we use the GPT architecture with 12 layers, 12 heads and 768 feedforward dimension.
Training details Following SQ- VAE (Takida et al., 2022),
Adam (Kingma & Ba, 2015)
optimizer with inital learning rate 0.0004, batch size 256, B1 = 0.9 and }3 = 0.99. We train 50K iterations and pick the best model by validating our models on 200 samples, which are randomly selected in advance from the training set. The learning rate will be decayed with y = 0.5 on 30,000 and 40,000 iterations. We have 512 codes with dimension 64 in the codebook, which is also the same as SQ-VAB (Takida et al., 2022). As suggested in Chorowski et al. (2019); Takida et al. (2022), we add a time-jitter regularization with a replacement probability of 0.5 during training, which is also used in SQ- SQ-VAE (Takida et al., 2022). The training takes ~ 12 hours using single GPU NVIDIA V100-16G.
During training, we clip the waveform with random start point to get a sequence of 100 tokens after the encoder of OT-VAE or SQ-VAE. Then we also performed next token generation with a SOS token concatenated at the beginning. We train 100k iterations with Adam W optimizer with the learning rate of 2e 4, and the batch size of 32. For generation, we start with the SOS token and 15 warm-up tokens obtained from the encdoer, and autoregressively generated the next 85 tokens with top-k (k = 5) multinomial sampling. Then the 100 tokens in total can be fed into the decoder of OT/SQ- VAE to obtain the generated log-mel spectrograms.
To evaluate the quality of generated audio, we follow Binkowski et al. (2019) to compute Frechet distance and the Maximum Mean Discrepancy between the generated data and the real data, which are also named FDSD and KDSD respectively in the reference paper. As suggested in Binkowski al. (2019), we obtained the features by feeding audio clips of 2x-length into the pre-trained Deep- Speech2 model (Amodei et al., 2016) from NVIDIA OpenSeq2Seq library (Kuchaiev et al., 2018) and averaged along the temporal dimension. Since the pre-trained model takes the raw waveforms as input, we also trained a HiFi-GAN model (Kong et al., 2020) on ZeroSpeech dataset to transform the mel-spectrogram into raw waveform. Due the time limit, we used the simple V3 configura- tion, trained with batch size of 128 on 8 V-100 gpus and took the checkpoint at 180k to perform the waveform generation.
Considering that in the test set (train parallel voice), there are 1559 out of 190 audio sequences that are longer than 2s, we cropped 63 chunks with random start point for each audio, which resulted in a total of 10,017 audio clips for reasonable estimates.
chttps://toroa Fidel wtw.tt....ddddsss....
17

---

Under review as a conference paper at ICLR 2023
Myluurequency spectragram am
Nei Prequenyyycy spectrogrem
5:00
08
70ch
2010
15
1021
3024
-00
10
-20 AB
0.3
15
0.5
355
Mel- Prequeeyyyycyy
Nel Rellfrreeeccy spectrogram
-15 02
25 03
30
20 de
>>
0.5
Thmss
Timee
Mel-troguacc spactrag am
Mel-hroquency spectrogram
4003
4220
6915
2310
-15 m
2
1024
2014
20
312
322
50 de
Figure 6: Generated audio spectrogram on ZeroSpeech2019 (Dunbar al., 2019) with the proposed OT-VAE.
Evaluation metrics Here we provide more details for the computation of Fréchet distance (i.e. FDSD) and Maximum Mean Discrepancy (i.e. KDSD). Given features from the target data X E mm and the features from the natural data Y Raxd we have:
FDSD(X, Y)
||xx
1x + Tr Ex + +EY 2(Ex 2x
(13)
KDSD(X,X) Y)
Σ
A(X, n(n -
Σ IsijSnitj
K(Y,, Y,)
m(m 1)
1sijSm,itj
k(X,, Y,),
(14)
where px, By and Ex, Ey indicate the means and variances X and Y respectively, whereas k: Rdd x Rd is positive definite kernel function. Follow Binkowski et al. (2019; 2018),
18

---

Under review as a conference paper at ICLR 2023
choose the polymial kernal:
k< y) (0)""++))
(15)
To evaluate the distribution of natural data, split the test set in half and compute these two metrics between the two subsets.
F DETAILS OF EXPERIMENTS ON 3D DANCE GENERATION
Our baseline (Li et al., 2022) Based on VQ-,AE, Li et al. (2022) is two-stage framework for music to dance generation. In the first stage, trains two VQ- to reconstruct 3D dance motions for upper and lower bodies separately. The training of VQ-VAE includes common heuristics such as: the commitment and embedding loss with stop-gradien operator Eqqation 1), exponential moving average (EMA) and codebook reset. In the second stage of training, GPT is trained to generate 3D motion conditional on music sequences.
Architecture We follow the official code of Li et al. (2022), which is released at https: //github. com/ 1 i yailyo21/BBa lando. The encoder-decoder architecture is composed of dilated 1D convolutions with kernel 4 x 4 and residual blocks. Convolutions with stride 2 are used in the encoder to reduce temporal resolution while transposed convolution with stride 2 are used in the decoder to increase the temporal resolution. The temporal downsampling rate is 8. Both codebooks for upper and lower bodies are with dimension 512 X 512. The GPT is composed transformer with 12 layers and heads. The embedding size in GPT is 768.
Training details Denoted a dance motion as X E (Jx33),), where T is the temporal length and J = 24 is the total number of joints in SMPL (Loper et al., 2015), we learn two codebooks for upper and lower bodies separately. We use the same reconstruction loss Li al. (2022) consisting the reconstruction between the input and the reconstructed motion as well as their the velocity V(.) and acceleration A(), i.e. Cre = (X, GZ)) LV (V(),, ((G())) + L1 (A()), (G))) We follow similar training schema as Li et al. (2022): first train the motion without the root transla- tion for 200K iteration then freeze the encoder, codebook and decoder to train an extra decoder for another 200K iteration optimize the root velocity. We Adam (Kingma & Ba, 2015) optimizer with learning rate 2e-4, batch size 128, T= 64, B1 = 0.9, B3 = 0.99.
We further train Generative Pre-trained Transformer (GPT) for music dance generation. We use music feature as Li et al. (2022). We encoder the motion into codebook index using our OT-VAE and train the GPT predict the codebook index in auto-regressive manner with cross-entropy loss. For inference, the motion is generated in an auto-regressive manner. The GPT is trained for 600K iteration with Adam (Kingma & Ba, 2015) optimizer. The learning rate set 2c-4 then decrease to le-5 at 400K iterations. We use batch size 32, = 256, 81 = 0.9, 82 = 0.99.
19


Reviews:
[{"review_id": "ULSxz0jvJ6", "reviewer": "Reviewer_cXiL", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "In this article the author seek to develop a method to solve the code collapsing problem which is common in VQ related methods. The nature of the proposed method is to regularize the code selection with a uniform prior so that different codes will be evenly used. The author conducted experiments in different settings including visual images, speech etc. ", "summary_of_the_review": "This manuscript present a sophisticated regularized version VQVAE and put the algorithm in OT framework. The manuscript is well written , the math is clean and the experiments are well organized. The largest concern I have is the novelty. ", "strength_and_weaknesses": "Strength: the authors put the regularization of code usage into a detailed and solid mathematic framework of optimal transport and then use an Gumbel softmax max to sample the code in a one-hot manner. This strategy allows for a lot of flexibity and stochasticity of temperature applied. The experimental design is solid \n\nweakness, not sure of the advanced framework is needed as the core of the problem is bias in code usage by regularizing the codebook usage. Therefore, I am not sure about the real novelty ", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "IU_FrF8t_L", "reviewer": "Reviewer_MP1B", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": " - As one key challenge of existing vector quantization methods comes from codebook collapse, this work proposes OT-VAE, which regularizes the quantization by explicitly assigning equal number of samples to each code. \n\n - The proposed method enforces the full utilization of the codebook while not requiring any heuristics. \n\n - Across three different data modalities, the authors empirically demonstrate that OT-VAE shows better reconstruction compared to other VQ-VAE variants. ", "summary_of_the_review": "As discussed in Strength and Weakness, I think this work is well-motivated. Meanwhile, with detailed figures and detailed explanations, the authors present OT-VAE clearly. Existing weak points about this work include: 1. some slightly-incorrect statement, such as the benefits over previous heuristics,  2. lack of some relevant discussions and comaprisons.", "strength_and_weaknesses": "#### **Strength**\n\n  - The problem of Vector Quantization is timely and important, with increasing number of applications. The method is also well-motivated, as existing key challenges in VQ lie in low codebook usage. \n\n - The approach is well-represented. The authors make the tables and figures easy to follow and nicely-illustrated.\n\n - Extensive experiments across three different modalities are conducted. Notably, OT-VAE achieves state-of-the-art results on 3D dance generation.\n\n\n#### **Weaknesses**\n\n - The authors stated that OT-VAE does not need to include other heuristics, such as stop-gradient, EMA, codebook reset. However, according to Table 5, the method is relatively sensitive towards regularization weight and initial temperature. Therefore, additional burdens on tuning these parameters exist.\n\n - The technique of L2 normalization has been adopted in ViT-VQGAN, which is not discussed here.\n\n - Quantitative improvements seem to be limited, especially for image modality, as there exists no qualitative comparisons except for 3D dance generation.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "0m0ScZLzSBL", "reviewer": "Reviewer_5Hey", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes to regularize the codebook learning in a VQ-VAE\ntowards a uniform utilization of the codebook. The approach, termed\nOT-VAE, uses a clustering equipartion objective in the latent space to\nattain this goal. Experiments on small-scale datasets of three\ndifferent modalities show the approach improves the reconstruction\nerror and codebook utilization w.r.t. a VQ-VAE baseline, and improves\n3d motion generation for one dataset.\n", "summary_of_the_review": "The paper presents an interesting exploration into regularizing the\nlearned codebook of a VQ-VAE with an OT objective. The empirical\nevidence for the impact of the proposed technique in learning the\ndiscrete prior (e.g. for compression or generation), is very limited.\nThis is a weakness being these some of the main applications of\nVQ-VAEs.\n", "strength_and_weaknesses": "**Strengths**\n- The paper is well written and the ideas are clearly explained.\n- The experimental results show the method is able to achieve\n  improvements in reconstruction accross multiple modalities (images,\n  speech, 3D dance motion), and improvements in the task of 3D dance\n  motion generation.\n\n**Weaknesses**\n- The experimental section is limited to small datasets. The most\n  interesting application of VQ-VAEs is possibly  generation\n  (learning the discrete latent prior). However, the impact of\n  the proposed OT-VAE for generation is only demonstrated for one\n  small (40 evaluation samples) 3D motion generation dataset, which in my opinion is\n  too weak a signal for validating the improvements of the method for\n  VQ-VAE-based generation.\n- Unless I missed it, the comparison should also show the\n  extra computation required to compute the $\\Gamma$ matrix (9). \n- I would suggest revising the general claim that the method \"does not\n  require heuristics\" (abstract, page 2), by clearly specifying the\n  kind of heuristic this is referring to.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject