Momentum Contrastive Autoencoder: Using Contrastive Learning for Latent Space Distribution Matching in WAE
Authors: Devansh Arpit
Aadyot Bhatnagar
Huan Wang
Caiming Xiong
Venue: ICLR 2022 Submitted (2022)
Keywords: Wasserstein autoencoder
contrastive learning

Abstract: Wasserstein autoencoder (WAE) shows that matching two distributions is equivalent to minimizing a simple autoencoder (AE) loss under the constraint that the latent space of this AE matches a pre-specified prior distribution. This latent space distribution matching is a core component of WAE, and a challenging task. In this paper, we propose to use the contrastive learning framework that has been shown to be effective for self-supervised representation learning, as a means to resolve this problem. We do so by exploiting the fact that contrastive learning objectives optimize the latent space distribution to be uniform over the unit hyper-sphere, which can be easily sampled from. 
We show that using the contrastive learning framework to optimize the WAE loss achieves faster convergence and more stable optimization compared with existing popular algorithms for WAE. This is also reflected in the FID scores on CelebA and CIFAR-10 datasets, and the realistic generated image quality on the CelebA-HQ dataset.

Under review as a conference paper at ICLR 2022
MOMENTUM CONTRASTIVE AUTOENCODER: USING CONTRASTIVE LEARNING FOR LATENT SPACE DISTRI- BUTION MATCHING IN WAE
Anonymous authors Paper under double-blind review
ABSTRACT
Wasserstein autoencoder (WAE) shows that matching two distributions is equivalent to minimizing simple autoencoder (AE) loss under the constraint that the latent space of this AE matches pre-specified prior distribution. This latent space distribution matching is core component of WAF, and challenging task. In this paper, we propose to use the contrastive learning framework that has been shown means to resolve
to be effective for self- supe-sueee representation learning, this problem. We do so by exploiting the fact that contrastive learning objectives optimize the latent space distribution to be uniform over the unit hyper-sphere, which can be easily sampled from. We show that using the contrastive learning framework to optimize the WAE loss achieves faster convergence and more stable optimization compared with existing popular algorithms for WAE. This is also reflected the FID scores on CelebA and CIFAR-10 datasets, and the realistic generated image quality on the CelebA-HQ dataset.
INTRODUCTION
The main goal of generative modeling is to learn good approximation of the underlying data distribution from finite data samples, while facilitating an efficient way to draw samples. Popular algorithms such as variational autoencoders (VAE, Kingma Welling (2013); Rezende al. (2014)) and generative adversarial networks (GAN, Goodfellow et al. (2014)) are theoretically-grounde models designed to meet this goal. However, they come with some challenges. For instance, VAFs suffer from the posterior collapse problem (Chen et al., 2016; Zhao et al., 2017; Van Den Oord et al., 2017), and a mismatch between the posterior and prior distribution (Kingma et al., 2016; Tomczak & Welling, 2018; Dai & Wipf, 2019; Bauer & Mnih, 2019). GANs are known to have the mode collapse problem (Che et al., 2016; Dumoulin et al., 2016; Donahue et al., 2016) and optimization instability (Arjovsky & Bottou, 2017) due to their saddle point problem formulation.
Wasserstein autoencoder (WAE) Tolstikhin et al. (2017) proposes a general theoretical framework that can potentially avoid some of these challenges. They show that the divergence between two distributions is equivalent to the minimum reconstruction error, under the constraint that the marginal distribution of the latent space is identical a prior distribution. The core challenge this framework is to match the latent space distribution to a prior distribution that is easy to sample from. Tolstikhin al. (2017) investigate GANs and maximum mean discrepancy (MMD, Gretton et al. (2012)) for this task and empirically find that the GAN-based approach yields better performance despite its instability. Existing research has tried address this challenge (Kolouri et al., 2018; Knop et al., 2018) (see section 2 for discussion).
This paper aims to design generative model address the latent space distribution matching problem of WAEs. To do so, we make a simple observation that allows us to use the contrastive learning framework. Contrastive learning achieves state-of-the-art results in self-supervised representation learning tasks (He et al., 2020; Chen et al., 2020) by forcing the latent representations to be I) augmentation invariant; 2) distinct for different data samples. It has been shown that the contrastive learning objective corresponding maximum entropy over the unit hyper-sphere (Wang & Isola, 2020). We observe that applying this contrastive loss term to the latent representation of an AE therefore matches it to the uniform
the latter goal pushes the leamne representations to achieve

---

Under review as a conference paper at ICLR 2022
distribution over the unit hyper-sphere. Due to the use of the contrastive learning approach, we call our algorithm Momentum Contrastive Autoencoder (MoCA). Our contributions are as follows:
1. we address the fundamental algorithmic challenge of Wasserstein auto-encoders (WAE), viz the latent space distribution matching problem, which involves matching the marginal distribution the latent space to a prior distribution. We achieve this by making the observation that the contrastive term in the recent contrastive learning framework implicitly achieves this precise goal. This is also our novelty.
2. we show that our proposal of using the contrastive learning framework to optimize the WAE loss achieves faster convergence and more stable optimization compared with existing popular algorithms for WAE.
we perform thorough ablation analysis of the impact of the dyper-parameee introduced by the contrastive learning framework, on the performance and behavior of WAE.
2 RELATED WORK
There has been a considerable amount of research on autoencoder based generative modeling. In this paper focus Wasserstein autoencoders (WAE). Nonetheless, we discuss other autoencoder methods for completeness, and then focus on prior work that aim at achieving the WAE objective.
AE based generative models: One of the earliest model in this category is the de-noising autoencoder (Vincent et al., 2008). Bengio et al. (2013b) show that training an autoencoder to de-noise a corrupted input leads the learning of Markov chain whose stationary distribution is the original data distribution it is trained on. However, this results in inefficient sampling and mode mixing problems (Bengio et al., 2013b; Alain & Bengio, 2014).
Variational autoencoders (VAE) (Kingma & Welling, 2013; Rezende et al., 2014) overcome these challenges by maximizing a variational lower bound of the data likelihood, which involves a KL term minimizing the divergence between the latent's posterior distribution and a prior distribution. This allows for efficient approximate likelihood estimation as well as posterior inference through ancestral sampling once the model is trained. Despite these advantages, followup works have identified few important drawbacks of VAEs. The VAE objective is at the risk posterior collapse learning a latent space distribution which is independent of the input distribution if the KL term dominates the reconstruction term (Chen et al., 2016; Zhao et al., 2017; Van Den Oord et al., 2017). The poor sample qualities of VAE has been attributed to mismatch between the prior (which is used for drawing samples) and the posterior (Kingma et al., 2016; Tomezak & Welling, 2018; Dai & Wipf, 2019; Bauer & Mnih, 2019). Dai & Wipf (2019) claim that this happens due mismatch between the AE latent space dimension and the intrinsic dimensionality of the data manifold (which is typically unknown), and propose two stage VAE to remedy this problem. VQ-VAE (Oord et al., 2017) take different approach and propose a discrete latent space as an inductive bias in VAE.
Ghosh et al. (2019) observe that VAEs can be interpreted as deterministic autoencoders with noise injected in the latent space as a form of regularization. Based on this observation, they introduce deterministic autoencoders and empirically investigate various other regularizations.
Similar to our work, the recently proposed DC-VAE (Parmar et al., 2021) also uses contrastive loss for generative modeling. However, the objective resulting from their version of instance discrimination estimates the log likelihood function rather than the WAE objective (which estimates Wasserstein distance). Also, they integrate the GAN loss to the instance discrimination version of VAE loss.
There has been research on AEs with hyperspherical latent space, that we use in our paper. Davidson et al. (2018) 8) propose to replace the Gaussian prior used in VAE with the Von Mises-Fisher (vMF) distribution, which is analogous the Gaussian distribution but on the unit hypersphere.
WAE: Tolstikhin et al. (2017) make the observation that the optimal transport problem can be equivalently framed as autoencoder objective under the constraint that the latent space distribution matches a prior distribution. They experiment with two alternatives to satisfy this constraint in the form of penalty- MMD (Gretton et al., 2012) and GAN (Goodfellow et al., 2014)) loss, and they find that the latter works better in practice. Training an autoencoder with an adversarial loss was also proposed earlier in adversarial autoencoders (Makhzani et al., 2015).
2

