Asymptotics of smoothed Wasserstein distances in the small noise regime
Anonymous Author(s)
AffiliationAddress email
Abstract
We study the behavior of the Wasserstein- 2 distance between discrete measuresµ and ν in R d when both measures are smoothed by small amounts of Gaussiannoise.
This procedure, known as Gaussian-smoothed optimal transport , has recently attracted attention as a statistically attractive alternative to the unregularized
Wasserstein distance.
We give precise bounds on the approximation properties ofthis proposal in the small noise regime, and establish the existence of a phasetransition: we show that, if the optimal transport plan from µ to ν is unique and aperfect matching, there exists a critical threshold such that the difference between
W 2 ( µ, ν ) and the Gaussian-smoothed OT distance W 2 ( µ ∗ N σ , ν ∗ N σ ) scaleslike exp( − c/σ 2 ) for σ below the threshold, and scales like σ above it.
Theseresults establish that for σ sufficiently small, the smoothed Wasserstein distanceapproximates the unregularized distance exponentially well
Introduction: optimal transport
Optimal Transport (OT) has seen a recent surge of applications in machine learning, in areas such asgenerative modeling [2, 16], image processing [13, 27, 31], and domain adaptation [7, 8].
A naturalstatistical question raised by these applications is to estimate the OT distances with samples.
Thesedistances, known as the Wasserstein distances, are defined by
[Equation]
where Π( µ, ν ) denotes the set of joint measures with marginals µ and ν , known as transport plans .
It is well known that plug-in estimators for this quantity, obtained by replacing µ and ν with empirical measures consisting of i.i.d.
samples, have performance in high dimensions, with rates ofconvergence typically of order n − 2 /d
[3, 11, 12, 15, 22] when d > 2 p .
Moreover, minimax lowerbounds show that this curse of dimensionality is unavoidable in general [25, 32].
The existence of the curse of dimensionality for OT has led to a series of proposals to obtainbetter rates of convergence by imposing additional structural assumptions—such as latent lowdimensionality [25] or smoothness [24, 33]—or by replacing W p by a better-behaved surrogate,such as an entropy-regularized version with much better statistical and computational properties [1, 9, 17, 23, 28].
A particularly intriguing option, developed by [18], consists in smoothing the Wasserstein distanceby adding Gaussian noise.
The following result shows the statistical benefits of this approach.
Proposition 1.1 ([21]) .
For d > 1 and σ > 0 , denote by N σ the centered Gaussian measure on R dwith covariance σ 2 I d .
For any compactly supported probability measure µ in R d , let x 1 , x 2 , . . .
, x n be i.i.d.
samples from µ , and define the empirical measure
[Equation]
Then there exists a constant c
= c ( µ, σ, d ) such that
[Equation]
[18] call this framework Gaussian-smoothed optimal transport (GOT), and follow up work hasshown that it possesses significant statistical benefits, with fast rates of convergence and clean limitlaws [19, 20, 38].
To leverage the beneficial properties of the GOT framework, it is necessary to understand how wellthe smoothed distance W 2 ( µ ∗N σ , ν ∗N σ ) approximates the standard Wasserstein distance W 2 ( µ, ν ) .
An application of the triangle inequality [18, Lemma 1] shows that
[Equation]
Indeed, the triangle inequality implies | W 2 ( µ ∗N σ , ν [Equation] ) − [Equation]
+ [Equation] ∗ N σ ) and the latter two terms are of order at most σ .
In general, this upper bound is unimprovable,as we show below.
On the other hand, it can also be very loose: if µ is a translation of ν , then
W 2 ( µ ∗ N σ , ν ∗ N σ )
=
W 2 ( µ, ν ) for all σ ≥ 0 .
These examples raise a natural question: how welldoes W 2 ( µ ∗ N σ , ν ∗ N σ ) approximate W 2 ( µ, ν ) when σ is small, and how does the answer to thisquestion depend on the measures µ
and ν ?
The main goal if this paper is to give a sharp answer to this question for finitely supported measures.
We focus on the finite support case for two reasons.
First, when µ and ν are finitely supported,µ ∗ N σ and ν ∗ N σ are each finite mixtures of Gaussians, and the behavior of Wasserstein distancesfor such measures is a topic of active research [6, 10].
Second, as our results indicate, the behaviorof this quantity for finitely supported measures is unexpectedly rich, with a sharp dichotomy in ratesdepending on the structure of the optimal transport plan between µ and ν : we show that when theunique optimal transport plan between µ and ν is a perfect matching , then there exist positive σ ∗and c such that
[Equation]
In other words, for sufficiently small σ , the GOT distance approximates the standard W 2 distanceexponentially well, substantially sharpening (1).
More strikingly, we establish the existence of aphase transition: for σ < σ ∗ , the gap is exponentially small, whereas for σ > σ ∗ , the gap scaleslinearly.
By contrast, if the optimal transport plan between µ and ν is not unique or is not a perfectmatching, then no phase transition appears: the upper bound of (1) is tight even in a neighborhoodof σ = 0 .
Our work provides a precise understanding on how GOT resembles vanilla OT in the vanishing noise( σ ↓ 0 ) regime.
These results complement those recently obtained by [5] in the large noise regime,who show that if µ and ν have n matching moments, n ≥ 1 , then W 2 ( µ ∗ N σ , ν ∗ N σ ) =
O ( σ − n )as σ → ∞ .
Along with results in [5], our work completes the limiting picture of the Euclideanheat semigroup acting on atomic measures under the Wasserstein distance.
All the relevant rates arepresented in Table 1.
We note that our work leaves open the question of characterizing the rates for non-atomic measures.
It is possible to show that, for general measures, there are measures exhibiting polynomial rates intermediate between σ and e − c/σ 2 ; however, these rates appear to depend delicately on the geometryof the measures and their support.
Giving a full characterization of the rate for general probabilitymeasures is an attractive open question.
Preliminaries and main results
We are concerned with the optimal transport problem between discrete measures
[Equation]
[Table] µ and ν agree up to n th moment σ − n [5] in the space R d , equipped with the squared Euclidean cost function c ( x, y ) = ∥ x − y ∥ 2 . Here{ α i } ki =1 and { β j } ℓj =1 are positive numbers such that (cid:80) ki
=1 α i = (cid:80) ℓj =1 β j = 1 , and the sets X :=
{ x i } and Y : = { y j } consist of distinct elements of R d . (Note that we do not require that X ∩Y = ∅ .) Explicitly, we may write
[Equation]
[Table]
[Equation]
[Table]
We call a minimizer in (2) an optimal coupling .
We shall show that the behavior of the quantity W 2 ( µ ∗ N σ , ν ∗ N σ ) as σ → 0 depends strongly onthe structure of the optimal couplings between µ and ν .
We rely on the following definition.
Definition 2.1.
The measures µ
and ν possess a unique perfect matching if there exists a uniquesolution π ∗ to (2) , and if the support of this unique solution is a perfect matching, i.e., the set { x ∈
X , y ∈ Y : π ∗ ( x, y ) > 0 } is a bijection between X and Y .
Our main results (Theorems 4.1 and 5.1) show that the GOT distance approximates the Wasserstein distance exponentially well for small σ
if and only if the measures possess a unique perfectmatching.
To obtain these bounds, we show that if σ is small and µ and ν possess a unique perfectmatching, then the optimal plan for µ and ν is also approximately optimal, in an appropriate sense,for the convolved measures µ ∗ N σ and ν ∗ N σ .
By contrast, if µ and ν do not possess a uniqueperfect matching, then we explicitly exhibit an alternate coupling between µ ∗ N σ and ν ∗ N σ withsignificantly smaller cost, therefore showing that W 2 ( µ ∗ N σ , ν ∗ N σ ) is smaller than W 2 ( µ, ν ) byan amount that scales linearly in σ .
To identify the range of σ for which the exponential error bound holds, we introduce a robust versionof optimality for π ∗ in the perfect matching case.
This definition, strong cyclical monotonicity , extends the classical cyclical monotonicity criterion from optimal transport [see, e.g. 37], and captureshow sensitive the optimal plan is to perturbations of the source and target measure.
We show thatthis notion is closely related to the strong convexity and smoothness of the dual optimal solutions tothe optimal transport problem, known as “potentials.”
The strong convexity of these potentials haspreviously been explored in computational and statistical contexts [26, 36], but to our knowledge itsconnection to the stability of discrete optimal tranport plans is new.
Strong cyclical monotonicity
In this section, we consider transport plans in the form of a perfect matching between { x i } and { y i } .
By relabeling the points, we may assume without loss of generality that the optimal transport planbetween µ and ν is the matching given by
[Equation]
In this section, we develop a robust notion of optimality for Γ .
This notion is based on a strengthening of the classic optimality condition for optimal transport, based on cyclical monotonicity. Werecall the following definition.
Definition 3.1 (See, e.g., 30) .
A set S ⊆ R d × R d is cyclically monotone if for any( a 1 , b 1 ) , . . .
, ( a n , b n ) ∈ S , we have
[Equation]
where we set b n +1 :=
b 1 .
The significance of this notion is the following fundamental result.
Theorem 3.2 (See 37, Theorem 5.10) .
If π ∈ Π( µ, ν ) has cyclically monotone support, then it isan optimal transport plan between µ and ν .
Our main definition strengthens this characterization by requiring the inequalities in the definitionof cyclical monotonicity to be strict.
Definition 3.3.
We say f : [ k ] × [ k ] → R ≥ 0 is a positive residual function on [ k ] , if f ( i, i ) = 0 ,f
( i, j ) > {+
+} for i ̸
= j , and f ( i, j ) =
f ( j, i ) for all i, j ∈ [ k ] .
Definition 3.4 (Strong cyclical monotonicity) .
For a positive residual function f on [ k ] , we say thatΓ is f -strongly cyclically monotone, if for any 1 ≤ n ≤ k and distinct τ (1) , τ (2) , . . .
, τ ( n ) ∈
[ k ](with the convention τ ( n + 1) =
τ (1) ), we have
[Equation] [Equation]
+ 1)) , or equivalently,
[Equation]
Strong cyclical monotonicity indicates that the optimal plan with support Γ is superior to any otherplan by a positive margin in its transport cost. The importance of Definition 3.4 is that is equivalentto robustness of the optimality of Γ under small perturbations of the points { x i } and { y i } .
To makethis connection precise, we make the following definition.
Definition 3.5. For ϵ ≥ 0 , we say Γ is ϵ -robust, if for any distinct τ (1) , τ (2) , . . . , τ ( n ) ∈
[ k ] , andany α τ (1) , α τ (2) , . . , α τ ( n ) ∈ R d such that [Equation] , there holdsn (cid:88)
[Equation]
We write
[Equation]
The quantity R (Γ) , which we call “robustness of optimality,” captures the behavior of the optimalplan when the supports of µ and ν are slightly perturbed.
As the following proposition indicates,robustness in this sense is equivalent to strong cyclical monotonicity.
Proposition 3.6.
Γ is strongly cyclically monotone
if and only if R (Γ) > 0 .
Proposition 3.6 may be viewed as a robust analogue to Theorem 3.2—if π ∈ Π( µ, ν ) has strongcyclically monotone support, then it is a robustly optimal transport plan between µ and ν , in thesense that it remains an optimal plan even when µ and ν are corrupted by noise.
As we establish in
Section 4, this observation is central to the analysis of GOT.
3.1 Implementability and explicit bounds on R (Γ)
Despite the mathematical simplicity of Definition 3.4, it is not clear how to verify it for a particularset Γ , nor how to establish that this property holds for an optimal plan between µ and ν .
To thisend, we propose another condition, strong implementability , which is equivalent to strong cyclicalmonotonicity but is more amenable to analysis.
We also show that both conditions are equivalent toµ and ν possessing a unique perfect matching.[29] introduced the notion implementability and established it as an equivalent condition of cyclicalmonotonicity.
In parallel to the results in [29], we also introduce the following stronger condition ofimplementability.
Definition 3.7 (Strong implementability) .
For a positive residual function f on [ k ] , we say that Γ isf -strongly implementable, if there exists a potential function φ , such that for any i, j ∈ [ k ] , we have
[Equation]
Analogous to the equivalence result in [29], we show that strong cyclical monotonicity and strongimplementability are both equivalent to the uniqueness and optimality of Γ .
Proposition 3.8.
The following three statements are equivalent:
(i) Γ is f -strongly cyclically monotone for some f ; (ii) Γ is f -strongly implementable; (iii) Γ is the unique optimal transport plan from { x i } to { y i } . Remark 3.9.
We can imply from the direction (i) to (iii) of Proposition 3.8 that, if the directedbipartite graph with vertex set { x 1 , x 2 , . . .
, x k , y 1 , y 2 , . .
. , y k } and arcs between x i and y j withweight ∥ x i − y j ∥ 2 does not possess an alternating cycle of zero total cost, then Γ is the uniqueoptimal transport plan from { x i } to { y i } . One may use this sufficient condition to verify uniquenessof an optimal transport plan Γ in practice. The equivalence in Proposition 3.8 holds for any positive residual function f ; however, in the contextof optimal transport with the squared Euclidean cost, it is most natural to focus on the quadratic case. The positive residual function constructed in the equivalence between (iii) and (i) in Proposition 3.8is of the form f
( i, j ) =
[Equation] 2 ∥
y i − y j ∥ 2 for some λ > 0 , in which case the implementability conditionreads
[Equation]
Quadratic residual functions are closely connected to convex analysis and to the theory of optimaltransport.
This condition is equivalent to the existence of a λ -strongly convex potential φ satisfying∇ φ ( y i )
= x i for all i ∈
[ k ] [35], or, equivalently, the existence of a Lipschitz Brenier map from µto ν [4]. The regularity of Brenier maps is a deep question in analysis [see, e.g. 14].
Proposition 3.8establishes that, in the finite-support case, this question is equivalent to the uniqueness of the optimaltransport plan for µ and ν .
More generally, we have the following theorem characterizing the properties of strongly implementable plans with residual functions of quadratic type.
Theorem 3.10.
The following conditions are equivalent:
(i) For some 0 ≤ α < β , there exists a potential function φ : R d → R d which is α -strongly convex and β -smooth, such that x i = ∇ φ ( y i ) for all i ∈ [ k ] .
(ii) Γ is strongly implementable for
[Equation]
or equivalently, there exists { ˜ φ
( y i ) } ki =1 ⊂ R d , such that for all i, j ∈ [ k ] ( i ̸ = j ),
[Equation]
Proof.
This is a direct application of Theorem 4 in [35].
The condition (i) in Theorem 3.10 is equivalent to the set { ( y i , x i , φ ( y i )) } being F α,β -interpolable in Definition 2 of [35], and the condition(ii) is equivalent to equation (4) in Theorem 4 of [35].
Theorem 3.10 also formally encompasses the choice β = + ∞ when φ is strongly convex but notsmooth, in which case (3) reads
[Equation]
which recovers the positive residual function used in the proof of Proposition 3.6.
Remark 3.11.
We should emphasize that the f defined in Theorem 3.10 is indeed a positive residualfunction given α < β , since Cauchy-Schwartz gives
[Equation]
As a direct consequence of the direction (ii) to (i) in Theorem 3.10, if Γ is strongly implementablefor a positive residual function f which is quadratic in y i − y j and x i − x j , there exists a smoothand strongly convex potential function verifying implementability.
Corollary 3.12.
Suppose Γ is strongly implementable for
[Equation]
where λ xx , λ xy and λ yy are nonnegative numbers which satisfy λ 2 xy + λ xy
= λ
xx λ yy .
Then thereexists a potential function φ : R d → R d which is λ
xy λ xx -strongly convex
and λ yy λ xy -smooth, such thatx i = ∇ φ ( y i ) for all i ∈ [ k ] .
By the Smith–Knott optimality criterion for optimal transport [34], the conclusion that x i = ∇ φ ( y i )for all i ∈ [ k ] is equivalent to the potential function φ solving the following dual version of (2):
[Equation]
(6) where ϕ ∗ denotes the Legendre conjugate.
Finally, we show that the above characterizations give rise to lower bounds on R (Γ) , which are easyto compute in O ( k 2 ) time. The following bound gives a quantitative link between robustness ofoptimality and the residual function f .
Proposition 3.13.
Suppose Γ is strongly implementable for a positive residual function f .
Then Γis ϵ -robust for
[Equation]
This implies that
[Equation]
By combining this bound with Theorem 3.10, we obtain a simple lower bound when f is of quadratictype.
Proposition 3.14.
When the equivalence in Theorem 3.10 holds, Γ is ϵ -robust for
[Equation]
This implies that
[Equation]
Remark 3.15.
When condition (i) in Theorem 3.10 holds, α -strong convexity and β -smoothnessimplies
[Equation]
Thus the condition (8) may be replaced by the bound
[Equation]
which is more interpretable.
We should emphasize that, to check the ϵ -robustness of T with either (8)or (9) requires prior knowledge on the parameters α and β , which are inherent to the optimaltransport plan Γ .
Exponential rates for unique perfect matchings
Our main results show that the robustness of optimality R (Γ) controls the gap between W 2 ( µ, ν )and W 2 ( µ ∗ N σ , ν ∗ N σ ) .
Theorem 4.1.
If σ ∗ =
R (Γ) > 0 , then for σ ∈ (0 , σ ∗ ) , [Equation]
The proof depends on the following simple lemma, which shows that robustness to optimality implies that the Wasserstein distance is unchanged when µ and ν are corrupted by small noise.
Lemma 4.2.
If σ ∗ =
R (Γ) > 0 , then for any measure Q in R d supported on B (0 , σ ∗ ) ,
[Equation]
Proofs of these results appear in the supplementary material. In the regime where σ does not exceed R (Γ) , the above theorem tells that the GOT distance is anexcellent approximation of the OT distance.
Our second main result is a converse to that statement,showing that if σ goes beyond R (Γ) , the loss W 2 ( µ, ν ) − W 2 ( µ ∗ N σ , ν ∗ N σ ) is bounded belowby a linear function of σ .
We start with the following proposition, which quantifies a “violation ofcyclical monotonicity” under possibly large perturbations in the sources and targets.
Proposition 4.3.
If Γ is an optimal transport plan, for any M ≥ 0 , denote
G ( M ) :
[Equation]
Then G ( M ) is a concave function of M for M ∈ [0 , + ∞ ) .
Note that G ( M ) vanishes for M < σ ∗ .
The next theorem shows that as long as G ( M ) is notnegligible for M ≳ σ ∗ , the approximation loss for σ ≥ σ ∗ is linear in σ .
Theorem 4.4.
If σ ∗ =
R (Γ) > 0 , then
[Equation]
Here G ( M ) is defined as in Proposition 4.3.
In particular, if G (3 σ ∗ ) ≥ c 0 σ ∗ for an absoluteconstant c 0 > 0 , then there exists a constant C = C ( c 0 ) > 0 such that for σ ∈ (0 , 2 σ ∗ ) ,
[Equation]
The proof of Theorem 4.1, Lemma 4.2, Proposition 4.3, and Theorem 4.4 can be found in thesupplementary material.
Beyond perfect matchings
In the case that R (Γ) = 0 , or equivalently by Proposition 3.8 and Proposition 3.6 that the optimaltransport map between µ and ν is not a perfect matching, Theorems 4.1 and 4.4 are not applicable.
In this situation, we are able to show that the approximation error is linear, even in a neighborhoodof zero.
In fact, this holds whenever there exists an optimal transport plan between µ and ν which isnot a perfect matching.
To analyze this case, we return to the setting of general discrete measures:
[Equation]
(cid:88)
[Equation]
Theorem 5.1.
Let µ and ν be as in (10) . If µ and ν do not possess a unique perfect matching, thenthere exists c 0 > 0 such that for σ ∈ (0 , c 0 ) ,
[Equation]
[Equation]
Theorem 5.1 tells that, unless the optimal transport plan between µ and ν is unique and a perfectmatching, the loss from approximating the OT distance with the GOT distance is at least linear inσ . We derive Theorem 5.1 from the following lemma, which shows that Theorem 5.1 holds in thespecial case where µ is a single point mass, and ν is uniform on two points.
Lemma 5.2.
Let x, y 1 and y 2 be different points in R d .
For µ 0 :
=
δ ( x ) and ν 0 :
= 12 δ ( y 1 )+ 12 δ
( y 2 ) ,there exists c 0 > 0 , such that for σ ∈ (0 , c 0 ) , we have
[Equation]
We obtain the full strength of Theorem 5.1 by reducing to the special case of Lemma 5.2 on aparticular subset of the support of µ and ν .
Full details appear in the supplementary material.
Numerical example
In this section, we present a numerical example to demonstrate different regimes of the rate
W 2 ( µ, ν ) − W 2 ( µ ∗ N σ , ν ∗ N σ ) , in respect of Theorem 4.1 and Theorem 4.4.
For the sake ofclarity, we consider atomic measures µ and ν both defined on R 2 .
One of the simplest cases wherea coupling Γ has R (Γ) = 0 is
[Equation]
It is easy to see that the optimal transport plan from µ to ν is not unique, which is also a consequenceof Proposition 3.8, Proposition 3.6 and the fact that R (Γ)
=
for the map
[Equation]
= { (( [Equation] 1 , − 1) , ( − 1 , 1)) , ((1 , 1) , (1 , − 1)) } that achieves the optimal cost.
We also consider the family of perturbed measures
[Equation]
The source and target distributions corresponding to p = k/ 10 for k = 1 , 2 , 3 , 4 are depicted in Figure 1.
For each k , the unique optimal transport plan from µ k = µ ( k/ 10) to ν is given by
[Equation]
For each of these GOT tasks, we draw 200 samples from the source distribution µ k ∗ N σand target distribution ν ∗ N σ , and use the empirical W 2 distance as an estimate of the true W 2 ( µ k ∗ N σ , ν ∗ N σ ) .
We repeat the process 100 times and report the mean, as shown in thefollowing figures.
By Theorem 4.1 and Theorem 4.4, we expect W 22 ( µ k , ν )
− W 22 ( µ k ∗ N σ , ν ∗ N σ ) to be of scalee − c/σ 2 for σ ∈ (0 , R (Γ k )) , and W 22 ( µ k , ν ) − W 22 ( µ k ∗ N σ , ν ∗ N σ ) ≳ σ for σ ≥ R (Γ k ) .
Thistransition from exponential to linear is visible in Figure 2.
Using Proposition 3.14, we obtain a lower bound on R (Γ k ) , which we plot with a vertical dashed linein Figure 2. Exponential decay is visible to the left of the dashed lines, as anticipated. Figure 3 showsthis behavior more clearly on a logarithmic scale, where we observe that log( − log( W 2 ( µ, ν ) − W 2 ( µ ∗ N σ , ν ∗ N σ ))) is linear in log( σ ) for small σ .
This paper develops approximation results for Gaussian-smoothed optimal transport, showing that GOT approximates the Wasserstein distance exponentially well for small σ when µ and ν possess aunique perfect matching. By contrast, if µ and ν do not possess a unique perfect matching, then thegap between W 2 ( µ, ν ) and W 2 ( µ ∗ N σ , ν ∗ N σ ) is linear in σ as σ → 0 .
The difference betweenthese two behaviors can be traced to the fact that if µ and ν possess a unique perfect matching, then
[Equation]
the optimal transport plan between them is stable under small perturbations of µ and ν .
In particular,for noise distributions Q with sufficiently small support, W 2 ( µ, ν ) =
W 2 ( µ ∗ Q, ν ∗ Q ) .
On theother hand, if µ and ν do not possess a unique perfect matching, then their optimal plan is not stable,and can change even under infinitesimal perturbation.
Our techniques are based on a new notion of stability for discrete optimal transport plans, which wecall robustness of optimality and characterize by developing strong variants of the classic notions ofcyclical monotonicity and implementability for optimal transport plans.
Just as cyclical monotonicity is closely connected to the convexity of the dual potentials for the optimal transport problem, weshow that strong cyclical montonicity is equivalent to strong convexity.
This characterization showsthat a discrete optimal transport plan is robust to perturbations of the source and target measuresif and only if its support lies in the subdifferential of a strongly convex function.
We anticipatethat this characterization will have statistical and computational implications for the estimation ofoptimal transport plans between discrete measures.