Momentum Contrastive Autoencoder: Using Contrastive Learning for Latent Space Distribution Matching in WAE

1 INTRODUCTION . The main goal of generative modeling is to learn a 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 et al . ( 2014 ) ) 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 . 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 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 . 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 . Existing research has tried to address this challenge ( Kolouri et al. , 2018 ; Knop et al. , 2018 ) ( see section 2 for a discussion ) . This paper aims to design a generative model to 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 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 . 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 of 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 . 3. we perform a thorough ablation analysis of the impact of the hyper-parameters 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 we focus on 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 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 ; 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 a few important drawbacks of VAEs . The VAE objective is 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 ) . 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 ) . Dai & Wipf ( 2019 ) claim that this happens due to mismatch between the AE latent space dimension and the intrinsic dimensionality of the data manifold ( which is typically unknown ) , and propose a two stage VAE to remedy this problem . VQ-VAE ( Oord et al. , 2017 ) take a 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 ) propose to replace the Gaussian prior used in VAE with the Von Mises-Fisher ( vMF ) distribution , which is analogous to 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 an 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 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 ) . There has been research on making use 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 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 ) . 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 it for backpropagation ( which is computationally expensive ) ; though their choice of Sinkhorn algorithm for latent space distribution matching allows their approach to be general . 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 Z ∼ 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 ) 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 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 ) : = 1K ∑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 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: 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.