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: This paper presents an extension of Wasserstein autoencoder (WAE) by modifying the regularization term in learning objective of variational autoencoder. This term measures the divergence between the distribution of the encoded training samples and the samplable prior distribution. The modification is based on the sliced-Wasserstein distance where the distance between two distributions is measured through slicing or projecting the high-dimensional distributions into one-dimensional marginal distributions. As a result, a closed-form solution to the integral in Eq. (9) is obtained via a numerical method. The adversarial learning in WAE, designed to fulfill the calculation of high-dimensional distance, can be avoided. In general, this is an interesting work by introducing new idea of sliced-Wasserstein distance.