---

Under review as a conference paper at ICLR 2022
There has been research making of sliced distances to achieve the WAE objective. For instance, Kolouri et al. (2018) observe that Wasserstein distance for one dimensional distributions have a closed form solution. Motivated by this, they propose to use Warreddtessr distance, which involves a large number of projections of the high dimensional distribution onto one dimensional spaces which allows approximating the original Wasserstein distance with the average of one dimensional Wasserstein distances. A similar idea using the sliced-Crand Cramer distance is introduced in Knop et al. (2018).
Patrini et al. (2020) on the other hand propose a more general framework which allows for matching the posterior of the autoencoder to any arbitrary prior of choice (which is a challenging task) through the use of the Sinkhorn algorithm (Cuturi, 2013). However, it requires differentiating through the Sinkhorn iterations and unrolling for hackpropagation (which computationally expensive); though their choice of Sinkhorn algorithm for latent space distribution matching allows their approach to be general.
MOMENTUM CONTRASTIVE AUTOENCODER
We present the proposed algorithm in this section. We begin by restating the WAE theorem that connects the autoencoder loss with the Wasserstein distance between two distributions. Let X ~ Px be random variable sampled from the real data distribution on X. Z Q() be its latent representation in Z S IRd, and X = g(2) be its reconstruction by a deterministic decoder/generator g: Z + X. Note that the encoder Q can also be deterministic in the WAE framework, and we let f(x) dist Q(ZIX) for some deterministic f - Z.
Lei Theorem 1. (Bousquet et al., 2017; Tolstikhin et al., 2017) Let Pz be a prior distribution on Z, let
Pg = g# Pz be the porw-forward of Pz under g (i.e. the distribution of X = g(Z) when ~ Pz). and let Qz = f+ be the push-forward of Px under f. Then,
inf Wo(Px,Po) Q:Qz=P2 Z~-(Z(x)
E
[c(X,(())]]
E [c(X (((X))(((X) f:1f#xx)=P2 X
inf
(1)
X~Px
where We denotes the Wasserstein distance for some measurable cost function c.
The above theorem states that the Wasserstein distance between the true (Px) and generated (Pg)
data distributions can be equivalently computed by finding the minimum (w.r.t. f) reconstruction loss, under the constraint that the marginal distribution of the latent variable Qz matches the prior distribution Pz. Thus the Wasserstein distance itself can be minimized by jointly minimizing the reconstruction loss w.r.t. both f (encoder) and g (decoder/generaatr)) as long as the constraint met.
In this work, we parameterize the encoder network ; X - Idd such that latent variable 2 (() has unit & norm. Our goal is then to match the distribution of this Z to the uniform distribution over the unit hyper-sphere S. {2 Rd ||=|| 1). To do so, we study the so-called "negative sampling" component of the contrastive loss used in self-supervised learning,
K
Long CnangU:TT. JJT,K)
E PN
log
1 K
f(x) f(x, )/T
Here, f X + Sa is neural network whose output has unit lg norm, is the temperature hyper- parameter, and K is the number of samples (another hyper-parameter). Theorem of Wang & Isola (2020) shows that for any fixed i, when K- DC.
((x) f(x >/r] [/r]
lim (Ineg (f: K) log K)- = E K-000
log
E TT
(3)
Crucially, this limit is minimized exactly when the mush-forward f (i.e. the distribution of the latent random variable Z = =fX when X ~ Px) is uniform on Sa. Moreover, even the Monte Carlo approximation of Eq. (with mini-batch size B and some K such that B < K < oo)
B
K ( ((;))/T
I MC (J:T, K, B) = B neg
log K
(4)
3

---

Under review as a conference paper at ICLR 2022
Algorithm PyTorch-like pseudocode of Momentum Contrastive Autoencoder algorithm
Enc_9r Ene_k: encoder netwerks for quezy and key, Their oucputs are 12 nermalized bea: doceder network
dictionary aa eueue ot 3 randomly initialized keys (dxR)
m: momentum
lambda: regularization ccefficient for entropy meximization tav: logit terperature
for in jata_loader,
load minibatch with samples 2_9 Rxa_g(x) queries: Rxd
2_x - Ene_k(t) deetaca keys: Bxd, no gradlent through keys X_eec (ec reconstructed input
positive logita: BX1 _pss = bmm (B,, (R..
negative logits: BaK
__ng = mmlr_..view (3, k))
logits: [x+1+
logits = cat [[]_pos, negg] dim=1)
compute loss
labals = zeros (B) pasitive elements are in the 0-th index
ILoon CrossEntropytoss tagata/ labels)
maximi zing entropy 2_9
Ierec [X_ree x) vx .)mm) / B reconstruction loss
L L_re lamada con # mementum contrastive autoenzoder loss
update Enc_g
networks
L. b.baakward update (En____aaaa_)))
update (Dec.params)
update Enc_k Enc k.p_ans
Bac_k.params (1-m) Eac_gg
update dictionazy
enqueue (Q, 2_k] enqueue current minibatch
doqueue (0) dequeue the carliest minibatch
fom:: batch matrix multiplication; ma: matrix maltiplication: cat: concatenation. enqueue appends Q keys 2x R fxd the current batch; dequeue removes oldest keys from Q
is consistent estimator (up to a constant) of the entropy of f#Px called the redistribution esti-
ef(.)² mate (Ahmad Lin, 1976). This follows if notice that k(0;T, T, K) := 1 شッ ) f(2))/T the un-normalized kernel density estimate of f(x;) using the i.i.d. samples {0y) so -L (J:T, K, B) = - Σe log k(z;;; T, K) (Wang & Isola, 2020). So minimizing Tong (and importantly I maximizes the entropy of ##
Tolstikhin et al. (2017) attempted to enforce the constraint that f#Px and Pz matching distributions by regularizing the reconstruction loss with the MMD or a GAN-based estimate of the divergence between /#Px and Pz. By letting Pz be the uniform distribution the unit hyper-sphere Sd, the insights above allow us to instead minimize the much simpler regularized loss
B
L.f.
x - ALM (f:t,K,B
(5)
B
Training: For simplicity, we will now use the notation Enc(-) and Dec(-) respectively denote the encoder and decoder network of the autoencoder. Further, the d-dimensional output of Enc(-) Vx. Based on the theory above, we aim minimize the loss L(Enc, Dec; X, B, K). where is the regularization coefficient, 7 is the temperature hyper- parameter, B is the mini-batch size, and K 2 B the number of samples used estimate Lneg-
le normalized, i.e., ||Ene(x)]l2
In practice, we propose to use the momentum contrast (MoCo, He et al. (2020)) framework to implement Iiney- Let Enc; be parameterized by 0, at step of training. Then, we Enc; be the same encoder parameterized by the exponential moving average 0, = (1 - m) Ei= m² Letting TK be the K most recent training examples, and letting ((j) = I Li/B] be the time
4

---

Under review as a conference paper at ICLR 2022
which xy appeared in training mini-batch, we replace L Lun at time step t with
B
K
"End
B Exc)(x;) (xnc)
LMaCo
B log
Σ
exp
1
(6)
B
This approach allows us use the latent vectors of inputs outside the current mini-batch without re-computing them, offering substantial computational advantages over other contrastive learning frameworks such as SimCLR (Chen et al., 2020). Foreing the parameters of Encc to evolve according to an exponential moving average is necessary for training stability, as is the second term encouraging the similarity of (x,)(x,) and Enc, (z:) (so-called "positive samples" in the terminology of con- trastive learning). Note that we do not use any data augmentations in our algorithm, but this similarity term is still non-trivial since the networks Enc, and Enc, are not identical. Pseudo-code of our final algorithm, which we call Momentum Contrastive Autoencoder (MoCA), is shown in Algorithm I (pseudo-code style adapted from He et al. (2020)). Finally, in all our experiments, inspired by Grill et al. (2020) we set the exponential moving average parameter m for updating the Enc' network at (1 mo) (cos(s1/T) + 1)/2, where T the total number of training
the 4th iteration as m iterations, and mg is the base momentum hyper-parameter.
Once the model is trained, the marginal distribution of the latent space (i.e. the push- forward Enc#Pxx) should be close to a uniform distribution over the unit hyper-sphere. We can therefore draw samples from the learned distribution as follows: we first sample 2 ~ ~V from the standard multivariate normal distribution in Fd and then generate sample Ig = Dec(z/|[<112).
Inference:
4 EXPERIMENTS
We conduct experiments on CelebA (Liu al. 2015), CIFAR-10 (Krizhevsky et al., 2009), CelebA- HQ (Karras et al., 2018) datasets, and synthetic datasets using our proposed algorithm as well as existing algorithms that implement the WAE objective. Unless specified otherwise, for CIFAR-10 and CelebA datasets, we use two architectures: Al: the architecture from Tolstikhin et al. (2017), which is commonly used as means to fairly compare against existing methods; and A2: ResNet-18 18 based architecture with much fewer parameters. For CelebA-HQ. we use a variant of ResNet- 18 with 6 residual blocks instead of 4 for both the encoder and decoder. The remaining architecture and optimization details are provided in appendix A.
4.1 CONVERGENCE SPEED AND OBJECTIVE ESTIMATION
Latent Space: We try to evaluate how well the contrastive term in objective addresses the problem of matching the marginal distribution Qz = Enc# of the autoencoder latent space to the prior distribution Pz, viz, the uniform distribution on the unit hyper-sphere. To systematically study the approximation quality and convergence rate in the latent space, we design a synthetic task where we remove the reconstruction loss, and set the objective exclusively match the latent distribution with the prior distribution (thus the decoder network is ignored). Specifically, for all the methods we investigate, the reconstruction error term is removed and only the regularization term aimed latent space distribution matching is used. Note that this is still a non-trivial objective.
To this end, we generate 100 dimensional synthetic data points as our dataset (1000 samples) which are fed through a 2 layer MLP to get their corresponding latent representations (128 dimensional). We then use different algorithms for matching the latent space distribution with a prior distribution (chosen to be the uniform distribution over the unit hyper-sphere in 128 dimensions). We use the Sinkhorn loss, Sliced Wasserstein Distance (SWD) and Maximum Mean Discrepancy (MMD) as baselines. For all algorithms, we use Adam optimizer with learning rate Ie 3 and train for 80 epochs on the synthetic data. In order to study how well the different algorithms match the latent space marginal distribution to the prior, during the training process, we measure the distance between the latent representations of the synthetic input data and random points sampled from the uniform distribution over the unit hyper-sphere in 128 dimensions. We use Sliced Wasserstein Distance (SWD) to measure this distance These SWD estimates are shown in figure 1. As evident, the contrastive objective converges significantly faster.
Note that SWD both an evaluation metrie, well as of the haseline loss function for distribution matching in this experiment. To explain why SWD as objective performs worse than some of
5

---

Under review as a conference paper at ICLR 2022
Slied Wavvessenn Distance
0.005
Contrastive MMD
Sinkcham
8000
McCA =
WAE-MMD J. =
0.004
MoCA
WAE-MMD
McCA 10
WAE-MMD 10
0.003
4000
MoCA 100
WAE-MMD
McCA 1000
WAE-MMD 1000
0.001
2000
0.000
50
60
Epachs
Epochs
Figure I: Systematic analysis of convergence speed for various objectives for matching the marginal distribution of the latent space to prior, without the reconstruction term. Con- trastive loss is faster than existing methods at this task. SWD measured in latent space.
Figure 2: Convergence speed and quality of objective approximation for MoCA vs WAE-MMD (baseline) for various values of their respective regularization coefficient on CIFAR-10. SWD measured in image space.
We measure how well our algorithm approximates the WAE objective compared with baseline. Since the MMD loss outperforms other baseline methods in the above experiments in the latent space, we only compare our algorithm with WAE-MMD (original WAE algorithm from Tolstikhin al. (2017)), and empirically study the comparative convergence rate of these two algorithms. We train MoCA and WAF-MMD with their respective regularization coefficients e {0, 1, 10, 100, 1000} on CIFAR-10 using architecture A2. Starting at initialization (epoch=0), at every 5 epochs during training, we estimate the Sliced Wasserstein distance (SWD) between the test set images and the samples generated by each algorithm at the corresponding epoch. We use the code provided in to compute SWD, whose implementation is aimed specifically at measuring the SWD between 2 image datasets. The results are shown in figure (all hyper-parameters were chosen to be identical for both algorithms for faimesss). The figure shows that larger À values in general result smaller SWD, and SWD estimates decrease with epochs for MoCA, confirming that proposed algorithm indeed minimizes the Wasserstein Distance. Additionally, looking the comparative convergence rate of MoCA and WAE-MMD, we can make three observations:
Image Space:
1. Faster convergence: at epoch 5, the SWD for MoCA is already below 1800 compared with that of WAE-MMD (above 2300) across all values of the regularization coefficient X. This shows that MoCA converges faster than WAE-MMD.
2. Better approximation: the final SWD values of MoCA are significantly lower than those of WAF-MMD across the corresponding À values.
3. Stable optimization: larger values of À result in smaller values SWD for MoCA but this is not always true for WAE-MMD (see À 1000 for WAE-MMD).
4.2 LATENT SPACE BEHAVIOR
Isotropy: We qualitatively investigate the behavior of MoCA in terms of latent space distribution matching, i.e., how closely Qz Enc# Px of the autoencoder latent space matches the prior distribution Pz. Since the encoder is parameterized to output unit Eg norm vectors, we only need to evaluate how close Q(Z) to being isotropic. As computationally efficient proxy, we compute the singular value decomposition (SVD) of the latent representation corresponding 10,000 randomly sampled images from the training set. We use SVD because the spread of singular values of the latent space samples is indicative of how isotropic the latent space is. If the singular values of all the singular vectors are close to each other, the latent space distribution is more isotropic.
In this experiment, train MoCA CIFAR-10. We compute SVD of latent representation for models trained with different values of the regularization coefficient À E {0, 500, 1000, 3000) ) Larger is designed increase the entropy of the latent space to better match the uniform distribution. As Figure 3 shows, for models trained with larger X, the singular values (and therefore the latent
the other objectives, we hypothesize that its gradients are worse. distant, but related example classification tasks, where we care about accuracy, but accuracy as a loss function does not provide any gradients.
hhtpps: thub. com/ koshi h an awd pytorch

---

Under review as a conference paper at ICLR 2022
Singular value
30
A D. FID 64.34
= 500, RDD 61.00
110
d-54
d2128
20
1000, 53.33
100
d5256
= 2000, HDD = 46.27
90
10
70
60
80
100
120
10a
107 10a 10
Singular value index
A
Figure 4: The interplay between the latent dimension of the MoCA network and the regularization coefficient X. Larger d requires significantly larger to achieve comparable FID scores for the generated sample quality,
Figure 3: SVD of latent representation E I228 for models trained with various values uniform singu-
of X. Larger À results in lar values, i.c., closer to uniform distribution, and lower (better) FID for generated samples.
distribution) become more uniform. Corresponding FID scores on generated samples reflect this effect since models trained with larger have lower (better) FID scores.
Latent space dimensionality d and regularization coefficient À Real data often lies on a low dimensional (dg < n1) manifold that is embedded in a high dimensional (n) space (Bengio al., 2013a). Autoencoders attempt to map the probability distribution of the data to a designated prior distribution in latent space of of dimension d, and vice versa. However. if there is a significant mismatch in the dimension do of the true data manifold and the latent space's 's dimension d, learning mapping between the two becomes impossible (Dai & Wipf, 2019). This results many "holes" in the learned latent space which do not correspond to the training data distribution.
Given the importance of this problem, we study how the latent dimension d influences the quality of samples generated by used by our model MoCA. We also simultaneously analyze the influence of the regularization coefficient X, since the value of À enforces how much we want the mapped latent distribution to be close to the uniform distribution on the unit hyper-sphere.
For this experiment, we use d E {64, 128, 256, 512} and study the value of À on a wide range on the log scale between 100 and 144,000 (in some cases). Due to the large number of experiments in this analysis, we train each configuration until the epoch reconstruction loss (mean squared error) reaches 50. For this reason the FID scores much higher than the fully trained models reported in other experiments (where reconstruction loss reaches ~255). The results are shown in Figure 4. We find that for d 64 the performance is quite stable across different À values. However, set to larger values, a significantly larger value of X is required to reach a similar FID score. We hypothesize that this happens because Eq. 6, the dot product of the two encoding vectors is more likely to be orthogonal in a higher dimensional space, which suggests that the value of the contrastive regularization term becomes smaller compared with the reconstruction term in Eq. 5 (considering all other factors fixed). To compensate for this, larger regularization coefficient would be required.
is
t-SNE Visualization of Latent Space: We present qualitative comparison between the latent representation representations learned by Hyperspherical VAE (S-VAE, Davidson et al. (2018)) and MoCA since both algorithms are aimed at learning latent representation that is embedded on the unit hypersphere. For this experiment, we train both algorithms on the training set of the MNIST dataset. Once the models are trained, we compute the latent representation for the test set using both models. These latent representations are projected to 2 dimensional space using T-SNE for visualization. The projections are shown in Figure in appendix. We find that despite its simplicity, MoCA learns perceptually distinguishable class clusters similar to S-VAE. Experimental details provided in appendix B.
4.3 IMAGE GENERATION QUALITY
Qualitatively, we visualize random samples from our trained models on all the datasets. Figure (rows 1-2) contains random samples from CelebA-HQ. The faces in these images look reasonably realistic. Figure 6 (rows 5-6) contains random samples from CIFAR- 10 and CelebA, as well as reconstructions (rows 1-2) of images from these datasets. Most generated samples look realistic
7

---

Under review as a conference paper at ICLR 2022
Figure 5: Random samples (rows 1-2) from model trained with MoCA on CelebA-HQ, and that model's interpolations (rows 3-4) between images in latent space. The leftmost and rightmost columns of rows 3-4 are the original images from the test set of CelebA-HQ which we interpolating.
Figure 6: Left: CIFAR-I0. Right: CelebA. Rows 1-2 show original image (odd column) and its reconstruction (even column). Rows 3-4 show model's interpolation between two test images in latent space. The leftmost and rightmost columns of rows 3-4 are the original images from the corresponding test set. Rows 5-6 show random samples drawn from trained model.
(especially for the two CelebA datasets), and the reconstructions perceptually match the original input in most cases.
We also present latent space interpolations between images from the test set of CelebA-HQ in Figure (rows 3-4). We present latent space interpolations for CIFAR-10 and CelebA in Figure (rows 3-4). For these interpolations, images x and z', let 2, = az + (1 a)2²' for some 0 < a < 1, and then generate the interpolated image a. Dec (=a/ll =. 2l2. These latent space interpolations show that algorithm causes the generator to learn a smooth function from the unit
compute the latent vectors
Enc(z) and z' Enc(x') (x') for two
S3 to image space, and moreover,
almost all intermediate samples look quite realistic.
We present additional random samples, image reconstructions, and latent space interpolations for these three datasets in appendix E.
8

---

Under review as a conference paper at ICLR 2022
DatalModel
WAE-MMD
WAE-GAN
Sinhorn AE
SWAE
CWAE
MoCA-AI
MoCA-A2
CelebA
55
42
56
79
49.69
48.43
44.59
Table I: Comparison of MoCA with existing baselines using FID (lower is better) on CelebA dataset. MoCA-A1 uses architecture similar to the one used in WAE (Tolstikhin et al., 2017), while MoCA-A2 uses a ResNet-18 based architecture.
4.4 QUANTITATIVE COMPARISON
For quantitative analysis we report the Fréchet Inception Distance (FID) score (Heusel et al., 2017). We compare the performance of MoCA with other existing algorithms that try to achieve the WAE objective and AE algorithms with hyper-spherical latent space, because these two class of algorithms are directly related to our proposal. Specifically, we compare with WAE-MMD and WAE-GAN (Tolstikhin et al., 2017), Sinkhorn autoencoder (Patrini et al., 2020) (with hyperspherical latent space), sliced Wasserstein autoencoder (SWAE, Kolouri et al. (2018)) and Cramer- Wold autoencoder (CWAE, Knop al. (2018)). All numbers cited from their respective papers.
Results are shown in table 1. We find that MoCA outperforms all of the methods except WAE-GAN, which uses GAN objective for matching the latent space distribution to the prior. Experiments comparing MoCA with WAE-MMD on CIFAR-10 can be found in appendix C.
4.5 ABLATION ANALYSIS AND IMPACT OF PAPERRPARAMETERS INTRODUCED BY THE CONTRASTIVE LEARNING LOSS
Unlike most existing autoencoder based generative models, our proposal of using the contrastive learn- ing framework, specifically momentum contrastive learning (He et al., 2020) due to its computational efficiency compared to its competitor Chen et (2020), introduces a number of ayper-parameters in addition to the regularization coefficient X. Therefore, it is important to shed light on their behavior during the training process of MoCA. This section explores how these various hyper-parameters impact the quality of generated samples. To keep the analysis tractable and quantitative, we use the Fréchet Inception Distance (FID) score to evaluate the performance of the trained models.
Due to lack of space, discuss these experiments in appendix D. and summarize the results here. In appendix D.I, we study the impact of the regularization coefficient À used in MoCA, on the reconstruction loss achieved at the end of training, and find that larger values of À help reduce the reconstruction loss even more. In appendix D.2, we discuss how the regularization coefficient should be scaled depending on the input image size, and find that its optimal value scales linearly with input size. In appendix D.3, we discuss the impact of the momentum hyper-parameter used in the contrastive algorithm, and show that its value must be closer appendix D.4, we discuss the impact of the temperature hyper-parameee algorithm, and show that its value must be smaller than I for good performance, and discuss possible reasons. Finally, in appendix D.5, we study the impact of the size of dictionary used in the contrastive algorithm, and find the the performance remains stable across wide range of sizes.
for good performance. In used in the contrastive
CONCLUSION
We propose a novel algorithm for addressing the latent space distribution matching problem of Wasserstein autoencoders (WAE) called Momentum Contrastive Autoencoder (MoCA). The main idea behind MoCA to use the contrastive learning framework to match the autoencoder's latent space marginal distribution with the uniform distribution on the unit hyper-sphere. We show that using the contrastive learning framework to estimate the WAE objective achieves faster convergence and more stable optimization compared with existing popular algorithms for WAE. We perform thorough ablation analysis and study the impact of the hyper-parameter introduced in the WAE framework due to our proposal of using the contrastive learning, and discuss how to set their values.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Ibrahim Ahmad and Pi-Erh Lin. A nonparametric estimation of the entropy for absolutely continuous distributions (corresp.). IEEE Transactions on Information Theory, 22(3):372-375, 1976.
Guillaume Alain and Yoshua Bengio. What regularized auto-encoders learn from the data-generating distribution. The Journal of Machine Learning Research, 15(1):3563-3593, 2014.
Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adversarial networks. arXiv preprint 170.002,,. 2017.
Matthias Bauer and Andriy Mnih. Resampled priors for variational autoencoders. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 66-75. PMLR, 2019.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798-188 2013a.
Yoshua Bengio, Li Yao, Guillaume Alain, and Pascal Vincent. Generalized denoising auto-encoders as generative models. Advances in neural information processing systems, pp. 8098899999771 907,907,22 2013b.
Olivier Bousquet, Sylvain Gelly, Ilya Tolstikhin, Carl-Johann Schoelkopf. From optimal transport to generative modeling: the vegan cookbook. arXiv preprint arXiv:1705.07642, 2017.
and Bernhard
Tong Che, Yanran Li, Athul Paul Jacob, Yoshua Bengio, and Wenjie Li. Mode regularized generative adversarial networks. arXiv preprint 16xiv:66/2.021336 2016.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint 2xXiv:2002.05709 2020.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint arXiv: 1.1.03.6,,.0 2016.
Marco Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In NIP,, volume 2, pp. 4, 2013.
Bin Dai and David Wipf. Diagnosing and enhancing vae models. arXiv preprint arXiv:190..05789 2019.
Tim Davidson, Luca Falorsi, Nicola De Cao, Thomas Kipf, and Jakub Tomczak. Hyperspherical variational auto-encoders. arXiv preprint arXiv: 00891, 2018.
Jeff Donahue, Philipp Krähenbühh, and Trevor Darrell. Adversarial feature learning. arXiv preprint 175.976605 2016.
Vincent Dumoulin, Ishmacl Belghazi. Ben Poole, Olivier Mastropietro, Alex Lamb, Martin Arjovsky, and Aaron Courville. Adversarially learned inference. arXiv preprint arXiv:1606.0077 2016.
Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational deterministic autoencoders. arXiv preprint arXiv: 12436, 2019.
Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural informa- tion processing systems, pp. 2672-2680, 2014.
kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012. Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf. and Alexander Smola. A
Jean-Basee Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: new approach to self-supervised learning. arXiv preprint arXiv:2006. 07733, 2020.
10

---

Under review as a conference paper at ICLR 2022
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 9729-9738, 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, pp. 6626-6637, 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv: 2015.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation. In International Conference on Learning Representations, 2018.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes. arXiv preprint
Durk Kingma, Tim Salimans, Rafal Jozefowicz. Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in neural information processing systems, pp. 4743 2016.
Szymon Knop, Jacek Tabor, Przemystaw Spurek, Igor Podolak, Marcin Mazur, and Stanislaw Jastrzebski. Cramer-wold autoencoder. arXiv preprint arXiv: 1805.09235, 2018.
Soheil Kolouri, Phillip E Pope, Charles E Martin, and Gustavo K Rohde. Sliced wasserstein auto- encoders. In International Conference Learning Representations, 2018
Alex Krizhevsky et al. Learning multiple layers of features from tiny images. 2009.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of the IEEE international conference computer vision, pp. 3730-3738, 2015,
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow. and Brendan Frey. Adversarial autoencoders. arXiv preprini arXiv: /5/ 1.05644, 2015.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. arXiv preprint /1t::/711.0099 2017.
Gaurav Parmar, Dacheng Li, Kwonjoon Lee, and Zhuowen Tu. Dual contradistinctive generative autoencoder. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 823 832, 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vetioo, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- tighpeerformncc deep learning library. Advances in Neural Information Processing Systems 32, pp. 8024-8035. 2019.
Giorgio Patrini, Rianne van den Berg. Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence, pp. 733-743. PMLR, 2020.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning. pp. 1278-1286. PMLR, 2014.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. arXiv preprint arXiv: / 1.015558 2017.
11

---

Under review as a conference paper at ICLR 2022
Jakub Tomezak and Max Welling. Vae with vampprior. In International Conference Artificial Intelligence and Statistics, pp. 2214-1223, 2018.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, pp. 6306-6315, 2017.
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning, pp. 1096-1103, 2008.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity the hypersphere. arXiv preprint arXiv:2005./024: 2020.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Towards deeper understanding of variational autoencoding models. arXiv preprint arXiv:1702.08658, 2017.
12

---

Under review as a conference paper at ICLR 2022
APPENDIX
A TRAINING AND EVALUATION DETAILS
We ran all our experiments in Pytorch 1,5.0 (Paszke et al., 2019). We used 8 V100 GPUs on google could platform for our experiments.
Datasets:
CIFAR-10 contains 50k training images and 10k test images of size 32 x 32.
CelebA dataset contains a total of 2003 64 x 64 images divided into ~ 1800 training images and 200 test images. The images were pre-processed by first taking 140x 140 center crops and then resizing to 64x64 resolution.
CelebA-HQ contains a total of 330 1024 1024 images. We resized these images to 256 256 and split the dataset into ~27k training images and ~-3 test images.
Architecture and optimization:
The network architecture AI is identical to the CNN architecture used in Tolstikhin et al. (2017) except that we use batch norm (loffe Szegedy, 2015) in every layer (similar Ghosh et al. (2019)), and the latent dimension is 128 for the CelebA dataset. This architecture roughly has around 38 million parameters. We found that using 64 dimensions for CelebA in this architecture prevented the reconstruction loss from reaching small values.
The encoder of the network architecture A2 is modification of the standard ResNet- 18 architecture He et al. (2016) in that the first convolutional layer has filters of size 3 x 3, and the final fully connected layer has latent dimension 128. The decoder architecture is a mirrored version of the encoder with upsampling instead of downsampling. Additionally, the final convolutional layer uses an upscaling factor of for CIFAR- 10 and 2 for CelebA. The architecture roughly has around 24 million parameters.
Both Al and A2 were trained on CIFAR-10 with MoCA hyperparameters K = 30000, 7 0.05, mo 0.99. A1 used À 2000 while A2 used 100. A1 was trained for 40 epochs while A2 was trained for 100 epochs using the Adam optimizer with batch size 64, and learning rate 0.001. The learning rate was decayed by factor of 2 after 60 epochs for A2.
Both Al and A2 were trained on CelebA with MoCA hyperparameters K = 60000, 7 0.05, mg 0.99. A1 used À = 1000 while A2 used 100. Both models were trained for 200 epochs using the Adam optimizer with batch size 64, and learning rate 0.001 decayed by factor of 2 every 60 epochs.
During our experiments, found that the choice of hyper-parameters 7 and was stable across the two architectures and datasets and they were chosen based on our ablation studies. The value of K was decided based on the size of the dataset (CelebA being larger than CIFAR-IO in our case). Finally, we found that the value of X was generally subjective to the dataset and architecture being used. We typically ran a grid search over À {100, 1000, 3000, 6000).
The images in Figure 5 (CelebA-HQ 256 x 256) were generated using variant of the ResNet-18 architecture. The base ResNet-18 18 architecture has residual blocks, each containing convolutional layers and an additional convolutional layer which spatially downsamples its input by a factor of 2 x 2. For the encoder, we use the same architecture, but with 6 blocks (to downsample a 256 x 256 image to 4 x 4, which we then flatten and project into the latent space). The decoder is mirrored version of the encoder, but with de-convolution upsampling layers instead of downsampling layers. The latent space of this architecture is 128 dimensional. We train this model with MoCA hyperparameters 20000, K = 30000, T = 1, mg = 0.99. The model was trained for 1000 epochs using the Adam optimizer with batch size 64, and learning rate 0.002 (decayed by factor of every 40 epochs until epoch 400).
Quantitative evaluation:
In all our experiments, FID was always computed using the test set of the corresponding dataset. We always use 10,000 samples for computing FII..
13

---

Under review as a conference paper at ICLR 2022
Figure 7: Left: MoCA. Right: Hyperspherical VAE (S-VAE). T-SNE projections of autoencoder la- tent representation of MNIST test set. Both algorithms learn a hyperspherical latent space embedding. The colors correspond to the 10 digit classes in MNIST. Despite the simplicity of MoCA over S- S-VAE, both algorithms are comparable at implicitly learning perceptually distinguishable class clusters.
DatalModel
WAE-MMD
MoCA-A1
MoCA-A2
CIFAR-10
80.9
61.49
54.36
Table 2: Comparison of MoCA with WAE-MMD on CIFAR- 10 dataset using FID (lower is better).
B EXPERIMENTAL DETAILS FOR EXPERIMENTS COMPARING MOCA WITH S-VAE
For the T-SNE experiment, we train both algorithms using Adam optimizer for 25 epochs with fixed learning rate 0.001 (other hyper-parameters take the default Pytorch Adam values). For MoCA. we used À 5, K 10, 000, т 0.99. Both algorithms use the same MLP architecture with hidden layer encoder and hidden layer decoder, and ReLU activations. both models, the hidden layer dimensions was 128 for all layers and the latent space dimension was 6. The latent space visualizations are shown in figure
C ADDITIONAL QUANTITATIVE RESULTS
Experiments comparing MoCA with WAE-MMD on CIFAR-10 can be found in table
D ABLATION ANALYSIS
For this section, unless specified otherwise, we the CelebA dataset with the ResNet-18 autoencoder architecture (architecture A2 in the previous section), K 60000. For optimization, we use Adam with learning rate 0.001, batch size 64, drop this learning rate by half every 60 epochs, and train for total of 200 epochs. All other optimization hyper-parameters are set to the default Pytorch values.
1, mo 0.999, d = 128, X = 3000,
0
0.9
0.99
0.999
size
32
32
64 x 64
256 256
FID
86.57
98.97
47.96
47.46
x
100
2000
20000
Table 4: Smaller base momentum mg causes model performance to degrade significantly. Performance is measured using the FID score (lower is better). Note that we use the cosine schedule described in section 3 in all these experiments.
Table 3: The optimal value of À scales lin- early with input size. We consider À between 100 and 50000 and report the value x that achieves the best FID score for each input size.
14

---

Under review as a conference paper at ICLR 2022
Momentum Contrastive Autoencoder
40
Reconstruction loss
38
36
34
32
10
10a
102
10a
Figure 8: Impact of regularization coefficient on reconstruction loss at the end of training. Larger values of Ả do not interfere with the reconstruction loss, rather help achieve lower reconstruction loss.
D.1 RECONSTRUCTION LOSS VS
We study the impact of the regularization coefficient À used during the training of MoCA, on reconstruction loss achieved at the end training. To provide some context, in VAEs, a trade-off emerges between the reconstruction loss and the KL term (regularization) because if KL is fully minimized, the posterior p(=[2; becomes a Normal distribution and hence independent of the input Therefore, strong regularization worsens the reconstruction loss in VAEs. However, such trade-off does not exist in WAE, in which the regularization aims the make the marginal distribution of the latent space close the prior. We see this effect in the experiments shown in figure 8 for CIFAR- 10 dataset using the A2 architecture. We find that larger values of À in fact help reduce the reconstruction loss even more.
D.2 CHOOSING GIVEN INPUT SIZE
An important consideration when selecting the regularization coefficient is the relative scale reconstruction loss ||2; (ff(x:) (r:)) ))1 and contrastive loss.
We show that the optimal value of the regularization weight À scales linearly with input size. For this experiment, we downscale CelebA-HQ (256 256) to 64 x 64 and 32 x 32, and we study the impact of À for the different input sizes. We construct the models for generating d images by removing log, (256/d) /d) of the residual blocks from the encoder/decoder of the base model used for CelebA-HQ (256 x 256) (described appendix A), as each block downsamples/tpsaapphe the We train all models for 400 epochs using the Adam optimizer with batch
image factor of size 64 and learning rate 0.002 (decayed by a factor of 2 every 40 epochs).
For each case, we report the value of À {100, 200, 500, 1000, 2000, 5000, 10000, 20000, 50000) that achieves the best FID score for each image size. We find that the optimal value of À is roughly proportional the number of pixels in the input (Table 3).
The full data for this experiment (which support the claims of Table 3) can be found in Figure 9, We consider {100, 200, 500, 1000, 2000, 5000, 10000. 20000. 256 x 256, as quality rapidly deteriorates when À < 5000). Please note that the absolute FID scores are not comparable between different image sizes! Rather, we emphasize the relative trends.
50000) 20000,50000) for each image size (except
D.3 IMPORTANCE OF MOMENTUM m
We study the impact of the contrastive learning syper-parameter m on the generated sample quality. is the exponential moving average momentum encoder network Enc,- m typically kept to be close to for training stability (He al., 2020). We confirm this intuition for generative model well in Table 4. We use the base
used for updating the parameters of the
15

---

Under review as a conference paper at ICLR 2022
FID vs. À for Different Image Sizes
32 x 32 Images
55
50
100
200
500
1000
2000
5000
10000
20000
50000
64 x 64 Images
56
55
FID
54
53
100
200
500
1000
2000
5000
10000
20000
50000
256 x 256 Images
120
100
80
100
200
500
1000
2000
5000
10000
20000
50000
Figure 9: Impact of regularization weight on FID score based on input size. Optimal values X (labeled with red star) scale linearly with input size. Note that absolute FID scores comparable between different image sizes! This figure focuses on relative trends.
100
T=1.0
FID
0.05
75
2000
4000
6000
8000
10000
A
Figure 10: Impact of optimal choice of À and best FID for generated samples. Optimal À is lower for lower т. Best FID is better for lower т. This suggests that entropy is maximized more accurately when 7 is smaller.
value mo 0.9.
0.9. 0.999). Note that we use the cosine schedule to compute the value of m every iteration (as discussed in section 3), which makes m increase from the base value mg to over the course of training. We find that FID scores much worse when is not close to 1.
D.4 IMPORTANCE OF TEMPERATURE
7
Based on the discussion below Eq. 3, the negative term in the contrastive loss essentially estimates the entropy of the latent space distribution due to its equivalent kernel density estimation (KDE) interpretation (Eq. 4). Therefore, the temperature the smoothing parameter of this KDE and controls the granularity of the estimated distribution. Thus for larger temperature, the estimated distribution becomes smoother and the entropy estimation becomes poor, which should result in poor quality of generated samples. Additionally, for larger т,
7 used in the contrastive loss acts
16

