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

Abstract:
The Sliced-Wasserstein distance (SW) is being increasingly used in machine learning applications as an alternative to the Wasserstein distance and offers significant computational and statistical benefits. Since it is defined as an expectation over random projections, SW is commonly approximated by Monte Carlo. We adopt a new perspective to approximate SW by making use of the concentration of measure phenomenon: under mild assumptions, one-dimensional projections of a highdimensional random vector are approximately Gaussian. Based on this observation, we develop a simple deterministic approximation for SW. Our method does not require sampling a number of random projections, and is therefore both accurate and easy to use compared to the usual Monte Carlo approximation. We derive nonasymptotical guarantees for our approach, and show that the approximation error goes to zero as the dimension increases, under a weak dependence condition on the data distribution. We validate our theoretical findings on synthetic datasets, and illustrate the proposed approximation on a generative modeling problem.

heading: 1 Introduction
text: 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.

heading: 2 Background
text: 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.

heading: 2.1 Optimal transport distances
text: 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].

heading: 2.2 Central limit theorems for random projections
text: There is a rich literature on the typical behavior of one-dimensional random projections of highdimensional vectors. To be more specific, let (θi)i∈N∗ be i.i.d. standard one-dimensional Gaussian random variables and (Xi)i∈N∗ be a sequence of one-dimensional random variables. Denote for any d ∈ N∗, θ1:d = {θi}di=1 and X1:d = {Xi}di=1. Several central limits theorems ensure that, under relatively mild conditions, the sequence of distributions of d−1/2 〈θ1:d, X1:d〉 ∈ R given θ1:d ∈ Rd converges in distribution to a Gaussian random variable in probability. This line of work goes back to [24, 25], whose contributions have then been sharpened and generalized in [31–38]. In particular, a recent study [26] gives a quantitative version of this phenomenon. More precisely, denote for any d ∈ N∗ by µXd , the distribution of X1:d (i.e., the joint distribution of X1, X2, . . . , Xd) and γd the zero-mean Gaussian distribution with covariance matrix (1/d)Id. Assume that for any d ∈ N∗, µXd ∈ P2(Rd). Then, [26, Theorem 1] shows that there exists a universal constant C ≥ 0 such that∫
Rd W22
( θ?]µ X d ,N ( 0, d−1m2(µ X d ) )) dγd(θ) ≤ CΞd(µXd ) , with (6)
Ξd(µ X d ) = d −1{α(µXd ) + ( m2(µ X d )β1(µ X d ) )1/2 + m2(µ X d ) 1/5β2(µ X d ) 4/5} , (7)
m2(µ X d ) = E [ ‖X1:d‖2 ] , α(µXd ) = E [∣∣∣‖X1:d‖2 − m2(µXd )∣∣∣] , βq(µXd ) = E 1q [|〈X1:d, X ′1:d〉|q] , (8)
where q ∈ {1, 2} and (X ′i)i∈N∗ is an independent copy of (Xi)i∈N∗ . A formal statement of this result is also given for completeness in the supplement.

heading: 3 Approximate Sliced-Wasserstein distance based on concentration of random projections
text: We develop a novel method to approximate the Sliced-Wasserstein distance of order 2, by extending the bound in (6) and deriving novel properties for SW. We then derive nonasymptotical guarantees of the corresponding approximation error, which ensure that our estimate is accurate for highdimensional data under a weak dependence condition.

heading: 3.1 Sliced-Wasserstein distance with Gaussian projections
text: First, to enable the use of (6) for the analysis of SW, we introduce a variant of SWp (4) whose projections are drawn from the Gaussian distribution considered in (6), instead of uniformly on the sphere. The Sliced-Wasserstein distance of order p ∈ [1,+∞) based on Gaussian projections is defined for any µ, ν ∈ Pp(Rd) as
S̃W p
p(µ, ν) = ∫ Rd Wpp(θ ? ]µ, θ ? ] ν)dγd(θ) . (9)
In the next proposition, we establish a simple mathematical relation between traditional SW and the newly introduced one: we prove that S̃Wp is equal to SWp up to a proportionality constant that only depends on the data dimension d and the order p.
Proposition 1. Let p ∈ [1,+∞). Then, S̃Wp (9) is related to SWp (4) as follows: for any µ, ν ∈ Pp(Rd), S̃Wp(µ, ν) = (2/d)1/2 { Γ(d/2 + p/2) / Γ(d/2) }1/p SWp(µ, ν), where Γ is the Gamma function.
Since (6) only applies to the Wasserstein distance of order 2, we will focus on SW of that same order in the rest of the paper. In this case, SW with Gaussian projections is equal to the original SW. Indeed, we can show that the constant (2/d)1/2 {Γ(d/2 + p/2) / Γ(d/2)}1/p defined in Proposition 1 is equal to 1 when p = 2, by using the property Γ(d/2 + 1) = (d/2)Γ(d/2).

