Fast Approximation of the Sliced-Wasserstein Distance Using Concentration of Random Projections

1 Introduction . Recent years have witnessed the emergence of numerical methods inspired by optimal transport ( OT ) to solve machine learning problems . In particular , Wasserstein distances are a core ingredient of OT and define metrics between probability measures . Despite their nice theoretical properties , they are in general computationally expensive in large-scale settings . Several workarounds that scale better to large problems have been developed , and include the Sliced-Wasserstein distance ( SW , [ 1 , 2 ] ) . The SW metric is a computationally cheaper alternative to Wasserstein as it exploits the analytical form of the Wasserstein distance between univariate distributions . More precisely , consider two random variables X and Y in Rd with respective distributions µ and ν , and denote by θ ? ] µ , θ ? ] ν the univariate distributions of the projections of X , Y along θ ∈ Rd . SW then compares µ and ν by computing E [ W ( θ ? ] µ , θ ? ] ν ) ] , where the expectation E is taken with respect to θ uniformly distributed on the unit sphere , and W is the Wasserstein distance . In practice , the expectation is typically estimated by Monte Carlo : one uniformly draws L projection directions { θl } Ll=1 , and approximates SW with L−1 ∑L i=1 W ( ( θl ) ? ] µ , ( θl ) ? ] ν ) . Since the Wasserstein ∗Corresponding author : kimia.nadjahi @ telecom-paris.fr 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . distance between univariate distributions can easily be computed in closed-form , this scheme leads to significant computational benefits as compared to the Wasserstein distance , provided that L is not chosen too large . SW has then been successfully applied in several practical tasks , such as classification [ 3 , 4 ] , Bayesian inference [ 5 ] , the computation of barycenters of measures [ 1 , 6 ] , and implicit generative modeling [ 7–12 ] . Besides , SW has been shown to offer nice theoretical properties as well . Indeed , it satisfies the metric axioms [ 13 ] , the estimators obtained by minimizing SW are asymptotically consistent [ 14 ] , the convergence in SW is equivalent to the convergence in Wasserstein [ 14 , 15 ] , and even though the sample complexity of Wasserstein grows exponentially with the data dimension [ 16–18 ] , the sample complexity of SW does not depend on the dimension [ 19 ] . However , the latter study also demonstrated with a theoretical error bound , that the quality of the Monte Carlo estimate of SW depends on the number of projections and the variance of the one-dimensional Wasserstein distances [ 19 , Theorem 6 ] . In other words , to ensure that the induced approximation error is reasonably small , one might need to choose a large value for L , which inevitably increases the computational complexity of SW . Alternative approaches have been proposed to overcome this issue , and mainly consist in picking more “ informative ” projection directions : e.g. , SW based on orthogonal projections [ 10 , 20 ] , maximum SW [ 21 ] , generalized SW distances [ 22 ] and distributional SW distances [ 23 ] . In this paper , we adopt a different perspective and leverage concentration results on random projections to approximate SW : previous work showed that , under relatively mild conditions , the typical distribution of low-dimensional projections of high-dimensional random variables is close to some Gaussian law [ 24 , 25 ] . Recently , this phenomenon has been illustrated with a bound in terms of the Wasserstein distance [ 26 ] : let { Xi } di=1 be a sequence of real random variables with distribution µd , such that X1 , . . . , Xd are independent with finite fourth-order moments ; then , E [ W ( θ ? ] µd , Nµd ) 2 ] goes to zero as d increases , where Nµd is a univariate Gaussian distribution whose variance depends on µd and the expectation is taken with respect to a Gaussian variable θ . This result has very recently been used to bound the “ maximum-sliced distance ” between any probability measure and its Gaussian approximation [ 27 ] . In our work , we use it to design a novel technique that estimates SW with a simple deterministic formula . As opposed to Monte Carlo , our method does not depend on a finite set of random projections , therefore it eliminates the need of tuning the hyperparameter L and can lead to a significant computational time reduction . Besides , our proposal is quite different from the aforementioned variants of SW which consist in selecting “ informative ” projection directions : these alternatives are defined as optimization problems whose resolution is challenging ( e.g. , [ 23 , Section 3.2 ] ) and are then computed by finding an approximate solution . This incurs an additional computational cost and estimation error , while our method directly approximates SW ( thus , does not define an alternative distance ) via simple deterministic operations , does not rely on any hyperparameters , and comes with theoretical guarantees on its induced error . The important steps to formulate our approximate SW are summarized as follows . We first define an alternative SW whose projection directions are drawn from the same Gaussian distribution as in [ 26 ] , instead of uniformly on the unit sphere , and establish its relation with the original SW. By combining this property with [ 26 , Theorem 1 ] , we bound the absolute difference between SW applied to any two probability measures µd , νd on Rd , and the Wasserstein distance between the univariate Gaussians Nµd , Nνd . Then , we explain why the mean parameters of µd and νd should be zero for the approximation error to decrease as d grows . Nevertheless , we show that it is not a limiting factor , by exploiting the following decomposition of SW : SW between µd , νd can be equivalently written as the sum of the difference between their means and the SW between the centered versions of µd , νd . Our approach then consists in estimating SW between the centered versions with the Wasserstein term between Gaussian approximations to meet the zero-means condition , and recover SW between the original measures via the aforementioned property . Since the Wasserstein distance between Gaussian distributions admits a closed-form solution , our approximate SW is very easy to compute , and faster than the Monte Carlo estimate obtained with a large number of projections . We derive nonasymptotical guarantees on the error induced by our approach . Specifically , we define a weak dependence condition under which the error is shown to go to zero with increasing d. Our theoretical results are then validated with experiments conducted on synthetic data . Finally , we leverage our theoretical insights to design a novel adversarial framework for a typical generative modeling problem in machine learning , and illustrate its advantages in terms of accuracy and computational time , over generative models based on the Monte Carlo estimate of SW. Our empirical results can be reproduced with our open source code2 . 2 Background . We first give some background on optimal transport distances and concentration of measure for random projections . All random variables are defined on a probability space ( Ω , F , P ) with associated expectation operator E. We denote N∗ = N \ { 0 } , and for d ∈ N∗ , P ( Rd ) is the set of probability measures on Rd . 2.1 Optimal transport distances . Let p ∈ [ 1 , +∞ ) and Pp ( Rd ) = { µ ∈ P ( Rd ) : ∫ Rd ‖x‖ p dµ ( x ) < +∞ } be the set of probability measures on Rd with finite moment of order p. The Wasserstein distance of order p between any µ , ν ∈ Pp ( Rd ) is defined as Wpp ( µ , ν ) = inf π∈Π ( µ , ν ) ∫ Rd×Rd ‖x− y‖p dπ ( x , y ) , ( 1 ) where ‖ · ‖ denotes the Euclidean norm , and Π ( µ , ν ) the set of probability measures on Rd × Rd whose marginals with respect to the first and second variables are given by µ and ν respectively . In some particular settings , Wp is relatively easy to compute since the optimization problem in ( 1 ) admits a closed-form solution : we give two examples that will be useful in the rest of the paper . Gaussian distributions . Denote by N ( m , Σ ) the Gaussian distribution on Rd with mean m ∈ Rd and covariance matrix Σ ∈ Rd×d symmetric positive-definite . The Wasserstein distance between two Gaussian distributions , also known as the Wasserstein-Bures metric , is given by [ 28 ] W22 { N ( m1 , Σ1 ) , N ( m2 , Σ2 ) } = ‖m1 −m2‖2 + Tr [ Σ1 + Σ2 − 2 ( Σ 1/2 1 Σ2Σ 1/2 1 ) 1/2 ] , ( 2 ) where Tr is the trace operator . Univariate distributions . Consider µ , ν ∈ Pp ( R ) , and denote by F−1µ and F−1ν the quantile functions of µ and ν respectively . By [ 29 , Theorem 3.1.2 . ( a ) ] , Wpp ( µ , ν ) = ∫ 1 0 ∣∣F−1µ ( t ) − F−1ν ( t ) ∣∣p dt . ( 3 ) If µ = n−1 ∑n i=1 δxi and ν = n −1∑n i=1 δyi , with { xi } ni=1 , { yi } ni=1 ⊂ Rn and δz the Dirac distribution with mass on z , ( 3 ) can simply be calculated by sorting { xi } ni=1 and { yi } ni=1 as x ( 1 ) ≤ . . . ≤ x ( n ) and y ( 1 ) ≤ . . . ≤ y ( n ) . Indeed , in this case , Wpp ( µ , ν ) = n−1 ∑n i=1 |x ( i ) − y ( i ) |p . However , when the empirical distributions are multivariate , Wp ( µ , ν ) is not analytically available in general , so its computation is expensive : the standard methods used to solve the linear program in ( 1 ) have a worst-case computational complexity in O ( n3 log n ) , and tend to have a super-cubic cost in practice [ 30 , Chapter 3 ] . The Sliced-Wasserstein distance [ 1 , 2 ] defines a practical alternative metric by leveraging the computational efficiency of Wp for univariate distributions . Let Sd−1 be the d-dimensional unit sphere and σ the uniform distribution on Sd−1 . For θ ∈ Sd−1 , θ ? : Rd → R denotes the linear form x 7→ 〈θ , x〉 with 〈· , ·〉 the Euclidean inner-product . Then , SW of order p ∈ [ 1 , ∞ ) between µ , ν ∈ Pp ( Rd ) is SWpp ( µ , ν ) = ∫ Sd−1 Wpp ( θ ? ] µ , θ ? ] ν ) dσ ( θ ) , ( 4 ) where for any measurable function f : Rd → R and ξ ∈ P ( Rd ) , f ] ξ is the push-forward measure of ξ by f : for any Borel set A in R , f ] ζ ( A ) = ζ ( f−1 ( A ) ) , with f−1 ( A ) = { x ∈ Rd : f ( x ) ∈ A } . Since θ ? ] µ , θ ? ] ν are univariate distributions , the Wasserstein distances in ( 4 ) are conveniently computed using ( 3 ) . Besides , in practical applications , the expected value in ( 4 ) is typically approximated with a standard Monte Carlo method : 2See https : //github.com/kimiandj/fast_sw SWpp , L ( µ , ν ) = ( 1/L ) L∑ l=1 Wpp ( ( θl ) ? ] µ , ( θl ) ? ] ν ) , with { θl } Ll=1 i.i.d . from σ . ( 5 ) Computing SWp , L between two empirical distributions then amounts to projecting sets of n observations in Rd along L directions , and sorting the projected data . The resulting computational complexity is O ( Ldn + Ln log n ) , which is more efficient than Wp in general . This complexity means that the Monte Carlo estimate is more expensive when d , n and L increase , and it is often unclear how L should be chosen in order to control the approximation error ; see [ 19 , Theorem 6 ] .

Summary: This paper develops a novel deterministic approximation for the Sliced-Wasserstein Distance (SWD), based on the fact that a random linear projection of a high-dimensional vector has an approximately Gaussian distribution. In many cases, the approximation involves no tuning parameters and can be computed much more quickly than the usual Monte Carlo estimate of SWD, especially for very high-dimensional data. If the $d$ dimensions of the data are not too strongly correlated, then the approximation converges no slower than $d^{-1/8}$ as $d \to \infty$. Experiments with synthetic data validate this theoretical result, as well as the computational benefits of the proposed approximation. Finally, the paper demonstrates an application to image generation.