arXiv:2605.20122v1 [stat.ML] 19 May 2026
Optimizing Computational-Statistical Runtime for Wasserstein Distance Estimation Peter Matthew Jacobs Department of Statistics University of Wisconsin-Madison Madison, WI 53706 [email protected]
Jeff M. Phillips Kahlert School of Computing University of Utah Salt Lake City, UT 84101 [email protected]
Abstract Squared Wasserstein distance is a frequently used tool to measure discrepancy between probability distributions. This distance is typically computed between empirical measures of size n from two underlying random samples. Unfortunately, even in lower dimensional Euclidean space problems (d ∈ {2, 3}), algorithms for Wasserstein distance computation with approximate or exact precision guarantees scale poorly in the runtime as a function of n and the desired precision. In response, we consider the computational-statistical runtime, where the goal is to estimate from samples the Wasserstein distance between potentially smooth measures up to ε-additive error in expectation with respect to the sampling; we allow O(1) computational cost for collecting a sample. Towards this, we develop a SampleSketch-Solve paradigm where we introduce a regular cartesian grid sketch of the samples. We show that (especially under α-Hölder smooth distributions) this can compress the data without increasing asymptotic error, and also regularizes the structure which enables faster exact algorithms. Ultimately, we approximate d+1+o(1) W22 (P, Q) within ε error in ε− max(2, 1+α ) time for 0 < α < 1 Hölder smooth distributions P, Q on (0, 1)d ; an optimal Θ(ε−2 ) for α > 1/2 when d = 2 and nearly optimal as α → 1 when d = 3.
1
Introduction
In data analysis and machine learning tasks such as parameter estimation via Minimum Distance Estimation (see for example Bernton et al. [2019]) or non-parametric two-sample hypothesis testing (see for example Ramdas et al. [2017]), require comparison of underlying unknown probability distributions P, Q ∈ P(Rd ). The squared Wasserstein distance, W22 (P, Q), is one increasingly popular choice for this task. When P, Q are supported on Rd with finite second moments it is defined as W22 (P, Q) := inf E(X,Y )∼γ ∥X − Y ∥22 , γ∈Π(P,Q)
where Π(P, Q) is the set of couplings of P, Q (joint distributions over Rd × Rd with marginals P and Q respectively). A standard approach for estimating W22 (P, Q) is to collect n random samples from P and Q, construct empirical measures P̂n , Q̂n , and use the plug-in estimator W22 (P̂n , Q̂n ). When n is large this is computationally burdensome. The best known computational algorithm can estimate 2 W22 (P̂n , Q̂n ) (even in the non uniformly discrete case) with additive error ε in runtime O(n ∆ε∞ ) [Lahn et al., 2019] where ∆∞ is the maximum squared distance between pairs of points across the two point sets. In this paper we take the position that analysis should start with a distributional assumption on the input P, Q, and that it should aim for computationally efficient algorithms. While there is value in Preprint.
considering algorithms for computation of Wasserstein distance between discrete point sets (e.g., Pn , Qn of size n), we argue that often Pn and Qn should be considered as empirical measures from unknown distributions P, Q. This means that the target accuracy of computing W22 (Pn , Qn ) should aim to match that induced by the sample. Finally, towards connecting this all together, we assume it takes O(1) time to create one additional sample from P or Q. The average loss of a randomized algorithm now takes into consideration both statistical error, and computational cost. A randomized b which takes as input a precision parameter ε, and randomly algorithm in this context is a procedure A returns an output Â(ε) in time O(f (ε)), which on average is ε close to W22 (P, Q). We summarize this in the next definition. Definition 1.1. (Computational-Statistical Runtime (CSR)) Suppose P and Q are two probability distributions on Rd for which we can access random samples at O(1) cost each draw. We say a b | ε > 0} achieves the O(f (ε)) Computational-Statistical Runtime if randomized algorithm {A(ε) b for every ε > 0, A(ε) has runtime O(f (ε)) and achieves b − W 2 (P, Q)| ≤ ε E|A(ε) 2 where the expectation is with respect to all sources of randomness induced by the algorithm. Crucially, to achieve a CSR of O(f (ε)), the algorithm itself must not only choose how to use samples, but also how many of them to gather in the first place. Ultimately, this aligns with the typical challenge faced by the practitioner: choosing an algorithm to estimate the underlying quantity (in this case W22 (P, Q)) up to ε error while minimizing the time to do so, including the cost of collecting samples. Or alternatively, if one already has n samples Pn , Qn , this perspective advocates to treat them as estimates P̂n , Q̂n of some unknown true distribution P, Q, and then the goal is to solve W22 (P̂n , Q̂n ) up to ε error, where ε is determined by E|W22 (P, Q) − W22 (P̂n , Q̂n )| = ε. As error ε scales with the diameters of the supports of P and Q, in order to prove ε-additive error, we need an assumption that ∆∞ = maxx,x′ ∈dom(P )∪dom(Q) ∥x − x′ ∥2 is bounded where dom(P ) denotes the support of the measure P . We henceforth assume the dimension d is constant, and P, Q are defined on [0, 1]d , so ∆∞ = O(1). Alternatively, we could bound error in terms of ε∆∞ . Table 1: ε-approximate W22 on [0, 1]d Related Work. A long line of works focuses purely on Notes Õ(·) Time the computational problem, exactly or ε-approximately, Ref. 2 d Kuhn [1956] uw n3 computing the W2 between discrete measures on [0, 1] and uw,2d n3/2 /ε3/2 of size n; see Table 1. The Hungarian algorithm [Kuhn, Agarwal Phillips [2006] 1956] achieves O(n3 ) runtime in the case that the dis- Altschuler et al. n2 /ε2 crete measures are uniform. The unweighted (uw) match- [2017] et al. n5/2 /ε ing problem was improved in R2 (2d) to Õ((n/ε)3/2 ) by Dvurechensky [2018] Agarwal and Phillips [2006] and then Õ(n5/4 poly(1/ε)) Lahn et al. [2019] n2 /ε+n/ε2 n/ε2 by Lahn and Raghvendra [2021] (where Õ denotes absorp- Lahn et al. [2019] 2d Lahn and Raghven- uw,2d n5/4 poly(1/ε) tion of logarithmic factors in the asymptotic notation). For dra [2021] general W22 computation, the Greenkhorn algorithm was originally shown to achieve Õ(n2 /ε3 ) [Altschuler et al., 2017] and then an improved analysis in Lin et al. [2019] yielded Õ(n2 /ε2 ). Dvurechensky et al. [2018] uses accelerated gradient descent which Lin et al. [2019] shows achieves Õ(n5/2 /ε). Lin et al. [2019] √separately provide an adaptive primal dual accelerated mirror descent algorithm achieving O(n2 γ/ϵ) where γ = O(n), in line with the worst case bound of Dvurechensky et al. [2018]. Lahn et al. [2019] then uses graph-scaling algorithms to achieve time O(n2 /ε + n/ε2 ) for general Optimal Transport computation and which improves to Õ(n/ε2 ) in R2 . But these works do not address statistical accuracy for the underlying quantity W22 (P, Q); they assume the n locations signify ground truth. On the statistical side, the dyadic partitioning argument presented by Weed and Bach [2019] plus triangle inequality for W2 implies that when P̂n and Q̂n are empirical measures from P, Q respectively, E|W22 (P̂n , Q̂n ) − W22 (P, Q)| is O(n−1/4 ) when d < 4 and and Õ(n−1/d ) for d ≥ 4, (bounded ∆∞ allows to convert W2 closeness to W22 closeness). Chizat et al. [2020] improves upon this, showing that E|W22 (P̂n , Q̂n ) − W22 (P, Q)| is O(n−1/2 ) when d < 4 and Õ(n−2/d ) for d ≥ 4. Thus ε error needs n = Õ(1/εmax(2,d/2) ). Alternatives to the traditional empirical measure plug-in approach rely on Besov smoothness [Niles-Weed and Berthet, 2022] of P, Q. 2
Table 2: Computational-Statistical Runtimes (CSR) for ε-additive W22 (P, Q) estimation on (0, 1)d , where d′ = 2⌊ d2 ⌋. For conciseness we omit not only poly(log(1/ε)) terms but also ε−o(1) terms. Note Our Result [Theorem 4.1] achieves ε−2 for α > 1/2 in R2 and and ε−2 in R3 when α → 1. Method
Notes
in R2
in R3
′
ε−7.5
ε−7.5
− max(6,d+2)
ε ε−(max(d,4)+1−I(d=2))
ε−6.5 ε−6 ε−4
ε−6.5 ε−6 ε−5
ε− max(2,d+1) d+1 ε− max(2, 1+α )
ε−3 ε−2
ε−4 ε−2
ε−(d +5.5)
Sinkhorn Chizat et al. [2020]
(d′ +11) − 2
ε
SH + Richardson Extrap Chizat et al. [2020] Entropic OT Chizat et al. [2020] Empirical Plugin + Lahn et al. [2019] Our Result [Corollary 4.1.3] Our Result [Theorem 4.1]
CSR in Rd
α-Höl
For hardness results, Niles-Weed and Rigollet [2022] imply that the number of samples needed to estimate W22 (P, Q) up to ε in mean absolute error is Ω(1/ε2 ). A simple, well-known (folklore) algorithm which achieves the CSR of O(1/ε2 ) in the 1 dimensional case works by collecting n = 1/ε2 , and then sorting the sample from P (p(1) < p(2) < . . . p(n) ) and the samples from Pn Q (q(1) < q(2) < · · · < q(n) ), and computing n1 j=1 (p(j) − q(j) )2 (this is simply the squared Wasserstein distance between the empirical measures). The algorithm takes O(n log n) time given the samples, and the CSR is Õ(ε−2 ) (by for example Chizat et al. [2020] Theorem 2). Our work will show that even for d = 2 or d = 3, it is possible to nearly achieve the optimal Õ(ε−2 ) CSR as long as the underlying distributions are sufficiently smooth. Chizat et al. [2020] also explicitly considers the joint computational-statistical runtime. They provide CSRs using Entropic Optimal Transport, the Sinkhorn-Divergence, and a Richardson extrapolation of ′
(d′ +11)
the Sinkhorn Divergence, respectively achieving Õ(ε− max(6,d+2) ), Õ(ε−(d +5.5) ) and Õ(ε− 2 ) where d′ := 2⌊ d2 ⌋. However, these rates are all inferior to the rate achieved by using the plug-in empirical measure approach with the O(nε−2 ) algorithm of Lahn et al. [2019] in R2 , with the tight statistical bound for the plug-in empirical measure approach from Chizat et al. [2020] mentioned above. These bounds are summarized in Table 2, along with our results which we outline next. Contributions. We introduce what we call the Sample-Sketch-Solve paradigm for achieving fasts CSRs for W22 (P, Q) with P, Q supported on (0, 1)d , and the distributions possess Hölder smooth densities. The paradigm operates in three phases: (1) it samples n points from distributions P and Q; (2) it sketches these samples into a compressed representation; and (3) it solves for the (approximately) optimal transportation plan on the compressed representation. In particular, our compressed representation is a regular Cartesian grid on [0, 1]d , with L divisions on each dimension. Each sample within a grid cell is mapped to the center so that there are max(n, Ld ) 1 total atoms. We set L = Θ(ε− 1+α ) when P and Q are α-Hölder smooth. This regular grid structure is useful in two ways: for achieving a high level of precision in approximating W22 (P, Q) when P and Q are smooth using the sketched measures, and for efficiently solving for the Optimal Transport cost between the sketched measures. In Section 2, we present new bounds on the error in approximating W22 (P, Q) between continuous probability distributions using this grid based discretization. We consider Holder smoothness constraints on the probability densities of P, Q, and rely on a novel, direct coupling based argument via a recent result of Collins and Tong [2025] on the regularity of the Brenier map. Still, our result here is not trivially obtainable using the triangle inequality, and is instead a result of a midpoint rule type cancellation of first order terms. Thus we crucially rely on mapping to the center point of the bounded diameter grid cells of our sketch. Then in Section 3 we outline a new approach for efficient exact computation of W22 (P, Q) that leverages (1) P, Q are supported on a regular grid with L splits per axis, and (2) each atom has normalized weight a multiple of 1/n. In this setting we can use the regular grid structure and the decomposable nature of W22 to leverage a result of Auricchio et al. [2018] to transform the problem into a (d + 1)-partite graph min-cost flow problem. Then given the control on the weights and 3
the bounded grid structure, we control the capacities of this flow problem, and in turn apply a fast algorithm of van den Brand et al. [2023]. This results in a Ld+1+o(1) exact algorithm for Wasserstein distance computation on our sketch. Finally, in Section 4 we put these pieces together, for α-Hölder distributions, to obtain the O(1/ε2 ) CSR algorithm for ε-additive error for W22 (P, Q) in R2 when α > 1/2, nearly O(1/ε2 ) in R3 when d+1+o(1)
α → 1, and in general O(ε− max(2, 1+α ) ) in Rd for 0 < α < 1 ; we also separately analyze the general non-smooth case. Table 2 displays the CSRs for the general case and α Hölder (α-Höl) case. 1.1
Preliminaries
(x)−f (y)∥ For a function f : A ⊆ Rd → Rz , |f |C 0,α (A) := supx,y∈A,x̸=y ∥f ∥x−y∥ denotes the α-Hölder α 0,α semi-norm of f where 0 < α ≤ 1. When |f |C (A) < ∞, we call f α-Hölder smooth [Tsybakov, 2008, Singh et al., 2018]. The higher the value of α, the higher the level of smoothness. The α = 1 case is referred to as Lipchitz smoothness. Note that on bounded domains, by the Mean Value Theorem if f is differentiable with upper bounded derivative, it is Lipchitz smooth. For 0 < α ≤ 1, f ∈ C 1,α (A) means f is differentiable and its gradient ∇f is α-Hölder smooth. When we say a probability distribution P supported on (0, 1)d is α-Hölder smooth for some 0 < α ≤ 1, we mean d it possesses a probability density function µ which R is α-Hölder smooth on (0, 1) . That is, for any Borel measurable subset of (0, 1)d , A, P (A) = x∈A µ(x)dx, and for some C > 0, |µ|C 0,α ((0,1)d ) =
= C < ∞. We encapsulate these common parameterizations in the supx,y∈(0,1)d ,x̸=y |µ(x)−µ(y)| ∥x−y∥α following assumption: Assumption 1.2. P, Q are probability measures on (0, 1)d with densities µ and ν respectively satisfying for some C > 1 and for some 0 < α < 1, 1. P and Q are α-Hölder smooth: |µ|C 0,α ((0,1)d ) + |ν|C 0,α ((0,1)d ) ≤ C 2. P and Q bounded above and below: C −1 ≤ µ(x), ν(x) ≤ C for x ∈ (0, 1)d When P, Q satisfy both parts of this assumption we say they are (α, C)-controlled. Gridding. Consider h > 0 such that L = h1 is an integer, Gh denotes the partition of (0, 1)d constructed by partitioning each axis into intervals of width h. Let Ḡdh ⊂ (0, 1)d be the 1/hd points at the center of these grid cells. Gh : (0, 1)d → Ḡdh is the function that projects a point to the center of the grid cell in Gh containing it. Pushforward. If f : (0, 1)d → Z for some set Z is a Borel measurable map, and P is a Borel probability measure on (0, 1)d , the notation f #P denotes the pushforward. Specifically for measurable A ⊆ Z, f #P (A) = P (x ∈ (0, 1)d : f (x) ∈ A). IdA refers to the identity map defined on A.
2
Grid Based Discretization Error Bounds
In this section we provide an improved upper bound on the error in approximating W22 (P, Q) by constructing discrete measures Gh #P and Gh #Q and using the plug-in approach W22 (Gh #P, Gh #Q). Note that Gh #P (resp. Gh #Q) are the measures constructed from P (resp. Q) by collapsing all mass within each grid cell in each of the O(hd ) grid cells in Gh onto the center of each cell. For Entropic Optimal transport this subject was addressed by Chizat et al. [2020] but the case with λ = 0 (i.e no entropic regularization, which corresponds to W22 ) has not been considered. We show a continuum of approximation depending on the smoothness level of the probability measures P, Q. Our proof relies on regularity results guaranteeing the smoothness of the optimal transport map which deterministically pushes forward P into Q. We first remind the reader of the famous Theorem due to Brenier [1991] guaranteeing the existence of such a map. Existing Result 2.1 ((Brenier) Theorem 1.16 of Chewi et al. [2024]). If P, Q are defined on Rd with bounded second moments, and P possesses a probability density, then there exists a convex function u : Rd → R and corresponding function Z : Rd → Rd × Rd , defined by Z(x) := (x, ∇u(x)) such that Z#P is the unique coupling of P and Q achieving the infima in the definition of W22 (P, Q). 4
By Result 2.1, (∇u)#P = Q and T := ∇u is thus referred to as the optimal R transport map. Note that by a change of measure (when Brenier’s Theorem applies), W22 (P, Q) = x∈Rd ∥x − T (x)∥22 P (dx). Caffarelli [1992, 1996] initiated the study of the regularity of the optimal transport map. We use a new contribution in transport map regularity theory due to Collins and Tong [2025] which establishes the global α-Hölder smoothness of ∇u when P and Q have α-Hölder densities themselves. Existing Result 2.2 (Global Regularity of Brenier Map for Convex Domains via Collins and Tong [2025] Theorem 5.1). Suppose P and Q satisfy Assumption 1.2 for some C > 1, 0 < α < 1. Then the convex functions u, v of Brenier’s Theorem satisfy u ∈ C 1,1−ξ ((0, 1)d ) and v ∈ C 1,1−ξ ((0, 1)d ) and |∇u|C 0,1−ξ ((0,1)d ) + |∇v|C 0,1−ξ ((0,1)d ) ≤ M for any 0 < ξ < 1 where M depends only on d, α, C, ξ. In particular the Optimal Transport maps ∇u and ∇v are (1 − ξ)-Hölder continuous for any 0 < ξ < 1. Before we proceed to the main result, we need a corollary. Corollary 2.2.1. Let P, Q be (α, C)-controlled for some C > 1 and 0 < α < 1, with optimal transportation map Tµ from P to Q (resp. Tν from Q to P ), and displacement function ℓµ (x) = x − Tµ (x) (resp. ℓν (x) = x − Tν (x)). Then ∥ℓµ (x) − ℓµ (Gh (x))∥ ≤ Kh1−ξ (resp. ∥ℓν (x) − ℓν (Gh (x))∥ ≤ Kh1−ξ ) for any 0 < ξ < 1, and where K is a constant that depends only on d,C, α, ξ. Proof. Since P, Q are (α, C)-controlled, Existing Result 2.2 holds. This means the OT map, Tµ , (which is equivalently ∇u where u is the convex function from Existing Result 2.2), is (1 − ξ)-Hölder continuous for any 0 < ξ < 1. Since the identity function, f (x) = x, is Lipchitz continuous, it is also (1 − ξ)-Hölder continuous; also the sum of two (1 − ξ)-Hölder continuous functions is also (1 − ξ)-Hölder. Hence the difference between x and Tµ (x), ℓµ (x), is (1 − ξ)-Hölder for any ξ ∈ (0, 1). The argument is analogous for ℓν . Now we state our main contribution of this section. Theorem 2.3. If P and Q are (α, C)-controlled (Assumption 1.2) for some C > 1, 0 < α < 1, then for some constant K depending only on d, C, α, and every h > 0 W22 (Gh #P, Gh #Q) − W22 (P, Q) ≤ Kh1+α Here we show the proof of the upper bound W22 (Gh #P, Gh #Q) − W22 (P, Q) ≤ Kh1+α . The argument for the lower bound W22 (Gh #P, Gh #Q)−W22 (P, Q) ≥ −Kh1+α is similar (see appendix Section A for complete details). In the argument below note that we repeatedly use the general constant K to indicate a constant depending only on d, C and α. proof of upper bound. The proof will proceed in three phases. First, we will define a transportation plan π which couples Gh #P with Gh #Q and show it is indeed a valid coupling. Second, we will use this coupling to upper bound W22 (Gh #P, Gh #Q) by W22 (P, Q) + remainder. Third, we leverage properties of the scale of the grid cells, and the α-Hölderness of both µ, ν and the Brenier maps’ to bound each term in the remainder by either either Kh2 or Kh1+α . Phase 1: Consider the probability measure π := S#γ where S : (0, 1)d ×(0, 1)d → (0, 1)d ×(0, 1)d is defined as S(x, y) = (Gh (x), Gh (y)) and γ is an optimal coupling of P, Q under squared Euclidean cost. π projects the mass of γ onto the grid cell centers, and thus since γ couples P and Q, π is a valid coupling for Gh #Pµ and Gh #Pν . We now formalize this by proving that the marginals are preserved. That is, for any A ⊆ (0, 1)d , π(A × (0, 1)d ) = γ(S −1 (A × (0, 1)d )) d = γ((x, y) ∈ (0, 1)2d : x ∈ G−1 h (A), y ∈ (0, 1) )
= P (G−1 h (A)) = Gh #P (A) and similarly π((0, 1)d × A) = Gh #Q(A). Phase 2: We can now state W22 (Gh #P, Gh #Q) ≤ E(X,Y )∼π ∥X − Y ∥22 =
5
Z x∈(0,1)d ,y∈(0,1)d
∥Gh (x) − Gh (y)∥22 dγ(x, y).
Applying Brenier’s theorem we can write γ = (Id(0,1)d , Tµ )#P where Tµ is the OT map (pushing P to Q) and so the above integral upper bound becomes (after change of measure) Z W22 (Gh #P, Gh #Q) ≤ ∥Gh (x) − Gh (Tµ (x))∥22 µ(x)dx. x∈(0,1)d
To isolate the contribution of W22 (P, Q) to the integral on RHS, we denote ∆z = z − Gh (z) for z ∈ (0, 1)d so that Gh (x) − Gh (Tµ (x)) = (x − Tµ (x)) − (∆x − ∆Tµ (x)). Thus using Brenier’s Theorem and expanding the square in the integral using the above identity, and R using that W22 (P, Q) = x (x − Tµ (x))2 µ(x)dx, we have Z W22 (Gh #P, Gh #Q) − W22 (P, Q) ≤ ∥∆x − ∆Tµ (x)∥22 − 2⟨x − Tµ (x), ∆x − ∆Tµ (x)⟩µ(x)dx. x
Phase 3: We will use three properties: [P1] For all x ∈ (0, 1)d , ∥∆x∥2 = ∥x − Gh (x)∥ ≤
√
dh .
[P2] The Brenier displacement ℓµ (x) := x − Tµ (x) takes values in the bounded domain (−1, 1)d . [P3] Since P, Q are (α, C)-controlled, by Corollary 2.2.1, the Brenier displacement ℓµ (x) = x − Tµ (x) satisfies ∥ℓ(x) − ℓ(Gh (x))∥ ≤ Khα (by setting ξ = 1 − α). [P4] Since P, Q are (α, C)-controlled, µ(x) ≤ C = O(1). Additionally, as µ is α-Hölder, |µ(x) − µ(Gh (x))| ≤ Khα . R The first term in the remainder, x ∥∆x − ∆Tµ (x)∥22 dx, scales as Kh2 by triangle inequality yielding ∥∆x − ∆Tµ (x)∥ ≤ ∥∆x∥ + ∥∆Tµ (x)∥ and then application of [P1]. We are left with Z W22 (Gh #Pµ , Gh #Pν )−W22 (Pµ , Pν ) ≤ Kh2 −2 ⟨x−Tµ (x), ∆x−∆Tµ (x)⟩µ(x)dx. (1) x∈(0,1)d
It remains to bound the absolute value of the remaining integral. We have Z | ⟨x − Tµ (x), ∆x − ∆Tµ (x)⟩µ(x)dx| ≤ x Z Z | ⟨x − Tµ (x), ∆x⟩µ(x)dx| + | ⟨x − Tµ (x), ∆Tµ (x)⟩µ(x)dx|. x
x
The second of these terms we can show is symmetric to the first by using from Brenier’s Theorem that the optimal coupling is unique. In particular, we can re-write γ = (Id(0,1)d , Tµ )#P from above in the other direction as γ = (Tν , Id(0,1)d )#Q (where Tν is the OT map pushing Q to P ). Then we can rewrite this by first integrating over the coupling measure, and then swapping the direction: Z Z ⟨x − Tµ (x), ∆Tµ (x)⟩µ(x)dx = ⟨x − y, ∆y⟩dγ(x, y) x∈(0,1)d (x,y)∈(0,1)d ×(0,1)d Z = ⟨Tν (y) − y, ∆y⟩ν(y)dy. y∈(0,1)d
R
R Now the two terms we need to bound | x∈(0,1)d ⟨x − Tµ (x), ∆x⟩µ(x)dx| and | y∈(0,1)d ⟨Tν (y) − y, ∆y⟩ν(y)dy| are symmetric to each other, and just show how to bound the first term. We arbitrarily label the cells C1 , . . . , C(1/h)d (without loss of generality assuming 1/h R is an integer), and the corresponding cell centers as c1 , . . . , c(1/h)d . Observe that for every j, Cj ⟨ℓµ (cj ), x − cj ⟩µ(cj ) = 0 since Cj is the grid cell centered at cj . We can now rewrite this first term, and subtract 6
this 0 expression: Z | ⟨x − Tµ (x), ∆x⟩µ(x)|dx x∈(0,1)d
Z
Z
x∈(0,1)d
≤
XZ
x∈(0,1)d
I(x ∈ Cj )⟨ℓµ (cj ), x − cj ⟩µ(cj )dx
j
|⟨ℓµ (x), x − cj ⟩µ(x) − ⟨ℓµ (cj ), x − cj ⟩µ(cj )| dx
(by Tri Ineq)
x∈Cj
j
≤
X
⟨ℓµ (x), ∆x⟩µ(x)dx −
=
XZ
|⟨ℓµ (x), x − cj ⟩µ(x) − ⟨ℓµ (x), x − cj ⟩µ(cj ) + ⟨ℓµ (x), x − cj ⟩µ(cj ) − ⟨ℓµ (cj ), x − cj ⟩µ(cj )|dx
x∈Cj
j
by ±⟨ℓµ (x), x − cj ⟩µ(cj ) ≤
XZ
Z |⟨ℓµ (x), x − cj ⟩(µ(x) − µ(cj ))|dx +
Cj
j
|⟨ℓµ (x) − ℓµ (cj ), x − cj ⟩µ(cj )|dx Cj
by Tri Ineq ≤
XZ
Z ∥ℓµ (x)∥2 ∥x − cj ∥2 |µ(x) − µ(cj )|dx +
Cj
j
∥ℓµ (x) − ℓµ (cj )∥2 ∥x − cj ∥2 |µ(cj )|dx Cj
by Cauchy-Schwarz XZ ≤K j
≤K
X
Z h|µ(x) − µ(cj )|dx +
Cj
h∥ℓµ (x) − ℓµ (cj )∥dx
by [P2] and [P4]
Cj
V ol(Cj )h1+α = Kh1+α
by [P3] and [P4]
j
R Thus | x∈(0,1)d ⟨x − Tµ (x), ∆(x)⟩µ(x)dx| ≤ Kh1+α . An analagous argument yields R 1+α | y∈(0,1)d ⟨Tν (y)−y, ∆y⟩ν(y)dy| ≤ Kh . The upper bound W22 (Gh #P, Gh #Q)−W22 (P, Q) ≤ 1+α Kh follows from these and Equation 1. 2
2
−y | Just by using that |x − y| ≤ |xx+y , we additionally have the following Lemma.
Lemma 2.4. Assuming P, Q are (α, C)-controlled for some C > 1 and 0 < α < 1 and W2 (P, Q) > ε > 0 then for some constant K depending only on d, c, α and every h > 0, |W2 (Gh #P, Gh #Q) − W2 (P, Q)| ≤
Kh1+α ε
Proof. Using Theorem 2.3 |W2 (Gh #P, Gh #Q) − W2 (P, Q)| ≤
2.1
Kh1+α |W22 (Gh #P, Gh #Q) − W22 (P, Q)| ≤ W2 (Gh #P, Gh #Q) + W2 (P, Q) ε
Non-Smooth Case
Note that the α → 0 bound of O(h) holds for any distributions P and Q, regardless of their level of smoothness. We prove this below. Lemma 2.5. If P and Q are probability measures on (0, 1)d , then for h > 0 √ |W2 (Gh #P, Gh #Q) − W2 (P, Q)| ≤ h 2 d . 7
Proof. Let πP = (Id(0,1)d , Gh )#P and πQ = (Id(0,1)d , Gh )#Q. Then πP is a coupling of P and Gh #P and πQR is a coupling of Q and Gh #Q. a change of measure we thus have R P Using W22 (Gh #P, P ) ≤ x∈(0,1)d ∥x−Gh (x)∥22 P (dx) = j x∈Cj ∥x−cj ∥22 P (dx) ≤ dh2 , and likewise W22 (Gh #Q, Q) ≤ dh2 . By triangle inequality for W2 the result follows. It is not hard to show that W2 (Gh #P, P ) = Θ(h) when P is for example the uniform distribution on (0, 1)d . And thus the brute force best bound one can obtain using the triangle inequality for W2 is |W2 (Gh #P, Gh #Q) − W2 (P, Q)| = O(h).
3
Efficient Exact Computation on a Regular Grid
In this section we isolate a structured regime in which the squared Wasserstein distance admits a substantially faster exact algorithm. Let L, k ∈ N, and recall that Ḡd1/L are the centers of a regular grid on (0, 1)d . Let X X P = p x δx , Q= qx δx x∈Ḡd 1/L
x∈Ḡd 1/L
be probability measures supported on Ḡd1/L , with px , qx ∈ k1 Z≥0 for every grid point x. After aggregating repeated support points, the input can be viewed as a pair of d-dimensional histograms on regular Ld axis-aligned bins of (0, 1)d . The squared Euclidean ground cost is separable on this grid. Indeed, if 2a1 + 1 2ad + 1 2bd + 1 2b1 + 1 x= ,..., ,..., , y= , 2L 2L 2L 2L with a1 , . . . , ad , b1 , . . . , bd ∈ {0, 1, . . . , L − 1}, then d
∥x − y∥22 =
1 X (aℓ − bℓ )2 . L2 ℓ=1
Therefore the cost has the separable form required by the exact reduction of Auricchio et al. [2018], which maps optimal transport between d-dimensional histograms with separable ground cost to an equivalent uncapacitated minimum-cost flow problem on a (d + 1)-partite graph. Instantiating their construction with Ld histogram bins yields a graph with (d + 1)Ld vertices and m := dLd+1 arcs. Theorem 3.1. Assume L, k ∈ N. Let P and Q be probability measures supported on the cell-center grid Ḡd1/L ⊂ [0, 1]d and assume every mass is an integer multiple of 1/k. Then W22 (P, Q), for every fixed dimension d, can be computed exactly in time (dLd+1 )1+o(1) log L log k. For for every fixed dimension d, this is Õ(Ld+1+o(1) ) where Õ(·) hides logarithmic factors of L, k. Proof. By the exact reduction of Auricchio et al. [2018], it suffices to solve the resulting uncapacitated minimum-cost flow instance on a (d + 1)-partite graph with m = dLd+1 arcs. In this reduction, each −bℓ )2 arc cost is one coordinate contribution to the squared Euclidean cost, hence is of the form (aℓ L 2 for some coordinate index ℓ and some integers aℓ , bℓ ∈ {0, 1, . . . , L − 1}. Therefore every arc cost is a multiple of 1/L2 . Next scale every supply and demand by k. Because each histogram mass lies in k1 Z≥0 , all node imbalances become integral; moreover, since P and Q are probability measures, the total amount of flow becomes exactly k. Now also scale every arc cost by L2 . After this rescaling, all costs become integral, and every cost is at most max
0≤aℓ ,bℓ ≤L−1
(aℓ − bℓ )2 ≤ (L − 1)2 < L2 .
Since the total flow is k, assigning capacity k to every formerly uncapacitated arc does not change the feasible region. Thus we obtain a directed minimum-cost flow instance with integral demands, 8
capacities, and costs, where the maximum edge capacity satisfies U ≤ k and the maximum edge cost satisfies C < L2 . We may therefore apply the deterministic exact minimum-cost flow algorithm of van den Brand et al. [2023], which computes an exact optimum in time m1+o(1) log U log C, for instances with integral demands, capacities, and costs. Substituting U ≤ k and C < L2 yields a running time of m1+o(1) log k log L. [ denote the optimum of the Finally, let OPT denote the original value of W22 (P, Q), and let OPT scaled min-cost flow instance. Scaling the flow values by k multiplies the objective by k, and scaling the arc costs by an additional factor of L2 multiplies the objective by another factor of L2 . Hence [ = kL2 OPT. OPT Dividing by kL2 therefore recovers the exact value of W22 (P, Q). Corollary 3.1.1. Under assumptions of Theorem 3.1, when k = poly(L) computing W22 (P, Q) • in d = 2 takes time L3+o(1) • in d = 3 takes time L4+o(1) .
4
Faster CSR Algorithm for W22
In this section we present a fast Computational-Statistical Runtime (CSR) algorithm for additive ε approximation of the W22 between two measures P, Q supported on (0, 1)2 . This combines our two above results: first applying the statistical result about Grid-based discretization (Theorem 2.3) to compress a sample, and second applying the exact algorithm when our samples have been discretized to a grid (Theorem 3.1) for computing W22 . Theorem 4.1. Let P and Q be (α, C)-controlled distributions (with C > 1, 0 < α < 1) on (0, 1)d , for which we can draw random samples in O(1) time each. For constant dimension d and ε > 0, we can compute Â(ε) so that in expectation |Â(ε) − W22 (P, Q)| ≤ ε, in time Õ((1/ε)max(2,
d+1+o(1) 1+α
) ).
Proof. The proof follows the three phases of the sample-sketch-solve paradigm. We start by sampling n = O(1/εmax(2,d/2) ) samples from both P and Q to produce discrete measures P̂n , Q̂n , respectively. 1 Next we apply a grid-based discretization with an edge length h = O(ε 1+α ) to obtain Gh #P̂n and Gh #Q̂n . By Theorem 2.3 we have |W22 (P, Q) − W22 (Gh #P, Gh #Q)| ≤ Kh1+α = ε/2. Since Gh #P̂n is an empirical measure of n independent samples of Gh #P and Gh #Q̂n is an empirical measure of n independent samples of Gh #Q, by Chizat et al. [2020] Theorem 2 we have E|W22 (Gh #P, Gh #Q) − W22 (Gh #P̂n , Gh #Q̂n )| ≤ ε/2 where the expectation is with respect to the sampling from P and Q. Thus with Â(ε) := W22 (Gh #P̂n , Gh #Q̂n ) and applying triangle inequality, E|W22 (P, Q) − Â(ε)| = E|W22 (P, Q) − W22 (Gh #P̂n , Gh #Q̂n )| ≤ |W22 (P, Q) − W22 (Gh #P, Gh #Q)| + E|W22 (Gh #P, Gh #Q) − W22 (Gh #P̂n , Gh #Q̂n )| ≤ ε. Finally, now the distributions Gh #P̂n , Gh #Q̂n lie on a regular Cartesian grid Ḡd1/L with L = 1/h = O(1/ε1/(1+α) ) and each atom has mass that is an integer multiple of 1/n with n = O(1/εmax(2,d/2) ). Thus we can invoke Theorem 3.1 to exactly compute W22 (Gh #P̂n , Gh #Q̂n ) in Ld+1+o(1) log n = 1/ε(d+1+o(1))/(1+α) time. Adding in the Õ(1/εmax(2,d/2) ) time for n samples completes the proof.
We now provide the following corollaries about how this analysis is exactly or nearly optimal for d = 2, 3 dimensions. 9
Corollary 4.1.1. Let P and Q be (α, C)-controlled distributions on (0, 1)d for d = 2, for which we can draw random samples in O(1) time each. For α > 1/2 and ε > 0, we can compute Â(ε) so that in expectation |Â(ε) − W22 (P, Q)| ≤ ε, in time Õ(1/ε2 ). Corollary 4.1.2. Let P and Q be (α, C)-controlled distributions on (0, 1)d for d = 3, for which we can draw random samples in O(1) time each. For ε > 0, we can compute Â(ε) so that in expectation 4+o(1)
|Â(ε) − W22 (P, Q)| ≤ ε, in time ε− 1+α , which approaches ε−(2+o(1)) as α → 1. 4.1
Non-Smooth Case
When we do not have a bound on the smoothness of P, Q, then we can use Lemma 2.5, and as a result Theorem 4.1 behaves as with α = 0 with a CSR of ε− max(2,d+1+o(1)) . Corollary 4.1.3. Let P and Q be distributions on (0, 1)d , for which we can draw random samples in O(1) time each. For constant dimension d and ε > 0, we can compute Â(ε) so that in expectation |Â(ε) − W22 (P, Q)| ≤ ε, in time Õ(ε− max 2, d+1+o(1) ). For d = 2 the CSR is 1/ε3+o(1) and for d = 3 it is 1/ε4+o(1) .
10
References Pankaj K Agarwal and Jeff M Phillips. On bipartite matching under the rms distance. In Proceedings of the Eighteenth Canadian Conference on Computational Geometry’, Kingston, Canada, pages 143–146. Citeseer, 2006. Jason Altschuler, Jonathan Niles-Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration. Advances in neural information processing systems, 30, 2017. Gennaro Auricchio, Federico Bassetti, Stefano Gualandi, and Marco Veneroni. Computing kantorovich-wasserstein distances on d-dimensional histograms using (d + 1)-partite graphs. In Advances in Neural Information Processing Systems 31, 2018. Espen Bernton, Pierre E Jacob, Mathieu Gerber, and Christian P Robert. On parameter estimation with the wasserstein distance. Information and Inference: A Journal of the IMA, 8(4):657–676, 2019. Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375–417, 1991. Luis A Caffarelli. The regularity of mappings with a convex potential. Journal of the American Mathematical Society, 5(1):99–104, 1992. Luis A Caffarelli. Boundary regularity of maps with convex potentials–ii. Annals of mathematics, 144(3):453–496, 1996. Sinho Chewi, Jonathan Niles-Weed, and Philippe Rigollet. Statistical optimal transport. arXiv preprint arXiv:2407.18163, 3, 2024. Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavier Vialard, and Gabriel Peyré. Faster wasserstein distance estimation with the sinkhorn divergence. Advances in Neural Information Processing Systems, 33:2257–2269, 2020. Tristan C Collins and Freid Tong. Boundary regularity of optimal transport maps on convex domains. arXiv preprint arXiv:2507.05395, 2025. Pavel Dvurechensky, Alexander Gasnikov, and Alexey Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by sinkhorn’s algorithm. In International conference on machine learning, pages 1367–1376. PMLR, 2018. Harold W Kuhn. Variants of the hungarian method for assignment problems. Naval research logistics quarterly, 3(4):253–258, 1956. Nathaniel Lahn and Sharath Raghvendra. An o (n 5/4) time -approximation algorithm for rms matching in a plane. In Proceedings of the 2021 ACM-SIAM Symposium on Discrete Algorithms (SODA), pages 869–888. SIAM, 2021. Nathaniel Lahn, Deepika Mulchandani, and Sharath Raghvendra. A graph theoretic additive approximation of optimal transport. Advances in Neural Information Processing Systems, 32, 2019. Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International Conference on Machine Learning, pages 3982–3991. PMLR, 2019. Jonathan Niles-Weed and Quentin Berthet. Minimax estimation of smooth densities in wasserstein distance. The Annals of Statistics, 50(3):1519–1540, 2022. Jonathan Niles-Weed and Philippe Rigollet. Estimation of wasserstein distances in the spiked transport model. Bernoulli, 28(4):2663–2688, 2022. Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. On wasserstein two-sample testing and related families of nonparametric tests. Entropy, 19(2):47, 2017. 11
Shashank Singh, Ananya Uppal, Boyue Li, Chun-Liang Li, Manzil Zaheer, and Barnabás Póczos. Nonparametric density estimation under adversarial losses. Advances in Neural Information Processing Systems, 31, 2018. Alexandre B Tsybakov. Nonparametric estimators. In Introduction to Nonparametric Estimation, pages 1–76. Springer, 2008. Jan van den Brand, Li Chen, Rasmus Kyng, Yang P. Liu, Richard Peng, Maximilian Probst Gutenberg, Sushant Sachdeva, and Aaron Sidford. A deterministic almost-linear time algorithm for minimumcost flow. FOCS, 2023. Jonathan Weed and Francis Bach. Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance. Bernoulli, 25(4A):2620–2648, 2019.
12
A
Discretization Error for W22
In this section we prove the following theorem: Theorem 2.3. If P and Q are (α, C)-controlled (Assumption 1.2) for some C > 1, 0 < α < 1, then for some constant K depending only on d, C, α, and every h > 0 W22 (Gh #P, Gh #Q) − W22 (P, Q) ≤ Kh1+α In the main body of the paper, we already proved the upper bound W22 (Gh #P, Gh #Q) − W22 (P, Q) ≤ Kh1+α . It remains to prove the corresponding lower bound W22 (Gh #P, Gh #Q) − W22 (P, Q) ≥ −Kh1+α , and we take this up in the next Subsection. A.1
Proof of Discretization Error Lower Bound
As we did in the upper bound proof, we continue to assume without loss of generality that h1 is an integer. Lemma A.1 (Discretization Error Lower Bound). If P and Q are (α, C) controlled (Assumption 1.2) for some C > 1, 0 < α < 1, then for some constant K depending only on d, C, α, and every h > 0, W22 (P, Q) ≤ W22 (Gh #P, Gh #Q) + Kh1+α . Note that the below argument only relies on smoothness of the densities, not at all on the transport maps. Existing Result 2.2, which guarantees the smoothness of the transport maps, uses that the densities are bounded below. Since we do not use transport map smoothness in the lower bound argument, we do not need the densities to be bounded below. In particular, only for the sake of consistency with the assumption of the upper bound does Theorem A.1 say P, Q are (α, C) controlled. This is stronger than what is nescessary for the lower bound argument. The lower bound argument uses only the smoothness of the density µ associated to P and ν associated to Q. This argument has a similar flavor to the upper bound argument. We find a coupling for P and Q, call it Γ, derived from an optimal coupling π ∗ for Gh P and Gh Q. Then we extract the cost W22 (Gh #P, Gh #Q) from the cost of Γ, and upper bound the leftover quantity utilizing the smoothness of the densities. Proof. Let π ∗ be an optimal coupling between the discrete measures Gh #P and Gh #Q. We define a probability measure Γ on (0, 1)2d by defining the probability density function γ : (0, 1)2d → R as γ(x, y) =
XX i
I(x ∈ Ci )I(y ∈ Cj )
j
π ∗ (ci , cj ) µ(x)ν(y). P (Ci )Q(Cj )
(2)
R and setting, for A ⊆ (0, 1)2d , Γ(A) = (x,y)∈A γ(x, y)d(x, y). Within Ci × Cj , Γ retains the shape of the independent coupling P × Q, but rescales the measure to be the amount of measure transferred between the center of Ci and Cj under the optimal coupling between the measures snapped onto the grid cell centers. We now show Γ is indeed a coupling of P and Q, and then we extract the quantity W22 (Gh #P, Gh #Q) from the cost of Γ. But first, note that Z Z Z γ(x, y)d(x, y) = γ(x, y)dydx = (x,y)∈(0,1)2d
=
x∈(0,1)d
y∈(0,1)d
XXZ
Z
i
=
j
XX i
x∈Ci ∗
π ∗ (ci , cj ) µ(x)ν(y)dydx y∈Cj P (Ci )Q(Cj )
π (ci , cj )
j
= 1. 13
(3)
where in the last inequality we used that π ∗ is a probabilty measure supported on (ci )i∈hd × (cj )j∈hd . Thus γ is indeed a probability density on (0, 1)2d . Additionally, for A ⊆ (0, 1)d , we have that X Γ(A × (0, 1)d ) = Γ((A ∩ Ci ) × (0, 1)d ) i
=
XXZ i
= =
j
X P (A ∩ Ci ) X P (Ci )
X P (A ∩ Ci ) P (Ci )
i
=
x∈A∩Ci
X X π ∗ (ci , cj ) P (A ∩ Ci )Q(Cj ) P (Ci )Q(Cj ) i j
i
=
π ∗ (ci , cj ) µ(x)ν(y)dydx y∈Cj P (Ci )Q(Cj )
Z
X P (A ∩ Ci ) P (Ci )
i
(4)
π ∗ (ci , cj )
j
Gh #P (ci ) P (Ci )
= P (A). where we used that π is a coupling between Gh #P and Gh #Q and then the definition of Gh #P . By a symmetric argument, for A ⊆ (0, 1)d , ∗
Γ((0, 1)d × A) = Q(A). In particular Γ is a coupling of P and Q. Thus we have that Z W22 (P, Q) ≤ ∥x − y∥22 γ(x, y)d(x, y) = XXZ i
j
x∈Ci
Z
(x,y)∈(0,1)2d π ∗ (ci , cj )
y∈Cj P (Ci )Q(Cj )
(5) ∥x − y∥22 µ(x)ν(y)dydx
Now we extract the cost W22 (Gh #P, Gh #Q) from the cost of Γ. Using that for each i ∈ [(1/h)d ], j ∈ [(1/h)d ], x, y ∈ (0, 1)d , x − y = (x − ci ) + (cj − y) + (ci − cj ), we have that ∥x − y∥22 = ∥x − ci ∥22 + ∥cj − y∥22 + ∥ci − cj ∥22 + 2⟨x − ci , cj − y⟩ + 2⟨x − ci , ci − cj ⟩+ (6) 2⟨cj − y, ci − cj ⟩. Expanding ∥x − y∥22 in this way in equation 5, we have that W22 (P, Q) ≤
π ∗ (ci , cj ) ∥x − ci ∥22 µ(x)dx+ P (C ) i x∈Ci i j XXZ π ∗ (ci , cj ) ∥y − cj ∥22 ν(y)dy+ Q(C ) j y∈Cj i j XX π ∗ (ci , cj )∥ci − cj ∥22 +
XXZ
i
j
Z X X π ∗ (ci , cj ) Z 2 (⟨x − ci , cj − y⟩ + ⟨x − ci , ci − cj ⟩ + ⟨cj − y, ci − cj ⟩) µ(x)ν(y)dydx. P (Ci )Q(Cj ) x∈Ci y∈Cj i j (7) ∥x − ci ∥22 ≤ Kh2 for x ∈ Ci and likewise ∥y − cj ∥22 ≤ Kh2 for y ∈ Cj for general constant K depending only on d. Thus the total of the first two terms on the RHS of equation 7 is bounded by Kh2 . Since π ∗ is an optimal coupling between Gh #P and Gh #Q, the third term on RHS of equation 7 is exactly W22 (Gh #P, Gh #Q). Thus we have that W22 (P, Q) ≤ W22 (Gh #P, Gh #Q) + Kh2 + Z Z X X π ∗ (ci , cj ) 2 (⟨x − ci , cj − y⟩ + ⟨x − ci , ci − cj ⟩ + ⟨cj − y, ci − cj ⟩) µ(x)ν(y)dydx. P (Ci )Q(Cj ) x∈Ci y∈Cj i j (8) 14
By Cauchy-Schwarz, and again using that the diameter of Ci , Cj is at most Kh, we have that Z X X π ∗ (ci , cj ) Z 2 (⟨x − ci , cj − y⟩) ≤ P (Ci )Q(Cj ) x∈Ci y∈Cj i j Z X X π ∗ (ci , cj ) Z 2 ∥x − ci ∥2 ∥cj − y∥2 µ(x)ν(y)dydx ≤ P (Ci )Q(Cj ) x∈Ci y∈Cj i j ! Z X X π ∗ (ci , cj ) Z 2 ∥x − ci ∥2 µ(x)dx ∥cj − y∥2 ν(y)dy ≤ P (Ci )Q(Cj ) x∈Ci y∈Cj i j XX Kh2 π ∗ (ci , cj ) ≤ i
(9)
j
Kh2 . By Equations 8 and 9, we have that W22 (P, Q) ≤ W22 (Gh #P, Gh #Q) + Kh2 + ∗
2
X X π (ci , cj ) P (Ci )Q(Cj ) i j
Z
Z (⟨x − ci , ci − cj ⟩ + ⟨cj − y, ci − cj ⟩) µ(x)ν(y)dydx.
x∈Ci
(10)
y∈Cj
To handle the first remaining integrand, we express µ(x) = µ(x) − µ(ci ) + µ(ci ). Then we have that Z X X π ∗ (ci , cj ) Z (⟨x − ci , ci − cj ⟩) µ(x)ν(y)dydx = 2 P (Ci )Q(Cj ) x∈Ci y∈Cj i j X X π ∗ (ci , cj ) Z 2 ⟨x − ci , ci − cj ⟩µ(cj )dx+ (11) P (Ci ) x∈Ci i j X X π ∗ (ci , cj ) Z 2 ⟨x − ci , ci − cj ⟩(µ(x) − µ(ci ))dx P (Ci ) x∈Ci i j The first integrand above is zero due to symmetry of the linear term about ci . The second term in absolute value is handled using Cauchy-Schwarz and the α-Hölder continuity of µ, yielding for this term an upper bound of Kh1+α where K depens only on d, C. Here we have used that ∥ci − cj ∥2 is upper bounded for all i, j by a constant depending only on d. Hence 2
Z X X π ∗ (ci , cj ) Z (⟨x − ci , ci − cj ⟩) µ(x)ν(y)dydx ≤ Kh1+α . P (C )Q(C ) i j x∈C y∈C i j i j
By an analgous argument using the α-H"older continuity of ν, we have that 2
Z X X π ∗ (ci , cj ) Z (⟨cj − y, ci − cj ⟩) µ(x)ν(y)dydx ≤ Kh1+α . P (C )Q(C ) i j x∈C y∈C i j i j
Using these and Equation 10, we conclude that for some K depending only d, C, W22 (P, Q) ≤ W22 (Gh #P, Gh #Q) + Kh1+α .
LLM Use Disclosure In addition to standard help with formatting or phrasing, we used LLM to search literature. This included the help in producing a complete related work, and contextualizing results. Crucially it also helped us identify the [Auricchio et al., 2018] paper which transforms the sketched representation into integral flow problem which can be solved efficiently. We of course reviewed all literature and techniques before integrating into the paper. 15