heading: 3.2 Approximate Sliced-Wasserstein distance
text: Our next result is an easy consequence of (6) and Proposition 1, and shows that the absolute difference between SW2(µd, νd) and W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} for any µd, νd ∈ P2(Rd), is bounded from above by Ξd(µd) + Ξd(νd) (7).
Theorem 1. There exists a universal constant C > 0 such that for any µd, νd ∈ P2(Rd),∣∣SW2(µd, νd)−W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))}∣∣ ≤ C(Ξd(µd) + Ξd(νd))1/2 , (10) where, for ξd ∈ {µd, νd}, Ξd(ξd) and m2(ξd) are defined in (7) and (8) respectively.
Since W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} has a closed-form solution by (2), it provides a computationally efficient approximation of SW2(µd, νd) whose accuracy is quantified by Theorem 1. Next, we identify settings where this approximation is accurate, by analyzing the error Ξd(µd) + Ξd(νd).
Our first observation is that µd and νd should have zero means for the error to go to zero as d→ +∞, and we develop a novel approximation of SW that takes into account this constraint. Going back to the definition of Ξd(µXd ) in (7), setting X̄i = Xi − E[Xi] and X̄ ′i = X ′i − E[X ′i], we get
m2(µ X d ) = E[‖X̄1:d‖2] + ‖E[X1:d]‖ 2 (11)
β22(µ X d ) = E [〈 X̄1:d, X̄ ′ 1:d 〉2] + 4E [〈 E[X1:d], X̄1:d 〉2] + ‖E[X1:d]‖4 . (12)
By Equations (11) and (12), since in practice the norm of the mean E[X1:d] is expected to increase linearly with d1/2 at least, so are m2(µXd ) and β2(µ X d ) as functions of d. As a consequence, Ξd(µ X d ) cannot be shown to converge to 0 as d → ∞ in this setting, but only to be bounded. However, if the data are centered, the norm of the mean is zero, thus Ξd(µXd ) might be decreasing. Therefore, we derive a convenient formula to compute SW2(µd, νd) from SW2(µ̄d, ν̄d) where for any ξd ∈ P2(Rd), ξ̄d is the centered version of ξ, i.e. the pushforward measure of ξd by x 7→ x−mξd with mξd = ∫ Rd y dξd(y). This result is the last ingredient to formulate our approximation of SW.
Proposition 2. Let µd, νd ∈ P2(Rd) with respective means mµd ,mνd . Then, the Sliced-Wasserstein distance of order 2 can be decomposed as
SW22(µd, νd) = SW 2 2(µ̄d, ν̄d) + (1/d)‖mµd −mνd‖2 . (13)
Based on Proposition 2, instead of estimating SW2(µd, νd) with W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} directly, we propose approximating SW2(µ̄d, ν̄d) with W2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))} and then using (13). This strategy yields our final approximation of SW, which is defined for any µd, νd ∈ P2(Rd) as
ŜW 2 2(µd, νd) = W 2 2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))}+ (1/d)‖mµd −mνd‖2 , (14)
where for ξd ∈ {µ̄d, ν̄d}, m2(ξd) is defined in (8). Note that (14) can be simplified since by (2), W22{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))} = d−1(m2(µ̄d)1/2 − m2(ν̄d)1/2)2. Besides, if µd and νd are both supported on a finite set of points, ŜW2(µd, νd) has a closed-form expression: given ξd = n−1 ∑n j=1 δx(j) ∈ P2(Rd) with x(j) ∈ Rd for j ∈ {1, . . . , n}, we then have mξd = n−1 ∑n j=1 x
(j), and m2(ξd) = n−1 ∑n j=1 ‖x(j)‖2. The associated computational complexity is therefore in O(dn).
Hence, we introduced an alternative technique to estimate SW which does not rely on a finite set of random projections, as opposed to the commonly used Monte Carlo technique (5). Our approach thus eliminates the need for practitioners to tune the number of projections L, but also to sort the projected data. As a consequence, it is more efficient to compute ŜW2(µd, νd) than SW2,L(µd, νd) for large L. We illustrate this latter point with empirical results in Section 4.

