Sliced Wasserstein Auto-Encoders

Scalable generative models that capture the rich and often nonlinear distribution of high-dimensional data , ( i.e. , image , video , and audio ) , play a central role in various applications of machine learning , including transfer learning Isola et al . ( 2017 ) ; Murez et al . ( 2018 ) , super-resolution Ledig et al . ( 2016 ) ; Kolouri & Rohde ( 2015 ) , image inpainting and completion Yeh et al . ( 2017 ) , and image retrieval Creswell & Bharath ( 2016 ) , among many others . The recent parametric generative models , including Generative Adversarial Networks ( GANs ) Goodfellow et al . ( 2014 ) ; Radford et al . ( 2015 ) ; Arjovsky et al . ( 2017 ) ; Berthelot et al . ( 2017 ) and Variational auto-encoders ( VAE ) Kingma & Welling ( 2013 ) ; Mescheder et al . ( 2017 ) ; Bousquet et al . ( 2017 ) enable an unsupervised and end-to-end modeling of the high-dimensional distribution of the training data . Learning such generative models boils down to minimizing a dissimilarity measure between the data distribution and the output distribution of the generative model . To this end , and following the work of Arjovsky et al . ( 2017 ) and Bousquet et al . ( 2017 ) , we approach the problem of generative modeling from the optimal transport point of view . The optimal transport problem Villani ( 2008 ) ; Kolouri et al . ( 2017 ) provides a way to measure the distances between probability distributions by transporting ( i.e. , morphing ) one distribution into another . Moreover , and as opposed to the common information theoretic dissimilarity measures ( e.g. , f -divergences ) , the p-Wasserstein dissimilarity measures that arise from the optimal transport problem : 1 ) are true distances , and 2 ) metrize a weak convergence of probability measures ( at least on compact spaces ) . Wasserstein distances have recently attracted a lot of interest in the learning community Frogner et al . ( 2015 ) ; Gulrajani et al . ( 2017 ) ; Bousquet et al . ( 2017 ) ; Arjovsky et al . ( 2017 ) ; Kolouri et al . ( 2017 ) due to their exquisite geometric characteristics Santambrogio ( 2015 ) . See the supplementary material for an intuitive example showing the benefit of the Wasserstein distance over commonly used f -divergences . In this paper , we introduce a new type of auto-encoders for generative modeling ( Algorithm 1 ) , which we call Sliced-Wasserstein auto-encoders ( SWAE ) , that minimize the sliced-Wasserstein distance between the distribution of the encoded samples and a samplable prior distribution . Our work is most closely related to the recent work by Bousquet et al . ( 2017 ) and more specifically the follow-up work by Tolstikhin et al . ( 2017 ) . However , our approach avoids the need to perform adversarial training in the encoding space and is not restricted to closed-form distributions , while still benefiting from a Wasserstein-like distance measure in the latent space . Calculating the Wasserstein distance can be computationally expensive , but our approach permits a simple numerical solution to the problem . Finally , we note that there has been several concurrent papers , including the work by Deshpande et al . ( 2018 ) and Şimşekli et al . ( 2018 ) , that also looked into the application of sliced-Wasserstein distance in generative modeling . Regardless of the concurrent nature of these papers , our work remains novel and is distinguished from these methods . Deshpande et al . ( 2018 ) use the sliced-Wasserstein distance to match the distributions of high-dimensional reconstructed images , which require large number of slices , O ( 104 ) , while in our method and due to the distribution matching in the latent space we only need O ( 10 ) slices . We also note that Deshpande et al . ( 2018 ) proposed to learn discriminative slices to mitigate the need for a very large number of random projections that is in essence similar to the adversarial training used in GANs , which contradicts with our goal of not using adversarial training . Şimşekli et al . ( 2018 ) , on the other hand , take an interesting but different approach of parameter-free generative modeling via sliced-Wasserstein flows . 1 NOTATION AND PRELIMINARIES . Let X denote the compact domain of a manifold in Euclidean space and let xn ∈ X denote an individual input data point . Furthermore , let ρX be a Borel probability measure defined on X . We define the probability density function pX ( x ) for input data x to be : dρX ( x ) = pX ( x ) dx Let φ : X → Z denote a deterministic parametric mapping from the input space to a latent space Z ( e.g. , a neural network encoder ) . To obtain the density of the push forward of ρX with respect to φ , i.e. , ρZ = φ∗ ( ρX ) , we use Random Variable Transformation ( RVT ) Gillespie ( 1983 ) ) . In short , the probability density function of the encoded samples z can be expressed in terms of φ and pX by : pZ ( z ) = ∫ X pX ( x ) δ ( z − φ ( x ) ) dx , ( 1 ) where δ denotes the Dirac distribution function . Similar to variational Auto-Encoders ( VAEs ) Kingma & Welling ( 2013 ) and the Wasserstein Auto-Encoders ( WAE ) Tolstikhin et al . ( 2017 ) , our main objective is to encode the input data points x ∈ X into latent codes z ∈ Z such that : 1 ) x can be recovered/approximated from z , and 2 ) the probability density function of the encoded samples , pZ , follows a prior distribution qZ . Let ψ : Z → X be the decoder that maps the latent codes back to the original space such that pY ( y ) = ∫ X pX ( x ) δ ( y − ψ ( φ ( x ) ) ) dx , ( 2 ) where y denotes the decoded samples . It is straightforward to see that when ψ = φ−1 ( i.e . ψ ( φ ( · ) ) = id ( · ) ) , the distribution of the decoder pY and the input distribution pX are identical . Hence , in its most general form , the objective of such auto-encoders simplifies to learning φ and ψ , so that they minimize a dissimilarity measure between pY and pX , and between pZ and qZ . In what follows , we briefly review the existing dissimilarity measures for these distributions . 1.1 MINIMIZING DISSIMILARITY BETWEEN pX AND pY We first emphasize that the VAE often assumes stochastic encoders and decoders Kingma & Welling ( 2013 ) , while we consider the case of only deterministic mappings . Although , we note that , similar to WAE , SWAE can also be formulated with stochastic encoders . Different measures have been used previously to compute the dissimilarity between pX and pY . Most notably , Nowozin et al . ( 2016 ) showed that for the general family of f -divergences , Df ( pX , pY ) , ( including the KL-divergence , JensenShannon , etc . ) , using the Fenchel conjugate of the convex function f and minimizing Df ( pX , pY ) leads to a min-max problem that is equivalent to the adversarial training widely used in the generative modeling literature Goodfellow et al . ( 2014 ) ; Makhzani et al . ( 2015 ) ; Mescheder et al . ( 2017 ) . Others have utilized the rich mathematical foundation of the OT problem and Wasserstein distances Arjovsky et al . ( 2017 ) ; Gulrajani et al . ( 2017 ) ; Bousquet et al . ( 2017 ) ; Tolstikhin et al . ( 2017 ) to define a distance between pX and pY . In Wasserstein-GAN , Arjovsky et al . ( 2017 ) utilized the Kantorovich-Rubinstein duality for the 1-Wasserstein distance , W1 ( pX , pY ) , and reformulated the problem as a min-max optimization that is solved through an adversarial training scheme . Inspired by the work of Bousquet et al . ( 2017 ) and Tolstikhin et al . ( 2017 ) , it can be shown that ( see supplementary material for a proof ) : Wc ( pX , pY ) ≤W ‡c ( pX , pY ) : = EpX ( c ( x , ψ ( φ ( x ) ) ) ) ( 3 ) = ∫ X c ( x , ψ ( φ ( x ) ) ) pX ( x ) dx , Furthermore , the r.h.s . of equation 3 supports a simple implementation where for i.i.d samples of the input distribution , { xn } Nn=1 , the upper bound can be approximated as : W ‡c ( pX , pY ) ≈ 1 N N∑ n=1 c ( xn , ψ ( φ ( xn ) ) ) ( 4 ) The r.h.s of equation 3 and equation 4 take advantage of the existence of pairs xn and yn = ψ ( φ ( xn ) ) , which make f ( · ) = ψ ( φ ( · ) ) a transport map between pX and pY ( but not necessarily the optimal transport map ) . In this paper , we minimize W ‡c ( pX , pY ) following equation 4 to minimize the discrepancy between pX and pY . Next , we focus on the discrepancy measures between pZ and qZ . 1.2 MINIMIZING DISSIMILARITY BETWEEN pZ AND qZ If qZ is a known distribution with an explicit formulation ( e.g . Normal distribution ) the most straightforward approach for measuring the ( dis ) similarity between pZ and qZ is the log-likelihood of z = φ ( x ) with respect to qZ , formally : supφ ∫ X pX ( x ) log ( qZ ( φ ( x ) ) ) dx ( 5 ) maximizing the log-likelihood is equivalent to minimizing the KL-divergence between pZ and qZ , DKL ( pZ , qZ ) ( see supplementary material for more details and derivation of Equation equation 5 ) . This approach has two major limitations : 1 ) The KL-Divergence and in general f -divergences do not provide meaningful dissimilarity measures for distributions supported on non-overlapping lowdimensional manifolds Arjovsky et al . ( 2017 ) ; Kolouri et al . ( 2018 ) ( see supplementary material ) , which is common in hidden layers of neural networks , and therefore they do not provide informative gradients for training φ , and 2 ) we are limited to distributions qZ that have known explicit formulations , which is restrictive as it eliminates the ability to use the much broader class of samplable distributions . Various alternatives exist in the literature to address the above-mentioned limitations . These methods often sample Z̃ = { z̃j } Nj=1 from qZ and Z = { zn = φ ( xn ) } Nn=1 from pX and measure the discrepancy between these sets ( i.e . point clouds ) . Note that there are no one-to-one correspondences between z̃js and zns . In their influential WAE paper , Tolstikhin et al . ( 2017 ) proposed two different approaches for measuring the discrepancy between Z̃ and Z , namely the GAN-based and the maximum mean discrepancy ( MMD ) -based approaches . The GAN-based approach proposed in Tolstikhin et al . ( 2017 ) defines a discriminator network , DZ ( pZ , qZ ) , to classify z̃js and zns as coming from ‘ true ’ and ‘ fake ’ distributions correspondingly , and proposes a min-max adversarial optimization for learning φ and DZ . The MMD-based approach , utilizes a positive-definite reproducing kernel k : Z × Z → R to measure the discrepancy between Z̃ and Z . The choice of the kernel and its parameterization , however , remain a data-dependent design parameter . An interesting alternative approach is to use the Wasserstein distance between pZ and qZ . Following the work of Arjovsky et al . ( 2017 ) , this can be accomplished utilizing the Kantorovich-Rubinstein duality and through introducing a min-max problem , which leads to yet another adversarial training scheme similar to the GAN-based method in Tolstikhin et al . ( 2017 ) . Note that , since elements of Z̃ and Z are not paired , an approach similar to equation 4 could not be used to minimize the discrepancy . In this paper , we propose to use the sliced-Wasserstein metric , Rabin & Peyré ( 2011 ) ; Rabin et al . ( 2011 ) ; Bonneel et al . ( 2015 ) ; Kolouri et al . ( 2016b ) ; Carriere et al . ( 2017 ) ; Kolouri et al . ( 2018 ) , to measure the discrepancy between pZ and qZ . We show that using the sliced-Wasserstein distance ameliorates the need for training an adversary network or choosing a data-dependent kernel ( as in WAE-MMD ) , and provides an efficient , stable , and simple numerical implementation . Before explaining our proposed approach , it is worthwhile to point out the major difference between learning auto-encoders as generative models and GANs . In GANs , one needs to minimize a distance between { ψ ( z̃j ) |z̃j ∼ qZ } Mj=1 and { xn } Mn=1 , which are high-dimensional point clouds for which there are no correspondences between ψ ( z̃j ) s and xns . For the auto-encoders , on the other hand , there exists correspondences between the high-dimensional point clouds { xn } Mn=1 and { yn = ψ ( φ ( xn ) ) } Mn=1 , and the problem simplifies to matching the lower-dimensional point clouds { φ ( xn ) } Mn=1 and { z̃j ∼ qZ } Mj=1 . In other words , the encoder performs a nonlinear dimensionality reduction , that enables us to solve a simpler problem compared to GANs . Next we introduce the details of our approach .

Summary: The authors propose a new autoencoding algorithm for the unsupervised generative modeling which they call Sliced Wasserstein Autoencoders (SWAE). SWAE minimizes a reconstruction cost (measured with respect to the non-negative cost function c(x,x') defined for pairs of input images x, x'), regularized by a penalty measuring a discrepancy between the prior distribution over the latent space qz and the push-forward pz of the unknown data distribution through the deterministic encoder. The authors present an extensive theoretical argument supporting the choice of this objective and a number of empirical results performed on MNIST, LSUN bedrooms, and Celeba.