Momentum Contrastive Autoencoder
Authors: Devansh Arpit
Aadyot Bhatnagar
Huan Wang
Caiming Xiong
Venue:  (2021)
Keywords: generative model
contrastive learning
autoencoder
Wasserstein autoencoder

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 in WAE, and is in itself 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. This results in a simple and scalable algorithm that avoids many of the optimization challenges of existing generative models, while retaining the advantage of efficient sampling. Quantitatively, we show that our algorithm achieves a new state-of-the-art FID of 54.36 on CIFAR-10, and performs competitively with existing models on CelebA in terms of FID score. We also show qualitative results on CelebA-HQ in addition to these datasets, confirming that our algorithm can generate realistic images at multiple resolutions.

Under review as a conference paper at ICLR 2021
MOMENTUM CONTRASTIVE AUTOENCODER
Anonymous authors Paper under double-blind review
ABSTRACT
Wasserstein autoencoder (WAH) shows that matching two distributions is equivalent 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 in WAE, and is in itself a challenging task. In this paper, we propose to use the contrastive learning framework that has been shown to be effective for self-suuprviied 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. This results in simple and scalable algorithm that avoids many of the optimization challenges of existing generative models, while retaining the advantage of efficient sampling. Quantitatively, we show that our algorithm achieves a new 10, and performs competitively with existing models CelebA in terms of FID score. We also show qualitative results on CelebA-HQ in addition to these datasets, confirming that our algorithm can generate realistic images at multiple resolutions.
FID of 54.36 on CIFAR-
INTRODUCTION
The main goal of generative modeling is to learn given data distribution while facilitating an efficient way to draw samples from them. Popular algorithms such as variational autoencoders (VAE, Kingma & Welling (2013)) and generative adversarial networks (GAN, Goodfellow al. (2014)) are theoretically-grumnee
models designed to meet this goal. However, they come with some challenges. For instance, VAEs suffer from the posterior collapse problem (Chen et al., 2016; Zhao al., 2017; Van Den Oord et al., 2017), and mismatch between the posterior and prior distribution (Kingma et al., 2016; Tomezak & Welling, 2018; Dai Wipf, 2019; Bauer & Mnih, 2019). GANs are known to have the mode collapse problem (Che et al., 2016; 6; Dumoulin et al., 2016; Donahue et al., 2016) and optimization instability (Arjovsky & Bottou, 2017) due to their saddle point problem formulation.
With the Wasserstein autoencoder (WAE), Tolstikhin et al. (2017) propose a general theoretical framework that can potentially avoid 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 to prior distribution. The core challenge of this framework is match the latent space distribution to prior distribution that is easy to sample from. If this challenge is addressed appropriately, WAE can avoid many of the aforementioned challenges of VAE and GANs. Tolstikhin et 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. Others have proposed solutions overcome this challenge (Kolouri et al., 2018; Knop et al., 2018), but they come with their own pitfalls (see Section 2).
This paper aims to design a generative model that avoids the aforementioned challenges of existing approaches. To do so, we build on the WAE framework. In order to tackle the latent space distribution matching problem, we make simple observation that allows us to use the contrastive learning framework to solve this problem. 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 to the latter goal pushes the learned representations to achieve 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 distribution over the unit hyper-sphere. This approach avoids the aforementioned

---

Under review as a conference paper at ICLR 2021
optimization challenges of existing methods, thus resulting in simple and scalable algorithm for generative modeling that we call Momentum Contrastive Autoencoder (MoCA).
2 RELATED WORK
There are many autoencoder based generative models in existing literature. 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 to 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) overcome these challenges by maxi- mizing a variational lower bound of the data likelihood, which involves a KL term minimizing the divergence between the latent's 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 a few important drawbacks of VAFs. 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). The VAE objective is also at the risk of posterior collapse- learning a latent space distribution which is independent of the input distribution the KL term dominates the reconstruction term (Chen et al., 2016; Zhao et al., 2017; Van Den Oord et al., 2017).
Dai & Wipf (2019) claim that the reason behind poor sample quality of VAEs is a mismatch between the prior and posterior, arising from the latent space dimension of the autoencoder being different from the intrinsic dimensionality of the data manifold (which is typically unknown). To overcome this mismatch, they propose to learn a two stage VAE in which the second stage learns a VAE on the latent space samples of the first. They show that this two stage training and sampling significantly improves the quality of generated samples. However, training a second VAE is computationally expensive and introduces some of the same challenges mentioned above.
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. The further introduce a post-hoc density estimation for the latent space since the autoencoding step does not match it prior. In this context, one can view our proposed algorithm as way to regularize deterministic autoencoders while simultaneously learning latent space distribution which can be easily sampled from.
Tolstikhin et al. (2017) make the observation that the optimal transport problem can be equivalently framed as an autoencoder objective (WAE) 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 a 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). Our algorithm builds on the aforementioned WAE theoretical framework due to its theoretical advantages.
There has been research that aims at avoiding the latent space distribution matching problem all together by making use of sliced distances. 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 sliced- Wasserstein distance, which involves 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-Cramer distance is introduced in Knop et al. (2018). However, the number of required random projections becomes prohibitively high when the data lives on low dimensional manifold in a high dimensional space, making this approach computationally inefficient or otherwise inaccurate (Liutkus et al., 2019).
3 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
2

