Subspace Robust Wasserstein Distances

Francois-Pierre Paty 1 Marco Cuturi 2 1

AbstractMaking sense of Wasserstein distances betweendiscrete measures in high-dimensional settings re-mains a challenge. Recent work has advocateda two-step approach to improve robustness andfacilitate the computation of optimal transport, us-ing for instance projections on random real lines,or a preliminary quantization of the measures toreduce the size of their support. We proposein this work a “max-min” robust variant of theWasserstein distance by considering the maximalpossible distance that can be realized betweentwo measures, assuming they can be projected or-thogonally on a lower k-dimensional subspace.Alternatively, we show that the corresponding“min-max” OT problem has a tight convex relax-ation which can be cast as that of finding an opti-mal transport plan with a low transportation cost,where the cost is alternatively defined as the sumof the k largest eigenvalues of the second ordermoment matrix of the displacements (or match-ings) corresponding to that plan (the usual OT def-inition only considers the trace of that matrix). Weshow that both quantities inherit several favorableproperties from the OT geometry. We proposetwo algorithms to compute the latter formulationusing entropic regularization, and illustrate theinterest of this approach empirically.

1. IntroductionThe optimal transport (OT) toolbox (Villani, 2009) is gain-ing popularity in machine learning, with several applicationsto data science outlined in the recent review paper (Peyre &Cuturi, 2019). When using OT on high-dimensional data,practitioners are often confronted to the intrinsic instabilityof OT with respect to input measures. A well known resultstates for instance that the sample complexity of Wasserstein

1CREST-ENSAE, Palaiseau, France 2Google Brain, Paris,France. Correspondence to: Francois-Pierre Paty <fran-cois.pierre.paty@ensae.fr>.

Proceedings of the 36 th International Conference on MachineLearning, Long Beach, California, PMLR 97, 2019. Copyright2019 by the author(s).

