Title: Fast Optimal Transport through Sliced Generalized Wasserstein Geodesics
Venue: NeurIPS 2023
Accepted: yes
Keywords: Optimal Transport

Wasserstein distance

Generalized Geodesics

Sliced Wasserstein
Abstract: Wasserstein distance (WD) and the associated optimal transport plan have been proven useful in many applications where probability measures are at stake. In this paper, we propose a new proxy of the squared WD, coined $\textnormal{min-SWGG}$, that is based on the transport map induced by an optimal one-dimensional projection of the two input distributions. We draw connections between  $\textnormal{min-SWGG}$, and Wasserstein generalized geodesics in which the pivot measure is supported on a line. We notably provide a new closed form for the exact Wasserstein distance in the particular case of one of the distributions supported on a line allowing us to derive a fast computational scheme that is amenable to gradient descent optimization. We show that  $\textnormal{min-SWGG}$, is an upper bound of WD and that it has a complexity similar to as Sliced-Wasserstein, with the additional feature of providing an associated transport plan. We also investigate some theoretical properties such as metricity, weak convergence, computational and topological properties. Empirical evidences support the benefits of  $\textnormal{min-SWGG}$, in various contexts, from gradient flows, shape matching and image colorization, among others.

Fast Optimal Transport through Sliced Wasserstein Generalized Geodesics

Abstract

Wasserstein distance (WD) and the associated optimal transport plan have proven useful in many applications where probability measures are at stake. In this paper, we propose a new proxy for the squared WD, coined min-SWGG, which relies on the transport map induced by an optimal one-dimensional projection of the two input distributions. We draw connections between min-SWGG and Wasserstein generalized geodesics with a pivot measure supported on a line. We notably provide a new closed form of the Wasserstein distance in the particular case where one of the distributions is supported on a line, allowing us to derive a fast computational scheme that is amenable to gradient descent optimization. We show that min-SWGG is an upper bound of WD and that it has a complexity similar to that of Sliced-Wasserstein, with the additional feature of providing an associated transport plan. We also investigate some theoretical properties such as metricity, weak convergence, computational and topological properties. Empirical evidences support the benefits of min-SWGG in various contexts, from gradient flows, shape matching and image colorization, among others.

1 Introduction

Gaspard Monge, in his seminal work on Optimal Transport (OT) [42], studied the following problem: how to move with minimum cost the probability mass of a source measure to a target one, for a given transfer cost function? At the heart of OT is the optimal map that describes the optimal displacement as the Monge problem can be reformulated as an assignment problem. It has been relaxed by [33] by finding a plan that describes the amount of mass moving from the source to the target. Beyond this optimal plan, an interest of OT is that it defines a distance between probability measures: the Wasserstein distance (WD).

Recently, OT has been successfully employed in a wide range of machine learning applications, in which the Wasserstein distance is estimated from the data, such as supervised learning [30], natural

language processing [38] or generative modelling [5]. Its capacity to provide meaningful distances between empirical distributions is at the core of distance-based algorithms such as kernel-based methods [60] or k -nearest neighbors [6]. The optimal transport plan has also been used successfully in many applications where a matching between empirical samples is sought such as color transfer [55], domain adaptation [19] and positive-unlabeled learning [15].

Solving the OT problem is computationally intensive; the most common algorithmic tools to solve the discrete OT problem are borrowed from combinatorial optimization and linear programming, leading to a cubic complexity with the number of samples that prevents its use in large scale applications [53]. To reduce the computation burden, regularizing the OT problem with e.g. an entropic term has led to solvers with a quadratic complexity [23]. Other methods based on the existence of a closed form of OT have also been devised to efficiently compute a proxy for WD, as outlined below.

Projections-based OT. The Sliced-Wasserstein distance (SWD) [56, 10] leverages 1D-projections of distributions to provide a lower approximation of the Wasserstein distance, relying on the closed form of OT for 1D probability distributions. Computation of SWD leads to a linearithmic time complexity. While SWD averages WDs computed over several 1D projections, max-SWD [24] keeps only the most informative projection. These frameworks provide efficient algorithms that can handle millions of samples and have similar topological properties as WD [45]. Other works restrain SWD and max-SWD to projections onto low dimensional subspaces [52, 40] to provide more robust estimation of those OT metrics. Although effective as proxies for WD, those methods do not provide a transport plan in the original space R d . To overcome this limitation, [44] aims to compute transport plans in a subspace which are extrapolated to the original space.

Pivot measure-based OT. Other research works rely on a pivot, yet intermediate measure. They decompose the OT metric into Wasserstein distances between each input measure and the considered pivot measure. They exhibit better properties such as statistical sample complexity or computational efficiency [29, 65]. Even though the OT problems are split, they are still expensive when dealing with large sample size distributions, notably when only two distributions are involved.

Contributions. We introduce a new proxy for the squared WD that exploits the principles of aforementioned approximations of OT metric. The original idea is to rely on projections and onedimensional assignment of the projected distributions to compute the new proxy. The approach is well-grounded as it hinges on the notion of Wasserstein generalized geodesics [4] with pivot measure supported on a line. The main features of the method are as: i) its computational complexity is on par with SW, ii) it provides an optimal transport plan through the 1D assignment problem, iii) it acts as an upper bound of WD, and iv) is amenable to optimization to find the optimal pivot measure. As an additional contribution, we establish a closed form of the WD when an input measure is supported on a line.

Outline. Section 2 presents some background of OT. Section 3 formulates our new WD proxy, provides some of its topological properties and a numerical computation scheme. Section 4 builds upon the concept of Wasserstein generalized geodesics to reformulate our OT metric approximation as the Sliced Wasserstein Generalized Geodesics (SWGG) along its optimal variant coined min-SWGG. This reformulation allows deriving additional topological properties and an optimization scheme. Finally, Section 5 provides experimental evaluations.

Notations. Let ⟨· , ·⟩ be the Euclidean inner product on R d and let S d -1 = { u ∈ R d s.t. ∥ u ∥ 2 = 1 } , the unit sphere. We denote P ( R d ) the set of probability measures on R d endowed with the σ -algebra of Borel set and P 2 ( R d ) ⊂ P ( R d ) those with finite second-order moment i.e. P 2 ( R d ) = { µ ∈ P ( R d ) s.t. ∫ R d ∥ x ∥ 2 2 dµ ( x ) < ∞} . Let P n 2 ( R d ) be the subspace of P 2 ( R d ) defined by empirical measures with n -atoms and uniform masses. For any measurable function f : R d → R d , we denote f # its push forward, namely for µ ∈ P 2 ( R d ) and for any measurable set A ∈ R d , f # µ ( A ) = µ ( f -1 ( A )) , with f -1 ( A ) = { x ∈ R d s.t. f ( x ) ∈ A } .

2 Background on Optimal Transport

Definition 2.1 (Wasserstein distance) . The squared WD [63] between µ 1 , µ 2 ∈ P 2 ( R d ) is defined as:

with Π( µ 1 , µ 2 ) = { π ∈ P 2 ( R d × R d ) s.t. π ( A × R d ) = µ 1 ( A ) and π ( R d × A ) = µ 2 ( A ) , ∀ A measurable set of R d } .

Endowed with the WD, the space P 2 ( R d ) is a geodesic space. Indeed, since there exists a Monge map T 1 → 2 between µ 1 and µ 2 , one can define a geodesic curve µ 1 → 2 : [0 , 1] →P 2 ( R d ) [31] as:

The arg min of Eq. (1) is referred to as the optimal transport plan. Denoted π ∗ , it expresses how to move the probability mass from µ 1 to µ 2 with minimum cost. In some cases, π ∗ is of the form ( Id, T ) # µ 1 for a measurable map T : R d → R d , i.e. there is no mass splitting during the transport. This map is called a Monge map and is denoted T µ 1 → µ 2 (or shortly T 1 → 2 ). Thus, one has W 2 2 ( µ 1 , µ 2 ) = inf T s.t. T # µ 1 = µ 2 ∫ R d ∥ x -T ( x ) ∥ 2 2 dµ 1 ( x ) . This occurs, for instance, when µ 1 has a density w.r.t. the Lebesgue measure [12] or when µ 1 and µ 2 are in P n 2 ( R d ) [58].

which represents the shortest path w.r.t. Wasserstein distance in P 2 ( R d ) between µ 1 and µ 2 . The Wasserstein mean between µ 1 and µ 2 corresponds to t = 0 . 5 and we simply write µ 1 → 2 .

This notion of geodesic allows the study of the curvature of the Wasserstein space [1]. Indeed, the Wasserstein space is of positive curvature [51], i.e. it respects the following inequality:

for all pivot measures ν ∈ P 2 ( R d ) .

Solving and approximating Optimal Transport. The Wasserstein distance between empirical measures µ 1 , µ 2 with n -atoms can be computed in O ( n 3 log n ) , preventing from the use of OT for large scale applications [11]. Several algorithms have been proposed to lower this complexity, for example the Sinkhorn algorithm [23] that provides an approximation in near O ( n 2 ) complexity [2].

Notably, when µ 1 = 1 n ∑ n i =1 δ x i and µ 2 = 1 n ∑ n i =1 δ y i are 1D distributions, computing the WD can be done by matching the sorted empirical samples, leading to an overall complexity of O ( n log n ) . More precisely, let σ and τ two permutation operators s.t. x σ (1) ≤ x σ (2) ≤ ... ≤ x σ ( n ) and y τ (1) ≤ y τ (2) ≤ ... ≤ y τ ( n ) . Then, the 1D Wasserstein distance is given by:

Sliced WD. The Sliced-Wasserstein distance (SWD) [56] aims to scale up the computation of OT by leveraging the closed form expression (4) of the Wasserstein distance for 1D distributions. It is defined as the expectation of 1D-WD computed along projection directions θ ∈ S d -1 over the unit sphere:

where P θ # µ 1 and P θ # µ 2 are projections onto the direction θ ∈ S d -1 with P θ : R d → R , x ↦→⟨ x , θ ⟩ and where ω is the uniform distribution over S d -1 .

Since the integral in Eq. (5) is intractable, one resorts, in practice, to Monte-Carlo estimation to approximate the SWD.

Its computation only involves projections and permutations. For L directions, the computational complexity is O ( dLn + Ln log n ) and the memory complexity is O ( Ld + Ln ) . However, in high dimension, several projections are necessary to approximate accurately the SWD and many projections lead to 1D-WD close to 0. This issue is well known in the SW community [68], where different ways of performing effective sampling have been proposed [49, 46, 50] such as distributional or hierarchical slicing. In particular, this motivates the definition of max-Sliced-Wasserstein [24] which keeps only the most informative slice:

While being a non convex problem, it can be optimized efficiently using a gradient ascent scheme.

The SW-like distances are attractive since they are fast to compute and enjoy theoretical properties: they are proper metrics and metricize the weak convergence. However, they do not provide an OT plan.

Projected WD. Another quantity of interest based on the 1D-WD is the projected Wasserstein distance (PWD) [57]. It leverages the permutations of the projected distributions in 1D in order to derive couplings between the original distributions.

Let µ 1 = 1 n ∑ n i =1 δ x i and µ 2 = 1 n ∑ n i =1 δ y i in P n 2 ( R d ) . The PWD is defined as:

where σ θ , τ θ are the permutations obtained by sorting P θ # µ 1 and P θ # µ 2 .

As some permutations are not optimal, we straightforwardly have W 2 2 ≤ PWD 2 2 . Note that some permutations can appear highly irrelevant in the original space, leading to an overestimation of W 2 2 (typically when the distributions are multi-modal or with support lying in a low dimensional manifold, see Supp. 7.1 for a discussion).

In this paper, we restrict ourselves to empirical distributions with the same number of samples. They are defined as µ 1 = 1 n ∑ n i =1 δ x i and µ 2 = 1 n ∑ n i =1 δ y i in P n 2 ( R d ) . Note that the results presented therein can be extended to any discrete measures by mainly using quantile functions instead of permutations and transport plans instead of transport maps (see Supp. 7.2).

3 Definition and properties of min-SWGG

The fact that the PWD overestimates W 2 2 motivates the introduction of our new loss function coined min-SWGG which keeps only the most informative permutation. Afterwards, we derive a property of distance and grant an estimation of min-SWGG via random search of the directions.

Definition 3.1 (SWGG and min-SWGG) . Let µ 1 , µ 2 ∈ P n 2 ( R d ) and θ ∈ S d -1 . Denote by σ θ and τ θ the permutations obtained by sorting the 1D projections P θ # µ 1 and P θ # µ 2 . We define respectively SWGG and min-SWGG as:

One shall remark that the function SWGG corresponds to the building block of PWD in eq. (7).

One main feature of min-SWGG is that it comes with a transport map. Let θ ∗ ∈ argmin SWGG 2 2 ( µ 1 , µ 2 , θ ) be the optimal projection direction. The associated transport map is:

In Supp. 7.6 we give several examples of such transport plan. These examples show that the overall structure of the optimal transport plan is respected by the transport plan obtained via min-SWGG.

We now give some theoretical properties of the quantities min-SWGG and SWGG. Their proofs are given in Supp. 7.3.

Proposition 3.2 (Distance and Upper bound) . Let θ ∈ S d -1 . SWGG 2 ( · , · , θ ) defines a distance on P n 2 ( R d ) . Moreover, min-SWGG is an upper bound of W 2 2 , and W 2 2 ≤ min-SWGG 2 2 ≤ PWD 2 2 , with equality between W 2 2 and min-SWGG 2 2 when d > 2 n .

Remark 3.3. Similarly to max-SW, min-SWGG retains only one optimal direction θ ∗ ∈ S d -1 . However, the two distances strongly differ: i) min-SWGG is an upper bound and max-SW a lower bound of W 2 2 , ii) the optimal θ ∗ may differ (see Supp. 7.4 for an illustration), and iii) max-SW does not provide a transport plan between µ 1 and µ 2 .

Solving Eq. (9) can be achieved using a random search, by sampling L directions θ ∈ S d -1 and keeping only the one leading to the lowest value of SWGG.