---

Under review as a conference paper at ICLR 2021
be random variable sampled from the real data distribution on X, Z Q(Z()) be its latent
representation in 2 C TR and X =g( (()) be its reconstruction by deterministic decoder/generrator Z + X. Note that the encoder (Z) can also be deterministic in the WAE framework, and
we let f(X) dist ( for some deterministic f - 2.
Theorem 1. (Bousquet et al., 2017; Tolstikhin al., 2017) Let Pz be a prior distribution on Z, let Pgg = g# 2 be the push-forward of Pz under (i.e. the distribution of X =(()) when g ~ Pz). and let Qz - Px be the push-forward of Px under f. Then,
Wo(Px.P,) = inf
E
[c(X) =
inf f: X
[c(X ((f()(
(1)
Q:Q2=P2
X-Px Z~Q(Z(X)
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 (Pa) 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/generator) as long as the constraint met.
In this work, we parameterize the encoder network f X - IR such that latent variable Z = f(X) has unit & norm. Our goal then match the distribution of this Z the uniform distribution over the unit hyper-sphere S₃ = {2 E R | = 1}. To do so, we study the so-called "negative sampling" component of the contrastive loss used self-supervised learning.
K
Σ J()))) K
1
Inegl (f:T,K)
E Px
log
(2)
(x
Here, f: X' + Sa a neural network whose output has unit l2 norm, T 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 f, when K - oc.
lim (Imeg K) log K) K-o0
E x~Px
log
E T~~Px
(x)) f(x s(/)/ [r]
Crucially, this limit is minimized exactly when the push-forwan ## (i.e. the distribution of the latent random variable Z == =(X) 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 <0)
L MC (f:T,K,B)
K,
log
1
slag)(T
a consistent estimator (up to a constant) of the entropy of [#Px called the redistribution esti- mate (Ahmad & Lin, 1976). This follows if we notice that k( r, K) = Σx (a.)) f(t,))((t,)/ f(2))/r the un-normalized kernel density estimate of f(x,) using the i..d. samples {2y} so -LMG (f:T, K, B) = - Emn log kor;( K) (Wang & Isola, 2020). So minimizing Lneg (and importantly I.MC) maximizes the entropy of P*
Tolstikhin et al. (2017) attempted to enforce the constraint that f#Px and Pz were matching distributions by regularizing the reconstruction loss with the MMD GAN-based estimate of the divergence between f Px and Pz. By letting Pz be the uniform distribution over the unit hyper-sphere Sa, the insights above allow to instead minimize the much simpler regularized loss
(...g.A,,,BBKK. 1
M x o(f(x)))) + AL (f:T,K, K)
(5)
Training: For simplicity, will now the notation Enc(-) and Dec(-) to respectively denote the encoder and decoder network of the autoencoder. Further, the d-dimensional output of Enc(-) is la normalized, i.e., ||Enc(=)]|2 Vz. Based the theory above, aim minimize the
3

---

Under review as a conference paper at ICLR 2021
Algorithm PyTorch-like pseudocode of Momentum Contrastive Autoencoder algorithm
Enc_9, Ene_k: encoder netwerks for quezy and key, Their oucputs are L2 nermalized
bec: doceder network
diccionary aa eueue ot 3 randomly initialized keys (dxK)
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(z_g) reconstructed input
positive logita: BX1 _pss = bmm (B,, (R..
negative logits: BxK
__ng = mml (3, k))
logits: [x+1+
logits = cat [[]_pos, nong] dim=1)
compute loss
labals = zeros (B) pasitive elements are in the 0-th index
ILoon CrossEntropptoss I flog'ta/tay 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. L.backward update (En__..aaaa))))
update (Dec.parama
update Enc_k Enc k.p_ans
Bac_k.params (1-m) pne_g.
update dictionazy
enqueue (Q, 2_k] enqueue the current minibatch
doqueue (0) dequeue the carliest minibatch
tom:: batch matrix multiplication; ma: matrix maltiplication: cat: concatenation. enqueue appends Q with keys 2x R fxd from the current batch; dequeue removes oldest keys from Q
loss L(Enc, Dec; T, B, K), where À is the regularization coefficient, T is the temperature hyper- parameter, B is the mini-batch size, and K Z B the number of samples used estimate Lingg
In practice, we propose to use the momentum contrast (MoCo, He et al. (2020)) framework to implement Lneg- Let Enc, be parameterized by 0 at step of training. Then, we let Enc, be the m² Letting same encoder parameterized by the exponential moving average 0, = (1 m) Σt i= [j/B| be the time at . 2K be the K most recent training examples, and letting t(j) = t
which 2, appeared training mini-batch, we replace Lag at time step t with
B log
K
Encx(x,)² "e''Encc (I;)
B
Enc, (x;) Enc, (
LMOC = B
Exp
(6)
K
T
B
j=1
This approach allows re-computing them, offering substantial computational advantages over other contrastive learning frameworks such as SimCLR (Chen et al., 2020). Foreing the parameters of Enc' to evolve according to an exponential moving average is necessary for training stability, as is the second term encouraging the similarity of (I,,(x,) and Enc/ (x;) (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 noi-trivial since the networks Enc; and Enc; are identical. Pseudo-code of our final algorithm, which we call Momentum Contrastive Autoencoder (MoCA), is shown in Algorithm (pseudo-code style adapted from He et al. (2020)). Finally, all our experiments, inspired by Grill et al. (2020) we set the exponential moving average parameter for updating the Enc' network at the gth iteration as m = - (1 - mo) (cos(Tt/T)) + 1)/2, where T is the total number training iterations, and my is the base momentum hyper-parameter.
the latent vectors of inputs outside the current mini-batch without
Inference:
Once the model is trained, the marginal distribution of the latent space (i.e. the push- forward Enc#Px) should be close to uniform distribution over the unit hyper-sphere. We can therefore draw samples from the learned distribution as follows: we first sample 2 ~ ~N from the standard multivariate normal distribution in Rd and then generate sample xy = Dec(z/
4

---

Under review as a conference paper at ICLR 2021
Figure 1: Random samples (rows 1-2) from model trained with MoCA 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 are interpolating.
Figure 2: Left: CIFAR-10. 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.
4 EXPERIMENTS
We present quantitative and qualitative evaluation of the samples generated by Momentum Con- trastive Networks trained CelebA (Liu al., 2015), CIFAR-10 (Krizhevsky et al., 2009), and CelebA-HQ (Karras et al., 2018). For all datasets except CelebA-HQ, we use two architectures: Al: the architecture from Tolstikhin et al. (2017), which commonly used as means fairly compare against existing methods; and A2: a ResNet- 18 based architecture with much fewer parameters. For CelebA-HQ we use a variant of ResNet- 18 with residual blocks instead of for both the encoder and decoder. The remaining details are provided in Appendix A.
For quantitative analysis we report the Fréchet Inception Distance (FID) score (Heusel et al., 2017). In Table 1, we compare the performance of MoCA with VAE (Kingma & Welling, 2013), WAE (Tolstikhin et al., 2017), two stage VAE (Dai Wipf, 2019), and spectral normalized regularized autoencoder (RAE-SN, Ghosh et al. (2019)). The numbers for 2-Stage VAE and RAE-SN are cited from their respective papers. The CelebA numbers for VAE are WAE are cited from Tolstikhin et al. (2017), while those for CIFAR-10 1-10 cited from Dai & Wipf (2019). We achieve state-of-the-arr on CIFAR-10 using our model with the A2 architecture, and we achieve competitive performance CelebA.
Qualitatively, we visualize random (not -icked) samples from our trained models on all the datasets. Figure (rows 1-2) contains random samples from CelebA-HQ. Figure (rows 5-6) contains
5

---

Under review as a conference paper at ICLR 2021
FID\Model
VAE
WAE-MMD
WAE-GAN
2-Stage VAE
RAE-SN
MoCA-A1
MoCA-A2
CelebA
63
55
42
34
40.95
48.43
44.59
CIFAR-10
106
80.9
72.9
75.30
77.49
54.36
Table 1: Evaluation of MoCA against existing baselines using FID (lower better, best models in bold). MoCA-AI uses an architecture similar to the one used in WAE (Tolstikhin et al., 2017), while MoCA-A2 uses ResNet- 18 based architecture.
random samples from CIFAR- and CelebA, as well as reconstructions (rows 1-2) of images from these datasets. Most generated samples look realistic, even across multiple resolutions.
We also present latent space interpolations between images from the test set of CelebA-HQ in Figure 1 (rows 3-4). We present latent space interpolations for CIFAR-10 10 and CelebA in Figure 2 (rows 3-4). For these interpolations, we compute the latent vectors Enc(z) (z) and 2' Enc(x') for two images and xz, let 2,
(1 a)=²' for some 0 < S 1, and then generate the interpolated image T.. = Dec (za/||=a ll). These latent space interpolations show that our algorithm causes the generator to learn smooth function from the unit hyper-sphere Sg 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 B.
5 ABLATION ANALYSIS
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 its computational efficiency compared to its competitor Chen et al. (2020), introduces number of important to shed light on their behavior
in
addition to the regularization coefficient X. Therefore, during the training process of MoCA. This section explores how these various hyper parameters impact the quality of generated samples. We also study the impact of the latent dimension of the autoencoder under MoCA training. To keep the analysis tractable and quantitative, we use the Fréchet Inception Distance (FID) score to evaluate the performance of the trained models..
For this section, unless specified otherwise, architecture (architecture A2 in the previous section), = 1, m0g 0.999, Adam with learning rate 0.001, batch size 64, drop this
the CelebA dataset with the ResNet- 18 autoencoder
128, = 3000,
K 60000. For optimization, learning rate by half every 60 epochs, and train for a total 200 epochs. All other optimization hyper-parameters set to the default Pytorch values.
110
128
100
d52
d-512
90
80
10a
10
Figure 3: The interplay between the latent dimension d of the MoCA network and the regularization coefficient X. Larger requires significantly larger À to achieve comparable FID scores for the generated sample quality. See text for more details.
5.1 LATENT SPACE DIMENSIONALITY d AND REGULARIZATION COFFFICIENT À
Real data often lies on low dimensional (dg < n)) manifold that is embedded in a high dimensional (n) space (Bengio et al., 2013a). Autoencoders attempt to map the probability distribution the data to designated prior distribution in latent space of of dimension d, and vice versa. However, if there is significant mismatch in the dimension dy 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 in many "holes" " in the learned latent space which do not correspond to the training data distribution.

---

Under review as a conference paper at ICLR 2021
size 32 x 32
64 x 64
256 x 256
Mp
0
0.9
0.99
0.999
100
2000
20000
FID
86.57
98.97
47.96
47.46
Table 2: The optimal value of À scales linearly with input size, We consider between 100 and 50000 and report the value X" that achieves the best FID score for each input size. See Appendix C for the full data.
Table 3: Smaller base momentum causes model performance to degrade significantly. Per- formance is measured using the FID score (lower better). Note that we use the cosine schedule described in section 3 in all these experiments.
Given the importance of this problem, we study how the latent dimension d influences the quality of samples generated by our used by our model MoCA. We also simultaneously analyze the influence of the regularization coefficient A, 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 25). The results are shown in Figure 3. We find that for d 64 the performance is quite stable across different À values. However, set to larger values, we find that significantly larger value of À is required to reach similar FID score. We hypothesize that this is because a larger forces the "meaningful regions" be more uniformly distributed in the latent space. Therefore, even though there are "holes" in the latent space (due to dg < d), random sample from uniform distribution is more likely to be close to one of these meaningful regions.
is
5.2 CHOOSING À GIVEN INPUT SIZE
An important consideration when selecting the regularization coefficient is the relative scale of the reconstruction loss ||z, (x)(x)))) and contrastive loss. We downsample CelebA- HQ to 32 x 32, 64 x 64, and 256 x 256 for this experiment, and we report the value of A E {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 to the number of pixels in the input (Table 2). See Appendix C for more detailed discussion.
5.3 IMPORTANCE OF MOMENTUM
We study the impact of the contrastive learning m is the exponential moving average typer-paramen used for updating the parameters of the momentum encoder network Encç- m is typically kept to be close to 1 for training stability (He al., 2020). We confirm this intuition for our generative model as well in Table We use the base value mo {0, 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 1.
m on the generated sample quality.
5.4 IMPORTANCE OF TEMPERATURE
Based on the discussion below Eq. 3, the negative term in the contrastive loss essentially estimates the entropy of the latent space distribution due its equivalent kernel density estimation (KDE) interpretation (Eq. 4). Therefore, the temperature hyper-parameter used in the contrastive loss acts as 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 poor quality of generated samples. Additionally, for larger 7, intuitively, larger À should be needed in order to push the KDE samples apart from one another. We confirm these intuitions in Figure Note that due to the large number of experiments this analysis, we train each configuration until 50 epochs (explaining the inferior FID values).
5.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
7

---

Under review as a conference paper at ICLR 2021
Singulas valuee
30
100
T-1.0
T=0.05
20
A= 1080, FID
53.33
A 3000, FID 46.27
75
10
2000
4000
6000
8000
10000
À
20
60
100
120
Singular value value index
Figure 4: Impact of on 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.
Figure 5: SVD of latent representation E R for models trained with various values of Larger À results in ues, i.c., closer to uniform distribution, and lower (better) FID for generated samples.
uniform singular val-
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 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.
negative samples in the contrastive framework which are used to push the latent representations away from one another, encouraging the latent space to be more uniformly distributed. We therefore expect a small K would be bad for achieving this goal. Our experiments in Table 4 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 K chosen, except for K 100, for which FID much worse.
5.6 MISMATCH BETWEEN Qz AND Pz
Finally, we now try to evaluate how well the contrastive term in our objective addresses the problem of matching the marginal distribution Qz = Enc# Px of the autoencoder latent space to the prior distribution Pz. viz, the uniform distribution on the unit hyper-sphere. Since the encoder parameterized output unit & norm vectors, we only need evaluate how close QZZ) to being isotropic. As computationally efficient proxy, compute the singular value decomposition (SVD) of the latent representation corresponding to 10,000 randomly sampled images from training set. We do this for models trained with different values of the regularization coefficient E {0, 500, 1000, 3000). Larger is designed to increase the entropy of the latent space to better match it to the uniform distribution. As Figure 5 shows, for models trained with larger À, the singular values (and therefore the latent distribution) become more uniform. Corresponding FID scores on generated samples reflect this effect since models trained with larger À have lower (better) FID scores.
CONCLUSION AND FUTURE WORK
We propose novel algorithm for learning generative model called Momentum Contrastive Autoen- coders (MoCA). The main idea behind MoCA is to use the contrastive learning framework to match the autoencoder's latent space marginal distribution with the uniform distribution on the unit hyper- sphere. Our objective has theoretical connections with Wasserstein autoencoder, but our algorithm avoids many of the optimization challenges of existing autoencoder-based generative models. We demonstrate that algorithm can generate samples that competitive with or better than existing state-of-the the-art algorithms. Finally, we present a thorough analysis the various algorithm introduced due to the contrastive learning framework, and how they impact learning.
in
We note that contrastive learning currently yields state-of- of-the-art performance in self-supervised learning tasks. Since use it as part of our learning algorithm, natural question is whether can jointly perform representation learning and generative modeling using a single objective. We leave this research direction future work.
8

---

Under review as a conference paper at ICLR 2021
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-generatin 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 1xXiv:1011.44000 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. 899-899997777 00099907,220 907, 2013b.
Olivier Bousquet, Sylvain Gelly, Ilya Tolstikhin, Carl-Johann Simon-Gabriel and Bernhard Schoelkopf. From optimal transport to generative modeling: the vegan cookbook. arXiv preprint arXiv: 2017.
Tong Che, Yanran Li, Athul Paul Jacob, Yoshua Bengio, and Wenjie Li. Mode regularized generative adversarial networks. arXiv preprint drXiv:1612.02216. 2016.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2202.055: 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:1611.02733 2016.
Bin Dai and David Wipf. Diagnosing and enhancing vae models. arXiv preprint 1903.0559, 2019.
Jeff Donahue, Philipp Krähenbühl, and Trevor Darrell. Adversarial feature learning. arXiv preprint 17Xvv:1605.0978. 2016.
Vincent Dumoulin, Ishmael Belghazi, Ben Poole, Olivier Mastropietro, Alex Lamb, Martin Arjovsky, and Aaron Courville, Adversarially learned inference. arXiv preprint arXiv:/600 0000000704, 2016.
Partha Ghosh, Mchdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From 2019.
variational to deterministic autoencoders. arXiv preprint
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural informa- tion processing systems, pp. 2672-2680, 2014.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
Jean-Basien -an-BBstien 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: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07733 2020.
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.

---

Under review as a conference paper at ICLR 2021
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 on Computer Vision and Pattern 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 to a local nash equilibrium. In Advances in neural information processing systems, pp. 6626 6667-663, 2017.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.0310 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 P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. Advances neural information processing systems, pp. 4743-4751, 2016.
Szymon Knop, Jacek Tabor, Przemystaw Spurek, Igor Podolak, Marcin Mazur, and Stanistaw Jastrzebski. Cramer-wold autoencoder. arXiv preprint 1xXivv8805.00 1805.09235 2018.
Soheil Kolouri, Phillip E Pope, Charles E Martin, and Gustavo K Rohde. Sliced wasserstein auto- encoders. In International Conference on 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 in the wild. In Proceedings of the IEEE international conference on computer vision, pp. 3730-3738, 2015.
Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter. Sliced- wasserseessn flows: Nonparametric generative modeling via optimal transport and diffusions. In International Conference on Machine Learning, pp. 4104-4113, PMLR, 2019.
Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial autoencoders. arXiv preprini arXiv: 1.05644, 2015.
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 Veti,o Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performanne deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. 2019.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein anto-encoders. arXiv preprint arXiv: 17110017 2017.
Jakub Tomezak and Max Welling. Vae with vampprior. In International Conference Artificial Intelligence and Statistics, pp. 1214-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 composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning, pp. 1096-1103, 103, 2008
Manzagol. Extracting and
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity the hypersphere. arXiv preprint 4rX1v:2005.11222,,0000 2020.
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Towards deeper understanding of variational autoencoding models. arXiv preprint 1702,0880 2017.
10

---

Under review as a conference paper at ICLR 2021
APPPENDIX
A TRAINING AND EVALUATION DETAILS
We ran all our experiments in Pytorch 1.5.0 (Paszke et al., 2019).
Datasets:
CIFAR-10 contains 32 32 images of 50k training samples and 10k test samples.
CelebA dataset contains a total of ~203 64 x 64 images of divided into 180k training images and 200,000 test images. The images were pre-processed by first taking 140x140 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 ~-3k 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 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, my 0.99. Al used À 1000 while A2 used À 100. Both models were trained for 100 epochs using the Adam optimizer with batch size 64, and learning rate 0.001 decayed by a factor of every 60 epochs.
Both Al and A2 were trained on CelebA with MoCA hyperparameters K 60000, 7 = 0.05, my = 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 1 (CelebA-HQ 256 x 256) were generated using variant of the ResNet-18 architecture. The base ResNet-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, = 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..
11

---

Under review as a conference paper at ICLR 2021
B ADDITIONAL QUALITATIVE RESULTS
Figures 6, 7, and respectively depict additional randomly sampled images, reconstructions, and latent space interpolations for CelebA-HQ 256 x 256. Figures and are generated by the same model used to generate Figure 1, while Figure is generated by an earlier checkpoint of that model (selected for best visual quality).
Figures 9, 11, and 13 respectively depict additional randomly sampled images, reconstructions, and latent space interpolations CIFAR-10 using the same model that achieved the FIID score of 54.36 in table 1.
Figures 10, 12, 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.
Figure 6: Random samples generated by model trained (as described in Appendix A on CelebA-HQ 256 x 256 for 850 epochs. Model checkpoint picked based on best visual quality.
C ADDITIONAL DATA ON CHOOSING À BASED ON INPUT SIZE
In Section 5.2, show that the optimal value of the regularization weight À scales linearly with input size. For this experiment, we downscale CelebA-HO (256 x 256) to 64 64 and 32 x 32, and study the impact of À for the different input sizes. We construct the models for generating d images by removing log, (256/d) dl of the residual blocks from the encoder/decoder of the base model used for CelebA-HQ (256 256) (described in Appendix A), each block hownsamples/upsamplss
12

---

Under review as a conference paper at ICLR 2021
n
STR
IP
PO
Figure 7: Image reconstructions by a 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.
the image by factor of 2 x 2. We train all models for 400 epochs using the Adam optimizer with batch size 64 and learning rate 0.002 (decayed by a factor of 2 every 40 epochs).
The full data for this experiment (which support the claims of Table 2) can be found in Figure 15. We consider (100, 200, 500, 1000, 2000, 5000, 10000, 20000, 20000,50000) for each image size (except 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.
13

---

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

---

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

---

Under review as a conference paper at ICLR 2021
FID vs. for Different Image Sizes
32 x 32 Images
55
FID
100
1000
2000
5000
10000
20000
50000
64 x 64 Images
55
FID
53
100
200
500
1000
2000
5000
10000
20000
50000
256 x 256 Images
120
FID
100
80
100
200
500
1000
2000
5000
10000
20000
50000
A
Figure 15: Impact of regularization weight on FID score based on input size. Optimal values X (labeled with a red star) scale linearly with input size. Note that absolute FID scores are not comparable between different image sizes! This figure focuses on relative trends.
16


Reviews:
[{"review_id": "aZhxHe6EbGU", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Description: The authors propose momentum contrastive Wasserstein autoencoders (MoCA), which is an extension of the Wasserstein autoencoder (WAE) that aims to match the prior p(z) and aggregate variational posterior q(z) through the use of contrastive learning, as opposed to earlier proposed techniques (MMD, GAN). The use of contrastive learning here is theoretically motivated by the fact that -- as shown in Wang & Isola (2020) -- in the limit of infinitely many negative samples, the distribution induced by the contrastive encoder is uniform on the hypersphere. Therefore, if p(z) is the unit + uniform hypersphere, then we can leverage contrastive learning to drive the aggregate variational posterior q(z) to be as close to it as possible. This provides a principled way to train WAEs since its corresponding optimisation assumes that P_Z == Q_Z.\n\nAdvantages:\n- A theoretical connection is made between contrastive learning and the training of WAEs.\n- Ablations explore the intricacies of combining MoCA with WAE.\n- Paper is written clearly.\n\nDisadvantages:\n- The empirical results in Table 1 suggest that MoCA is at least competitive with CelebA, though (1)  these should have uncertainty estimates and (2) the WAE-GAN result for CIFAR10 is missing. You should include uncertainty estimates over multiple runs. This may also involve you having to run WAE-GAN / WAE-MMD yourself, since I assume that you were not able to quote uncertainty estimates from their paper. Even if uncertainty estimates were in the cited papers, I *strongly recommend* to run these experiments yourself to remove any confounding variables from the fact that these numbers were quoted from the literature (since the experimental setup could be vastly different compared to yours). I'm not suggesting you do this for all the methods in Table 1, but I highly suggest you do so for the WAE experiments. Otherwise I have little confidence in the results you have presented.\n\nComments / suggestions / typos:\n\n- Theorem 1, \"...when g ~ P_z\". Don't you mean \"...when z ~ P_z\"?\n- Use subfigures for Figure 2 to make the presentation more clear (distinguishing between reconstructions, interpolations, and samples from the prior)\n- For the FID calculation in Table 1, what is the reference dataset here? (I see in the supplementary material you say it's the test set, ok, maybe mention this in the caption for Table 1 since it's important)\n- Generally, VAE papers also report the ELBO on the test set to examine generalisation performance. Can a similar metric be derived for your case? For instance, ELBO surgery [1] rewrites the vanilla KL term (the KL between q(z_i|x) and p(z)) as the KL between the aggregate posterior and prior (KL[ q(z) || p(z) ]) + an extra mutual information term (also mentioned in the \"Related work\" section of the original WAE paper). In that case, if you use a WAE-GAN, the discriminator estimates KL[ q(z) || p(z) ], and you could use that to obtain a rough lower bound on log p(x) on the test set. In fact, that in some sense makes WAE-GAN slightly more appealing to use here than what you have presented. What are your thoughts on this?\n- It's not clear to me whether using MoCA would be preferable from a more computationally 'economic' perspective than WAE-GAN. In the cases of both, hyperparameters are necessary, and GAN training has been made easier in recent times with the right 'tricks' (say, careful normalisation of the discriminator and using WGAN/JSGAN). Can something be said about the memory/computation time tradeoffs between using WAE-GAN and MoCA? For instance, if you're using contrastive learning, you either need a large batch size (for SimCLR) or a large queue (MoCA), whereas with GAN training you don't need this.\nThe main contribution is the nice theoretical connection between WAE training and contrastive learning. However, I am concerned with the statistical significance of some of the results.\n\nJustification of rating: my main concern is in the statistical significance of the results in Table 1. I am open to raising my score if my concerns have been addressed.\n\nReferences:\n- [1]: Hoffman, M. D., & Johnson, M. J. (2016, December). Elbo surgery: yet another way to carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS (Vol. 1, p. 2).", "extra_scores": {}, "extra_text": {}}, {"review_id": "BoAASPJLR1I", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nThe authors aim at constructing another generative autoencoder: reconstruction loss plus matching-aggregate-posterior-with-prior penalty. For this they use a constrastive loss in latent space.\n\nWithout paying attention to writing style and clarity, etc., the paper should be rejected to be published at ICLR already for the following reasons:\n\n- lack of originality: any paper that follows the above recipe of minimizing reconstruction loss plus a penalty that matches (aggregate) posterior with prior needs more content than just swapping the divergence in latent space. Since contrastive losses are well known their invention can't be contributed to this paper.\nThis brings us to the following point:\n- lack of theoretical guarantees: There are no new theoretical guarantees or insights in this paper. Empirically lowering the FID score on cifar by a few points without any guarantees is imho not good enough anymore for research in generative autoencoders.\nFurthermore:\n- missing citations: the authors didn't keep up with the literature. among many others they miss e.g. [PB].\n\n[PB] G Patrini, R van den Berg, et al.: Sinkhorn AutoEncoders. UAI 2019.", "extra_scores": {}, "extra_text": {}}, {"review_id": "AT3VdMjjCt", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is fairly easy to follow and does not have serious presentation problems. Especially, the discussion on the background (WAE) is clear and concise.\n\nNone the less, the two weaknesses are\n1. The idea of having a hyperspherical latent space [1] is already been studied (and is a famous work), yet the author does not include any discussion with it. Even the author has discussed it, I can't see any difference between using the methodology presented in prior work and the submission. \n2. There is nothing new I can learn from this submission. Precisely, I can't find the presented approach has any benefits over existing variational auto-encoder variants. \n\nSuggestion:\n1. I can't agree that momentum encoder makes contrastive learning \"more computational efficient\". It requires extra handing on the momentum. I think what the author likes to claim is that the momentum encoder leverages an asynchronous update and hence the effective batch size is larger. \n2. In section 5.4, I can't agree with the statement from the author \"for larger temperature, the estimated distribution becomes smoother and the entropy estimation becomes poor, which should result in a poor quality of generated samples.\" This statement is very vague and is incorrect. \n\nTypo:\n1. Theorem 1: should be when Z~P_Z not when g ~ P_Z\n\n[1] Hyperspherical Variational Auto-Encoders, Davidson et al., UAI 2018.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "1leGDR4fvdS", "reviewer": "AnonReviewer3", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThis paper considers training autoencoders with a hyperspherical latent space distribution. The encoder maps inputs to latent variables with a unit norm, and a contrastive loss with momentum is used to encourage the latent variables to be distributed uniformly over the surface of the unit-hypersphere. Generations of this autoencoder model are compared against wasserstein auto-encoders and several VAEs and evaluated with FID scores on CIFAR10 and CelebA. On CelebA the model is outperformed by both WAE-GAN (Wasserstein auto-encoders trained with a GAN objective to match the latent distributions) and a 2-stage VAE. On CIFAR10 the model outperforms related work, but only when using a different encoder/decoder architecture.\nSamples, reconstructions and interpolations are also depicted for Cifar10, CelebA and CelebA-HQ. The influence of hyperparameter settings are studied through various ablation studies. \n\n**Pros**\n* The paper is well written.\n* The proposed idea is clearly presented and well motivated.\n* Hyperparameter settings and their influence on performance are extensively studied.\n\n**Cons**\n* The claims on performance gains on both Cifar10 and CelebA are overstated. The abstract states that state of the art FID scores on CIFAR10 (54.36) are obtained, which is actually far from the current state of the art on CIFAR10 (which is around 3.17, achieved by [1]). A more representative claim would be one that states that the proposed method outperforms related methods on CIFAR10 when making some architecture changes. On CelebA the statement is that the model performs competitively with respect to related work, even though related work such as the 2-stage VAE has an FID score of about 10 points lower. I understand that compute resources are not equally distributed among researchers and that lots of resources are required to scale to massive models in order to reach state of the art performance. Comparisons to related work on smaller models with comparable architectures is also a perfectly valid way to demonstrate the workings of a new algorithm/model. However, it would be more transparent if the results would then be framed in this context. \n* In table 1 with the FID scores, it seems like most of the performance gains of the proposed model are actually obtained by changing the architecture of the encoder/decoder that is used in related work to a ResNet-18-based architecture with much fewer parameters. Importantly, with an architecture comparable to related work the proposed method does not improve upon related work on CIFAR10. To disentangle the performance gains due to the architecture change and the proposed method, at least some of the related methods should be evaluated with this different ResNet-18-based architecture. With the current results table it is hard to conclude if there are any benefits to using this contrastive auto-encoder framework compared to the other related methods.\n* In section 5.1 the interplay of the latent dimensionality ($d$) and the hyperparameter ($\\lambda$) that controls the contribution of the latent contrastive loss compared to the reconstruction loss is discussed. The observation is made that a larger dimensionality $d$ requires a higher value of $\\lambda$ in order to obtain better FID scores (Fig 3). A hypothesis for this is given along the lines of there being holes in the latent space when the intrinsic data dimensionality is lower than that of the latent dimensionality, and that larger $\\lambda$ will make sure to distribute the latent representations of datapoints more uniformly. However, wouldn’t the following argument also explain the need for a larger $\\lambda$ for a higher $d$:  The probability of two random unit-norm encodings being orthogonal to each other increases with the dimensionality. This can lead to a larger probability of obtaining inner products with small magnitudes in the exponent of eq 6 and hence a contrastive loss that is closer to zero. Therefore, to rescale the contribution of the contrastive loss with respect to the reconstruction loss (which is independent of latent dimensionality) one needs to increase the $\\lambda$. Could this also be the reason?\n\n\n**Minor comments/questions**\n* There are various other related works on auto-encoders with hyperspherical latent spaces that are not mentioned, such as [2, 3]. Please discuss more related work with hyperspherical priors.\n* Please add a citation to [4] when referring to VAE’s in the second paragraph of the related work.\n* The sample quality of VAEs has improved rapidly recently, see for instance [5]. The related work section on the sample quality of VAEs could reflect this a bit more.\n* Typo in paragraph starting with Ghosh et al. in related work: “the further…” → “they further…”\n* In theorem 1 in the end of the second line, it states $g\\sim P_Z$. Should this be $Z\\sim P_Z$?\n* A factor of $1/K$ seems to be missing in the equation for $-L^{MC}_{neg}$ in the paragraph below eq. 4.\n\n[1] Ho et al., denoising diffusion probabilistic models, https://arxiv.org/abs/2006.11239\n[2] Patrini et al., Sinkhorn autoencoders, https://arxiv.org/abs/1810.01118\n[3] Davidson et al., Hyperspherical Variational Auto-encoders, https://arxiv.org/abs/1804.00891\n[4] Rezende et al., Stochastic backpropagation and approximate inference in deep generative models, https://arxiv.org/abs/1401.4082\n[5] Vahdat & Kautz, NVAE: A deep hierarchical variational autoencoder, https://arxiv.org/abs/2007.03898\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents an interesting approach for training generative autoencoders with a latent space that lies on a hyperspherical subspace. However, the reviewers have raised concerns regarding the similarity of this work with several prior works and have questioned the experimental setup. Without the authors' response, we cannot situate this paper among the prior work properly. Thus, I recommend Reject at this point. 
Decision: Reject