heading: 3.3 Error analysis under weak dependence
text: We have discussed why centering the data is important to ensure that the approximation error goes to zero with increasing d. Next, we introduce a weak dependence condition under which the error is guaranteed to decrease as d increases.
We first consider a setting mentioned in [26] where µd = µ(1)⊗· · ·⊗µ(d) and νd = ν(1)⊗· · ·⊗ν(d), ⊗ denoting the tensor product of measures, and µ(j), ν(j) ∈ P4(R) for j ∈ {1, . . . , d}. We prove in this case that W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} converges to SW2(µd, νd) at a rate of d−1/8. This result is reported in the supplementary document, and can be interpreted as an extension of [26, Corollary 3] for SW.
We emphasize that the assumptions of this first setting severely restrict the scope of application of our approximation method: in several statistical and machine learning tasks, the random variables of interest {Xi}di=1 are not independent from each other (e.g. for image data, each Xi typically represents the value of a pixel at a certain position, thus depends on the neighboring pixels). Therefore, we relax this independence condition by considering a concept of ‘weak dependence’ inspired by [39] and properly defined in Definition 1. Definition 1. Let (Xj)j∈N∗ be a stationary sequence of one-dimensional random variables with mean zero, i.e. Xi and Xj have the same distribution for any i, j ∈ N∗ and E[X1] = 0. We say that (Xj)j∈N∗ is fourth-order weakly dependent if there exist some constant K ≥ 0 and a nonincreasing sequence of real coefficients {ρ(n)}n∈N such that, for any i, j ∈ N∗, i ≤ j,
|Cov(X2i , X2j )| ≤ Kρ(j − i) , |Cov(Xi, Xj)| ≤ Kρ(j − i) . (15) In addition, the sequence {ρ(n)}n∈N satisfies ∑+∞ n=0 ρ(n) ≤ ρ∞ < +∞.
Intuitively, in practical applications, the weak dependence condition would essentially require the components of the observations not to exhibit strong correlations; yet, they are allowed to depend on each other. Furthermore, since our weak dependence condition is weaker than the one introduced in [39, Theorem 1], it is satisfied by the various examples of models described in [39, Section 5]. We present some of them below, to illustrate Definition 1 more clearly. 1) Gaussian processes and associated processes [40, Section 3.1], provided that they are stationary. 2) Bernoulli shifts: Xt = H(εt, . . . , εt−r) for t ∈ N∗, where H : Rr+1 → R is a measurable
function and (εi)i∈N∗ is a sequence of i.i.d. real random variables. A simple example of such process is given by moving-average models.
3) Autoregressive models, defined as Xt = f(Xt−1, . . . , Xt−r) + εt for t ∈ N∗, where (εi)i∈N∗ a sequence of i.i.d. real random variables with E |ε1| <∞, and |f(u1, . . . , ur)− f(v1, . . . , vr)| ≤∑r i=1 ai |ui − vi| for some a1, . . . , ar ≥ 0 such that (∑r i=1 ai )1/r < 1.
We then consider a sequence of fourth-order weakly dependent random variables (Xj)j∈N∗ , and prove that Ξd(µXd ) goes to zero as d → ∞, with a rate of convergence depending on {ρ(n)}n∈N. This result is given in the supplementary document, and helps us refine Theorem 1 under this weak dependence condition: the next corollary establishes that the error approaches 0 at a rate of d−1/8. Corollary 1. Let (Xj)j∈N∗ and (Yj)j∈N∗ be sequences of random variables which are fourth-order weakly dependent. Set for any d ∈ N∗, X1:d = {Xj}dj=1 and Y1:d = {Yj}dj=1, and denote by µd, νd the distributions of X1:d, Y1:d respectively. Then, there exists a universal constant C > 0 such that∣∣SW2(µd, νd)−W2(N(0, d−1m2(µd)),N(0, d−1m2(νd)))∣∣ ≤ Cd−1/8. Hence, by replacing the independence condition of the first setting with weak dependence, we broaden the scope of application whilst guaranteeing that the approximation error goes to zero as d increases. We finally note that in these two settings, the data are required to have zero mean, which is automatically verified with our approximation method since we estimate SW between the centered distributions (see eq. (14)).