---

Under review as a conference paper at ICLR 2022
K
100
5000
10,000
30,000
60,000
120,000
FID
84.93
48.45
46.56
47.31
46.68
49.94
Table 5: The effect of dictionary size K on the quality of samples measured using the FID score (lower is better). The performance is largely stable across different values of K unless K is too small.
intuitively, a larger À should be needed in order to push the KDE samples apart from one another. We confirm these intuitions Figure 10. Note that due to the large number of experiments this analysis, train each configuration until 50 epochs (explaining the inferior FID values).
D.5 EFFECT OF K
Since we use the momentum contrastive framework, would be useful to understand how the dictionary size K affects the quality of generative model learned. The dictionary Q contains the negative samples in the contrastive framework which are used to push the latent representations away from one another, encouraging the latent space be more uniformly distributed. We therefore expect small K would be bad for achieving this goal. Our experiments in Table 5 confirm this intuition. We use K E (100, 5000, 10000. 30000, 60000, 120000). We find that the FID score is stable and small across the various values of K chosen, except for K 100, for which FID is much worse.
E ADDITIONAL QUALITATIVE RESULTS
Figures 11, 12, and 13 respectively depict additional randomly sampled images, reconstructions, and latent space interpolations for CelebA-HQ 256 256. Figures 12 and are generated by the same model used to generate Figure 5, while Figure 11 is generated earlier checkpoint of that model (selected for best visual quality).
Figures 14, 16, and 18 respectively depict additional randomly sampled images, reconstructions, and latent space interpolations for CIFAR-10 using the same model that achieved the FID score of 54.36 in table 1.
Figures 15. 17, and respectively depict additional randomly sampled images, reconstructions, and latent space interpolations for CelebA using the same model that achieved the FID score of 44.59 in table 1.
17