This gives an overall computational complexity of O ( Ldn + Ln log n ) and a memory complexity of O ( dn ) . In low dimension, the random search estimation is effective: covering all possible

permutations through S d -1 can be done with a low number of directions. In high dimension, many more directions θ are needed to have a relevant approximation, typically O ( L d -1 ) . This motivates the design of gradient descent techniques for finding θ ∗ .

4 SWGGas minimizing along the Wasserstein generalized geodesics

Solving problem in Eq. (9) amounts to optimize over a set of admissible permutations. This problem is hard since SWGG is non convex w.r.t. θ and piecewise constant, thus not differentiable over S d -1 . Indeed, as long as the permutations remain the same for different directions θ , the value of SWGG remains constant. When the permutations change, the objective SWGG "jumps" as illustrated in Fig. 1.

In this section, we tackle this problem by providing an alternative formulation of min-SWGG that allows smoothing the different kinks of

µ

1

Q

θ

1

#

µ

2

µ

2

Q

θ

2

#

µ

1

Q

θ

1

#

µ

1

Q

θ

2

#

µ

2

0

π

/2

π

3

π

/2

0

7

SWGG value

SWGG, hence, making min-SWGG amenable to optimization. This formulation relies on Wasserstein generalized geodesics we introduce hereinafter.

We show that this alternative formulation brings in computational advantages and allows establishing some additional topological properties and deriving an efficient optimization scheme. We also provide a new closed form expression of the Wasserstein distance W 2 2 ( µ 1 , µ 2 ) when either µ 1 or µ 2 is supported on a line.

4.1 SWGGbased on Wasserstein Generalized Geodesics

Wasserstein generalized geodesics (see Supp. 8 for more details) were first introduced in [4] in order to ensure the convergence of Euler scheme for Wasserstein Gradient Flows. This concept has been used notably in [29, 44] to speed up some computations and to derive some theoretical properties. Generalized geodesic is also highly related with the idea of linearization of the Wasserstein distance via an L 2 space [65, 43], see Supp. 9 for more details on the related works.

Generalized geodesics lay down on a pivot measure ν ∈ P n 2 ( R d ) to transport the distribution µ 1 toward µ 2 . Indeed, one can leverage the optimal transport maps T ν → µ 1 and T ν → µ 2 to construct a curve t ↦→ µ 1 → 2 g ( t ) linking µ 1 to µ 2 as

The related generalized Wasserstein mean corresponds to t = 0 . 5 and is denoted µ 1 → 2 g .

Intuitively, the optimal transport maps between ν and µ i , i = 1 , 2 give rise to a sub-optimal transport map between µ 1 and µ 2 :

One can be interested in the cost induced by the transportation of µ 1 to µ 2 via the transport map T 1 → 2 ν , known as the ν -based Wasserstein distance [47] and defined as

Notably, the second part of Eq. (13) straddles the square Wasserstein distance with Eq. (3). Remarkably, the computation of W 2 ν can be efficient if the pivot measure ν is chosen appropriately. As established in Lemma 4.6, it is the case when ν is supported on a line. Based on these facts, we propose hereafter an alternative formulation of SWGG.

Definition 4.1 (Pivot measure) . Let µ 1 and µ 2 ∈ P n 2 ( R d ) . We restrict the pivot measure ν to be the Wasserstein mean of the measures Q θ # µ 1 and Q θ # µ 2 :

where θ ∈ S d -1 and Q θ : R d → R d , x ↦→ θ ⟨ x , θ ⟩ is the projection onto the subspace generated by θ . Moreover µ 1 → 2 θ is always defined as the middle of a geodesic as in Eq (2).

One shall notice that Q θ # µ 1 and Q θ # µ 2 are supported on the line defined by the direction θ , so is the pivot measure ν = µ 1 → 2 θ . We are now ready to reformulate the metric SWGG.

Proposition 4.2 (SWGGbased on generalized geodesics) . Let θ ∈ S d -1 , µ 1 , µ 2 ∈ P n 2 ( R d ) and µ 1 → 2 θ be the pivot measure. Let µ 1 → 2 g,θ be the generalized Wasserstein mean between µ 1 and µ 2 ∈ P n 2 ( R d ) with pivot measure µ 1 → 2 θ . Then,

The proof is in Supp.10.1. From Proposition 4.2, SWGG is the µ 1 → 2 θ -based Wasserstein distance between µ 1 and µ 2 . This alternative formulation allows establishing additional properties of min-SWGG.

4.2 Theoretical properties

Additionally to the properties derived in Section 3 (SWGG is a distance and min-SWGG is an upper bound of W 2 2 ), we provide below other theoretical guarantees.

Proposition 4.3 (Weak Convergence) . min-SWGG metricizes the weak convergence in P n 2 ( R d ) . In other words, let ( µ k ) k ∈ N be a sequence of measures in P n 2 ( R d ) and µ ∈ P n 2 ( R d ) . We have:

where L , 2 -→ stands for the weak convergence of measure i.e. ∫ R d fdµ k → ∫ R d fdµ for all continuous bounded functions f .

Beyond the weak convergence, min-SWGG possesses the translation property, i.e. the translations can be factored out as the Wasserstein distance does (see [53, remark 2.19] for a recall).

Proposition 4.4 (Translation) . Let T u (resp. T v ) be the map x ↦→ x -u (resp. x ↦→ x -v ), with u , v vectors of R d . We have:

This property is useful in some applications such as shape matching, in which translation invariances are sought.

where m 1 = ∫ R d x dµ 1 ( x ) and m 2 = ∫ R d x dµ 2 ( x ) are the means of µ 1 , µ 2 .

The proofs of the two Propositions are deferred to Supp. 10.2 and 10.3.

Remark 4.5 (Equality) . min-SWGG and W 2 2 are equal in different cases. First, [43] showed that it is the case whenever µ 1 is the shift and scaling of µ 2 (see Supp. 9.1 for a full discussion). In Lemma 4.6, we will state that it is also the case if one of the two distributions is supported on a line.

4.3 Efficient computation of SWGG

SWGG defined in Eq. (14) involves computing three WDs that are fast to compute, with an overall O ( dn + n log n ) complexity, as detailed below. Building on this result, we provide an optimization scheme that allows optimizing over θ with O ( sdn + sn log sn ) operations at each iteration, with s a (small) integer. We first start by giving a new closed form expression of the WD whenever one distribution is supported on a line, that proves useful for deriving an efficient computation scheme.

New closed form of the WD. The following lemma states that W 2 2 ( µ 1 , µ 2 ) admits a closed form whenever µ 2 is supported on a line.

This lemma leverages the computation of the WD between µ 2 and the orthogonal projection of µ 1 onto the linear subspace defined by the line. Additionally, it provides an explicit formulation for the optimal transport map T 1 → 2 .

Lemma 4.6. Let µ 1 , µ 2 in P n 2 ( R d ) with µ 2 supported on a line of direction θ ∈ S d -1 . We have:

with Q θ as in Def. 4.1. Note that W 2 2 ( µ 1 , Q θ # µ 1 ) = 1 n ∑ ∥ x i -Q θ ( x i ) ∥ 2 2 and W 2 2 ( Q θ # µ 1 , µ 2 ) = W 2 2 ( P θ # µ 1 , P θ # µ 2 ) are the WD between 1D distributions. Additionally, the optimal transport map is given by T 1 → 2 = T Q θ # µ 1 → µ 2 ◦ T µ 1 → Q θ # µ 1 = T Q θ # µ 1 → µ 2 ◦ Q θ . In particular, the map T 1 → 2 can be obtained via the permutations of the 1D distributions P θ # µ 1 and P θ # µ 2 . The proof is provided in Supp. 10.4.

Efficient computation of SWGG . Eq. (14) is defined as the Wasserstein distance between a distribution (either µ 1 or µ 2 or µ 1 → 2 g,θ ) and a distribution supported on a line ( µ 1 → 2 θ ). As detailed in Supp. 10.5, computation of Eq. (14) involves three Wasserstein distances between distributions and their projections: i) W 2 2 ( µ 1 , Q θ # µ 1 ) , ii) W 2 2 ( µ 2 , Q θ # µ 2 ) , iii) W 2 2 ( µ 1 → 2 g,θ , µ 1 → 2 θ ) , and a one dimensional Wasserstein distance W 2 2 ( P θ # µ 1 , P θ # µ 2 ) , resulting in a O ( dn + n log n ) complexity.

Optimization scheme for min-SWGG. The term W 2 2 ( µ 1 → 2 g,θ , µ 1 → 2 θ ) in Eq. (14) is not continuous w.r.t. θ . Indeed, the generalized mean µ 1 → 2 g,θ depends only on the transport maps T µ 1 → 2 θ → µ 1 and T µ 1 → 2 θ → µ 2 , which remain constant as long as different projection directions θ lead to the same permutations σ θ and τ θ . Hence, we rely on a smooth surrogate ˜ µ 1 → 2 g,θ of the generalized mean and we aim to minimize the following objective function:

To define ˜ µ 1 → 2 g,θ , one option would be to use entropic maps in Eq. (11) but at the price of a quadratic time complexity. We rather build upon the blurred Wasserstein distance [26] to define ˜ µ 1 → 2 g,θ as it can be seen as an efficient surrogate of entropic transport plans in 1D. In one dimensional setting, ˜ µ 1 → 2 g,θ can be approximated efficiently by adding an empirical Gaussian noise followed by a sorting pass. In our case, it resorts in making s copies of each sorted projection P θ ( x σ ( i ) ) and P θ ( y τ ( i ) ) respectively, to add an empirical Gaussian noise of deviation √ ϵ/ 2 and to compute averages of sorted blurred copies x s σ s , y s τ s . We finally have ( ˜ µ 1 → 2 g,θ ) i = 1 2 s ∑ is k =( i -1) s +1 x s σ s ( k ) + y s τ s ( k ) . [26] showed that this blurred WD has the same asymptotic properties as the Sinkhorn divergence.

The surrogate ˜ SWGG ( µ 1 , µ 2 , θ ) is smoother w.r.t. θ and can thus be optimized using gradient descent, converging towards a local minima. Once the optimal direction θ ∗ is found, min-SWGG resorts to be the solution provided by SWGG ( µ 1 , µ 2 , θ ∗ ) . Fig. 2 illustrates the effect of the smoothing on a toy example and more details are given in Supp. 10.6. The computation of ˜ SWGG ( µ 1 , µ 2 , θ ) is summarized in Alg. 1.

θ

1

θ

2

µ

1

µ

2

µ

1

→

2

g,θ

1

µ

1

→

2

g,θ

2

µ

1

→

2

g,θk

0

π

/2

π

3

π

/2

5

10

SWGG

2

SWGG

2

5 Experiments

Wehighlight that min-SWGG is fast to compute, gives an approximation of the WD and the associated transport plan. We start by comparing the random search and the gradient descent schemes for finding the optimal direction in subsection 5.1. Subsection 5.2 illustrates the weak convergence property of min-SWGG through a gradient flow application to match distributions. We then implement an efficient algorithm for colorization of gray scale images in 5.3, thanks to the new closed form expression of the WD. We finally evaluate min-SWGG in a shape matching context in subsection 5.4. When possible from the context, we compare min-SWGG with the main methods for approximating the WD namely SW, max-SW, Sinkhorn [23], factored coupling [29] and subspace robust WD (SRW) [52]. Supp. 11 provides additional results on the behavior of min-SWGG and experiments on other tasks such as color transfer or on data sets distance computation. All the code is available at 1

5.1 Computing min-SWGG

Let consider Gaussian distributions in dimensions d ∈ { 2 , 20 , 200 } . We first sample n = 1000 points from each distribution to define µ 1 and µ 2 . We then compute min-SWGG 2 2 ( µ 1 , µ 2 ) computed using different schemes, either by random search, by simulated annealing [54] or by gradient descent. We report the obtained results in Fig. 3 (left). For the random search scheme, we repeat each experiment 20 times and we plot the average value of min-SWGG ± 2 times the standard deviation.

For the gradient descent, we select a random initial θ . We observe that, in low dimension, all schemes provide similar values of min-SWGG. When the dimension increases, optimizing the direction θ yields a more accurate approximation of the true Wasserstein distance (see plots' title in Fig. 3). On Fig. 3 (right), we compare the empirical runtime evaluation for min-SWGG with different competitors for d = 3 and using n samples from Gaussian distributions, with n ∈ { 10 2 , 10 3 , 10 4 , 5 × 10 4 , 10 5 } . We observe that, as expected, min-SWGG with random search is as fast as SW with a super linear time complexity. With the optimization process, it is faster than SRW for a given number of samples. Wealso note that SRW is more demanding in memory and hence does not scale as well as min-SWGG. We give more details on this experimentation and a comparison with competitors in Supp. 11.2.

10

1000

33

.

0

33

.

5

34

.

0

34

.

5

35

.

0

Distance

W

2

2

=

32.4

10

1000

376

378

380

382

W

2

2

=

346.1

random search

optimized

annealing

10

1000

3945

3950

3955

3960

3965

W

2

2

=

3836.0

10

2

10

3

10

4

10

5

10

-

3

10

-

2

10

-

1

10

0

10

1

10

2

10

3

Seconds

SW,

L

=

200

min-SWGG,

L

=

200

min-SWGG optim

Factored Coupling

Sinkhorn

max-SW

SRW

Number of projections

Number of samples in each distribution

2

2

5.2 Gradient Flows

We highlight the weak convergence property of min-SWGG. Initiating from a random initial distribution, we aim to move the particles of a source distribution µ 1 towards a target one µ 2 by reducing the objective min-SWGG 2 2 ( µ 1 , µ 2 ) at each step. We compare both variants of min-SWGG against SW, max-SW and PWD, relying on the code provided in [37] for running the experiment; we report the results on Fig. 4. We consider several target distributions, representing diverse scenarios and fix n = 100 . We run each experiment 10 times and report the mean ± the standard deviation. In every case, one can see that µ 1 moves towards µ 2 and that all methods tend to have similar behavior. One can notice though that, for the distributions in d = 500 dimensional space, min-SWGG computed with the optimization scheme leads to the best alignment of the distributions.

0

250

500

750

1000

1250

1500

1750

2000

Number of iterations

-

4

-

3

-

2

-

1

0

Log

10

(

W

2

)

SW

max-SW

min-SWGG (random search)

min-SWGG (optim)

PWD

0

250

500

750

1000

1250

1500

1750

2000

-

4

-

3

-

2

-

1

0

0

250

500

750

1000

1250

1500

1750

2000

Number of iterations

-

3

-

2

-

1

0

Log

10

(

W

2

)

0

250

500

750

1000

1250

1500

1750

2000

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

Gaussian 500d

5.3 Gray scale image colorization

Lemma 4.6 states that the WD has a closed form when one of the 2 distributions is supported on a line, allowing us to compute the WD and the OT map with a complexity of O ( dn + n log n ) . This particular situation arises for instance with RBG images ( µ 1 , µ 2 ∈ P n 2 ( R 3 ) ), where black and white images are supported on a line (the line of grays). One can address the problem of image colorization through color transfer [25], where a black and white image is the source and a colorful image the target. Our fast procedure allows considering large images without sub-sampling with a reasonable computation time. Fig. 5 gives an example of colorization of an image of size 1280 × 1024 that was computed in less than 0.2 second, while being totally untractable for the O ( n 3 log n ) solver of WD.

1.0

Source Image

Target Image

Source Image

Colorized Image

Target Image

Blue

0.4:

0.1

DUD

00

2'0

Red

0.1

0.6

0.5

10

9.8

02

5.4 Point clouds registration

Iterative Closest Point (ICP) is an algorithm for aligning point clouds based on their geometries [7]. Roughly, its most popular version defines a one-to-one correspondence between point clouds, computes a rigid transformation (namely translation, rotation or reflection), moves the source point clouds using the transformation, and iterates the procedure until convergence. The rigid transformation is the solution of the Procrustes problem i.e. arg min (Ω ,t ) ∈ O ( d ) × R d ∥ Ω( X -t ) -Y ∥ 2 2 , where X , Y are the source and the target cloud points and O ( d ) the space of orthogonal matrices of dimension d . This Procrustes problem can be solved using a SVD [59] for instance.

We perform the ICP algorithm with different variants to compute the one-to-one correspondence: neareast neighbor (NN) correspondence, OT transport map (for small size datasets) and min-SWGG

transport map. Note that SW, PWD, SRW, factored coupling and Sinkhorn cannot be run in this context where a one-to-one correspondence is mandatory; subspace detours [44] are irrelevant in this context (see Supp. 11.5). We evaluate the results of the ICP algorithm in terms of: i) the quality of the final alignment, measured by the Sinkhorn divergence between the re-aligned and target point cloud; ii) the speed of the algorithm given by the running time until convergence. We consider 3 datasets of different sizes. The results are shown in Table 1 and more details about the setup, can be found in Supp. 11.5. In Supp. 11.5 we give a deeper analysis of the results, notably with different criteria for the final assignment, namely the Chamfer and the Frobenius distance. One can see that the assignment provided by OT-based methods is better than NN. min-SWGG allows working with large datasets, while OT fails to provide a solution for n = 150000 .