heading: 4 Experiments
text: Synthetic experiments. The goal of these experiments is to illustrate our theoretical results derived in Section 3. In each setting, we generate two sets of d-dimensional samples, denoted by {x(j)}nj=1
and {y(j)}nj=1 with n = 104 and x(j), y(j) ∈ Rd for j ∈ {1, . . . , n}. We then approximate SW between their empirical distributions in P2(Rd), given by µd = n−1 ∑n j=1 δx(j) and νd =
n−1 ∑n j=1 δy(j) .
First, we analyze the consequences of centering data. Here, {x(j)}nj=1, {y(j)}nj=1 are n independent samples from Gaussian or Gamma distributions: see the supplementary document for more details. We compute |W2{N(0, d−1m2(µd)),N(0, d−1m2(νd))} − SW2(µd, νd)| on the one hand, and |W2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))} − SW2(µ̄d, ν̄d)| on the other hand. In the Gaussian case, the exact value of SW2 is known (we report it in the supplementary document), while for the Gamma distributions, it is approximated with Monte Carlo based on 2 × 104 random projections. Figures 1(a) and 1(b) show that the error goes to zero as d increases if the data are centered. This confirms our analysis provided in Section 3.2 about the influence of the mean, and in Section 3.3 on sequences of independent random variables.
Next, we consider autoregressive processes of order one (AR(1)). An AR(1) process is defined as X1 = ε1 and, for t ∈ N∗, Xt = αXt−1 + εt, where α ∈ [0, 1] and (εi)i∈N∗ is an i.i.d. sequence of real random variables with E[ε1] = 0 and finite second-order moment. If α < 1, the process has a stationary distribution and (Xj)j∈N∗ satisfies the weak dependence condition in its stationary regime [41]. In practice, we generate a sample by using this recursion formula for 104 + d steps, and keeping the last d samples. The discarded samples correspond to a “burn-in” phase which helps reaching the stationary solution of the process. We generate {x(j)}nj=1 and {y(j)}nj=1 using the same distribution for the noise (either a Gaussian or Student’s t-distribution, as described in the supplementary document). This means that both datasets come from the same distribution, thus the exact value of SW2 is zero. We plot on Figures 1(c) and 1(d) the approximation error according to d ∈ [10, 103] for different values of α. The error converges to zero with increasing d, which is consistent with Corollary 1.
Note that Figure 1 exhibits rate of convergence that are better than the one in d−1/8 derived in Section 3.3: in Figure 1(b), the slope is approximately −0.45 (Gaussian) and −0.7 (Gamma), and in Figures 1(c) and 1(d), it is on average −0.35. This suggests that our theoretical bounds might be improved, and we further investigate this aspect for the Gaussian case: we consider the case where {x(j)}nj=1, {y(j)}nj=1 are n independent samples from Gaussian distributions with diagonal
covariance matrices, and we prove that E|W2{N(0, d−1m2(µ̄d)),N(0, d−1m2(ν̄d))}−SW2(µ̄d, ν̄d)| goes to 0 as dn→ +∞ with a convergence rate in d−1/2n−1/2. We provide the complete statement and formal proof in Section S3.1 (Proposition S3). This result is consistent with Figure 1(b), and is a first encouraging step towards the following research direction: we will study if our proofs and the ones in [26] can be refined when assuming additional structure on the distributions (e.g., sub-Gaussian and sub-exponential), in order to identify the settings under which our current bounds are tight or can be improved.
Finally, we compare our approximation scheme against the standard Monte Carlo estimation, in terms of accuracy and computation time. We use the same setting as in Figure 1(b), where the n samples are independently drawn from Gamma distributions. We compute ŜW2(µd, νd) (14) and SW2,L(µd, νd) (5) with L ∈ {100, 1000, 5000}, and we compare each approximation with SW2,2×104(µd, νd), which we consider as the exact value of SW. Figure 2 reports the approximation error and computation time of each scheme for d ∈ [10, 104], and shows that our method is more accurate and faster than Monte Carlo. In particular, when d = 103, the average computation time of our technique is 0.02s, while the second best approximation (Monte Carlo with L = 5000) takes more than 8s. Besides, we observe that Monte Carlo is very sensitive to the hyperparameters, since it loses accuracy when L decreases and gets slower as L and d increase. This observation is consistent with the computational complexity of SW2,L recalled in Section 2.1. On the other hand, our approximation scheme is extremely efficient even for large d and n, since it is based on a simple deterministic formula which does not require projecting and sorting data along random directions.
Image generation. Finally, we leverage our theoretical insights to design a novel method for a typical generative modeling application. The problem consists in tuning a neural network that takes as input k-dimensional samples from a reference distribution (e.g., uniform or Gaussian), to generate images of dimension d > k. During the training phase, the parameters of the network are updated by iteratively minimizing a dissimilarity measure between the dataset to fit and the generated images.
In [9], the dissimilarity measure is Monte Carlo SW approximated with 104 random projections, and the resulting generative model is called the Sliced-Wasserstein generator (SWG). This model performs well on moderately high-dimensional image datasets (e.g., 28 × 28 for MNIST images [42]). However, for very large dimensions (e.g., 64 × 64 × 3 for the CelebA dataset [43]), Monte Carlo SW requires more than 104 random projections to capture relevant information, which leads to very expensive training iterations and potential memory issues. To offer better scalability, SWG can be augmented with a discriminator network [9, Section 3.2] that aims at finding a lower-dimensional space in which the two projected datasets are clearly distinguishable. The intuition behind this heuristic is that the more distinct the two datasets are from each other, the fewer projection directions Monte Carlo SW requires to provide useful information. The training then consists in optimizing the generator’s and discriminator’s objective functions in an alternating fashion.
Our novel approach builds on SWG and modifies the saddle-point problem in [9, Section 3.2]: motivated by the gain in accuracy and time illustrated in Figure 2 on high-dimensional datasets, we propose to replace Monte Carlo SW with our approximate SW (14) in the generator’s objective; then, to make sure that our approximation is accurate, we regularize the discriminator’s objective:
max ψ
L(ψ) + λ1 ∥∥Cov[d′ψ(X)]∥∥2F + λ1∥∥Cov[d′ψ(gφ(Z))]∥∥2F (16)
+ λ2 E [ ‖d′ψ(X)‖−2 ] + λ2 E [ ‖d′ψ(gφ(Z))‖−2 ] (17)
where L is the discriminator’s loss used in SWG, gφ and d′ψ are the generator’s last layer and the discriminator’s penultimate layer respectively (parameterized by φ, ψ), X and Z are the random variables corresponding to the images to fit and the generator’s input, Cov denotes the covariance matrix, ‖ · ‖F the Frobenius norm, and λ1, λ2 ≥ 0. The regularization in (16) enforces the weak dependence condition (Corollary 1), while (17) prevents the network to converge to d′ψ = 0. We call this generative adversarial network regularized deterministic SWG (reg-det-SWG).
To investigate the consequences of (i) regularizing the discriminator, and (ii) replacing the Monte Carlo SW with our approximation, we design another model, called regularized SWG (reg-SWG): similarly to SWG, the generator minimizes SW2,104 , but the discriminator’s objective is regularized as in (16), (17). We then compare reg-det-SWG against SWG and reg-SWG, by training the models on MNIST and CelebA and measuring their respective training time and Fréchet Inception Distances (FID, [44]): see Table 1. We used the same network architectures for all methods, and tuned (λ1, λ2) via cross-validation: more details on the experimental setup are given in Section S4. First, we observe that the regularized models produce images of higher quality, since reg-SWG and reg-det-SWG return lower FID values than SWG. The FID of reg-SWG and reg-det-SWG are close for both datasets, thus the two models seem to yield similar performances. Hence, we report in Figure 3 the images generated by SWG and reg-det-SWG only.
The training process is more expensive when regularizing the discriminator: the average running time per epoch is higher for the regularized models. We also observe that reg-det-SWG is faster than reg-SWG, which is consistent with the fact that our approximation method is faster than Monte Carlo on high-dimensional settings. To further illustrate this point, we reported the average time spent in computing the generative loss per epoch, i.e. SW2,104 for SWG and reg-SWG, and ŜW2 for reg-det-SWG: see column TSW in Table 1. On GPU, reg-det-SWG is at least 15 times faster than SWG and reg-SWG on MNIST, and 6 times faster on CelebA. Note that the models were trained using PyTorch, thus Monte Carlo SW benefits from a GPU-accelerated implementation of the sorting operation (with the function torch.sort). We also reported the computation times when models are trained on CPU. In this case, computing ŜW2 takes at most less than 3s per epoch, whereas the Monte Carlo estimation executes in several minutes (e.g., approximately 45min on CelebA). As a result, the total training time is almost the same for reg-det-SWG and SWG on CelebA, and the lowest for reg-det-SWG on MNIST.

