Momentum Contrastive Autoencoder

1 INTRODUCTION . The main goal of generative modeling is to learn a 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 et al . ( 2014 ) ) are theoretically-grounded 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 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 . 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 a prior distribution . The core challenge of this framework is to match the latent space distribution to a 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 to 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 a simple observation that allows us to use the contrastive learning framework to solve this problem . Contrastive learning achieves state-of-the-art results in selfsupervised representation learning tasks ( He et al. , 2020 ; Chen et al. , 2020 ) by forcing the latent representations to be 1 ) 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 optimization challenges of existing methods , thus resulting in a 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 a 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 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 a few important drawbacks of VAEs . The poor sample qualities of VAE has been attributed to a mismatch between the prior ( which is used for drawing samples ) and the posterior ( Kingma et al. , 2016 ; Tomczak & 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 if 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 to a prior . In this context , one can view our proposed algorithm as a way to regularize deterministic autoencoders while simultaneously learning a 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 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-Cramer distance is introduced in Knop et al . ( 2018 ) . However , the number of required random projections becomes prohibitively high when the data lives on a 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 be a random variable sampled from the real data distribution on X , Z ∼ Q ( Z|X ) be its latent representation in Z ⊆ Rd , and X̂ = g ( Z ) be its reconstruction by a deterministic decoder/generator g : Z → X . Note that the encoder Q ( Z|X ) can also be deterministic in the WAE framework , and we let f ( X ) dist= Q ( Z|X ) for some deterministic f : X → Z . Theorem 1 . ( Bousquet et al. , 2017 ; Tolstikhin et al. , 2017 ) Let PZ be a prior distribution on Z , let Pg = g # PZ be the push-forward of PZ under g ( i.e . the distribution of X̂ = g ( Z ) when g ∼ PZ ) , and let QZ = f # PX be the push-forward of PX under f . Then , Wc ( PX , Pg ) = inf Q : QZ=PZ E X∼PX Z∼Q ( Z|X ) [ c ( X , g ( Z ) ) ] = inf f : f # PX=PZ E X∼PX [ c ( X , g ( f ( X ) ) ] ( 1 ) where Wc 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/generator ) as long as the constraint is met . In this work , we parameterize the encoder network f : X → Rd such that latent variable Z = f ( X ) has unit ` 2 norm . Our goal is then to match the distribution of this Z to the uniform distribution over the unit hyper-sphere Sd = { z ∈ Rd : ‖z‖2 = 1 } . To do so , we study the so-called “ negative sampling ” component of the contrastive loss used in self-supervised learning , Lneg ( f ; τ , K ) = E x∼PX { x−i } K i=1∼PX log 1 K K∑ j=1 ef ( x ) T f ( x−j ) /τ ( 2 ) Here , f : X → Sd is a neural network whose output has unit ` 2 norm , τ is the temperature hyperparameter , and K is the number of samples ( another hyper-parameter ) . Theorem 1 of Wang & Isola ( 2020 ) shows that for any fixed t , when K →∞ , lim K→∞ ( Lneg ( f ; τ , K ) − logK ) = E x∼PX [ log E x−∼PX [ ef ( x ) T f ( x− ) /τ ] ] ( 3 ) Crucially , this limit is minimized exactly when the push-forward f # PX ( i.e . the distribution of the latent random variable Z = f ( X ) when X ∼ PX ) is uniform on Sd . Moreover , even the Monte Carlo approximation of Eq . 2 ( with mini-batch size B and some K such that B ≤ K < ∞ ) LMCneg ( f ; τ , K , B ) = 1 B B∑ i=1 log 1 K K∑ j=1 ef ( xi ) T f ( xj ) /τ ( 4 ) is a consistent estimator ( up to a constant ) of the entropy of f # PX called the redistribution estimate ( Ahmad & Lin , 1976 ) . This follows if we notice that k ( xi ; τ , K ) : = ∑K j=1 e f ( xi ) T f ( xj ) /τ is the un-normalized kernel density estimate of f ( xi ) using the i.i.d . samples { xj } Kj=1 , so −LMCneg ( f ; τ , K , B ) = − 1B ∑B i=1 log k ( xi ; τ , K ) ( Wang & Isola , 2020 ) . So minimizing Lneg ( and importantly LMCneg ) maximizes the entropy of f # PX . 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 or a GAN-based estimate of the divergence between f # PX and PZ . By letting PZ be the uniform distribution over the unit hyper-sphere Sd , the insights above allow us to instead minimize the much simpler regularized loss L ( f , g ; λ , τ , B , K ) = 1 B B∑ i=1 ‖xi − g ( f ( xi ) ) ‖22 + λLMCneg ( f ; τ , K , B ) ( 5 ) Training : For simplicity , we will now use the notation Enc ( · ) and Dec ( · ) to respectively denote the encoder and decoder network of the autoencoder . Further , the d-dimensional output of Enc ( · ) is ` 2 normalized , i.e. , ‖Enc ( x ) ‖2 = 1 ∀x . Based on the theory above , we aim to minimize the Algorithm 1 PyTorch-like pseudocode of Momentum Contrastive Autoencoder algorithm # Enc_q , Enc_k : encoder networks for query and key . Their outputs are L2 normalized # Dec : decoder network # Q : dictionary as a queue of K randomly initialized keys ( dxK ) # m : momentum # lambda : regularization coefficient for entropy maximization # tau : logit temperature for x in data_loader : # load a minibatch x with B samples z_q = Enc_q ( x ) # queries : Bxd z_k = Enc_k ( x ) .detach ( ) # keys : Bxd , no gradient through keys x_rec = Dec ( z_q ) # reconstructed input # positive logits : Bx1 l_pos = bmm ( z_q.view ( B,1 , d ) , z_k.view ( B , d,1 ) ) # negative logits : BxK l_neg = mm ( z_q.view ( B , d ) , Q.view ( d , K ) ) # logits : Bx ( 1+K ) logits = cat ( [ l_pos , l_neg ] , dim=1 ) # compute loss labels = zeros ( B ) # positive elements are in the 0-th index L_con = CrossEntropyLoss ( logits/tau , labels ) # contrastive loss maximizing entropy of z_q L_rec = ( ( x_rec - x ) * * 2 ) .sum ( ) / B # reconstruction loss L = L_rec + lambda * L_con # momentum contrastive autoencoder loss # update Enc_q and Dec networks L.backward ( ) update ( Enc_q.params ) update ( Dec.params ) # update Enc_k Enc_k.params = m * Enc_k.params + ( 1-m ) * Enc_q.params # update dictionary enqueue ( Q , z_k ) # enqueue the current minibatch dequeue ( Q ) # dequeue the earliest minibatch bmm : batch matrix multiplication ; mm : matrix multiplication ; cat : concatenation . enqueue appends Q with the keys zk ∈ RB×d from the current batch ; dequeue removes the oldest B keys from Q loss L ( Enc , Dec ; λ , τ , B , K ) , where λ is the regularization coefficient , τ is the temperature hyperparameter , B is the mini-batch size , and K ≥ B is the number of samples used to estimate Lneg . In practice , we propose to use the momentum contrast ( MoCo , He et al . ( 2020 ) ) framework to implement Lneg . Let Enct be parameterized by θt at step t of training . Then , we let Enc′t be the same encoder parameterized by the exponential moving average θ̃t = ( 1−m ) ∑t i=1m t−iθi . Letting x1 , . . . , xK be the K most recent training examples , and letting t ( j ) = t − bj/Bc be the time at which xj appeared in a training mini-batch , we replace LMCneg at time step t with LMoCo = 1 B B∑ i=1 log 1 K K∑ j=1 exp ( Enct ( xi ) TEnc′t ( j ) ( xj ) τ ) − 1 B B∑ i=1 Enct ( xi ) TEnc′t ( xi ) τ ( 6 ) This approach allows us to 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 ) . Forcing 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 Enct ( xi ) and Enc′t ( xi ) ( so-called “ positive samples ” in the terminology of contrastive learning ) . Note that we do not use any data augmentations in our algorithm , but this similarity term is still non-trivial since the networks Enct and Enc′t are not identical . Pseudo-code of our final algorithm , which we call Momentum Contrastive Autoencoder ( MoCA ) , is shown in Algorithm 1 ( 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 the tth iteration as m = 1− ( 1−m0 ) · ( cos ( πt/T ) + 1 ) /2 , where T is the total number of training iterations , and m0 is the base momentum hyper-parameter . Inference : Once the model is trained , the marginal distribution of the latent space ( i.e . the pushforward Enc # PX ) 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 z ∼ N ( 0 , I ) from the standard multivariate normal distribution in Rd and then generate a sample xg : = Dec ( z/‖z‖2 ) .

Summary: This 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.