Source

Target

6 Conclusion

Table 1: Sinkhorn Divergence between final transformation on the source and the target. Timings in seconds are into parenthesis. Best values are boldfaced. An example of a point clouds ( n = 3000 ) is provided on the left.

In this paper, we hinge on the properties of sliced Wasserstein distance and on the Wasserstein generalized geodesics to define min-SWGG, a new upper bound of the Wasserstein distance that comes with an associated transport map. Topological properties of SWGG are provided, showing that it defines a metric and that min-SWGG metrizes the weak convergence of measure. We also propose two algorithms for computing min-SWGG, either through a random search scheme or a gradient descent procedure after smoothing the generalized geodesics definition of min-SWGG. We illustrate its behavior in several experimental setups, notably showcasing its interest in applications where a transport map is needed.

The set of permutations covered by min-SWGG is the one induced by projections and permutations on the line. It is a subset of the original Birkhoff polytope and it would be interesting to characterize how these two sets relates. In particular, in the case of empirical realizations of continuous distributions, the behavior of min-SWGG, when n grows, needs to be investigated. In addition, the fact that min-SWGG and WD coincide when d > 2 n calls for embedding the distributions in higher dimensional spaces to benefit from the greater expressive power of projection onto the line. Another important consideration is to establish a theoretical upper bound for min-SWGG.

Acknowledgments

The authors gratefully acknowledge the financial support of the French Agence Nationale de la Recherche (ANR), under grant ANR-20-CHIA-0021-01 (project RAIMO 2 ), grant ANR-20-CHIA0030 (project OTTOPIA) and grant ANR-18-CE23-0022-01 (project MULTISCALE). Clément Bonet is supported by the project DynaLearn from Labex CominLabs and Région Bretagne ARED DLearnMe.

References

Aleksandr Danilovich Alexandrov. A theorem on triangles in a metric space and some of its applications. Trudy Mat. Inst. Steklov. , 38:5-23, 1951.

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.

David Alvarez-Melis and Nicolo Fusi. Geometric dataset distances via optimal transport. Advances in Neural Information Processing Systems , 33:21428-21439, 2020.

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures . Springer Science & Business Media, 2005.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning , pages 214-223. PMLR, 2017.

Arturs Backurs, Yihe Dong, Piotr Indyk, Ilya Razenshteyn, and Tal Wagner. Scalable nearest neighbor search for optimal transport. In International Conference on Machine Learning , pages 497-506. PMLR, 2020.

Paul J Besl and Neil D McKay. Method for registration of 3-d shapes. In Sensor fusion IV: control paradigms and data structures , volume 1611, pages 586-606. Spie, 1992.

Garrett Birkhoff. Tres observaciones sobre el algebra lineal. Univ. Nac. Tucuman, Ser. A , 5:147-154, 1946.

Nicolas Bonneel and David Coeurjolly. Spot: sliced partial optimal transport. ACMTransactions on Graphics (TOG) , 38(4):1-13, 2019.

Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and radon wasserstein barycenters of measures. Journal of Mathematical Imaging and Vision , 51(1):22-45, 2015.

Nicolas Bonneel, Michiel Van De Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. In Proceedings of the 2011 SIGGRAPH Asia conference , pages 1-12, 2011.

Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics , 44(4):375-417, 1991.

Giuseppe Buttazzo, Guillaume Carlier, and Katharina Eichinger. Wasserstein interpolation with constraints and application to a parking problem. arXiv preprint arXiv:2207.14261 , 2022.

Luis A Caffarelli. Monotonicity properties of optimal transportation¶ and the fkg and related inequalities. Communications in Mathematical Physics , 214(3):547-563, 2000.

Laetitia Chapel, Mokhtar Z Alaya, and Gilles Gasso. Partial optimal tranport with applications on positive-unlabeled learning. In Advances in Neural Information Processing Systems , volume 33, pages 2903-2913, 2020.

Rui Chibante. Simulated annealing: theory with applications . BoD-Books on Demand, 2010.

Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature. In Neural Information Processing Systems Machine Learning for Creativity and Design Workshop , 2018.

Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. EMNIST: Extending MNIST to handwritten letters. In 2017 international joint conference on neural networks (IJCNN) , pages 2921-2926, 2017.

Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence , 39(9):1853-1865, 2017.

Thomas M Cover. The number of linearly inducible orderings of points in d-space. SIAM Journal on Applied Mathematics , 15(2):434-439, 1967.

Katy Craig. The exponential formula for the wasserstein metric. ESAIM: Control, Optimisation and Calculus of Variations , 22(1):169-187, 2016.

JA Cuesta-Albertos, C Matrán, and Araceli Tuero-Diaz. Optimal transportation plans and convergence in distribution. journal of multivariate analysis , 60(1):72-83, 1997.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Ishan Deshpande, Yuan-Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced wasserstein distance and its use for gans. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10648-10656, 2019.

Sira Ferradans, Nicolas Papadakis, Julien Rabin, Gabriel Peyré, and Jean-François Aujol. Regularized discrete optimal transport. In International Conference on Scale Space and Variational Methods in Computer Vision , pages 428-439. Springer, 2013.

Jean Feydy. Geometric data analysis, beyond convolutions . PhD thesis, Université Paris-Saclay Gif-sur-Yvette, France, 2020.

Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouve, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 2681-2690, 2019.

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python Optimal Transport. Journal of Machine Learning Research , 22(78):1-8, 2021.

Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet, Geoffrey Schiebinger, and Jonathan Weed. Statistical optimal transport via factored couplings. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 2454-2465. PMLR, 2019.

Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. Advances in neural information processing systems , 28, 2015.

Wilfrid Gangbo and Robert J McCann. The geometry of optimal transportation. Acta Mathematica , 177(2):113-161, 1996.

Jonathan J Hull. A database for handwritten text recognition research. IEEE Transactions on Pattern Analysis and Machine Intelligence , 16(5):550-554, 1994.

Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS) , volume 37, pages 199-201, 1942.

Tapas Kanungo, David M Mount, Nathan S Netanyahu, Christine D Piatko, Ruth Silverman, and Angela Y Wu. An efficient k-means clustering algorithm: Analysis and implementation. IEEE transactions on pattern analysis and machine intelligence , 24(7):881-892, 2002.

Young-Heon Kim, Brendan Pass, and David J Schneider. Optimal transport and barycenters for dendritic measures. Pure and Applied Analysis , 2(3):581-601, 2020.

Scott Kirkpatrick, C Daniel Gelatt Jr, and Mario P Vecchi. Optimization by simulated annealing. science , 220(4598):671-680, 1983.

Soheil Kolouri, Kimia Nadjahi, Umut Simsekli, Roland Badeau, and Gustavo Rohde. Generalized sliced wasserstein distances. Advances in neural information processing systems , 32, 2019.

Matt Kusner, Yu Sun, Nicholas Kolkin, and Kilian Weinberger. From word embeddings to document distances. In International conference on machine learning , pages 957-966. PMLR, 2015.

Yann LeCun. The MNIST database of handwritten digits. http://yann. lecun. com/exdb/mnist/ , 1998.

Tianyi Lin, Zeyu Zheng, Elynn Chen, Marco Cuturi, and Michael I Jordan. On projection robust optimal transport: Sample complexity and model misspecification. In International Conference on Artificial Intelligence and Statistics , pages 262-270, 2021.

Quentin Mérigot, Alex Delalande, and Frédéric Chazal. Quantitative stability of optimal transport maps and linearization of the 2-wasserstein space. In International Conference on Artificial Intelligence and Statistics , pages 3186-3196, 2020.

Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci. , pages 666-704, 1781.

Caroline Moosmüller and Alexander Cloninger. Linear optimal transport embedding: Provable wasserstein classification for certain rigid transformations and perturbations. Information and Inference: A Journal of the IMA , 12(1):363-389, 2023.

Boris Muzellec and Marco Cuturi. Subspace detours: Building transport plans that are optimal on subspace projections. Advances in Neural Information Processing Systems , 32, 2019.

Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli. Statistical and topological properties of sliced probability divergences. Advances in Neural Information Processing Systems , 33, 2020.

Kimia Nadjahi, Alain Durmus, Pierre E Jacob, Roland Badeau, and Umut Simsekli. Fast approximation of the sliced-wasserstein distance using concentration of random projections. Advances in Neural Information Processing Systems , 34:12411-12424, 2021.

Luca Nenna and Brendan Pass. Transport type metrics on the space of probability measures involving singular base measures. Applied Mathematics & Optimization , 87(2):28, 2023.

Peer Neubert and Peter Protzel. Compact watershed and preemptive slic: On improving tradeoffs of superpixel segmentation algorithms. In 2014 22nd international conference on pattern recognition , pages 996-1001. IEEE, 2014.

Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui. Distributional sliced-wasserstein and applications to generative modeling. In International Conference on Learning Representations , 2021.

Khai Nguyen, Tongzheng Ren, Huy Nguyen, Litu Rout, Tan Nguyen, and Nhat Ho. Hierarchical sliced wasserstein distance. arXiv preprint arXiv:2209.13570 , 2022.

Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. Communications in Partial Differential Equations , 26(1-2):101-174, 2001.

François-Pierre Paty and Marco Cuturi. Subspace robust wasserstein distances. In International conference on machine learning , pages 5072-5081, 2019.

Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning , 11(5-6):355-607, 2019.

Martin Pincus. A monte carlo method for the approximate solution of certain types of constrained optimization problems. Operations research , 18(6):1225-1228, 1970.

Julien Rabin, Sira Ferradans, and Nicolas Papadakis. Adaptive color transfer with relaxed optimal transport. In IEEE international conference on image processing (ICIP) , pages 48524856, 2014.

Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernot. Wasserstein barycenter and its application to texture mixing. In Alfred M. Bruckstein, Bart M. ter Haar Romeny, Alexander M. Bronstein, and Michael M. Bronstein, editors, Scale Space and Variational Methods in Computer Vision , pages 435-446, Berlin, Heidelberg, 2012. Springer Berlin Heidelberg.

Mark Rowland, Jiri Hron, Yunhao Tang, Krzysztof Choromanski, Tamas Sarlos, and Adrian Weller. Orthogonal estimation of wasserstein distances. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 186-195. PMLR, 2019.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(5863):94, 2015.

Peter H Schönemann. A generalized solution of the orthogonal procrustes problem. Psychometrika , 31(1):1-10, 1966.

Matteo Togninalli, Elisabetta Ghisu, Felipe Llinares-López, Bastian Rieck, and Karsten Borgwardt. Wasserstein weisfeiler-lehman graph kernels. Advances in Neural Information Processing Systems , 32, 2019.