heading: 5 Conclusion
text: In this work, we presented a novel method to approximate the Sliced-Wasserstein distance of order 2, which relies on the concentration-of-measure phenomenon for random projections. The resulting method computes SW with simple deterministic operations, which are computationally efficient even on high-dimensional settings and do not require any hyperparameters. We proved nonasymptotical guarantees showing that, under a weak dependence condition, the approximation error goes to zero as the dimension increases. Our theoretical findings are then illustrated with experiments on synthetic datasets. Motivated by the computational efficiency and accuracy of our approximate SW, we finally designed a novel approach for image generation that leverages our theoretical insights. As compared to generative models based on SW estimated with Monte Carlo, our framework produces images of higher quality with further computational benefits. This encourages the use of our approximate SW on other algorithms that rely on Monte Carlo SW, e.g. autoencoders [8] or normalizing flows [12].
The weak dependence condition can be inappropriate to describe the underlying geometry of real data in ML applications, and in that case, approximating SW with our method seems inadequate. To overcome this problem, we encourage practitioners to resort to models where real data are represented by features that can be made weakly dependent. This strategy has proven successful in our image generation experiment: the reg-det-SWG model uses our approximation to compare two sets of features (instead of the raw images) whose covariance matrices are regularized to enforce weak dependence. Since many ML techniques make use of features and regularizers, we believe that our methodology is not restrictive and can then be applied to other standard problems than image generation. Besides, our weak dependence condition in Definition 1 is weaker than the one in [39], which is a notion commonly used in statistics.
Our empirical results on synthetic data show that the approximation error goes to zero with a faster convergence rate than the one we proved. Then, the main current limitation of our framework is that our theoretical convergence rate in d−1/8 might be slower than necessary. We proved that the overall approximation error is upper-bounded by a term in d−1/2 when comparing Gaussians with diagonal covariance matrices, and the improvement of our error bounds for other specific distributions is left for future work. On the other hand, the extension of our methodology to variants of SW is another challenging future research direction. To the best of our knowledge, the literature on the concentration of measure phenomenon focuses on linear random projections, therefore the derivation of deterministic approximations for SW based on nonlinear projections seems highly nontrivial. A more promising direction would be to generalize our approach to SW based on k-dimensional linear projection by leveraging the bound in [26, Theorem 1] for k > 1.
Since this paper is focused on developing a theoretically-grounded novel method to estimate a distance between probability distributions, we believe it will not pose any negative societal or ethical consequence. On the other hand, as demonstrated in Section 4, our contribution provides tools to speed up existing machine learning algorithms on CPU, which is useful when powerful hardware resources are not available, or when their use is deliberately avoided for environmental purposes.
Acknowledgments and Disclosure of Funding
This work is partly supported by the industrial chair “Data Science & Artificial Intelligence for Digitalized Industry & Services” from Télécom Paris. Umut Şimşekli’s research is supported by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute). Pierre E. Jacob gratefully acknowledges support by the National Science Foundation through grant DMS1844695. Alain Durmus acknowledges support of the Lagrange Mathematical and Computing Research Center. Kimia Nadjahi is grateful to Pierre Colombo for his helpful advice on how to train the generative models in Section 4 on Télécom Paris’s GPUs.

Review summaries:
This paper introduces a (possibly) cheaper approximation of Sliced Wasserstein (SW) distance. The proposed technique is based on calculating 1D Wasserstein distance between Gaussian projection directions. This is motivated by the fact that Wasserstein distance between Gaussian distributions admits a closed-form solution, which leads to an easy compute of the proposed approximate SW and discard a large number of uniformly sampling random projections for the Monte Carlo approximation of the original SW distance. 

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.

The paper presents a method to estimate the Sliced-Wasserstein distance. The method is based on the Gaussian approximation of the projected data, thus the SW distance could be approximated through the SW distance between Gaussians which has a closed-form. Hence, the method is better than the Monte Carlo SW distance in term of computational speed. The authors also demonstrate their method on toy data set and real data sets, i.e. MNIST and Celeb A.

The paper proposes a new approach to approximate the sliced Wasserstein distance. In particular, the paper extends the central limit theorems for random projections to derive a deterministic approximation of sliced Wasserstein (based on the closed form of Wasserstein distance between two Gaussians) that is efficient (low error) in high dimension under centering assumption (weak dependence assumption). On the experimental side, the paper illustrates the derived theory on the synthetic data and compares the proposed approach to the conventional Monte Carlo approximation of the sliced Wasserstein distance.