SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Guillaume Mahey
INSA Rouen Normandie - Université Bretagne Sud
LITIS - IRISA
guillaume.mahey@insa-rouen.fr
Laetitia Chapel
Université Bretagne Sud - Institut Agro Rennes-Angers
IRISA
laetitia.chapel@irisa.fr
Gilles Gasso
INSA Rouen Normandie
LITIS
gilles.gasso@insa-rouen.fr
Clément Bonet
Université Bretagne Sud
LMBA
clement.bonet@univ-ubs.fr
Nicolas Courty
nicolas.courty@univ-ubs.fr
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.
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 linearistic 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  $\mathbb{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 one-dimensional 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  $\langle \cdot, \cdot \rangle$  be the Euclidean inner product on  $\mathbb{R}^d$  and let  $\mathbb{S}^{d-1} = \{\pmb{u} \in \mathbb{R}^d \text{ s.t. } \| \pmb{u} \|_2 = 1\}$ , the unit sphere. We denote  $\mathcal{P}(\mathbb{R}^d)$  the set of probability measures on  $\mathbb{R}^d$  endowed with the  $\sigma$ -algebra of Borel set and  $\mathcal{P}_2(\mathbb{R}^d) \subset \mathcal{P}(\mathbb{R}^d)$  those with finite second-order moment i.e.  $\mathcal{P}_2(\mathbb{R}^d) = \{\mu \in \mathcal{P}(\mathbb{R}^d) \text{ s.t. } \int_{\mathbb{R}^d} \| \pmb{x} \|_2^2 d\mu(\pmb{x}) < \infty\}$ . Let  $\mathcal{P}_2^n(\mathbb{R}^d)$  be the subspace of  $\mathcal{P}_2(\mathbb{R}^d)$  defined by empirical measures with  $n$ -atoms and uniform masses. For any measurable function  $f: \mathbb{R}^d \to \mathbb{R}^d$ , we denote  $f_\#$  its push forward, namely for  $\mu \in \mathcal{P}_2(\mathbb{R}^d)$  and for any measurable set  $A \in \mathbb{R}^d$ ,  $f_\# \mu(A) = \mu(f^{-1}(A))$ , with  $f^{-1}(A) = \{\pmb{x} \in \mathbb{R}^d \text{ s.t. } f(\pmb{x}) \in A\}$ .
Definition 2.1 (Wasserstein distance). The squared WD [63] between  $\mu_1, \mu_2 \in \mathcal{P}_2(\mathbb{R}^d)$  is defined as:
$$
W _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \stackrel {\text {d e f}} {=} \inf  _ {\pi \in \Pi \left(\mu_ {1}, \mu_ {2}\right)} \int_ {\mathbb {R} ^ {d} \times \mathbb {R} ^ {d}} \| \boldsymbol {x} - \boldsymbol {y} \| _ {2} ^ {2} d \pi (\boldsymbol {x}, \boldsymbol {y}) \tag {1}
$$
with  $\Pi(\mu_1, \mu_2) = \{\pi \in \mathcal{P}_2(\mathbb{R}^d \times \mathbb{R}^d) \text{ s.t. } \pi(A \times \mathbb{R}^d) = \mu_1(A)\text{ and } \pi(\mathbb{R}^d \times A) = \mu_2(A), \forall A$ . measurable set of  $\mathbb{R}^d\}$ .
The arg min of Eq. (1) is referred to as the optimal transport plan. Denoted  $\pi^{*}$ , it expresses how to move the probability mass from  $\mu_{1}$  to  $\mu_{2}$  with minimum cost. In some cases,  $\pi^{*}$  is of the form  $(Id,T)_{\#}\mu_{1}$  for a measurable map  $T:\mathbb{R}^d\to \mathbb{R}^d$ , i.e. there is no mass splitting during the transport. This map is called a Monge map and is denoted  $T^{\mu_1\rightarrow \mu_2}$  (or shortly  $T^{1\rightarrow 2}$ ). Thus, one has  $W_{2}^{2}(\mu_{1},\mu_{2}) = \inf_{T\mathrm{s.t.}T_{\#}\mu_{1} = \mu_{2}}\int_{\mathbb{R}^{d}}\| \pmb {x} - T(\pmb {x})\|_{2}^{2}d\mu_{1}(\pmb {x})$ . This occurs, for instance, when  $\mu_{1}$  has a density w.r.t. the Lebesgue measure [12] or when  $\mu_{1}$  and  $\mu_{2}$  are in  $\mathcal{P}_2^n (\mathbb{R}^d)$  [58].
Endowed with the WD, the space  $\mathcal{P}_2(\mathbb{R}^d)$  is a geodesic space. Indeed, since there exists a Monge map  $T^{1\to 2}$  between  $\mu_{1}$  and  $\mu_{2}$ , one can define a geodesic curve  $\mu^{1\rightarrow 2}:[0,1]\to \mathcal{P}_2(\mathbb{R}^d)$  [31] as:
$$
\forall t \in [ 0, 1 ], \mu^ {1 \rightarrow 2} (t) \stackrel {\text {d e f}} {=} (t T ^ {1 \rightarrow 2} + (1 - t) I d) _ {\#} \mu_ {1} \tag {2}
$$
which represents the shortest path w.r.t. Wasserstein distance in  $\mathcal{P}_2(\mathbb{R}^d)$  between  $\mu_1$  and  $\mu_2$ . The Wasserstein mean between  $\mu_1$  and  $\mu_2$  corresponds to  $t = 0.5$  and we simply write  $\mu^{1\rightarrow 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:
$$
W _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \geq 2 W _ {2} ^ {2} \left(\mu_ {1}, \nu\right) + 2 W _ {2} ^ {2} \left(\nu , \mu_ {2}\right) - 4 W _ {2} ^ {2} \left(\mu^ {1 \rightarrow 2}, \nu\right) \tag {3}
$$
for all pivot measures  $\nu \in \mathcal{P}_2(\mathbb{R}^d)$ .
Solving and approximating Optimal Transport. The Wasserstein distance between empirical measures  $\mu_1, \mu_2$  with  $n$ -atoms can be computed in  $\mathcal{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  $\mathcal{O}(n^2)$  complexity [2].
Notably, when  $\mu_{1} = \frac{1}{n}\sum_{i = 1}^{n}\delta_{x_{i}}$  and  $\mu_{2} = \frac{1}{n}\sum_{i = 1}^{n}\delta_{y_{i}}$  are 1D distributions, computing the WD can be done by matching the sorted empirical samples, leading to an overall complexity of  $\mathcal{O}(n\log n)$ . More precisely, let  $\sigma$  and  $\tau$  two permutation operators s.t.  $x_{\sigma (1)}\leq x_{\sigma (2)}\leq \ldots \leq x_{\sigma (n)}$  and  $y_{\tau (1)}\leq y_{\tau (2)}\leq \ldots \leq y_{\tau (n)}$ . Then, the 1D Wasserstein distance is given by:
$$
W _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) = \frac {1}{n} \sum_ {i = 1} ^ {n} \left(x _ {\sigma (i)} - y _ {\tau (i)}\right) ^ {2}. \tag {4}
$$
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  $\theta \in \mathbb{S}^{d - 1}$  over the unit sphere:
$$
\mathrm {S W} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \stackrel {\text {d e f}} {=} \int_ {\mathbb {S} ^ {d - 1}} W _ {2} ^ {2} \left(P _ {\#} ^ {\theta} \mu_ {1}, P _ {\#} ^ {\theta} \mu_ {2}\right) d \omega (\theta), \tag {5}
$$
where  $P_{\#}^{\theta}\mu_{1}$  and  $P_{\#}^{\theta}\mu_{2}$  are projections onto the direction  $\theta \in \mathbb{S}^{d - 1}$  with  $P^{\theta}:\mathbb{R}^{d}\to \mathbb{R},x\mapsto \langle x,\theta \rangle$  and where  $\omega$  is the uniform distribution over  $\mathbb{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  $\mathcal{O}(dLn + Ln\log n)$  and the memory complexity is  $\mathcal{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:
$$
\max  - \mathrm {S W} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \stackrel {\text {d e f}} {=} \max  _ {\theta \in \mathbb {S} ^ {d - 1}} W _ {2} ^ {2} \left(P _ {\#} ^ {\theta} \mu_ {1}, P _ {\#} ^ {\theta} \mu_ {2}\right). \tag {6}
$$
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  $\mu_{1} = \frac{1}{n}\sum_{i = 1}^{n}\delta_{\pmb{x}_{i}}$  and  $\mu_{2} = \frac{1}{n}\sum_{i = 1}^{n}\delta_{\pmb{y}_{i}}$  in  $\mathcal{P}_2^n (\mathbb{R}^d)$ . The PWD is defined as:
$$
\mathrm {P W D} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \stackrel {\text {d e f}} {=} \int_ {\mathbb {S} ^ {d - 1}} \frac {1}{n} \sum_ {i = 1} ^ {n} \| \boldsymbol {x} _ {\sigma_ {\theta} (i)} - \boldsymbol {y} _ {\tau_ {\theta} (i)} \| _ {2} ^ {2} d \omega (\theta), \tag {7}
$$
where  $\sigma_{\theta},\tau_{\theta}$  are the permutations obtained by sorting  $P_{\#}^{\theta}\mu_{1}$  and  $P_{\#}^{\theta}\mu_{2}$ .
As some permutations are not optimal, we straightforwardly have  $W_2^2 \leq \mathrm{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  $\mu_{1} = \frac{1}{n}\sum_{i = 1}^{n}\delta_{\pmb{x}_{i}}$  and  $\mu_{2} = \frac{1}{n}\sum_{i = 1}^{n}\delta_{\pmb{y}_{i}}$  in  $\mathcal{P}_2^n (\mathbb{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  $\mu_1, \mu_2 \in \mathcal{P}_2^n(\mathbb{R}^d)$  and  $\theta \in \mathbb{S}^{d-1}$ . Denote by  $\sigma_\theta$  and  $\tau_\theta$  the permutations obtained by sorting the 1D projections  $P_{\#}^{\theta} \mu_1$  and  $P_{\#}^{\theta} \mu_2$ . We define respectively SWGG and min-SWGG as:
$$
\operatorname {S W G G} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}, \theta\right) \stackrel {\text {d e f}} {=} \frac {1}{n} \sum_ {i = 1} ^ {n} \| \boldsymbol {x} _ {\sigma_ {\theta} (i)} - \boldsymbol {y} _ {\tau_ {\theta} (i)} \| _ {2} ^ {2}, \tag {8}
$$
$$
\min  - \operatorname {S W G G} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \stackrel {\text {d e f}} {=} \min  _ {\theta \in \mathbb {S} ^ {d - 1}} \operatorname {S W G G} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}, \theta\right). \tag {9}
$$
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  $\theta^{*} \in \mathrm{argmin}\, \mathrm{SWGG}_{2}^{2}(\mu_{1},\mu_{2},\theta)$  be the optimal projection direction. The associated transport map is:
$$
T \left(\boldsymbol {x} _ {i}\right) = \boldsymbol {y} _ {\tau_ {\theta^ {*}} ^ {- 1} \left(\sigma_ {\theta^ {*}} (i)\right)}, \quad \forall 1 \leq i \leq n. \tag {10}
$$
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  $\theta \in \mathbb{S}^{d - 1}$ .  $\mathrm{SWGG}_2(\cdot ,\cdot ,\theta)$  defines a distance on  $\mathcal{P}_2^n (\mathbb{R}^d)$ . Moreover, min-SWGG is an upper bound of  $W_{2}^{2}$ , and  $W_{2}^{2}\leq \min \text{-SWGG}_{2}^{2}\leq \mathrm{PWD}_{2}^{2}$ , with equality between  $W_{2}^{2}$  and min-SWGG $_2^2$  when  $d > 2n$ .
Remark 3.3. Similarly to max-SW, min-SWGG retains only one optimal direction  $\theta^{*} \in \mathbb{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  $\theta^{*}$  may differ (see Supp. 7.4 for an illustration), and iii) max-SW does not provide a transport plan between  $\mu_{1}$  and  $\mu_{2}$ .
Solving Eq. (9) can be achieved using a random search, by sampling  $L$  directions  $\theta \in \mathbb{S}^{d - 1}$  and keeping only the one leading to the lowest value of SWGG.
This gives an overall computational complexity of  $\mathcal{O}(Ldn + Ln\log n)$  and a memory complexity of  $\mathcal{O}(dn)$ . In low dimension, the random search estimation is effective: covering all possible
permutations through  $\mathbb{S}^{d - 1}$  can be done with a low number of directions. In high dimension, many more directions  $\theta$  are needed to have a relevant approximation, typically  $\mathcal{O}(L^{d - 1})$ . This motivates the design of gradient descent techniques for finding  $\theta^*$ .
4 SWGG as 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.  $\theta$  and piecewise constant, thus not differentiable over  $\mathbb{S}^{d - 1}$ . Indeed, as long as the permutations remain the same for different directions  $\theta$ , 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
SWGG, hence, making min-SWGG amenable to optimization. This formulation relies on Wasserstein generalized geodesics we introduce hereinafter.
Figure 1: (Left) Empirical distributions with examples of 2 sampled lines (Right) that lead to 2 possible values of SWGG when  $\theta \in [0,2\pi ]$
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(\mu_1, \mu_2)$  when either  $\mu_1$  or  $\mu_2$  is supported on a line.
4.1 SWGG based 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  $\nu \in \mathcal{P}_2^n (\mathbb{R}^d)$  to transport the distribution  $\mu_{1}$  toward  $\mu_{2}$ . Indeed, one can leverage the optimal transport maps  $T^{\nu \to \mu_1}$  and  $T^{\nu \to \mu_2}$  to construct a curve  $t\mapsto \mu_g^{1\to 2}(t)$  linking  $\mu_{1}$  to  $\mu_{2}$  as
$$
\mu_ {g} ^ {1 \rightarrow 2} (t) \stackrel {\text {d e f}} {=} ((1 - t) T ^ {\nu \rightarrow \mu_ {1}} + t T ^ {\nu \rightarrow \mu_ {2}}) _ {\#} \nu , \quad \forall t \in [ 0, 1 ]. \tag {11}
$$
The related generalized Wasserstein mean corresponds to  $t = 0.5$  and is denoted  $\mu_g^{1\rightarrow 2}$ .
Intuitively, the optimal transport maps between  $\nu$  and  $\mu_{i}, i = 1,2$  give rise to a sub-optimal transport map between  $\mu_{1}$  and  $\mu_{2}$ :
$$
T _ {\nu} ^ {1 \rightarrow 2} \stackrel {\text {d e f}} {=} T ^ {\nu \rightarrow \mu_ {2}} \circ T ^ {\mu_ {1} \rightarrow \nu} \quad \text {w i t h} \quad \left(T _ {\nu} ^ {1 \rightarrow 2}\right) _ {\#} \mu_ {1} = \mu_ {2}. \tag {12}
$$
One can be interested in the cost induced by the transportation of  $\mu_{1}$  to  $\mu_{2}$  via the transport map  $T_{\nu}^{1\to 2}$ , known as the  $\nu$ -based Wasserstein distance [47] and defined as
$$
W _ {\nu} ^ {2} \left(\mu_ {1}, \mu_ {2}\right) \stackrel {\text {d e f}} {=} \int_ {\mathbb {R} ^ {d}} \| \boldsymbol {x} - T _ {\nu} ^ {1 \rightarrow 2} (\boldsymbol {x}) \| _ {2} ^ {2} d \mu_ {1} (\boldsymbol {x}) = 2 W _ {2} ^ {2} \left(\mu_ {1}, \nu\right) + 2 W _ {2} ^ {2} \left(\nu , \mu_ {2}\right) - 4 W _ {2} ^ {2} \left(\mu_ {g} ^ {1 \rightarrow 2}, \nu\right). \tag {13}
$$
Notably, the second part of Eq. (13) straddles the square Wasserstein distance with Eq. (3). Remarkably, the computation of  $\bar{W}_{\nu}^{2}$  can be efficient if the pivot measure  $\nu$  is chosen appropriately. As established in Lemma 4.6, it is the case when  $\nu$  is supported on a line. Based on these facts, we propose hereafter an alternative formulation of SWGG.
Definition 4.1 (Pivot measure). Let  $\mu_{1}$  and  $\mu_{2} \in \mathcal{P}_{2}^{n}(\mathbb{R}^{d})$ . We restrict the pivot measure  $\nu$  to be the Wasserstein mean of the measures  $Q_{\#}^{\theta}\mu_{1}$  and  $Q_{\#}^{\theta}\mu_{2}$ :
$$
\mu_ {\theta} ^ {1 \rightarrow 2} \stackrel {{\mathrm {d e f}}} {{=}} \arg \min  _ {\mu \in \mathcal {P} _ {2} ^ {n} (\mathbb {R} ^ {d})} W _ {2} ^ {2} (Q _ {\#} ^ {\theta} \mu_ {1}, \mu) + W _ {2} ^ {2} (\mu , Q _ {\#} ^ {\theta} \mu_ {2}),
$$
where  $\theta \in \mathbb{S}^{d - 1}$  and  $Q^{\theta}:\mathbb{R}^{d}\to \mathbb{R}^{d},\pmb {x}\mapsto \theta \langle \pmb {x},\theta \rangle$  is the projection onto the subspace generated by  $\theta$ . Moreover  $\mu_{\theta}^{1\rightarrow 2}$  is always defined as the middle of a geodesic as in Eq (2).
One shall notice that  $Q_{\#}^{\theta}\mu_{1}$  and  $Q_{\#}^{\theta}\mu_{2}$  are supported on the line defined by the direction  $\theta$ , so is the pivot measure  $\nu = \mu_{\theta}^{1\rightarrow 2}$ . We are now ready to reformulate the metric SWGG.
Proposition 4.2 (SWGG based on generalized geodesics). Let  $\theta \in \mathbb{S}^{d - 1}$ ,  $\mu_1, \mu_2 \in \mathcal{P}_2^n(\mathbb{R}^d)$  and  $\mu_\theta^{1\rightarrow 2}$  be the pivot measure. Let  $\mu_{g,\theta}^{1\rightarrow 2}$  be the generalized Wasserstein mean between  $\mu_1$  and  $\mu_2 \in \mathcal{P}_2^n(\mathbb{R}^d)$  with pivot measure  $\mu_\theta^{1\rightarrow 2}$ . Then,
$$
\mathrm {S W G G} _ {2} ^ {2} \left(\mu_ {1}, \mu_ {2}, \theta\right) = 2 W _ {2} ^ {2} \left(\mu_ {1}, \mu_ {\theta} ^ {1 \rightarrow 2}\right) + 2 W _ {2} ^ {2} \left(\mu_ {\theta} ^ {1 \rightarrow 2}, \mu_ {2}\right) - 4 W _ {2} ^ {2} \left(\mu_ {g, \theta} ^ {1 \rightarrow 2}, \mu_ {\theta} ^ {1 \rightarrow 2}\right). \tag {14}
$$
The proof is in Supp.10.1. From Proposition 4.2, SWGG is the  $\mu_{\theta}^{1\rightarrow 2}$ -based Wasserstein distance between  $\mu_{1}$  and  $\mu_{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.