Stefan Van der Walt, Johannes L Schönberger, Juan Nunez-Iglesias, François Boulogne, Joshua D Warner, Neil Yager, Emmanuelle Gouillart, and Tony Yu. scikit-image: image processing in python. PeerJ , 2:e453, 2014.

Peter JM Van Laarhoven, Emile HL Aarts, Peter JM van Laarhoven, and Emile HL Aarts. Simulated annealing . Springer, 1987.

Cédric Villani. Optimal transport: old and new , volume 338. Springer, 2009.

Gemine Vivone, Mauro Dalla Mura, Andrea Garzelli, and Fabio Pacifici. A benchmarking protocol for pansharpening: Dataset, preprocessing, and quality assessment. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing , 14:6102-6118, 2021.

Wei Wang, Dejan Slepˇ cev, Saurav Basu, John A Ozolek, and Gustavo K Rohde. A linear optimal transportation framework for quantifying and visualizing variations in sets of images. International journal of computer vision , 101(2):254-269, 2013.

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.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747 , 2017.

Jingyi Zhang, Ping Ma, Wenxuan Zhong, and Cheng Meng. Projection-based techniques for high-dimensional optimal transport problems. Wiley Interdisciplinary Reviews: Computational Statistics , 15(2):e1587, 2023.

7 Proofs and supplementary results related to Section 3

7.1 Overestimation of WD by PWD

As stated in Section 2, the projected Wasserstein distance PWD (see Eq. 7) tends to overestimate the Wasserstein distance. This is due to the fact that some permutations σ θ and τ θ (with θ ∈ S d -1 ) involved in PWD computation may be irrelevant. Such situation occurs when the distributions are in high dimension but supported on a low dimensional manifold or when the distributions are multi-modal.

Let consider the distributions µ 1 and µ 2 lying on a low dimensional manifold. In high dimension, randomly sampled vectors θ tend to be orthogonal. Moreover, vectors orthogonal to the low dimensional manifold lead to 'collapsed' projected distributions P θ # µ 1 and P θ # µ 2 onto θ . Hence, such projection directions lead to permutations that can be randomly chosen. To empirically illustrate this behavior of PWD, we consider µ 1 and µ 2 as Gaussian distributions in R d , d = 10 but supported on the first two coordinates and we sample 200 points per distribution. Table 2 summarizes the obtained corresponding distances and shows that PWD overestimates the WD.

Now, let us consider two multimodal distributions µ 1 , µ 2 with K clusters such that each cluster of µ 1 has a close cluster from µ 2 (cyclical monotonicity assumption). Also we assume the same number of points in each cluster. OT plan will match the corresponding clusters and will lead to a relatively low value for W 2 2 (since cluster from µ 1 has a closely related cluster in µ 2 ). However as PWD may allow permutations that make correspondences between points from different clusters (since a source cluster and a target cluster can be far in the original space but very close when projected on 1D), the resulting distance will be much more larger, leading to an overestimation of the Wasserstein distance. Table 2 provides an illustration for K = 10 clusters and d = 2 .

7.2 Quantile version of SWGG

The main body of the paper expresses SWGG for empirical distributions µ 1 and µ 2 with the same number of points and uniform probability masses. In this section we derived SWGG in a more general setting of discrete distributions.

̸

Let remark that min-SWGG relies on solving a 1D optimal transport (OT) problem. So far, the 1D OT problem was derived for µ 1 , µ 2 ∈ P n 2 ( R ) and thus was expressed using the permutation operators τ and σ . In the general setting of distributions µ 1 ∈ P n 2 ( R ) and µ 2 ∈ P m 2 ( R ) with n = m , the 1D optimal transport is computed based on quantile functions. Hence, the expression of SWGG in the general setting of µ 1 ∈ P n 2 ( R ) and µ 2 ∈ P m 2 ( R ) hinges on quantile functions instead of permutations.

More formally, let µ ∈ P n 2 ( R ) ; its cumulative function is defined as:

and its quantile function (or pseudo inverse), is given by:

An important remark is that the quantile function is a step function with n (the number of atoms) discontinuities. Thus, it can be stored efficiently using two vectors of size n (one for the locations of the discontinuities and the other for the values of the discontinuities).

For µ 1 ∈ P n 2 ( R ) and µ 2 ∈ P m 2 ( R ) , we recover the Wasserstein distance through quantiles with:

Moreover, the optimal transport plan is given by:

where λ [0 , 1] is the Lebesgue measure on [0 , 1] . The transport plan can be stored efficiently using two vectors of size ( n + m -1) (see [53] Prop 3.4).

Following [53, Remark 9.6], one can define the quantile function related to the Wasserstein mean by :

Now, let µ 1 ∈ P n 2 ( R d ) and µ 2 ∈ P m 2 ( R d ) . Let µ 1 → 2 θ be the Wasserstein mean of the projected distributions on θ . Finally let π θ → 1 denote the transport plan from µ 1 → 2 θ to µ 1 and π θ → 2 be the transport plan from µ 1 → 2 θ to µ 2 . Following the construction of [4, Sec. 9.2], we shall introduce a multi marginal plan defined as:

where P 12 : ( R d ) 3 → ( R d ) 2 projects to the first two coordinates and P 13 projects to the coordinates 1 and 3. In particular, P 12 # π is the projection of π on its 2 first marginals and P 13 # π on the first and 3rd marginal. Similarly to the 2-marginal transport plan we defined Π( µ 1 → 2 θ , µ 1 , µ 2 ) = { π ∈ P 2 ( R d × R d × R d ) s.t. π ( A × R d × R d ) = µ 1 → 2 θ ( A ) , π ( R d × A, × R d ) = µ 1 ( A ) and π ( R d × R d × A ) = µ 2 ( A ) , ∀ A measurable set of R d } :

The generalized barycenter µ 1 → 2 g,θ is then defined as:

where P i is the projection on the i -th coordinate.

We finally have all the building blocks to compute SWGG in the general case. Let remark that the complexity goes from O ( dn + n log n ) in the P n 2 ( R d ) case to O ( d ( n + m ) + ( n + m ) log( n + m )) in the general case.

7.3 Proof of Proposition 3.2

We aim to prove that SWGG 2 2 ( µ 1 , µ 2 , θ ) is an upper bound of W 2 2 ( µ 1 , µ 2 ) and that SWGG ( µ 1 , µ 2 , θ ) is a distance ∀ θ ∈ S d -1 , µ i ∈ P n 2 ( R d ) , i = 1 , 2 .

Distance. Note that this proof will be derived for the alternative definition of SWGG in supp. 10.8. Let µ 1 = 1 n ∑ δ x i , µ 2 = 1 n ∑ δ y i , µ 3 = 1 n ∑ δ z i be in P 2 ( R d ) , let θ ∈ S d -1 . We note σ (resp. τ and π ) the permutation such that ⟨ x σ (1) , θ ⟩ ≤ ... ≤ ... ⟨ x σ ( n ) , θ ⟩ (resp. ⟨ y τ (1) , θ ⟩ ≤ ... ≤ ... y τ n , θ and z π (1) , θ ... ... z π ( n ) , θ ).

Non-negativity and finite value. From the ℓ 2 norm, it is derived

Identity property. From one side, µ 1 = µ 2 implies that ⟨ x i , θ ⟩ = ⟨ y i , θ ⟩ , ∀ 1 ≤ i ≤ n and that σ = τ , which implies SWGG 2 2 ( µ 1 , µ 2 , θ ) = 0 .

From the other side, SWGG 2 2 ( µ 1 , µ 2 , θ ) = 0 = ⇒ 1 n ∑ ∥ x σ ( i ) -y τ ( i ) ∥ 2 2 = 0 = ⇒ x σ ( i ) = y τ ( i ) , ∀ 1 ≤ i ≤ n = ⇒ µ 1 = µ 2 .

Equality The equality W 2 2 = min-SWGG 2 2 whenever d > 2 n comes from the fact that all the permutations are within the range of SWGG. In particular minimizing SWGG is equivalent to solve the Monge problem. We refer to Supp. 11.1 for more details.

Upper Bound The fact that min-SWGG 2 2 in an upper bound of W 2 2 comes from the sub-optimality of the permutations σ θ , τ θ . Indeed, they induce a one-to-one correspondence x σ θ ( i ) → y τ θ ( i ) ∀ 1 ≤ i ≤ n . This correspondence corresponds to a transport map T θ such that T θ # µ 1 = µ 2 . Since W 2 2 = inf T s.t. T # µ 1 = µ 2 1 n ∑ ∥ x -T ( x ) ∥ 2 2 we necessarily have W 2 2 ≤ min-SWGG 2 2 .

7.4 Difference between max-SW and min-SWGG

Herein, we give an example where the selected vectors θ for max-SW and min-SWGG differ.

Since these two distributions are far away on the x -coordinate, max-SW will catch this difference between the means by selecting θ ≈ ( 1 0 ) . Indeed, the projection on the x -coordinate represents the largest 1D WD.

Conversely, min-SWGG selects the pivot measure to be supported on θ ≈ ( 1 0 ) that separates the two distributions. Indeed, this direction better captures the geometry of the 2 distributions, delivering permutations that are well grounded to minimize the transport cost.

Fig. 6 illustrates that difference between max-SW and min-SWGG.

-

10

.

0

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

min-SWGG

Source

Target

Optimal direction

-

10

.

0

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

7

.

5

10

.

0

max-SW

Optimal direction

-

10

-

5

0

5

10

15

-

10

-

5

0

5

10

7.5 From permutations to transport map

In this section we provide the way of having a transport map from permutations.

Let µ 1 , µ 2 ∈ P n 2 ( R d ) , let θ ∗ ∈ arg min SWGG and let σ θ ∗ , τ θ ∗ the associated permutations. The associated map must be T ( x σ ( i ) ) = y τ ( i ) ∀ 1 ≤ i ≤ n . In the paper, we formulate the associated transport map as:

15

Moreover, the matrix representation of T is given by:

7.6 Examples of Transport Plan

Fig. 7 illustrates two instances of the transport plan obtained via min-SWGG. Even though these transport plans are not optimal, they were able to capture the overall structure of the true optimal transport plans.

2

70

8

6

4

2

0

2

4

6

8

8

6

4

2

0

2

4

6

Transport plan min-SWGG=11.46

Source

Target

8

6

4

2

0

2

4

6

8

8

6

4

2

0

2

4

6

W

2 =11.74

Source

Target

0

10

20

30

40

50

60

70

80

0

10

20

30

40

50

60

70

80

3

2

1

0

1

2

3

3

2

1

0

1

2

3

4

Transport plan min-SWGG=8.6

Source

Target

3

2

1

0

1

2

3

3

2

1

0

1

2

3

4

W

2 =7.8

Source

Target

0

10

20

30

40

50

60

70

80

0

10

20

30

40

50

60

70

80

0

10

20

30

40

50

60

70

80

0

10

20

30

40

50

60

80

0

10

20

30

40

50

60

70

80

0

10

20

30

40

50

60

2

70

80

The first example shows that the OT plan by min-SWGG exhibits a "block" structure, and thus approximates well the true Wasserstein distance. The second example shows that even in a context of superimposed distribution the "general transport direction" in min-SWGG is representative of that of the optimal transport map.

8 Background on Wasserstein Generalized Geodesics

We introduce some concepts related the Wasserstein generalized geodesics in Sec. 4.1. In this section, we provide more details about these geodesics in order to provide a wider view on this theory.

In the following definitions, we do not address the issue of uniqueness of the geodesics. However this is not a problem in our setup since we focus our study on pivot measure with n -atoms ν ∈ P n 2 ( R d ) . In this case, we have uniqueness of the ν -based Wasserstein distance [47].

Wasserstein generalized geodesics As mentioned in Sec. 4.1, Wasserstein generalized geodesics rely on a pivot measure ν ∈ P n 2 ( R d ) to transport µ 1 to µ 2 . Indeed, one can leverage the optimal transport maps T ν → µ 1 and T ν → µ 2 to construct a curve linking µ 1 to µ 2 . The generalized geodesic with pivot measure ν is defined as:

The generalized Wasserstein mean refers to the middle of the geodesic, i.e. when t = 0 . 5 and has been denoted µ 1 → 2 g .

Intuitively, the optimal transport maps between ν and µ i , i = 1 , 2 give rise to a sub-optimal transport map between µ 1 and µ 2 through:

T 1 → 2 ν links µ 1 to µ 2 via the generalized geodesic:

We recall here the ν -based Wasserstein distance induced by T 1 → 2 ν and introduced in Eq. (13). Definition 8.1. The ν -based Wasserstein distance [21, 47] is defined as:

Moreover, this new notion of geodesics comes with an inequality, which is of the opposite side to Eq. (3):

The parallelogram law is not respected but straddles with eq. (3) and eq. (31). We refer to Figure 8 for an intuition behind positive curvature [51], parallelogram law and generalized geodesics.

µ

1

µ

1

ν

ν

µ

2

µ

2

x

1

y

x

2

µ

1

→

2

(

t

)

µ

1

→

2

g

(

t

)

tx

1

+(1

-

t

)

x

2

Setting t = 0 . 5 in Eq. (31) and reordering the term gives:

Moreover one can remark that:

In particular situations W 2 ν and W 2 2 coincide. It is the case for 1D distributions where the Wasserstein space is known to be flat [4]. In that case, the Wasserstein mean and the generalized Wasserstein mean are the same.

