Distributional Sliced-Wasserstein and Applications to Generative Modeling

Sliced-Wasserstein distance ( SW ) and its variant , Max Sliced-Wasserstein distance ( Max-SW ) , have been used widely in the recent years due to their fast computation and scalability even when the probability measures lie in a very high dimensional space . However , SW requires many unnecessary projection samples to approximate its value while Max-SW only uses the most important projection , which ignores the information of other useful directions . In order to account for these weaknesses , we propose a novel distance , named Distributional Sliced-Wasserstein distance ( DSW ) , that finds an optimal distribution over projections that can balance between exploring distinctive projecting directions and the informativeness of projections themselves . We show that the DSW is a generalization of Max-SW , and it can be computed efficiently by searching for the optimal push-forward measure over a set of probability measures over the unit sphere satisfying certain regularizing constraints that favor distinct directions . Finally , we conduct extensive experiments with large-scale datasets to demonstrate the favorable performances of the proposed distances over the previous sliced-based distances in generative modeling applications . 1 INTRODUCTION . Optimal transport ( OT ) is a classical problem in mathematics and operation research . Due to its appealing theoretical properties and flexibility in practical applications , it has recently become an important tool in the machine learning and statistics community ; see for example , ( Courty et al. , 2017 ; Arjovsky et al. , 2017 ; Tolstikhin et al. , 2018 ; Gulrajani et al. , 2017 ) and references therein . The main usage of OT is to provide a distance named Wasserstein distance , to measure the discrepancy between two probability distributions . However , that distance suffers from expensive computational complexity , which is the main obstacle to using OT in practical applications . There have been two main approaches to overcome the high computational complexity problem : either approximate the value of OT or apply the OT adaptively to specific situations . The first approach was initiated by ( Cuturi , 2013 ) using an entropic regularizer to speed up the computation of the OT ( Sinkhorn , 1967 ; Knight , 2008 ) . The entropic regularization approach has demonstrated its usefulness in several application domains ( Courty et al. , 2014 ; Genevay et al. , 2018 ; Bunne et al. , 2019 ) . Along this direction , several works proposed efficient algorithms for solving the entropic OT ( Altschuler et al. , 2017 ; Lin et al. , 2019b ; a ) as well as methods to stabilize these algorithms ( Chizat et al. , 2018 ; Peyré & Cuturi , 2019 ; Chizat et al. , 2018 ; Schmitzer , 2019 ) . However , these algorithms have complexities of the order O ( k2 ) , where k is the number of supports . It is expensive when we need to compute the OT repeatedly , especially in learning the data distribution . ∗The work was finished when Nhat Ho worked at VinAI Research in the summer of 2020 . The second approach , known as `` slicing '' , takes a rather different perspective . It leverages two key ideas : the OT closed-form expression for two distributions in one-dimensional space , and the transformation of a distribution into a set of projected one-dimensional distributions by the Radon transform ( RT ) ( Helgason , 2010 ) . The popular proposal along this direction is Sliced-Wasserstein ( SW ) distance ( Bonneel et al. , 2015 ) , which samples the projecting directions uniformly over a unit sphere in the data ambient space and takes the expectation of the resulting one-dimensional OT distance . The SW distance hence requires a significantly lower computation cost than the original Wasserstein distance and is more scalable than the first approach . Due to its solid statistical guarantees and efficient computation , the SW distance has been successfully applied to a variety of practical tasks ( Deshpande et al. , 2018 ; Liutkus et al. , 2019 ; Kolouri et al. , 2018 ; Wu et al. , 2019 ; Deshpande et al. , 2019 ) where it has been shown to have comparative performances to other distances and divergences between probability distributions . However , there is an inevitable bottleneck of computing the SW distance . Specifically , the expectation with respect to the uniform distribution of projections in SW is intractable to compute ; therefore , the Monte Carlo method is employed to approximate it . Nevertheless , drawing from a uniform distribution of directions in high-dimension can result in an overwhelming number of irrelevant directions , especially when the actual data lies in a low-dimensional manifold . Hence , SW typically needs to have a large number of samples to yield an accurate estimation of the discrepancy . Alternatively , in the other extreme , Max Sliced-Wasserstein ( Max-SW ) distance ( Deshpande et al. , 2019 ) uses only one important direction to distinguish the probability distributions . However , other potentially relevant directions are ignored in Max-SW . Therefore , Max-SW can miss some important differences between the two distributions in high dimension . We note that the linear projections in the Radon transform can be replaced by non-linear projections resulting in the generalized sliced-Wasserstein distance and its variants ( Beylkin , 1984 ; Kolouri et al. , 2019 ) . Apart from these main directions , there are also few proposals that try either to modify them or to combine the advantages of the above-mentioned approaches . In particular , Paty & Cuturi ( 2019 ) extended the idea of the max-sliced distance to the max-subspace distance by considering finding an optimal orthogonal subspace . However , this approach is computationally expensive , since it could not exploit the closed-form of the one-dimensional Wasserstein distance . Another approach named the Projected Wasserstein distance ( PWD ) , which was proposed in ( Rowland et al. , 2019 ) , uses sliced decomposition to find multiple one-dimension optimal transport maps . Then , it computes the average cost of those maps equally in the original dimension . Our contributions . Our paper also follows the slicing approach . However , we address key friction in this general line of work : how to obtain a relatively small number of slices simultaneously to maintain the computational efficiency , but at the same time , cover the major differences between two high-dimensional distributions . We take a probabilistic view of slicing by using a probability measure on the unit sphere to represent how important each direction is . From this viewpoint , SW uses the uniform distribution while Max-SW searches for the best delta-Dirac distribution over the projections , both can be considered as special cases . In this paper , we propose to search for an optimal distribution of important directions . We regularize this distribution such that it prefers directions that are far away from one another , hence encouraging an efficient exploration of the space of directions . In the case of no regularization , our proposed method recovers max- ( generalized ) SW as a special case . In summary , our main contributions are two-fold : 1 . First , we introduce a novel distance , named Distributional Sliced-Wasserstein distance ( DSW ) , to account for the issues of previous sliced distances . Our main idea is to search for not just a single most important projection , but an optimal distribution over projections that could balance between an expansion of the area around important projections and the informativeness of projections themselves , i.e. , how well they can distinguish the two target probability measures . We show that DSW is a proper metric in the probability space and possesses appealing statistical and computational properties as the previous sliced distances . 2 . Second , we apply the DSW distance to generative modeling tasks based on the generative adversarial framework . The extensive experiments on real and large-scale datasets show that DSW distance significantly outperforms the SW and Max-SW distances under similar computational time on these tasks . Furthermore , the DSW distance helps model distribution converge to the data distribution faster and provides more realistic generated images than the SW and Max-SW distances . Organization . The remainder of the paper is organized as follows . In Section 2 , we provide backgrounds for Wasserstein distance and its slice-based versions . In Section 3 , we propose distributional ( generalized ) sliced-Wasserstein distance and analyze some of its theoretical properties . Section 4 includes extensive experiment results followed by discussions in Section 5 . Finally , we defer the proofs of key results and extra materials in the Appendices . Notation . For any θ , θ′ ∈ Rd , cos ( θ , θ′ ) = θ > θ′ ‖θ‖‖θ′‖ , where ‖.‖ is ` 2 norm . For any d ≥ 2 , S d−1 denotes the unit sphere in d dimension in ` 2 norm . Furthermore , δ denotes the Dirac delta function , and 〈· , ·〉 is the Euclidean inner-product . For any p ≥ 1 , Lp ( Rd ) is the set of real-valued functions on Rd with finite p-th moment . 2 BACKGROUND . In this section , we provide necessary backgrounds for the ( generalized ) Radon transform , the Wasserstein , and sliced-Wasserstein distances . 2.1 WASSERSTEIN DISTANCE . We start with a formal definition of Wasserstein distance . For any p ≥ 1 , we define Pp ( Rd ) as the set of Borel probability measures with finite p-th moment defined on a given metric space ( Rd , ‖.‖ ) . For any probability measures µ , ν defined on X , Y ⊆ Rd , we denote their corresponding probability density functions as Iµ and Iν . The Wasserstein distance of order p between µ and ν is given by ( Villani , 2008 ; Peyré & Cuturi , 2019 ) : Wp ( µ , ν ) : = ( inf π∈Π ( µ , ν ) ∫ X×Y ‖x− y‖pdπ ( x , y ) ) 1 p , where Π ( µ , ν ) is a set of all transportation plans π such that the marginal distributions of π are µ and ν , respectively . In order to simplify the presentation , we abuse the notation by using both Wp ( µ , ν ) and Wp ( Iµ , Iν ) interchangeably for the Wasserstein distance between µ and ν . When µ and ν are one-dimension measures , the Wasserstein distance between µ and ν has a closed-form expression Wp ( µ , ν ) = ( ∫ 1 0 |F−1µ ( z ) − F−1ν ( z ) |pdz ) 1/p where Fµ and Fν are the cumulative distribution function ( CDF ) of Iµ and Iν , respectively . 2.2 ( GENERALIZED ) RADON TRANSFORMS . Now , we review ( generalized ) Radon transform maps , which are key to the notion of ( generalized ) sliced-Wasserstein distance and its variants . The Radon transform ( RT ) maps a function I ∈ L1 ( Rd ) to the space of functions defined over space of lines in Rd . In particular , for any t ∈ R and direction θ ∈ Sd−1 , the RT is defined as follows ( Helgason , 2010 ) : RI ( t , θ ) : = ∫ Rd I ( x ) δ ( t− 〈x , θ〉 ) dx . The generalized Radon transform ( GRT ) ( Beylkin , 1984 ) extends the original one from integration over hyperplanes of Rd to integration over hypersurfaces . In particular , it is defined as : GI ( t , θ ) : =∫ Rd I ( x ) δ ( t − g ( x , θ ) ) dx , where t ∈ R and θ ∈ Ωθ . Here , Ωθ is a compact subset of R d and g : Rd × Sd−1 7→ R is a defining function ( cf . Assumptions H1-H4 in ( Kolouri et al. , 2019 ) for the definition of defining function ) inducing the hypersurfaces . When g ( x , θ ) = 〈x , θ〉 and Ωθ = Sd−1 , the generalized Radon transform becomes the standard Radon transform .

Summary: The paper presents a novel variant of the Sliced Wasserstein (SW) distance. Wasserstein distances have been used recently in lot of machine learning problems. One of the major problem is that, in its primal form, it is computationally expensive. In order to alleviate this problem, a class of methods, called sliced, leverage on the fact that Wasserstein has a closed form expression in 1D (which amount to sort the samples). It replaces the original Wasserstein distance by an expectation of 1D sub-problems over directions drawn uniformly on the unit hypersphere (akin to a Radon transform). Observing that not all the directions are meaningful, the authors propose a variant of SW where the expectation over all the directions is replaced by an expectation over a distribution of directions. The ‘extent’ of this distribution is controlled by an extra parameter. Interestingly, the authors show that this formulation is computationally tractable if one parametrizes this distribution by a measurable function, expressed as a neural network. This result is obtained by deriving the Lagrangian dual of the original problem. Comparisons with previous works are then given in two GAN scenarii: one on MNIST to explore the importance of the different parameters, and another on larger and more complicated datasets, where the FID score is exposed.