distances can grow exponentially in dimension (Dudley,1969; Fournier & Guillin, 2015), which means that an irre-alistic amount of samples from two continuous measures isneeded to approximate faithfully the true distance betweenthem. This result can be mitigated when data lives on lowerdimensional manifolds as shown in (Weed & Bach, 2017),but sample complexity bounds remain pessimistic even inthat case. From a computational point of view, that prob-lem can be interpreted as that of a lack of robustness andinstability of OT metrics with respect to their inputs. Thisfact was already a common concern of the community whenthese tools were first adopted, as can be seen in the use of`1 costs (Ling & Okada, 2007) or in the common practiceof thresholding cost matrices (Pele & Werman, 2009).

Regularization The idea to trade off a little optimalityin exchange for more regularity is by now considered acrucial ingredient to make OT work in data sciences. Aline of work initiated in (Cuturi, 2013) advocates adding anentropic penalty to the original OT problem, which resultsin faster and differentiable quantities, as well as improvedsample complexity bounds (Genevay et al., 2019). Follow-ing this, other regularizations (Dessein et al., 2018), notablyquadratic (Blondel et al., 2018), have also been investi-gated. Sticking to an entropic regularization, one can alsointerpret the recent proposal by Altschuler et al. (2018b)to approximate Gaussian kernel matrices appearing in theregularized OT problem with Nystrom-type factorizations(or exact features using a Taylor expansion (Cotter et al.,2011) as in (Altschuler et al., 2018a)), as robust approachesthat are willing to tradeoff yet a little more cost optimality inexchange for faster Sinkhorn iterations. In a different line ofwork, quantizing first the measures to be compared beforecarrying out OT on the resulting distriutions of centroids is afruitful alternative (Canas & Rosasco, 2012) which has beenrecently revisited in (Forrow et al., 2019). Another approachexploits the fact that the OT problem between two distribu-tions on the real line boils down to the direct comparison oftheir generalized quantile functions (Santambrogio, 2015,§2). Computing quantile functions only requires sortingvalues, with a mere log-linear complexity. The sliced ap-proximation of OT (Rabin et al., 2011) consists in projectingtwo probability distributions on a given line, compute theoptimal transport cost between these projected values, andthen repeat this procedure to average these distances over

Subspace Robust Wasserstein Distances

-1 -0.5 0 0.5 1

-1

-0.5

0

0.5

1

Optimal Vπ⋆

〈π⋆, C〉 = 1.24

λ1(Vπ⋆) = 1.04

-1 -0.5 0 0.5 1

-1

-0.5

0

0.5

1

Random feasible π

〈π, C〉 = 1.43

λ1(Vπ) = 1.15

0 50 100 150

θ in degrees0

0.5

1

1.5

2

Gap between P1 and S1

OT cost for θλ1(Vπ⋆(θ))λ2(Vπ⋆(θ))

-1 -0.5 0 0.5 1

-1

-0.5

0

0.5

1

1-Projection robust π⋆(θ∗)

〈π⋆(θ∗), C〉 = 1.26

λ1(Vπ⋆(θ∗)) = 0.853

Optimal projection θ∗

Figure 1. We consider two discrete measures (red and blue dots) on the plane. The left-most plot shows the optimal transport betweenthese points, in which the width of the segment is proportional to the mass transported between two locations. The total cost is displayedin the lower right part of the plot as 〈π?, C 〉, where C is the pairwise squared-Euclidean distance matrix. The largest eigenvalue of thecorresponding second order moment matrix Vπ? of displacements, see (1), is given below. As can be expected and seen in the second plot,choosing a random transportation plan yields a higher cost. The third plot displays the most robust projection direction (green line), thatupon which the OT cost of these point clouds is largest once projected. The maximal eigenvalue of the second order moment matrix(still in dimension 2) is smaller than that obtained with the initial OT plan. Finally, we plot as a function of the angle θ between (0, 180)the OT cost (which, in agreement with the third plot, is largest for the angle corresponding to the green line of the third plot) as well asthe corresponding maximal eigenvalue of the second order moment of the optimal plan corresponding to each of these angles θ. Themaximum of the red curve, as well as the minimum reached by the dark blue one, correspond respectively to the values of the projectionPk and subspace Sk robust Wasserstein distances described in §3. They happen to coincide in this example, but one may find examples inwhich they do not, as can be seen in Figure 11 (supplementary material). The smallest eigenvalue is given for illustrative purposes only.

several random lines. This approach can be used to definekernels (Kolouri et al., 2016), compute barycenters (Bon-neel et al., 2015) but also to train generative models (Kolouriet al., 2018; Deshpande et al., 2018). Beyond its practicalapplicability, this approach is based on a perhaps surpris-ing point-of-view: OT on the real line may be sufficient toextract geometric information from two high-dimensionaldistributions. Our work builds upon this idea, and morecandidly asks what can be extracted from a little more thana real line, namely a subspace of dimension k ≥ 2. Ratherthan project two measures on several lines, we consider inthis paper projecting them on a k-dimensional subspace thatmaximizes their transport cost. This results in optimizingthe Wasserstein distance over the ground metric, which wasalready considered for supervised learning (Cuturi & Avis,2014; Flamary et al., 2018).

Contributions This optimal projection translates into a“max-min” robust OT problem with desirable features. Al-though that formulation cannot be solved with convexsolvers, we show that the corresponding “min-max” prob-lem admits on the contrary a tight convex relaxation andalso has an intuitive interpretation. To see that, one can firstnotice that the usual 2-Wasserstein distance can be describedas the minimization of the trace of the second order momentmatrix of the displacements associated with a transport plan.We show that computing a maximally discriminating opti-mal k dimensional subspace in this “min-max” formulationcan be carried out by minimizing the sum of the k largesteigenvalues (instead of the entire trace) of that second order

moment matrix. A simple example summarizing the linkbetween these two “min-max” and “max-min” quantitiesis given in Figure 1. That figure considers a toy examplewhere points in dimension d = 2 are projected on linesk = 1, our idea is designed to work for larger k and d, asshown in §6.

Paper structure We start this paper with background ma-terial on Wasserstein distances in §2 and present an alter-native formulation for the 2-Wasserstein distance using thesecond order moment matrix of displacements described ina transport plan. We define in §3 our “max-min” and “min-max” formulations for, respectively, projection (PRW) andsubspace (SRW) robust Wasserstein distances. We study thegeodesic structure induced by the SRW distance on the spaceof probability measures in §4, as well as its dependence onthe dimension parameter k. We provide computational toolsto evaluate SRW using entropic regularization in §5. Weconclude the paper with experiments in §6 to validate andillustrate our claims, on both simulated and real datasets.

2. Background on Optimal TransportFor d ∈ N, we write JdK = 1, ..., d. Let P(Rd) be theset of Borel probability measures in Rd, and let

P2(Rd) =

µ ∈P(Rd)

∣∣∣∣ ∫ ‖x‖2 dµ(x) <∞.

Subspace Robust Wasserstein Distances

Monge and Kantorovich Formulations of OT Forµ, ν ∈P(Rd), we write Π(µ, ν) for the set of couplings

Π(µ, ν) = π ∈P(Rd × Rd) s.t. ∀A,B ⊂ Rd Borel,

π(A× Rd) = µ(A), π(Rd ×B) = ν(B),

and their 2-Wasserstein distance is defined as

W2(µ, ν) :=

(inf

π∈Π(µ,ν)

∫‖x− y‖2 dπ(x, y)

)1/2

.

Because we only consider quadratic costs in the remainderof this paper, we drop the subscript 2 in our notation andwill only useW to denote the 2-Wasserstein distance. ForBorel X ,Y ⊂ Rd, Borel T : X → Y and µ ∈ P(X ), wedenote by T#µ ∈ P(Y) the push-forward of µ by T , i.e.the measure such that for any Borel set A ⊂ Y ,

T#µ(A) = µ(T−1(A)

).

The Monge (1781) formulation of optimal transport is, whenthis minimization is feasible, equivalent to that of Kan-torovich, namely

W(µ, ν) =

(inf

T :T#µ=ν

∫‖x− T (x)‖2 dµ(x)

)1/2

.

W as Trace-minimization For any coupling π, we definethe d× d second order displacement matrix

Vπ :=

∫(x− y)(x− y)T dπ(x, y). (1)

Notice that when a coupling π corresponds to a Mongemap, namely π = (Id, T )#µ, then one can interpret evenmore naturally Vπ as the second order moment of all dis-placement (x − T (x))(x − T (x))T weighted by µ. Withthis convention, we remark that the total cost of a cou-pling π is equal to the trace of Vπ , using the simple identitytrace(x − y)(x − y)T = ‖x − y‖2 and the linearity ofthe integral sum. Computing theW distance can thereforebe interpreted as minimizing the trace of Vπ. This simpleobservation will play an important role in the next section,and more specifically the study of λl(Vπ), the l-th largesteigenvalue of Vπ .

3. Subspace Robust Wasserstein DistancesWith the conventions and notations provided in §2, we con-sider here different robust formulations of the Wassersteindistance. Consider first for k ∈ JdK, the Grassmannian ofk-dimensional subspaces of Rd :

Gk =E ⊂ Rd | dim(E) = k

.

For E ∈ Gk, we note PE the orthogonal projector ontoE. Given two measures µ, ν ∈ P2(Rd), a first attempt at

computing a robust version of W(µ, ν) is to consider theworst possible OT cost over all possible low dimensionalprojections:Definition 1. For k ∈ JdK, the k-dimensional projectionrobust 2-Wasserstein (PRW) distance between µ and ν is

Pk(µ, ν) = supE∈Gk

W(PE#µ, PE#ν

).

As we show in the supplementary material, this quantityis well posed and itself worthy of interest, yet difficult tocompute. In this paper, we focus our attention on the corre-sponding “min-max” problem, to define the k-dimensionalsubspace robust 2-Wasserstein (SRW) distance:Definition 2. For k ∈ JdK, the k-dimensional subspacerobust 2-Wasserstein distance between µ and ν is

Sk(µ, ν) = infπ∈Π(µ,ν)

supE∈Gk

[∫‖PE(x− y)‖2dπ(x, y)

]1/2

Remark 1. Both quantities Sk andPk can be interpreted asrobust variants of theW distance. By a simple applicationof weak duality we have that Pk(µ, ν) ≤ Sk(µ, ν).Lemma 1. Optimal solutions for Sk exist, i.e.

Sk(µ, ν) = minπ∈Π(µ,ν)

maxE∈Gk

[∫‖PE(x− y)‖2dπ(x, y)

]1/2

We show next that the SRW variant Sk can be elegantlyreformulated as a function of the eigendecomposition of thedisplacement second-order moment matrix Vπ (1):Lemma 2. For k ∈ JdK and µ, ν ∈P2(Rd), one has

S2k(µ, ν) = min

π∈Π(µ,ν)max

U∈Rk×dUUT=Ik

∫‖Ux− Uy‖2 dπ(x, y)

= minπ∈Π(µ,ν)

k∑l=1

λl(Vπ).

This characterization as a sum of eigenvalues will be cru-cial to study theoretical properties of Sk. Subspace robustWasserstein distances can in fact be interpreted as a convexrelaxation of projection robust Wasserstein distances: theycan be computed as the maximum of a concave functionover a convex set, which will make computations tractable.Theorem 1. For k ∈ JdK and µ, ν ∈P2(Rd),

S2k(µ, ν) = min

π∈Π(µ,ν)max

0ΩItrace(Ω)=k

∫d2

Ω dπ (2)

= max0ΩI

trace(Ω)=k

minπ∈Π(µ,ν)

∫d2

Ω dπ (3)

= max0ΩI

trace(Ω)=k

W2(

Ω1/2#µ,Ω

1/2#ν)

(4)

Subspace Robust Wasserstein Distances

where dΩ stands for the Mahalanobis distance

d2Ω(x, y) = (x− y)TΩ(x− y).

We can now prove that both PRW and SRW variants are,indeed, distances over P2(Rd).

Proposition 1. For k ∈ JdK, both Pk and Sk are distancesover P2(Rd).

Proof. Symmetry is clear for both objects, and for µ ∈P2(Rd), Sk(µ, µ) = Pk(µ, µ) = 0. Let µ, ν ∈ P2(Rd)such that Sk(µ, ν) = 0. Then Pk(µ, ν) = 0 and for anyE ∈ Gk, W(PE#µ, PE#ν) = 0, i.e. PE#µ = PE#ν.Lemma 7 (in the supplementary material) then shows thatµ = ν. For the triangle inequalities, let µ0, µ1, µ2 ∈P2(Rd). Let Ω? ∈ 0 Ω I, trace(Ω) = k be op-timal between µ0 and µ2. Using the triangle inequalities forthe Wasserstein distance,

Sk(µ0, µ2) =W[Ω

1/2? #µ0,Ω

1/2? #µ2

]≤ W

[Ω

1/2? #µ0,Ω

1/2? #µ1

]+W

[Ω

1/2? #µ1,Ω

1/2? #µ2

]≤ sup

0ΩItrace(Ω)=k

W[Ω1/2

#µ0,Ω1/2

#µ1

]+ sup

0ΩItrace(Ω)=k

W[Ω1/2

#µ1,Ω1/2

#µ2

]= Sk(µ0, µ1) + Sk(µ1, µ2).

The same argument, used this time with projections, yieldsthe triangle inequalities for Pk.

4. Geometry of Subspace Robust DistancesWe prove in this section that SRW distances share severalfundamental geometric properties with the Wasserstein dis-tance. The first one states that distances between Diracsmatch the ground metric:

Lemma 3. For x, y ∈ Rd and k ∈ JdK,

Sk(δx, δy) = ‖x− y‖.

Metric Equivalence. Subspace robust Wasserstein dis-tances Sk are equivalent to the Wasserstein distanceW:

Proposition 2. For k ∈ JdK, Sk is equivalent toW . Moreprecisely, for µ, ν ∈P2(Rd),√

k

dW(µ, ν) ≤ Sk(µ, ν) ≤ W(µ, ν).

Moreover, the constants are tight since

Sk(δx, δy) = W(δx, δy)

Sk(δ0, σ) =√

kdW(δ0, σ)

where δx, δy, δ0 are Dirac masses at points x, y, 0 ∈ Rd andσ is the uniform probability distribution over the centeredunit sphere in Rd.

Dependence on the dimension. We fix µ, ν ∈ P2(Rd)and we ask the following question : how does Sk(µ, ν)depend on the dimension k ∈ JdK ? The following lemmagives a result in terms of eigenvalues of Vπk , where πk ∈Π(µ, ν) is optimal for some dimension k, then we translatein Proposition 3 this result in terms of Sk.Lemma 4. Let µ, ν ∈P2(Rd). For any k ∈ Jd− 1K,

λk+1(Vπk+1) ≤ S2

k+1(µ, ν)− S2k(µ, ν)

≤ λk+1(Vπk)

where for L ∈ JdK, πL ∈ arg minπ∈Π(µ,ν)

L∑l=1

λl(Vπ).

Proposition 3. Let µ, ν ∈ P2(Rd). The sequence k 7→S2k(µ, ν) is increasing and concave. In particular, for k ∈

Jd− 1K,

S2k+1(µ, ν)− S2

k(µ, ν) ≥ W2(µ, ν)− S2

k(µ, ν)

d− k.

Moreover, for any k ∈ Jd− 1K,

Sk(µ, ν) ≤ Sk+1(µ, ν) ≤√k + 1

kSk(µ, ν).

Geodesics We have shown in Proposition 2 that for anyk ∈ JdK,

(P2(Rd),Sk

)is a metric space with the same

topology as that of the Wasserstein space(P2(Rd),W

).

We conclude this section by showing that(P2(Rd),Sk

)is

in fact a geodesic length space, and exhibits explicit constantspeed geodesics. This can be used to interpolate betweenmeasures in Sk sense.Proposition 4. Let µ, ν ∈P2(Rd) and k ∈ JdK. Take

π∗ ∈ arg minπ∈Π(µ,ν)

k∑l=1

λl(Vπ)

and let ft(x, y) = (1− t)x+ ty. Then the curve

t 7→ µt := ft#π∗

is a constant speed geodesic in(P2(Rd),Sk

)connecting

µ and ν. Consequently,(P2(Rd),Sk

)is a geodesic space.

Proof. We first show that for any s, t ∈ [0, 1],

Sk(µs, µt) = |t− s|Sk(µ, ν)

by computing the cost of the transport plan π(s, t) =(fs, ft)#π

∗ ∈ Π(µs, µt) and using the triangular inequality.Then the curve (µt) has constant speed

|µ′t| = limε→0

Sk(µt+ε, µt)

|ε|= Sk(µ, ν),

Subspace Robust Wasserstein Distances

and the length of the curve (µt) is

sup

n−1∑i=0

Sk(µti , µti+1)

∣∣∣∣ n ≥ 10 = t0 < ... < tn = 1

= Sk(µ, ν),

i.e. (µt) is a geodesic connecting µ and ν.

5. ComputationWe provide in this section algorithms to compute the saddlepoint solution of Sk. µ, ν are now discrete with respectivelyn and m points and weights a and b : µ :=

∑ni=1 aiδxi

and ν :=∑mj=1 biδyj . For k ∈ JdK, three different objects

are of interest: (i) the value of Sk(µ, ν), (ii) an optimalsubspace E∗ obtained through the relaxation for SRW, (iii)an optimal transport plan solving SRW. A subspace canbe used for dimensionality reduction, whereas an optimaltransport plan can be used to compute a geodesic, i.e. tointerpolate between µ and ν.

5.1. Computational challenges to approximate Sk

We observe that solving minπ∈Π(µ,ν)

∑kl=1 λl(Vπ) is chal-

lenging. Considering a direct projection onto the transporta-tion polytope

Π(µ, ν) =π ∈ Rn×m |π1m = a, πT1n = b

would result in a costly quadratic network flow problem.The Frank-Wolfe algorithm, which does not require suchprojections, cannot be used directly because the applicationπ 7→

∑kl=1 λl(Vπ) is not smooth.

On the other hand, thanks to Theorem 1, solving themaximization problem is easier. Indeed, we can projectonto the set of constraints R = Ω ∈ Rd×d | 0 Ω I ; trace(Ω) = k using Dykstra’s projection algo-rithm (Boyle & Dykstra, 1986). In this case, we will onlyget the value of Sk(µ, ν) and an optimal subspace, but notnecessarily the actual optimal transport plan due to the lackof uniqueness for OT plans in general.

Smoothing It is well known that saddle points are hardto compute for a bilinear objective (Hammond, 1984).Computations are greatly facilitated by adding smoothness,which allows the use of saddle point Frank-Wolfe algo-rithms (Gidel et al., 2017). Out of the two problems, themaximization problem is seemingly easier. Indeed, we canleverage the framework of regularized OT (Cuturi, 2013) tooutput, using Sinkhorn’s algorithm, a unique optimal trans-port plan π? at each inner loop of the maximization. Tosave time, we remark that initial iterations can be solvedwith a low accuracy by limiting the number of iterations,and benefit from warm starts, using the scalings computedat the previous iteration, see (Peyre & Cuturi, 2019, §4).

Algorithm 1 Projected supergradient method for SRWInput: Measures (xi, ai) and (yj , bj), dimension k,learning rate τ0π ← OT((x, a), (y, b), cost = ‖ · ‖2)U ← top k eigenvectors of VπInitialize Ω = UUT ∈ Rd×dfor t = 0 to max iter doπ ← OT((x, a), (y, b), cost = d2

Ω)τ = τ0/(t+ 1)Ω← ProjR [Ω + τVπ]

end forOutput: Ω, 〈Ω |Vπ〉

5.2. Projected Supergradient Method for SRW

In order to compute SRW and an optimal subspace, we cansolve equation (3) by maximizing the concave function

f : Ω 7→ minπ∈Π(µ,ν)

∑i,j

d2Ω(xi, yj)πi,j = min

π∈Π(µ,ν)〈Ω |Vπ〉

over the convex set R. Since f is not differentiable, butonly superdifferentiable, we can only use a projected super-gradient method. This algorithm is outlined in Algorithm 1.Note that by Danskin’s theorem, for any Ω ∈ R,

∂f(Ω) = Conv

Vπ∗

∣∣∣∣π∗ ∈ arg minπ∈Π(µ,ν)

〈Ω |Vπ〉

.

5.3. Frank-Wolfe using Entropy Regularization

Entropy-regularized optimal transport can be used to com-pute a unique optimal plan given a subspace. Let γ > 0be the regularization strength. In this case, we want tomaximize the concave function

fγ : Ω 7→ minπ∈Π(µ,ν)

〈Ω |Vπ〉+ γ∑i,j

πi,j [log(πi,j)− 1]

over the convex setR. Since for all Ω ∈ R, there is a uniqueπ∗ minimizing π 7→ 〈Ω |Vπ〉+ γ

∑i,j πi,j [log(πi,j)− 1],

fγ is differentiable. Instead of running a projected gradi-ent ascent on Ω ∈ R, we propose to use the Frank-Wolfealgorithm when the regularization strength is positive. In-deed, there is no need to tune a learning rate in Frank-Wolfe,making it easier to use. We only need to compute, for fixedπ ∈ Π(µ, ν), the maximum overR of Ω 7→ 〈Ω |Vπ〉:Lemma 5. For π ∈ Π(µ, ν), compute the eigendecompo-sition of Vπ = U diag (λ1, ..., λd)U

T with λ1 ≥ ... ≥ λd.Then for k ∈ JdK, Ω = U diag ([1k,0d−k])UT solves

max0ΩI

trace(Ω)=k

∫d2

Ω dπ.

This algorithm is outlined in algorithm 2.

Subspace Robust Wasserstein Distances

Algorithm 2 Frank-Wolfe algorithm for regularized SRWInput: Measures (xi, ai) and (yj , bj), dimension k, reg-ularization strength γ > 0, precision ε > 0π ← reg OT((x, a), (y, b), reg = γ, cost = ‖ · ‖2)U ← top k eigenvectors of VπInitialize Ω = UUT ∈ Rd×dfor t = 0 to max iter doπ ← reg OT((x, a), (y, b), reg = γ, cost = d2

Ω)U ← top k eigenvectors of Vπif∑kl=1 λl(Vπ)− 〈Ω |Vπ〉 ≤ ε〈Ω |Vπ〉 then

breakend ifΩ← U diag ([1k,0d−k])UT

τ = 2/(2 + t)

Ω← (1− τ)Ω + τ Ωend forOutput: Ω, π, 〈Ω |Vπ〉

5.4. Initialization and Stopping Criterion

We propose to initialize Algorithms 1 and 2 withΩ0 = UUT where U ∈ Rd×k is the matrix of the top keigenvectors (i.e. the eigenvectors associated with the topk eigenvalues) of Vπ∗ and π∗ is an optimal transport planbetween µ and ν. In other words, Ω0 is the projectionmatrix onto the k first principal components of thetransport-weighted displacement vectors. Note that Ω0

would be optimal is π∗ were optimal for the min-maxproblem, and that this initialization only costs the equivalentof one iteration.

When entropic regularization is used, Sinkhorn algorithm isrun at each iteration of Algorithms 1 and 2. We proposeto initialize the potentials in Sinkhorn algorithm with thelatest computed potentials, so that the number of iterationsin Sinkhorn algorithm should be small after a few iterationsof Algorithms 1 or 2.

We sometimes need to compute Sk(µ, ν) for all k ∈ JdK,for example to choose the optimal k with an “elbow” rule.To speed the computations up, we propose to compute thissequence iteratively from k = d to k = 1. At each iteration,i.e. for each dimension k, we initialize the algorithm withΩ0 = UUT , where U ∈ Rd×k is the matrix of the topk eigenvectors of Vπk+1

and πk+1 is the optimal transportplan for dimension k + 1. We also initialize the Sinkhornalgorithm with the latest computed potentials.

Instead of running a fixed number of iterations in Algo-rithm 2, we propose to stop the algorithm when the compu-tation error is smaller than a fixed threshold ε. The compu-

tation error at iteration t is:

|Sk(µ, ν)− Sk(t)|Sk(µ, ν)

≤ ∆(t)

Sk(t)

where Sk(t) is the computed “max-min” value and ∆(t)is the duality gap at iteration t. We stop as soon as∆(t)/Sk(t) ≤ ε.

6. ExperimentsWe first compare SRW with the experimental setup used toevaluate FactoredOT (Forrow et al., 2019). We then studythe ability of SRW distances to capture the dimension ofsampled measures by looking at their value for increasingdimensions k, as well as their robustness to noise.

6.1. Fragmented Hypercube

We first consider µ = U([−1, 1])d to be uniform over anhypercube, and ν = T#µ the pushforward of µ under themap T (x) = x+2 sign(x)(

∑k∗

k=1 ek), where sign is takenelementwise, k∗ ∈ JdK and (e1, ..., ed) is the canonical basisof Rd. The map T splits the hypercube into four differenthyperrectangles. T is a subgradient of a convex function, soby Brenier’s theorem (1991) it is an optimal transport mapbetween µ and ν = T#µ and

W2(µ, ν) =

∫‖x− T (x)‖2 dµ(x) = 4k∗.

Note that for any x, the displacement vector T (x)− x liesin the k∗-dimensional subspace spane1, ..., ek∗ ∈ Gk∗ ,which is optimal. This means that for k ≥ k∗, S2

k(µ, ν)is constant equal to 4k∗. We show the interest of plotting,based on two empirical distributions µ from µ and ν fromν, the sequence k 7→ S2

k(µ, ν), for different values of k∗.That sequence is increasing concave by proposition 3, andincreases more slowly after k = k∗, as can be seen onFigure 2. This is the case because the last d−k∗ dimensionsonly represent noise, but is recovered in our plot.

Figure 2. S2k(µ, ν) depending on the dimension k ∈ JdK, for k∗ ∈

2, 4, 7, 10, where µ, ν are empirical measures from µ and νrespectively with 100 points each. Each curve is the mean over100 samples, and shaded area show the min and max values.

We consider next k∗ = 2, and choose from the resultof Figure 2, k = 2. We look at the estimation error

Subspace Robust Wasserstein Distances

|W2(µ, ν)− S2k(µ, ν)| where µ, ν are empirical measures

from µ and ν respectively with n points each. In Figure 3,we plot this estimation error depending on the number ofpoints n. In Figure 4, we plot the subspace estimation er-ror ‖Ω∗ − Ω‖ depending on n, where Ω∗ is the optimalprojection matrix onto spane1, e2.

Figure 3. Mean estimation error over 500 random samples for npoints, n ∈ 25, 50, 100, 250, 500, 1000. The shaded areas rep-resent the 10%-90% and 25%-75% quantiles over the 500 samples.

Figure 4. Mean estimation of the subspace estimation error over500 samples, depending on n ∈ 25, 50, 100, 250, 500, 1000.The shaded areas represent the 10%-90% and 25%-75% quantilesover the 500 samples.

We also plot the optimal transport plan (in the sense ofW ,Figure 5 left) and the optimal transport plan (in the sense ofS2) between µ and ν (with n = 250 points each, Figure 5right).

Figure 5. Fragmented hypercube, n = 250, d = 30. Optimalmapping in the Wasserstein space (left) and in the SRW space(right). Geodesics in the SRW space are robust to statistical noise.

6.2. Robustness, with 20-D Gaussians

We consider µ = N (0,Σ1) and ν = N (0,Σ2), withΣ1,Σ2 ∈ Rd×d semidefinite positive of rank k. It meansthat the supports of µ and ν are k-dimensional subspaces ofRd. Although those two subspaces are k-dimensional, they

may be different. Since the union of two k-dimensionalsubspaces is included in a 2k-dimensional subspace, forany l ≥ 2k, S2

l (µ, ν) =W2(µ, ν).

For our experiment, we simulated 100 independent couplesof covariance matrices Σ1,Σ2 in dimension d = 20, eachhaving independently a Wishart distribution with k = 5degrees of freedom. For each couple of matrices, we drawn = 100 points from N (0,Σ1) and N (0,Σ2) and con-sidered µ and ν the empirical measures on those points.In Figure 6, we plot the mean (over the 100 samples) ofl 7→ S2

l (µ, ν)/W2(µ, ν). We plot the same curve for noisydata, where each point was added a N (0, I) random vector.With moderate noise, the data is only approximately on thetwo k = 5-dimensional subspaces, but the SRW does notvary too much.

Figure 6. Mean normalized SRW distance, depending on the di-mension. The shaded area show the 10%-90% and 25%-75%quantiles and the minimum and maximum values over the 100samples.

6.3. Sk is Robust to Noise

As in experiment 6.2, we consider 100 independent samplesof couples Σ1,Σ2 ∈ Rd×d, each following independentlya Wishart distribution with k = 5 degrees of freedom. Foreach couple, we draw n = 100 points from N (0,Σ1) andN (0,Σ2) and consider the empirical measures µ and νon those points. We then gradually add Gaussian noiseσN (0, I) to the points, giving measures µσ , νσ . In Figure 7,we plot the mean (over the 100 samples) of the relativeerrors

σ 7→ |S25 (µσ, νσ)− S2

5 (µ0, ν0)|S2

5 (µ0, ν0)

and

σ 7→ |W2(µσ, νσ)−W2(µ0, ν0)|W2(µ0, ν0)

.

Note that for small noise level, the imprecision in thecomputation of the SRW distance adds up to the errorcaused by the added noise. SRW distances seem morerobust to noise than the Wasserstein distance when the noisehas moderate to high variance.

Subspace Robust Wasserstein Distances

Figure 7. Mean SRW distance over 100 samples, depending onthe noise level. Shaded areas show the min-max values and the10%-90% quantiles.

6.4. Computation time

We consider the Fragmented Hypercube experiment, withincreasing dimension d and fixed k∗ = 2. Using k = 2 andAlgorithm 2 with γ = 0.1 and stopping threshold ε = 0.05,we plot in Figure 8 the mean computation time of both SRWand Wasserstein distances on GPU, over 100 random sam-plings with n = 100. It shows that SRW computation isquadratic in dimension d, because of the eigendecomposi-tion of matrix Vπ in Algorithm 2.

Figure 8. Mean computation times on GPU (log-log scale). Theshaded areas show the minimum and maximum values over the100 experiments.

6.5. Real Data Experiment

We consider the scripts of seven movies. Each script is trans-formed into a list of words, and using word2vec (Mikolovet al., 2018), into a measure over R300 where the weightscorrespond to the frequency of the words. We then computethe SRW distance between all pairs of films: see Figure 9for the SRW values. Movies with a same genre or thematictend to be closer to each other: this can be visualized byrunning a two-dimensional metric multidimensional scaling(mMDS) on the SRW distances, as shown in Figure 10 (left).

In Figure 10 (right), we display the projection of the twomeasures associated with films Kill Bill Vol.1 and Inter-stellar onto their optimal subspace. We compute the first(weighted) principal component of each projected measure,and find among the whole dictionary their 5 nearest neigh-bors in terms of cosine similarity. For Kill Bill Vol.1 , theseare: ’swords’, ’hull’, ’sword’, ’ice’, ’blade’. For Interstel-lar, they are: ’spacecraft’, ’planets’, ’satellites’, ’asteroids’,

D G I KB1 KB2 TM TD 0 0.186 0.186 0.195 0.203 0.186 0.171G 0.186 0 0.173 0.197 0.204 0.176 0.185I 0.186 0.173 0 0.196 0.203 0.171 0.181

KB1 0.195 0.197 0.196 0 0.165 0.190 0.180KB2 0.203 0.204 0.203 0.165 0 0.194 0.180TM 0.186 0.176 0.171 0.190 0.194 0 0.183T 0.171 0.185 0.181 0.180 0.180 0.183 0

Figure 9. S2k distances between different movie scripts. Bold

values correspond to the minimum of each line. D=Dunkirk,G=Gravity, I=Interstellar, KB1=Kill Bill Vol.1, KB2=Kill BillVol.2, TM=The Martian, T=Titanic.

’planet’. The optimal subspace recovers the semantic dis-similarities between the two films.

Figure 10. Left: Metric MDS projection for the distances of Fig-ure 9. Right: Optimal 2-dimensional projection between Kill BillVol.1 (red) and Interstellar (blue). Words appearing in both scriptsare displayed in violet. For clarity, only the 30 most frequent wordsof each script are displayed.

7. ConclusionWe have proposed in this paper a new family of OT distanceswith robust properties. These distances take a particular in-terest when used with a squared-Euclidean cost, in whichcase they have several properties, both theoretical and com-putational. These distances share important properties withthe 2-Wasserstein distance, yet seem far more robust torandom perturbation of the data and able to capture bettersignal. We have provided algorithmic tools to compute theseSRW distance. They come at a relatively modest overhead,given that they require using regularized OT as the innerloop of a FW type algorithm. Future work includes theinvestigation of even faster techniques to carry out thesecomputations, eventually automatic differentiation schemesas those currently benefitting the simple use of Sinkhorndivergences.

Acknowledegments. Both authors acknowledge the sup-port of a ”Chaire d’excellence de l’IDEX Paris Saclay”. Wethank P. Rigollet and J. Weed for their remarks and theirvaluable input.

Subspace Robust Wasserstein Distances

ReferencesAltschuler, J., Bach, F., Rudi, A., and Weed, J. Approximat-

ing the quadratic transportation metric in near-linear time.arXiv preprint arXiv:1810.10046, 2018a.

Altschuler, J., Bach, F., Rudi, A., and Weed, J. Mas-sively scalable sinkhorn distances via the nystrom method.arXiv preprint arXiv:1812.05189, 2018b.

Blondel, M., Seguy, V., and Rolet, A. Smooth and sparse op-timal transport. In International Conference on ArtificialIntelligence and Statistics, pp. 880–889, 2018.

Bonneel, N., Rabin, J., Peyre, G., and Pfister, H. Sliced andRadon Wasserstein barycenters of measures. Journal ofMathematical Imaging and Vision, 51(1):22–45, 2015.

Boyle, J. P. and Dykstra, R. L. A method for finding projec-tions onto the intersection of convex sets in hilbert spaces.In Advances in order restricted statistical inference, pp.28–47. Springer, 1986.

Brenier, Y. Polar factorization and monotone rearrangementof vector-valued functions. Communications on Pure andApplied Mathematics, 44(4):375–417, 1991.

Canas, G. and Rosasco, L. Learning probability measureswith respect to optimal transport metrics. In Pereira, F.,Burges, C. J. C., Bottou, L., and Weinberger, K. Q. (eds.),Advances in Neural Information Processing Systems 25,pp. 2492–2500. 2012.

Cotter, A., Keshet, J., and Srebro, N. Explicit ap-proximations of the gaussian kernel. arXiv preprintarXiv:1109.4603, 2011.

Cuturi, M. Sinkhorn distances: lightspeed computation ofoptimal transport. In Advances in Neural InformationProcessing Systems 26, pp. 2292–2300, 2013.

Cuturi, M. and Avis, D. Ground metric learning. TheJournal of Machine Learning Research, 15(1):533–564,2014.

Deshpande, I., Zhang, Z., and Schwing, A. G. Genera-tive modeling using the sliced wasserstein distance. InThe IEEE Conference on Computer Vision and PatternRecognition (CVPR), June 2018.

Dessein, A., Papadakis, N., and Rouas, J.-L. Regularizedoptimal transport and the rot mover’s distance. Journalof Machine Learning Research, 19(15):1–53, 2018.

Dudley, R. M. The speed of mean Glivenko-Cantelli conver-gence. Annals of Mathematical Statistics, 40(1):40–50,1969.

Fan, K. On a theorem of weyl concerning eigenvalues oflinear transformations i. Proceedings of the NationalAcademy of Sciences, 35(11):652–655, 1949.

Flamary, R., Cuturi, M., Courty, N., and Rakotomamonjy,A. Wasserstein discriminant analysis. Machine Learning,107(12):1923–1945, 2018.

Forrow, A., Hutter, J.-C., Nitzan, M., Rigollet, P.,Schiebinger, G., and Weed, J. Statistical optimal transportvia factored couplings. 2019.

Fournier, N. and Guillin, A. On the rate of convergence inWasserstein distance of the empirical measure. Probabil-ity Theory and Related Fields, 162(3-4):707–738, 2015.

Genevay, A., Chizat, L., Bach, F., Cuturi, M., and Peyre, G.Sample complexity of sinkhorn divergences. 2019.

Gidel, G., Jebara, T., and Lacoste-Julien, S. Frank-Wolfe al-gorithms for saddle point problems. In Proceedings of the20th International Conference on Artificial Intelligenceand Statistics (AISTATS), 2017.

Hammond, J. H. Solving asymmetric variational inequal-ity problems and systems of equations with generalizednonlinear programming algorithms. PhD thesis, Mas-sachusetts Institute of Technology, 1984.

Kolouri, S., Zou, Y., and Rohde, G. K. Sliced Wassersteinkernels for probability distributions. In Proceedings ofthe IEEE Conference on Computer Vision and PatternRecognition, pp. 5258–5267, 2016.

Kolouri, S., Martin, C. E., and Rohde, G. K. Sliced-wasserstein autoencoder: An embarrassingly simple gen-erative model. arXiv preprint arXiv:1804.01947, 2018.

Ling, H. and Okada, K. An efficient earth mover’s distancealgorithm for robust histogram comparison. IEEE Trans-actions on Pattern Analysis and Machine Intelligence, 29(5):840–853, 2007.

Mikolov, T., Grave, E., Bojanowski, P., Puhrsch, C., andJoulin, A. Advances in pre-training distributed wordrepresentations. In Proceedings of the International Con-ference on Language Resources and Evaluation (LREC2018), 2018.

Monge, G. Memoire sur la theorie des deblais et des rem-blais. Histoire de l’Academie Royale des Sciences, pp.666–704, 1781.

Overton, M. L. and Womersley, R. S. Optimality conditionsand duality theory for minimizing sums of the largesteigenvalues of symmetric matrices. Mathematical Pro-gramming, 62(1-3):321–357, 1993.

Subspace Robust Wasserstein Distances

Pele, O. and Werman, M. Fast and robust earth mover’sdistances. In IEEE 12th International Conference onComputer Vision, pp. 460–467, 2009.

Peyre, G. and Cuturi, M. Computational optimal transport.Foundations and Trends in Machine Learning, 11(5-6),2019.

Rabin, J., Peyre, G., Delon, J., and Bernot, M. Wassersteinbarycenter and its application to texture mixing. In In-ternational Conference on Scale Space and VariationalMethods in Computer Vision, pp. 435–446. Springer,2011.

Santambrogio, F. Optimal transport for applied mathemati-cians. Birkhauser, 2015.

Villani, C. Optimal Transport: Old and New, volume 338.Springer Verlag, 2009.

Weed, J. and Bach, F. Sharp asymptotic and finite-samplerates of convergence of empirical measures in Wassersteindistance. arXiv preprint arXiv:1707.00087, 2017.