Multi-marginal Another formulation of the ν -based Wasserstein distance is possible through the perspective of multi-marginal OT [4]. Let Π( µ 1 , µ 2 , ν ) = { π s.t. P 12 # π = π 1 → 2 , P 13 # π = π 1 → ν and P 23 # π = π 2 → ν } , where P ij is the projection onto the coordinates i, j . Let also Π ∗ ( µ i , ν ) be the space of optimal transport maps between µ i and ν . We have:

Equation (34) expresses the fact that we select the optimal plan from Π( µ 1 , µ 2 , ν ) which is already optimal for Π( µ i , ν ) . Mathematically, this minimization is not a multi-marginal problem, since the optimal plan is supposed to be already optimal for some coordinate.

The set { π ∈ Π( µ 1 , µ 2 , ν ) s.t. P i 3 # π ∈ Π ∗ ( µ i , ν ) i = 1 , 2 } is never empty, i.e. there is always existence of π 1 → 2 ν (thanks to the gluing lemma [63], page 23). Moreover, in situations where it is a singleton, there is uniqueness of π 1 → 2 ν . Uniqueness is an ingredient which overpasses the selection of a final coupling and comes with additional result.

Lemma 8.2 (Lemma 6 [47]) . Whenever { π ∈ Π( µ 1 , µ 2 , ν ) s.t. P i 3 # π ∈ Π ∗ ( µ i , ν ) i = 1 , 2 } is a singleton, W 2 ν is a proper distance. It is a semi-distance otherwise.

Notably, 1D pivot measure was studied in [35] to ensure a dendritic structure of the distributions along the geodesic.

9 Related Works

In this section we highlight the fact that several upper approximations of W 2 2 are in the framework of generalized geodesics. The differences lay in the choice of the pivot measure ν .

Factored Coupling. In [29], the authors impose a low rank structure on the transport plan by factorizing the couplings through a pivot measure ν expressed as the k -Wasserstein mean between µ 1 and µ 2 ( k ≤ n ). It is of particular interest since whenever the pivot distribution is the Wasserstein mean between µ 1 and µ 2 , W 2 ν and W 2 2 coincide.

Factored coupling results in a problem of computing the k -Wasserstein mean ( µ 1 → 2 ) followed by solving two OT problems between the clustered Wasserstein mean and the two input distributions ( W 2 2 ( µ 1 , µ 1 → 2 ) and W 2 2 ( µ 1 → 2 , µ 2 ) ). Even though the OT problems are smaller, they are still expensive in practice.

Moreover, in this scenario, the uniqueness of the OT plan T 1 → 2 ν is not ensured. It appears that [29] chooses the most entropic transport plan, i.e. simply T 1 → 2 ν = T µ 1 → 2 → µ 2 ◦ T µ 1 → µ 1 → 2 .

Subspace Detours. From a statistical point of view, it is beneficial to consider optimal transport on a lower dimensional manifold [66]. In [44], authors compute an optimal transport plan T µ E 1 → µ E 2 between projections on a lower linear subspace E of µ 1 and µ 2 , i.e. µ E i = P E # µ i , where P E is the linear projection on E . They aimed at leveraging T µ E 1 → µ E 2 to construct a sub-optimal map T 1 → 2 E between µ 1 and µ 2 .

The problem can be recast as a generalized geodesic problem with ν being the Wasserstein mean of µ E 1 and µ E 2 embedded in R d . Once again, uniqueness of T µ 1 → µ 2 ν is not guaranteed, authors provide two ways of selecting the map, namely Monge-Knothe and Monge-Independent lifting.

Subspace detours result in a problem where one needs to select a linear subspace E (which is a non convex procedure), compute an optimal transport between µ 1 and µ 2 (in O ( n 3 log n ) whenever dim( E ) > 1 ) and reconstruct T µ 1 → µ 2 E .

Linear Optimal Transport (LOT). Given a set of distributions ( µ i ) i m =1 ∈ P 2 ( R d ) m , LOT [65] embeds the set of distributions into the L 2 ( ν ) -space by computing the OT of each distribution to the pivot distribution. Mathematically, it computes T ν → µ i ∀ 1 ≤ i ≤ m and lies on estimating W 2 2 ( µ i , µ j ) with W 2 ν ( µ i , µ j ) through eq. (13).

In LOT, the pivot measure ν was chosen to be the average of the input measures [65], the Lebesgue measure on R d [41] or an isotropic Gaussian distribution [43].

Instead of computing ( m 2 ) expensive Wasserstein distances, it resorts only on m Wasserstein distances between ( µ i ) i m and ν . While significantly reducing the computational cost when several distributions are at stake, it does not allow speeding up the computation when only two distributions are involved.

9.1 Linear Optimal Transport with shift and scaling

In this section, we recall the result from [43]. The theorem states that the ν -based approximation is very close to WD whenever µ 1 , µ 2 are continuous distributions which are very close to be shift and scaling of each other. It can applies to a continuous version of SWGG, however it works with discrete measures in the particular case of equality between W 2 ν and W 2 2 .

Theorem 9.1 (Theorem 4.1 [43]) . Let Λ = { S a (shift) , a ∈ R d } ∪ { R c (scaling) , c ∈ R } , Λ µ,R = { h ∈ Λ s.t. ∥ h ∥ µ ≥ R } and G µ,R,ϵ = { g ∈ L 2 ( R d , µ ) s.t. ∃ h ∈ Λ µ,R s.t. ∥ g -h ∥ µ ≤ ϵ }

Let ν , µ ∈ P 2 ( R d ) , with µ, ν ≪ λ (the Lebesgue measure). Let R > 0 , ϵ > 0

For g 1 , g 2 ∈ G µ,R,ϵ and ν = λ on a convex compact subset of R d , we have:

If µ and ν satisfy the assumption of Caffarelli's regularity theorem [14], then for g 1 , g 2 ∈ G µ,R,ϵ , we have:

where C, C depdends on ν, µ and R .

10 Proofs and other results related to Section 4

10.1 Proof of Proposition 4.2: equivalence between the two formulations of SWGG

In this section, we prove that the two definitions of SWGG in Def. 3.1 and Prop. 4.2 are equivalent. Let θ ∈ S d -1 be fixed.

From one side in Def. 3.1, we have:

where σ θ and τ θ are the permutations obtained by sorting P θ # µ 1 and P θ # µ 2 .

From the other side we note D ( µ 1 , µ 2 , θ ) the quantity:

We want to prove that SWGG 2 2 ( µ 1 , µ 2 , θ ) = D ( µ 1 , µ 2 , θ ) , ∀ µ 1 , µ 2 ∈ P n 2 ( R d ) and θ ∈ S d -1 . Eq. (13) in the main paper states that D ( µ 1 , µ 2 , θ ) is equivalent to ∫ R d ∥ x -T 1 → 2 µ 1 → 2 θ ( x ) ∥ 2 2 dµ 1 ( x ) .

Finally, Lemma 4.6 states that the transport map T 1 → 2 µ 1 → 2 θ is fully determined by the permutations on the line: the projections part is a one-to-one correspondence between x and θ ⟨ x , θ ⟩ (resp. between y and θ ⟨ y , θ ⟩ ). More formally T 1 → 2 µ 1 → 2 θ ( x σ θ ( i ) ) = y τ θ ( i ) ∀ 1 ≤ i ≤ n . And thus we recover:

which concludes the proof.

10.2 Proof of Weak Convergence (Proposition 4.3)

We want to prove that, for a sequence of measures ( µ k ) k ∈ N ∈ P n 2 ( R d ) , we have:

The notation µ k L , 2 -→ µ stands for the weak convergence in P n 2 ( R d ) i.e. ∫ R d f ( x ) dµ ( k ) ( x ) → ∫ R d f ( x ) dµ ( x ) for all continuous bounded functions f and for the Euclidean distance f ( x ) = ∥ x 0 -x ∥ 2 2 for all x 0 ∈ R d .

From one side, if min-SWGG 2 2 ( µ k , µ ) → 0 = ⇒ W 2 2 ( µ k , µ ) → 0 = ⇒ µ k L , 2 -→ µ . The first implication is due to the fact that min-SWGG 2 2 is an upper-bounds of W 2 2 , the Wasserstein distance, and that WD metrizes the weak convergence.

From another side, assume µ k L , 2 -→ µ ; we have for any θ :

Let µ µ k → µ θ ∈ P n 2 ( R d ) stands for the Wasserstein mean of the projections Q θ # µ k and Q θ # µ and let µ µ → µ θ = Q θ # µ . We have µ µ k → µ θ converges towards (in law) to µ µ → µ θ , which implies that:

Since µ ∈ P n 2 ( R d ) , we have T µ µ k → µ θ → µ k -→ k T µ µ k → µ θ → µ (see [22], theorem 3.2). It implies that µ µ k → µ g,θ L -→ µ and particularly:

By combining the previous elements, we get:

The previous relation shows that µ k L , 2 -→ µ implies SWGG 2 2 ( µ k , µ, θ ) -→ k 0 for any θ . Hence, we can conclude that:

This concludes the proof.

Note that when µ 1 and µ 2 are continuous, [41] proved that when the distributions are smooth enough (i.e. respecting the Cafarelli theorem [14]), there is a bi-Holder equivalence between the ν -based Wasserstein distance and W 2 2 . Hence, it still holds for SWGG for any θ ∈ S d -1 :

where B depends on µ i , i ∈ { 1 , 2 } , θ and the dimension d . This bound is sufficient to prove that SWGG metrizes the weak convergence in this context. We refer to [41] for more details.

10.3 Proof of Translation property (Proposition 4.4)

We prove that min-SWGG 2 2 has the same behavior w.r.t. the translation as W 2 2 . This property is well known for Wasserstein and useful in applications such as shape matching.

Let µ 1 , µ 2 ∈ P n 2 ( R d ) , and let T u (resp. T v ) be the map x ↦→ x -u (resp. x ↦→ x -v ), with u , v vectors of R d .

To ease the notations, let define ˜ µ 1 = T u # µ 1 and ˜ µ 2 = T v # µ 2 .

Let remind that in the case of Wasserstein distance we have [53](Remark 2.19):

We aim to compute min-SWGG 2 2 (˜ µ 1 , ˜ µ 2 ) def = min-SWGG 2 2 ( T u # µ 1 , T v # µ 2 ) . Let express first with m 1 = ∫ R d x dµ 1 ( x ) and m 2 = ∫ R d x dµ 2 ( x ) .

where ˜ µ 1 → 2 θ is the Wasserstein mean of the projections along θ of the shifted measures ˜ µ 1 = T u # µ 1 and ˜ µ 2 = T v # µ 2 as in Proposition 2. The generalized Wasserstein mean ˜ µ 1 → 2 g,θ is defined accordingly (see also Proposition 11).

We have:

Similarly W 2 2 (˜ µ 2 , ˜ µ 1 → 2 θ ) = W 2 2 ( µ 2 , µ 1 → 2 θ ) -2 ⟨ v , m 2 -m 3 ⟩ + ∥ v ∥ 2 2 .

where m 3 = ∫ R d x d ˜ µ 1 → 2 θ ( x ) .

Let express now the third term in eq. (47). For that we require to define the generalized Wasserstein mean ˜ µ 1 → 2 g,θ with pivot measure ˜ µ 1 → 2 θ . By the virtue of eq. (11) in the main paper, we have:

Hence, the third term in (47) is:

Putting all together, we have:

10.4 Proof of the new closed form of the Wasserstein distance (Lemma 4.6)

We recall and prove the lemma that makes explicit a new closed form for WD. Let µ 1 , µ 2 be in P n 2 ( R d ) with µ 2 a distribution supported on a line whose direction is θ ∈ S d -1 . We have:

Moreover, the optimal map is given by T 1 → 2 = T Q θ # µ 1 → µ 2 T µ 1 → Q θ # µ 1 = T Q θ # µ 1 → µ 2 Q θ .

◦ ◦ Let µ 1 , µ 2 be in P n 2 ( R d ) with µ 2 a distribution supported on a line of direction θ . We have:

Moreover, the optimal map is given by:

Here Q θ is given in Def. 4.1 of the paper.

The proof of the Lemma was first inspired by [13](Proposition 2.3), where authors show that W 2 C ( µ 1 , µ 2 ) = W 2 C 1 ( µ 1 , µ ) + W 2 C 2 ( µ, µ 2 ) , with C 1 , C 2 and C some cost matrices with the constraints C ij = min s C 1 is + C 2 sj .

Let µ 1 = 1 n ∑ δ x i and µ 2 = 1 n δ y i be in P n 2 ( R d ) with µ 2 a distribution supported on a line with direction θ . Let Q θ # µ 1 = µ 1 = 1 n ∑ δ x i ∈ P n 2 ( R d ) . We emphasize here the fact that the atoms of µ 1 and µ 2 are supported on a line are denoted by the overline symbol.

From one side, we have:

Equation (61) is obtained thanks to the Pythagorean theorem since ⟨ x i , Q θ ( x i ) , y i ⟩ is a right triangle ∀ 1 ≤ i ≤ n . The equation (64) is obtained by taking the inf of the previous first term of the previous equation.

From the other side:

Where T 3 and T 4 are the optimal plan of W 2 2 ( µ 1 , µ 1 ) and + W 2 2 ( µ 1 , µ 2 ) . Similarly, (65) is obtained via the Pythagorean theorem. This concludes the proof.

We plot an illustration of the lemma in Figure 9.

6

4

2

0

2

4

6

-

4

-

2

0

2

4

6

Q θ

(

T

1

→

2

)

#

µ

1

θ

µ

1

µ

2

(

Q θ

)

#

µ

1

-

-

-

10.5 Details on the efficient computation of SWGG

We decompose the second formulation of SWGG. Let first remind that Q θ : R d → R d , x ↦→ θ ⟨ x , θ ⟩ and P θ : R d → R , x ↦→⟨ x , θ ⟩ are the projections on the subspace generated by θ .

We have:

First, by lemma 4.6,

as µ 1 → 2 θ 's support is on a line. Similarly,

and

We notice that 2 W 2 2 ( P θ # µ 1 , P θ # µ 1 → 2 θ ) + 2 W 2 2 ( P θ # µ 1 → 2 θ , P θ # µ 2 ) = W 2 2 ( P θ # µ 1 , P θ # µ 2 ) (as P θ # µ 1 → 2 θ is the Wasserstein mean between P θ # µ 1 and P θ # µ 2 ). We also notice that -4 W 2 2 ( P θ # µ 1 → 2 g,θ , P θ # µ 1 → 2 θ ) = 0 (it comes from the fact that the generalized Wasserstein mean is induced by the permutations on the line), we can put all together to have:

One can show that SWGG is divided into 3 Wasserstein distances between a distribution and its projections on a line and 1D Wasserstein problem. This results in a very fast computation of SWGG.

10.6 Smoothing of SWGG

In this section, we give details on the smoothing procedure of min-SWGG, an additional landscape of SWGG and its smooth counterpart ˜ SWGG and an empirical heuristic for setting hyperparameters s and ϵ .

Smoothing Procedure. A natural surrogate would be to add an entropic regularization within the definition of T µ 1 → 2 θ → µ i , i ∈ { 1 , 2 } and to solve an additional optimal transport problem. Nevertheless, it would lead to an algorithm with an O ( n 2 ) complexity. Instead, we build upon the blurred Wasserstein distance [26] between two distributions ν 1 and ν 2 :

where ∗ denotes the smoothing (convolution) operator and k ϵ/ 4 is the Gaussian kernel of deviation √ ϵ/ 2 . In our case, it resorts in making s copies of each sorted projections P θ ( x i ) and P θ ( y i ) respectively, to add a Gaussian noise of deviation √ ϵ/ 2 and to compute averages of sorted blurred copies x s σ s , y s τ s :

Further, we provide additional examples of the landscape of ˜ min-SWGG ( µ 1 , µ 2 ) and discuss how to choose empirically relevant s and ϵ values.

[26] has shown that the blurred WD has the same asymptotic properties as the Sinkhorn divergence, with parameter ϵ the strength of the blurring: it interpolates between WD (when ϵ → 0 ) and a degenerate constant value (when ϵ →∞ ).

To find a minimum of Eq. (16) in the paper (i.e. ˜ SWGG 2 2 ( µ 1 , µ 2 , θ ) ), we iterate over:

where η ∈ R + is the learning rate. This procedure converges towards a local minima with a complexity of O ( snd + sn log( sn )) for each iteration. Once the optimal direction θ ⋆ is found, the final solution resorts to be the solution provided by SWGG 2 2 ( µ 1 , µ 2 , θ ⋆ ) , where the induced optimal transport map is an unblurred matrix.

Heuristic for setting the hyperparameters of ˜ SWGG We here provide an heuristic for setting parameters s (number of copies of each points) and ϵ (strength of the blurring). We then give an example of the behavior of ˜ SWGG w.r.t. these hyper parameters.

s ∈ N + represents the number of copies of each sample. We observe empirically that the quantity sn should be large to provide a smooth landscape. It means that the s values can be small when n increases, allowing to keep a competitive algorithm (as the complexity depends on ns )

ϵ ∈ R + represents the variance of the blurred copies of each sample. Empirically, ϵ should depend on the variance of the distributions projected on the line. Indeed, an ϵ very close to zero will not smooth enough the discontinuities whereas a large ϵ will give a constant landscape.

As discussed in Section 4.3, finding an optimal θ ∈ S d -1 is a non convex problem and provides a discontinuous loss function. We give some examples of the landscape of ˜ SWGG w.r.t. different values of the hyperparameters in Fig. 10. The landscapes were computed with a set of projections θ regularly sampled with angles ∈ [0 , 2 π ] .

We observe that the larger s , the smoother ˜ SWGG. Additionally, raising ϵ tends to flatten ˜ SWGG w.r.t. θ (erasing local minima). Indeed similarly to Sinkhorn, a large ϵ blurred the transport plan and thus homogenize all the value of SWGG w.r.t. θ .

Moreover, we empirically observe that the number of samples for µ 1 and µ 2 enforces the continuity of SWGG. We then conjecture that the discontinuities of SWGG are due to artifact of the sampling and thus the smoothing operation erases this unwanted behavior. A full investigation of this assumption is left for future work.

7.5

5

2.5

0

2.5

5

7.5

8

6

4

2

0

2

4

6

8

Source

Target

0

50

100

150

200

250

300

350

Degree

20

40

60

80

100

Distance

SWGG

s=10 eps=0.2

s=10 eps=2.0

s=100 eps=0.2

s=100 eps=2.0

1

0.5

0

0.5

1

1

0.75

0.5

0.25

0

0.25

0.5

0.75

1

Source

Target

0

50

100

150

200

250

300

350

Degree

0.4

0.6

0.8

1

1.2

Distance

SWGG

s=10 eps=0.2

s=10 eps=2.0

s=100 eps=0.2

s=100 eps=2.0

10.7 Inconsequential of the pivot measure

Importantly, only the direction θ is of importance for the value of SWGG. Indeed, whenever ν ∈ P n 2 ( R d ) is supported on a line of direction θ , the position of the atoms is irrelevant for W ν and the associated transport plan whenever the atoms are distinct. Despite the fact that the pivot measure is inconsequential for the value of SWGG (at θ fixed), we choose it to be µ 1 → 2 θ . This choice is supported by the fact that µ 1 → 2 θ can be efficiently computed (as a 1D Wasserstein mean) and that some computation can be alleviated:

It is an important comment to derive the property of distance for SWGG; it also allows minimizing SWGG over θ ∈ S d -1 without consideration for ν , since any choice of ν supported on the subspace generated by θ give the same result for min-SWGG. This property of irrelevance comes from the

nature of the subspace where ν is supported, which is uni-dimensional. More formally we give the following proposition and its associated proof.

Proposition 10.1. Let µ 1 , µ 2 ∈ P n 2 ( R d ) . Let θ ∈ S d -1 . Let ν 1 , ν 2 ∈ P n 2 ( R d ) be two pivot measures supported on a line with direction θ , with disctincs atoms for each measure. We then have:

We give a proof of this proposition.

Thanks to lemma 4.6, we known that the transport map T 1 → 2 ν is fully induced by the transport plan T Q θ # µ 1 → Q θ # µ 2 ν . Let remind that T Q θ # µ 1 → Q θ # µ 2 ν is given by T ν → Q θ # µ 2 ◦ T Q θ # µ 1 → ν (see equation (12)). Moreover the two optimal transport plans are obtained via the ordering permutations, i.e. let σ, τ, π ∈ S ( n ) s.t:

With x i being the atoms of Q θ # µ 1 , y i the atoms of Q θ # µ 2 and z i being the atoms of Q θ # ν .

One have T µ 1 → ν ( x σ ( i ) ) = z π ( i ) (resp. T ν → µ 2 ( z π ( i ) ) = x τ ( i ) ) ∀ 1 ≤ i ≤ n . Composing these two identities gives:

The last equation shows that T 1 → 2 ν is in fact independent of π and thus of ν .

10.8 Proof that min-SWGG is a distance (generalized geodesic formulation)

This proof has already been established in 7.3. However we rephrase the proof in the context of generalized geodesics.

We aim to prove that SWGG 2 = √ 2 W 2 2 ( µ 1 , µ 1 → 2 θ ) + 2 W 2 2 ( µ 1 → 2 θ , µ 2 ) -4 W 2 2 ( µ 1 → 2 g,θ , µ 1 → 2 θ ) defines a metric.

Finite and non-negativity. Each term of SWGG 2 2 is finite thus the sum of the three terms is finite. Moreover, being an upper bound of WD makes it non-negative.

Symmetry. We have

Identity property.

From one side, when µ 1 = µ 2 = ⇒ T µ 1 → µ 1 → 2 θ = T µ 2 → µ 1 → 2 θ = Id , giving µ 1 → 2 g,θ = µ 1 = µ 2 . Thus:

From another side, SWGG 2 2 ( µ 1 , µ 2 , θ ) = 0 = ⇒ W 2 2 ( µ 1 , µ 2 ) = 0 = ⇒ µ 1 = µ 2 (by being an upper bound of WD).

Triangle Inequality. We have:

where, with an abuse of notation for clarity sake, T i θ is the optimal map between µ 1 → 2 θ and µ i and T g θ is the optimal map between µ 1 → 2 θ and µ 1 → 2 g,θ . The last line comes from the parallelogram rule of R d . Thanks to Proposition 10.1 we see that SWGG is simply the L 2 ( R d , ν ) square norm, i.e.:

with ν being any arbitrary pivot measure of P n 2 ( R d ) . And thus SWGG 2 is the L 2 ( R d , ν ) norm. This observation is enough to conclude that SWGG 2 is a proper distance for θ fixed.

11 Experiment details and additional results

WD, SW, Sinkhorn, Factored coupling are computed using the Python OT Toolbox [28] and our code is available at https://github.com/MaheyG/SWGG . The Sinkhorn divergence for the point cloud matching experiment was computed thanks to the Geomloss package [27].

11.1 Behavior of min-SWGG with the dimension and the number of points

In this section, we draw two experiments to study the behavior of min-SWGG w.r.t. the dimension and to the number of points.

Evolution with d In [20][Theorem of Section 2], authors aim at enumerate the number of permutations obtained via the projection of point clouds on a line. It appears that the number of permutations increases with the dimension. They even show that whenever d ≥ 2 n ( 2 n being the total number of points of the problem), all the possible permutations ( n ! ) are in the scope of a line. Fig. 11 depicts the number of obtainable permutations as a function of the dimension d , for n fixed. This theorem can be applied to min-SWGG to conclude that whenever d ≥ 2 n , we have min-SWGG 2 2 = W 2 2 .

It turns out empirically that the greater the dimension, the better the approximation of W 2 2 with min-SWGG (see Fig. 11) for a fixed n . More formally, the set of all possible transport maps is called the Birkhoff polytope and it is known that the minimum of the Monge problem is attained at the extremal points (which are exactly the set of permutations matrices, a set of n ! matrices in our context) [8]. The set of the transport maps in the scope of SWGG is a subset of the extremal points of the Birkhoff polytope (there are permutations matrices but not all possibilities are represented). Theoretically, the set of transport maps in the scope of SWGG is larger as d grows, giving a subset that is more and more tight with the extremal points of the Birkhoff polytope. This explains that min-SWGG can benefit from higher dimension.

We plot in Fig. 11 the evolution, over 50 repetitions, of the ratio min-SWGG ( µ 1 ,µ 2 ) W 2 2 ( µ 1 ,µ 2 ) with d , n = 50 and µ 1 ∼ N (1 R d , Id ) , µ 2 ∼ N ( -1 R d , Id ) .

Evolution with n Fig. 12 represents the evolution of W 2 2 ( µ 1 , µ 2 ) and min-SWGG 2 2 ( µ 1 , µ 2 ) for two distributions µ 1 ∼ N (1 R d , Id ) and µ 2 ∼ N ( -1 R d , Id ) , with d = 4 and a varying number of points. The results are averages over 10 repetitions.

0

50

100

150

200

250

300

Dimension

1

.

025

1

.

050

1

.

075

1

.

100

1

.

125

1

.

150

1

.

175

1

.

200

1

.

225

Ratio

min

-

SWGG

W

2

2

for

n

= 50

0

50

100

150

200

250

300

Dimension

0

200

400

600

800

1000

1200

1400

log

10

number of permutations

Number permutations

n

= 310

Number of permutation from a line

n

= 310

10

100

1000

Number of points

10

15

20

25

30

35

W

2

2

for

d

= 4

min-SWGG for

d

= 4

We observe that, when n is large enough, min-SWGG tends to stabilize around some constant value.

We conjecture that there may exist an upper bound for min-SWGG:

Where d ′ is the max of the dimensions of the distributions µ 1 , µ 2 [66], and ψ an unknown function.

11.2 Computing min-SWGG

We now provide here more details about the experimental setup of the experiments of Section 5.1.

Choosing the optimal θ We compare three variants for choosing the optimal direction θ : random search, simulated annealing and optimization (defined in Section 4.3). We choose to compare with simulated annealing since it is widely used in discrete problem (such as the travelling salesman) and known to perform well in high dimension [62] [16] [36]. We notice in Fig. 3 of the paper that the smooth version of min-SWGG is always (comparable or) better than the simulated annealing. In this experiment, we randomly sample 2 Gaussian distributions with different means and covariances matrices, whose parameters are chosen randomly. For optimizing min-SWGG, we use the Adam optimizer of Pytorch, with a fixed learning rate of 5 e -2 during 100 iterations, considering s = 10 and ϵ = 1 .

Fig. 13 provides the timings for computing the random search approximation, simulated annealing and the optimization scheme. In all cases, we recover the linear complexity of min-SWGG (blue curves) in a log space. For the computation timings we compute min-SWGG with random search

with L = 500 , simulated annealing (green curves) with 500 iterations with a temperature scheme (1 -k +1 500 ) 500 k =1 and the optimization scheme (considering s = 10 with a fixed number of iterations for the optimization scheme equals to 100).

100

10000

0

.

00

0

.

01

0

.

10

1

.

00

Time (sec)

100

10000

number of points

n

0

.

00

0

.

01

0

.

10

1

.

00

Annealing

100

10000

0

.

0

0

.

1

1

.

0

10

.

0

Additionally, we reproduce the same setup as in 5.1 for the SW, max-SW and PWD distance. For sake of readability we compared with min-SWGG optim and report the results in Fig. 14.

10

1000

0

10

20

30

40

Distance

W 2

2

=32.4

10

1000

Number of projections

0

100

200

300

W 2

2

=346.1

maxSW (optimized)

SWGG Optimized

10

1000

0

1000

2000

3000

4000

W 2

2

=3836.0

10

1000

33

33.5

34

34.5

35

Distance

W 2

2

=32.4

10

1000

Number of projections

376

378

380

382

W 2

2

=346.1

SWGG (Optimized)

10

1000

3945

3950

3955

3960

3965

W 2

2

=3836.0

Runtime Evaluation In the paper, on Fig. 3 (Right), we compare the empirical runtime evaluation on GPU for different methods. We consider Gaussian distributions in dimension d = 3 and we sample n points per distribution with n ∈ { 10 2 , 10 3 , 10 4 , 5 · 10 4 , 10 5 } . For SW Monte-Carlo and min-SWGG random search, we use L = 200 projections. For both max-SW and min-SWGG with optimization, we use 100 iterations with a learning rate of 1, and we fix s = 50 for min-SWGG. We use the official implementation of the Subspace Robust Wasserstein (SRW) with the Frank-Wolfe algorithm [52].

11.3 Gradient Flows

We rely on the code provided with [37] for running the experiment of Section 5.2.

We fix n = 100 , the source distribution is taken to be Gaussian and we consider four different target measures that represent several cases: i) a 2 dimensional Gaussian, ii) a 500 dimensional Gaussian (high dimensional case), iii) 8 Gaussians (multi-modal distribution) and iv) a two-moons distribution (non-linear case).

We fix a global learning rate of 5 e -3 with an Adam optimizer. For SW, PWD and SWGG (random search), we sample L = 100 directions. For the optimization methods max-SW, we set a learning rate of 1 e -3 with a number of 100 iterations for i), iii), and iv) and 200 iterations for ii). For min-SWGG (optimization), we took a learning rate of i) 1 e -1 , ii) 1 e -3 , iii) 5 e -2 , and iv) 1 e -3 . The hyper parameters for the optimization of min-SWGG are s = 10 and ϵ = 0 . 5 , except for the 500-dimensional Gaussian for which we pick ϵ = 10 .