---

Under review as a conference paper at ICLR 2022
Figure 11: Random samples generated by a model trained (as described in appendix on CelebA-HQ 256 x 256 for 850 epochs. Model checkpoint picked based on best visual quality.
18

---

Under review as a conference paper at ICLR 2022
HI
ADS
K
HE
STOP
IP
Kr
02
Oe
Figure 12: Image reconstructions by model trained (as described in appendix A) on CelebA-HQ
256 x 256. For each pair of columns, the left is the original image, and the right is the reconstruction.
19

---

Under review as a conference paper at ICLR 2022
A
RE
Figure 13: Latent space interpolations by model trained (as described in appendix A) on CelebA-HQ
256 x 256.
Figure 14: Randomly generated samples from the MoCA model trained CIFAR-10 with FID 54.36 in table I.
Figure 15: Randomly generated samples from the MoCA model trained on CelebA with FID 44.59 in table
20

---

Under review as a conference paper at ICLR 2022
Figure 16: Reconstructed test samples from the MoCA model trained on CIFAR-10 10 with FID 54.36 in table 1.
Figure 17: Reconstructed test samples from the MoCA model trained on CelebA with FID 44.59 in table 1.
Figure 18: Interpolation between two test im- ages in latent space for MoCA model trained on CIFAR-10 with FID 54.36 in table 1. The leftmost and rightmost columns are the original images from the corresponding test set.
Figure 19: Interpolation between two test im- ages latent space for MoCA model trained on CelebA with FID 44.59 in table 1. The leftmost and rightmost columns are the original images from the corresponding test set.
21


Reviews:
[{"review_id": "jPObqz4oNO8", "reviewer": "Reviewer_ZJcJ", "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": "Strengths:\n\n1. The paper is well-written and easy to follow. I do think that some of the notation such as \"push forward\" from measure theory, is really not needed or particularly useful here. Simpler terminology such as just using encoding and decoding functions would be more than sufficient.\n\n2. Some of the experiments are interesting and show the effects of the proposed regularization e.g. on the singular value distribution of the latent representation.\n\n3. Using a contrastive approach is a potentially effective way to match the prior and posterior distributions.\n\nWeaknesses: It is unclear that the proposed regularizer results in qualitatively better reconstructions than the baselines. FID is not a perfect measure and the samples from baselines should be shown side-by-side with the proposed approach to know whether there is indeed an improvement. I found the CIFAR-10 reconstruction results are somewhat poor. \n\nQuestion: the value of the lambda parameter is very large - what are the relative loss values during training/convergence (the reconstruction loss vs. regularizer loss)?", "summary_of_the_paper": "The paper presents a regularization technique for Wasserstein Auto-Encoders, based on contrastive learning. ", "summary_of_the_review": "The paper is interesting but has some shortcomings. I would like to see some results of the baselines to decide if the proposed regularizer does indeed improve results qualitatively. I do not believe that FID is a proper measure of quality (not just for this paper but for measurement of GAN sample quality, in general). I give the paper a slightly positive score based on the idea, but I am looking forward to some samples in the rebuttal to decide my final score.", "strength_and_weaknesses": "", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "bbqvGspY8gN", "reviewer": "Reviewer_nhFp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "I notice this is a re-submission from ICLR-2021. Thus some of my comments are based on the differences between two versions.\n\n## Strengths\n1. The paper is well written and well motivated.\n2. I think the idea of using contrastive learning to enforce a hypersphere prior for WAE is clever and neat.\n3. The authors provide extensive ablations on hyperparameters.\n\n\n## Weaknesses\n1. My main concern is the performance of the proposed method on CIFAR10 and CelebA. The interpolation, reconstruction, and samples in Figure 6 are very blurry, and hard to justify the benefit of using the proposed approach. The reported FID in Table 1 and 2 are very high. It would be nice to include a comparison of [1] (which has FID of 5.25 and 24.08 on CelebA and CIFAR10 respectively). Also, why is the two-stage VAE baseline in the previous version removed?\n2. It would be nice to include WAE-GAN in Figure 1 and 2, since it outperforms the proposed MoCA in Table 1.\n3. I think it would be interesting to see how to integrate the instance contrastive loss as in DC-VAE [2] into the proposed MoCA.\n\n\n[1] Aneja, Jyoti, et al. \"Ncp-vae: Variational autoencoders with noise contrastive priors.\" arXiv preprint arXiv:2010.02917 (2020).\n[2] Parmar, Gaurav, et al. \"Dual contradistinctive generative autoencoder.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.", "summary_of_the_paper": "This paper proposes a new approach to train Wasserstein auto-encoders (WAE) with contrastive learning techniques. Specifically, the paper proposes to enforce the marginal matching constraint of WAE by exploiting the fact that contrastive learning objectives optimize the latent space distribution to be uniform over the unit hypersphere.", "summary_of_the_review": "The main idea of the paper is well motivated. However, I still find the results on image datasets such as CIFAR10 and CelebA hard to justify the superiority of the proposed method over baselines. I lean towards weak rejection but am willing to amend my score if my concerns are addressed.", "strength_and_weaknesses": "", "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": "VCzPcyysrRJ", "reviewer": "Reviewer_pPeS", "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": "Strengths:\n- Overall, this paper is well organized.\n- The method proposed in this paper is a reasonable combination of existing state-of-the-art methods.\n- Experimental results show that the proposed method has stability and faster convergence, which is promising.\n\nWeaknesses:\n- The authors claim that MoCA can generate images with high quality. However, the experimental results in this paper do not show this very well. First of all, although the authors claim that the results in Figure 5 and Figure 6 are \"look realistic\", some of the face images seem to be collapsed, and the interpolation between the two images seems to be discontinuous. Since there is no qualitative comparison with existing methods, we cannot judge these methods as \"realistic\". In Table 1, the authors show the quantitative comparison results with the existing methods, but there are some puzzling points. First, why does MoCA-A2, which has fewer parameters, have higher performance? Also, why does WAE-GAN perform better than MoCA? Since Table 1 shows that WAE-GAN is better than MoCA, shouldn't it be compared with WAE-MMD? Furthermore, why are the quantitative results of CelebA-HQ not shown?\n- The authors show in their experiments that MoCA achieves faster convergence than existing methods, but they do not fully explain why MoCA shows such convergence. Why does the convergence of the proposed method become better when contrastive learning is included? In addition, Figure 1 shows the line graphs of only one training trial for each method, and the variance of each method is not shown. Therefore, I cannot judge whether the difference in results between methods is large or small.\n- In Section 4.1, the authors compare WAE-MMD with MoCA as the original WAE algorithm, but I do not understand why they do not compare it with WAE-GAN, which is also original. I think this should be done because Table 1 shows that WAE-GAN has better image generation performance than the proposed method.\n For example, the authors employ MoCo, but how effective is this in improving the performance of the proposed method? The authors do not seem to have verified such a thing.\n\nMinor comments:\n - The significant figures of the results of each method in Table 1 should be the same.\n - Section3: any fixed t -> any fixed \\tau", "summary_of_the_paper": "In this paper, the authors propose to use contrastive learning for matching in latent space in the Wasserstein autoencoder (WAE). In addition, they employ techniques such as momentum contrast in contrastive learning. Experimental results show that the proposed method, MoCA, is more stable and converges faster than existing methods. It is also capable of generating high-resolution images such as CelebA-HQ.\n", "summary_of_the_review": "In terms of convergence and stability, the proposed method is considered to be effective to a certain extent. Also, the idea of using contrastive learning for WAE is interesting. However, the explanation of the claim and the presentation of the results are insufficient.", "strength_and_weaknesses": "", "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": "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:

Meta-review: This paper presents a variant of the WAE which uses a contrastive criterion to enforce the marginal distribution matching constraint. Experiments show faster convergence in terms of Wasserstein distance, more visually appealing samples, and better FID scores compared with other WAE models.

The original WAE framework leaves open the choice of approximation for enforcing marginal distribution matching, and the original paper gives two such algorithms. Therefore, it's pretty natural to replace this approximation with something else (such as the contrastive criterion used here), so a submission would need to show evidence that it's significantly better than other approaches. Reviewers have expressed various concerns about the experiments. None of them are major problems, but overall the method doesn't seem consistently better than other WAE methods; e.g., the FID score is worse than that of WAE-GAN.

I encourage the authors to take the reviewers' comments into account in preparing the submission for future cycles.
Decision: Reject