Each experiment is run 10 times and shaded areas in Fig. 4 (see the main paper) represent the mean ± the standard deviation.

11.4 Gray scale image colorization

We now provide additional results on a pan-sharpening application to complete results provided in Section 5.3.

In pan-sharpening [64], one aims at constructing a super-resolution multi-chromatic satellite image with the help of a super-resolution mono-chromatic image (source) and low-resolution multichromatic image (target).

To realize this task, we choose to used a color transfer procedure, where the idea is to transfer the color palette from the target to the source image. This transfer is carry out by the optimal transport plan of the Wasserstein distance. More details on color transfer can be found in Supp. 11.6.

Additionally, we improve the relevance of the colorization by adding a proximity prior. For that, we used super pixels computed via the Watershed algorithm [48] thanks to the the scikit-image package [61]. Obtained high resolution colorized images of size 512 × 512 ( n = 262 144 ) are reported on Fig. 15.

Source

ColorizedSource

Target

Colonized Source

All pan-sharpening experiments were run on the PairMax data set [64]. The hyperparameters (markers and compactness) for the watershed super-pixels are: 500, 200, 200, 200 markers (an upper bound for the number of super pixel) for each image (by order of apparition) and compactness 1 e -8 (high values result in more regularly-shaped watershed basins) for all the images.

11.5 Point clouds registration

We here provide additional details and results about the experiments in Section 5.4.

Authors of [9] highlighted the relevance of OT in the point clouds registration context, plugged into an Iterative Closest Point (ICP) algorithm. They leveraged the 1D partial OT without consideration for the direction of the line. Our experiment shows the importance of θ : the smaller SWGG is, the better the registration.

In this experiment, having a one to one correspondence is mandatory: as such, we compare min-SWGG with a nearest neighbor assignment and the one provided by OT. Note that we do not compare min-SWGG with subspace detour [44], since: i) with empirical distributions, the reconstruction of the plan is degenerated (as it doesn't involve any computation), ii) the research of subspace can be intensive as no prior is provided.

To create the source distributions, we used random transformation (Ω , t ) ∈ O ( d ) × R d of the target domain. Ω was picked randomly from O ( d ) , the set of rotations and reflections, and t has random direction with ∥ t ∥ 2 = 5 . We also add a Gaussian noise N (0 , ϵId ) , with ϵ = 0 . 1 .

The ICP algorithm was run with 3 datasets with the following features: i) 500 points in 2D, ii) 3000 points in 3D, and iii) 150 000 points in 3D. min-SWGG was computed through the random search estimation with L = 100 . A stopping criterion was the maximum number of iterations of the algorithm, which varies with the dataset i.e. : i) 50, ii) 100, and iii) 200 respectively. The other stopping criterion is ∥ Ω -Id ∥ + ∥ t ∥ 2 ≤ ε with ε chosen respectively for the datasets as follows: i) 1 e -4 , ii) 1 e -2 , and iii) 1 e -2 , where (Ω , t ) ∈ O ( d ) × R d is the current transformation and ∥ · ∥ is the Frobenius norm. All these settings were run with 50 different seeds. Results are reported in Fig. 16.

Source

Target

0

2

4

6

8

0

10

20

30

40

50

60

Percent

0

20

40

60

80

100

Chamfer Distance

0

20

40

60

80

Percent

0

5000

10000

15000

20000

25000

Frobenius Distance

0

20

40

60

80

100

Percent

min-SWGG

Nearest

OT

0

0.2

0.4

Time in second

0

10

20

30

40

50

60

70

Percent

10

20

30

40

Number of Iterations

0

2.5

5

7.5

10

12.5

15

17.5

Percent

Source

Target

0

50

100

150

0

2.5

5

7.5

10

12.5

15

17.5

Percent

0.2

0.4

0.6

0.8

1

Chamfer Distance

0

10

20

30

40

50

60

Percent

0

20

40

60

80

100

120

Frobenius Distance

0

10

20

30

40

50

60

Percent

min-SWGG

Nearest

OT

0

25

50

75

100

125

150

Time in second

0

10

20

30

40

50

60

70

Percent

0

20

40

60

80

100

Number of Iterations

0

2.5

5

7.5

10

12.5

15

17.5

20

Percent

Target

25

100

25

min-SWGG

Nearest

20

25

80-

20

20-

20

15

15

60

15

IUBO

IUOOA

15

10-

40

10

10

10

20

10

Sinkhom Divergence

20

30

40

50

50

Chamfer Distance

100

150

200

250

20

Frobenius Distance

40

09

0S

Time in second

100

150

200

50

Number of lterations

100

150

200

From Fig. 16, one can see that for:

n = 500 : The registration obtained via OT is very powerful (it is fast to compute and converges toward a good solution). min-SWGG is slightly faster with better convergence result and an equivalent number of iterations. Finally the nearest neighbor does not converge to a solution closed to the target.

n = 3000 : registration by OT can converge poorly, moreover the timings are much higher than the competitors. min-SWGG shows efficient convergence results with an attractive computation time (order of fews seconds). We observe that the number of iterations can be very large and we conjecture that it is due to the fact that min-SWGG-based ICP can exit local minima. The nearest neighbor is fast but, most of the time, does not converge to global minima (i.e. the exact matching of the shapes).

n = 150000 : In this setting, OT is totally untractable (the cost matrix needs 180 GB in memory). min-SWGG shows good convergence and is most of the time very fast whenever

the number of iterations does not attain the stopping criterion. The nearest neighbor assignment is faster but only converges to local minima.

Note that, despite the fact that min-SWGG is slightly slower than the nearest neighbor, the overall algorithm can be faster due to the better quality of each iteration (min-SWGG can attain a minimum with less iterations).

In Table 3 we give additional results on the final distribution. We control the final convergence via the square Chamfer distance and the square Frobenius distance. The Square Chamfer distance is a sqaure distance between point cloud defined as:

and the Square Frobenius norm is a square distance between the transformation, defined as:

In both cases we can see that the final results for min-SWGG are much better than the results for NN and relatively closed to the results from OT.

Table 3: Square Chamfer distance (top) and Square Frobenius distance (bottom) between final transformation on the source and the target. Best values are boldfaced.

An other important aspect of ICP is that the algorithm tends to fall into local minima: the current solution is not good and further iterations do not allow a better convergence of the algorithm. We observed empirically that min-SWGG can avoid getting stuck on local minima when a reasonable number of directions θ is sampled ( L ∼ 100 ). We conjecture that the random search approximation is not always the ideal solution and hence may escape local minima. This may lead to a better convergence solution for min-SWGG-based ICP.

11.6 Color Transfer

In this section, we provide an additional experiment in a color transfer context.

We aim at adapting the color of an input image to match the color of a target one [25, 44]. This problem can be recast as an optimal transport problem where we aim at transporting the color of the source image X into the target Y . For that, usual methods lie down on the existence of a map T : X → Y . We challenge min-SWGG to this problem to highlight relevance of the obtained transport map.

Images are encoded as vector in R nm × 3 , where n and m are the size of the image and 3 corresponds to the number of channels (here RBG channels). We first compute a map T 0 : X 0 → Y 0 between a subsample of X and Y of size 20000 and secondly extend this mapping to the complete distributions T : X → Y using a nearest neighbor interpolation. The subsampling step is mandatory due to the size of the images but can deteriorate the quality of the transfer.

We compare the results obtained with maps obtained from Wasserstein distance, min-SWGG with random search (100 projections), subspace detour [44] and min-SWGG (optimized). Obtained images and the associated timings are provided in fig. 17.

Figure 17 shows that min-SWGG and W 2 2 provide visually equivalent solutions. Since, the quality of the color transfer is dependent on the size of the subsampling: using min-SWGG permits larger

Source

Full OT (runtime 62.73s)

SWGG (runtime 0.07s)

SWGG optim (runtime 0.10s)

Subspace Detour (runtime 0.181s)

Target

Source

Full OT (runtime 70.35s)

SWGG (runtime 0.08s)

SWGG optim (runtime 0.10s)

Subspace Detour (runtime 0.274s)

Target

subsamples than W 2 2 and thus improves the quality of the map T . Moreover one can note that min-SWGG (optimized) is the fastest to compute.

We now give more details about how to perform color transfer between two distributions. The first step is to encode n × m images as R nm × 3 vectors, with 3 channels in a RGB image. Note that m and n can differ for the source and target image. The second step consists of defining subsamples X 0 , Y 0 of X,Y , in our case we took X 0 , Y 0 ∈ R 20000 × d . We subsample the same number of points for the source and target image. In order to have a better subsampling of X and Y , it is common to perform a k -means [34] to derive X 0 and Y 0 ( X 0 , Y 0 are then taken as centroids of the k-means algorithm). The third step is to compute T 0 : X 0 → Y 0 . We set T as the optimal Monge map given by the Wasserstein distance and T as the optimal map given by min-SWGG. Finally, the fourth step deals with extending T 0 : X 0 → Y 0 to T : X → Y . ∀ x ∈ X . We compute the closest element x 0 ∈ X 0 and we pose:

More details on the overall procedure can be found in [25].

To perform the experiment, we took L = 100 projections for min-SWGG (random search). For min-SWGG (optimized), we fixed the following set of parameters for the gradient descent: learning rate 5 e -2 , number of iterations 20, number of copies s = 10 and ϵ = 1 . Regarding the subspace detour results, we used the code of [44] provided at https://github.com/BorisMuzellec/ SubspaceOT .

Additionally, we perform color transfer without sub-sampling with the help of min-SWGG (we the same hyperparameters). This procedure is totally untractable for either W 2 2 and subspace detours (due to memory issues). As we mentioned before, the subsampling phase can decrease the quality of the transfer and thus min-SWGG can deliver better result than before. Result are give in Fig. 18

Source

SWGG(runtime8.82s)

SWGG optim (runtime 6.55s)

Target

11.7 Data set distance

We finally evaluate min-SWGG in an other context: computing distances between datasets. Let D 1 = { ( x 1 i , y 1 i ) } n i =1 and D 2 = { ( x 2 i , y 2 i ) } n i =1 be source and target data sets such that x 1 i , x 2 i ∈ R d are samples and y 1 i , y 2 i are labels ∀ 1 ≤ i ≤ n . In [3], the authors compare those data sets using the Wasserstein distance with the entries of the cost matrix defined as:

and the corresponding distance as:

where α y is the distribution of all samples with label y , namely { x ∈ R d | ( x , y ) ∈ D} for D being either D 1 or D 2 and U is the Birkhoff polytope which encodes the marginal constraints. Notice that cost in Eq. (87) encompasses the ground distance and a label-to-label distance. This distance is appealing in transfer learning application since it is model-agnostic. However, it can be cumbersome to compute in practice since it lays down on solving multiple OT problems (to compute the cost matrix and the OTDD). To circumvent that, [3] proposed several methods to compute the cost matrix in Eq. (87). They used the Sinkhorn algorithm (in O ( n 2 ) ) or they assumed α y ∼ N ( m y , Σ y ) in order to get the WD through the Bures metric (that provides a closed form of OT for Gaussian distributions in O ( d 3 ) ), which is still prohibitive for high dimension. We challenge min-SWGG in this context.

In this experiment, we compare the following datasets: MNIST [39], EMNIST [18], FashionMNIST [67], KMNIST [17] and USPS [32]. We rely on the code of OTDD provided at https://github. com/microsoft/otdd . In order to make it compliant with the min-SWGG hypothesis, we require the empirical distributions α y to have the same number of atoms.

Fig. 19 provides results for a batch size of n = 40000 samples using the Sinkhorn divergence (with a regularisation parameter of 1 e -1 ) and for min-SWGG (optimized) on batch of size 40000. We report results for a learning rate of 1 e -5 , 20 iterations and s and ϵ to be 1 and 0.

1

1

0.8

1.1

1.2

1.2

1.3

1

1.3

1.2

1.3

0.8

1

1.2

1.2

1

1

1.3

1.2

1

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

0.9

0.9

0.7

0.9

1.1

1

1.2

0.9

1.2

1.3

1.2

0.8

0.8

1.1

1.1

0.9

0.9

1.2

1.1

0.9

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

We check that the orders of magnitude are preserved with min-SWGG. For example OTDD(MNIST,USPS) is smaller than OTDD(MNIST,FashionMNIST) for either Sinkhorn divergence or min-SWGG as distance between labels, this validate that min-SWGG is a meaningful distance in this case scenario. Moreover in our setup, the computation cost is more expensive for Sinkhorn than for min-SWGG and totally untractable for W 2 2 . On smaller batches (see Fig 20), the same observation can be made: min-SWGG is comparable (in term of magnitude) with W 2 2 , Sinkhorn and the Bures approximation.

We give additional results in Fig. 20 for batches of size of n = 2000 samples obtained with W 2 2 , Sinkhorn divergence (setting the entropic regularization parameter to 1 e -1 ), the Bures approximation, min-SWGG (random search with L = 1000 projections) and min-SWGG (optimized, with a learning rate of 5 e -1 , 50 iterations, s = 20 and ϵ = 0 . 5 ).

0.4

0.5

0.4

0.4

0.5

0.5

0.6

0.4

0.6

0.5

0.6

0.3

0.4

0.5

0.5

0.4

0.4

0.6

0.5

0.4

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

1

1

0.8

1

1.2

1.2

1.3

1

1.3

1.2

1.3

0.8

0.9

1.2

1.2

1

0.9

1.3

1.2

0.9

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

0.9

0.9

0.7

0.9

1

1

1.1

0.9

1.2

1.1

1.1

0.7

0.8

1.1

1

0.9

0.8

1.1

1

0.8

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

1

1.1

0.9

1.1

1.3

1.3

1.4

1.1

1.3

1.3

1.4

0.8

1

1.3

1.3

1.1

1

1.3

1.3

1.1

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

1

1.1

0.8

1.1

1.3

1.3

1.3

1.1

1.3

1.2

1.4

0.8

1

1.3

1.3

1.1

1

1.3

1.3

1

MNIST

EMNIST

Fashion

KMNIST

USPS

MNIST

EMNIST

Fashion

KMNIST

USPS

̸

Note that the Figs. 19 and 20 are not symmetric (OTDD(KMNIST,FashionMNIST) = OTDD(FashinMNIST,KMNIST)) because of the random aspect of batches.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a variant of slicing of the Wasserstein distance that keeps the optimal transport map property (in particular, existence of an optimal transport map). The typical formulation of slicing loses the Monge map as each slice comes with its own map and there is no natural way to combine maps coming from different slices. As in the Wasserstein case, the variant of the sliced optimal transport distance introduced here is efficient to compute, and performs similarly (both theoretically and in practice) to the sliced Wasserstein distance.\n\nSTRENGTHS:\nThe idea is very nice. A major advantage of sliced distances are their numerical efficiency but the big drawback is the lack of maps. By being able to define optimal transport maps the authors have removed a significant disadvantage, whilst keeping the main advantage.\n\nWEAKNESSES:\n(1) More theoretical background/justification, such as when do minimizers in Definition 4.1 exist would be helpful.\n(2) The distance defined in eq. (13), since it can be written as the square root of \\int_{R^d} \\| T^{\\nu\\to\\mu_1}(x) - T^{\\nu\\to\\mu_2}(x)\\|^2 d \\nu(x), is the linearized Wasserstein distance. The authors should connect to the literature on linearized optimal transport distances.\n(3) There were a MANY grammatical and spelling errors in the paper, including misspelling Wasserstein (line 48). The paper needs a very careful readthrough. This is my reason for scoring the paper as a 2 in presentation. I felt the explanations were otherwise quite good, and I would be happy to give a higher score once corrections have been made.\n\nQUESTIONS:\n(a) Connected to (1) in the weaknesses section: can you give sufficient and/or necessary conditions for existence of minimizers in Definition 4.1?\n(b) What would you do if minimizers do not exist?\n\nLIMITATIONS:\nThere was no discussion on limitations of the method, although a few directions for further work was proposed in the conclusions.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an approximation to the squared Wasserstein distance between probability measures (typically required in high dimensional space) based on the sliced Wasserstein paradigm. The general idea is to project points onto a line where computing Wasserstein distances becomes significantly simpler (a typical case involves only permutations and sorting). The authors build up their main construction - min SWGG, by leveraging the properties of prior works like Max SW and projected sliced Wasserstein distance (PWD) and derive a tighter upper bound for the squared Wasserstein distance (W2) in comparison to PWD leading to the definition of min-SWGG. To put it very simply min-SWGG is closely related to the PWD, where the minimum is chosen instead of the average over all possible angles.   \n\nThe authors then relate the min-SWGG formulation to Wasserstein Generalised Geodesics - which involves a third measure (called pivot) in order to approximate the actual Wasserstein distance. By choosing the pivot measure to lie on a line,  this connection enables a more tractable way to optimize for the best angle/plane to compute min-SWGG. \n\nExperiments are demonstrated for (1.) Image Color Transfer (2.) PointCloud Matching (3.) Empirical demonstration of min SWGG optimization. The results demonstrate a proof of concept (i.e. examples showing the min-SWGG can be computed with different options quite reasonably) and some cases of good accuracy and faster computation\n\nSTRENGTHS:\n- Overall, I found this paper to be very comprehensive in its core formulation of min-SWGG. This paper appears to rate highly in its theoretical contributions - (1.) The definition of a new OT approximation with interesting bounds to the actual WD and PWD (2.) Connecting min SWGG to the Generalized Wasserstein Geodesic using a pivot measure that lies on a line thereby showing a path for computational tractability (3.) Proofs of metricity, weak convergence, computational and topological properties etc.\n- The paper is written well with a good introduction and fairly self contained background\n\nWEAKNESSES:\n- The practical aspect of this paper is sort of underwhelming. It would have been very convincing to demonstrate on truly high dimensional applications. Most experiments are very synthetic, and despite a modest experimental section, I am still grasping for situations where min-SWGG outperforms SWD and perhaps WD in terms of computability, quality of distances and transport maps, barycenters etc.\n- To this aid, despite the commentary on the use of cumulative distributions in the supplementary (enabling transport plans instead of maps), I could not place any impactful demonstration. That being said almost all of the narrative is based on matching uniform distributions having equal number of points and it's unclear how to generalize from here. \n- Generally, most figures could do with a more descriptive caption. Especially Figure 4 which I am finding really hard to parse. What is the box with Gaussian 500d indicate? Which are the source and target distributions? What is the message of this experiment? (I suspect it is that with enough iterations, the optimization of min-SWGG provides lower W2 distances - please clarify)\n\nQUESTIONS:\n- I am curious to visualize the difference between the transport plan obtained using min-SWGG in comparison to the optimal transport plan (i.e. corresponding to WD). This could potentially be visualized on the Pointcloud example. \n- How do the other SW distances (SW, max-SW, PWD etc.) compare in figure 3 (left)?\n- What is the evaluation for some other metric in the point cloud example, for eg the chamfer distance, L2 error after alignment etc. Is there a specific reason to report only the sinkhorn divergence?\n\nLIMITATIONS:\nSee Weaknesses for technical limitations. I am not an expert, but I think this paper has no direct negative societal impact. \n\nOverall, I am inclined positively on this paper. The new definition and approximation are derived nicely and interesting connections are made to the Wasserstein Generalised Geodesics framework. However, the practical impact (even amongst competing optimal transport techniques) is not equally impressive, and at this point unclear if it is as widely applicable. Put together, I vote a weak acceptance, with an intention to re-assess more carefully after the rebuttal.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
This paper proposed a novel proxy of the squared Wasserstein distance called min-SWGG, which is based on Sliced-Wasserstein distance (SWD) and the projected Wasserstein distance (PWD). This quantity is also shown to be relevant to Wasserstein Generalized Geodesics with pivot measure, which makes it more interesting. The theoretical properties of min-SWGG has been discussed, including the topological properties and the relationships with Wasserstein Generalized Geodesics. The authors provide two efficient algorithms for computing min-SWGG.  The experimental results indicates better computational efficiency of the proposed algorithm compared to the state-of-the-art methods.

STRENGTHS:
The paper is well written and clear, the approach is well supported by the theoretical analysis. The experiment and application seem strong and nice. The authors also made comparisons to other state-of-the-art methods that makes the paper more complete.

WEAKNESSES:
- One question about the approach in line 140: In Remark 3.3 the authors mentioned about the overall computational complexity to calculate min-SWGG by random search for L times. I wonder if there is any upperbound for it when the dimension d is high? Or is there any results showing the possible order of relation of  L and dimention d, or other quntities?

- For the second row of Figure 2, I notice that for Gaussian 500d, min-SWGG (optim) performs the best but for the left one, min-SWGG (optim) may have inferior behaviors. What might be the issue? Could you maybe elaborate on it a little bit?

QUESTIONS:
See weakness.

LIMITATIONS:
I have not found the limitations of the approach to be addressed in the paper. Maybe the authors can elaborate more on these.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new formulation and algorithm to approximate Wasserstein 2-distance (i.e., the distance of two points is the squared Euclidean one) between two distributions. Though some of the development is general, the key parts and the algorithms apply only to distributions being the sums of delta functions with the same weights. The approach builds on approaches approximating the true WD by using 1-D slices of the distributions, the distance between which can be computed cheaply as pointwise sum of squared differences of sorted distributions.\n\nThe basic Sliced WD, which approximates the true WD by the integral of the distance between such 1-D projections over all directions (eq. (5)). Its approximation is Max-Sliced WD, which replaces the integral by maximization (eq. (6)), leading to a lower bound on the true WD.\n\nAnother existing sliced formulation is Projected WD, which approximates the true WD by the integral of the pointwise sum of squared differences between the two distributions in which the points are sorted as if the distributions were projected onto the direction, over all directions (eq. (7)). The approach proposed in the manuscript replaces the integral here by minimization (eq. (9)), which is called Min-SWGG. It yields an upper bound in the true WD. This problem is non-convex, in fact discontinuous because the objective jumps with every change of either permutation (which sort the points on the projections). Therefore, it is smoothed and at the same time reformulated using the concepts from differential geometry of WD, generalized geodesics. This leads to a smooth non-convex problem, approximating the initial min-WDGG formulation. Its objective for a given direction can be evaluated cheaply: its most expensive part is the WD between a general point cloud and a point cloud on a line, which can be computed cheaply by just sorting 1-D distributions. The obtained problem is solved by gradient descent (yielding thus a local minimum).\n\nThe experiments compare the Min-SWGG, optimized either by smoothing and gradient descent or by random search or by simulated annealing, with existing (approximate) WD methods on synthetic data. Then it is illustrated on colorization of gray-scale images by OT, where computing the exact OT is intractable. Then it is demonstrated on point cloud registration by the ICP algorithm, where the point correspondence in every ICP iteration is computed either by nearest-neighbor or OT.\n\nSTRENGTHS:\nThe problem addressed is useful in a number of applications.\n\nTechnically sound.\n\nClearly enough written.\n\nWEAKNESSES:\nThe formulation of min-SWGG (eq. (9)) is rather straightforward. Namely, approximating the integral (7) by the minimization (9) is analogous to approximating the integral (5) by the maximization (6). In this sense, I see the novelty of the paper as rather incremental.\n\nIt is inconvenient that a lot of non-trivial information is in the supplement (I wonder if this format of presentation is desirable for NIPS). E.g., in the main paper the problem is restricted for the distributions from $P_2^n(R^d)$, i.e., unweighted averages of delta-distributions (see line 116), so that the exact WD (formula (1)) reduces to a linear assignment problem with the costs being squared Euclidean distances, for which the transport plan is a permutation. I.e., given two clouds of points in R^d with n elements, we want to find a permutation of one point cloud such that the sum of squared Euclid. distances between the points is minimum. In this restricted formulation, the min-SWGG (9) just optimizes over a subset of all possible permutations, given by sorting the projections of the distributions onto some direction. It is mentioned that this formulation can be extended to arbitrary distributions, but the details of this extension are just sketched in the supplement (section A.2). It is not clear from A.2 what would be the complexity of the algorithm in this general setting. In particular, I do not see why the transport plan (eq. (4) in A.2) should always have only $n+m-1$ nonzeros. Moreover, the experiments are done only for the restricted setting.\n\nDetails of some experiments are also often given only in the supplement. E.g., Section 5.3 has too little details to me, so that I do not understand how exactly image colorization is formulated. The additional info in the supplement is not helpful.\n\nIn experiments, Table 1 (report on point cloud registration by ICP) might be unfair to the nearest-neighbor method because the table reports Sinkhorn divergence of the registered clouds, which is close to what is optimized in min-SWGG. A better criterion to report might be the difference between the estimated and the ground-truth transformation (assuming that the latter is available, which trivially is e.g. when the target cloud is constructed by transforming the source cloud). Moreover, NN might perform well in the regime when the point clouds are initially close to each other.\n\nMinor comments:\n- The formulas for $\\mu_1,\\mu_2$ on line 116 are repeated on lines 122. The text could be optimized such that they appear only once.\n- Denoting the distances (8) and (9) as (min-)SWGG is confusing because generalized geodesics (GG) are not needed for their definitions (they are needed only later to reformulate the definitions).\n\nQUESTIONS:\nQ1) In the restricted setting of unweighted discrete distributions (i.e., from $P_2^n(R^d)$), the problem (1) reduces to linear assignment problem with costs being squared Euclidean distances. This is a classical problem, for which a great number of exact and approximate algorithms have been proposed. Unfortunately, I am not that familiar with existing approximate algorithms for this problem. Did you double check that there is no existing approximate algorithm for this problem that would compare in efficiency with your algorithm but which would however not generalize to arbitrary distribution (e.g., various primal heuristics, moving between permutations by local changes)? In other words, is the strength of the proposed approach only in its generalizability to arbitrary distributions, or already for the distributions from $P_2^n(R^d)$?\n\nQ2) In particular, the experiments do not convince me that min-SWGG is more efficient than Sinkhorn. The right-most graph in Figure 3 does not make clear to what accuracy the individual algorithms were run - in other words, it mixes accuracy and runtime (it does not compare the runtimes for the same achieved accuracy, or vice versa). Note, Sinkhorn keeps only dual variables and it need not keep the cost matrix (n-by-n) explicitly, so its memory complexity scales up well with n and d. Please comment.\n\nQ3) Please comment on parallelizability of the algorithm. I guess that the gradient descent to minimize (9) is not easy to parallelize. Of course, random search is.\n\nQ4) A suggestion: for small $d$, it might be interesting to try optimizing (9) (non-smoothed) by derivative-free methods, such as Nelder-Mead.\n\nLIMITATIONS:\nThe experiments are done only for unweighted discrete distributions.\n\nThe experiments do not convincingly show that min-SWGG is more efficient than, e.g., Sinkhorn.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}