Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent
Authors: Jason Altschuler
Sinho Chewi
Patrik Robert Gerber
Austin J Stromme
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Bures-Wasserstein barycenter
dimension-free convergence
entropic regularization
first-order optimization
geometric median
non-convex optimization
Riemannian optimization

Abstract: We study first-order optimization algorithms for computing the barycenter of Gaussian distributions with respect to the optimal transport metric. Although the objective is geodesically non-convex, Riemannian gradient descent empirically converges rapidly, in fact faster than off-the-shelf methods such as Euclidean gradient descent and SDP solvers. This stands in stark contrast to the best-known theoretical results, which depend exponentially on the dimension. In this work, we prove new geodesic convexity results which provide stronger control of the iterates, yielding a dimension-free convergence rate. Our techniques also enable the analysis of two related notions of averaging, the entropically-regularized barycenter and the geometric median, providing the first convergence guarantees for these problems.

Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent
Jason M. Altschuler
Sinho Chewi MIT schewi cenriimiit..m edu
Patrik Gerber MIT prgerberomit.edu
Austin Stromme MIT astromme 0mi
MIT jasonalt@mit.co edu
Abstract
We study first-order optimization algorithms for computing the barycenter Gaussian distributions with respect the optimal transport metric. Although the objective geodesically non-convex, Riemannian GD empirically converges rapidly, in fact faster than off- fff-the-shelf methods such Euclidean GD and SDP solvers. This stands in stark contrast to the best-known theoretical results for Riemannian GD, which depend exponentially we prove new geodesic convexity results on auxiliary functionals; this provides strong control of the Riemannian GD iterates, ultimately yielding dimension-free convergence rate. Our techniques also enable the analysis two related notions of averaging, the peropicalll-regulaate barycenter and the geometric median, providing the first convergence guarantees for Riemannian GD for these problems.
the dimension. In this work,
1 Introduction
Averaging multiple data sources is among the most classical and fundamental subroutines in data science. However, modern challenge is that data often complicated than points in IR In this paper, we study the task of averaging probability distributions on Rd setting that commonly arises in machine learning and statistics (CD14 Ho+17 SLDI8 Dog+19], computer vision and graphics [Rab+1l] Sol+15) probability theory KS94; and signal processing [El++20] see
also the surveys PC+19 PZ19
and the references within.
The Wasserstein barycenter (ACI1] Rab+11] has emerged as a particularly canonical notion of average. Formally, let (RR) denote the space of probability measures Rd with finite second moment, let P be probability measure over (2(R)) and let W2 denote the 2- 2-Wasserstein distance (i.e. the standard optimal transport distance). Then, the Wasserstein burycenter of a solution of
/ (...
minimize (EP2(Rd)
related notion of average is the entropically-regularize Wasserstein barycenter of P, which is defined to be solution of
/
minimize (Ed)(Rd)
dP
where ent an entropic penalty which allows for incorporating prior knowledge into the average. Lastly, third related notion of average with better robustness properties (e.g., with breakdown point of 50% FFJ09) is the Wasserstein geometric median of P, which is defined to be solution of
f [Walb.)dPP
minimize (R
(3)
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Importantly, while these three notions of average can be defined using other metrics in lieu of W2, the Wasserstein distance critical for many applications since enables capturing geometric features of the distributions [CD14.
The many applications of Wasserstein barycenters and geometric medians (see e.g.. CETO
Rab+11
CD14, G Sol+15 BPCI6 SLD18 [LR20) have inspired significant research into their mathematical and statistical properties since their introduction roughly a decade ago [AC1]
Rab+11
For instance, on the mathematical side it is known that under mild conditions, the barycenter and
geometric median exist, are unique, and admit dual formulations related to multimarginal optimal
transport problems [CETO] ACTI; COO1S]. And on the statistical side, |C AC17 [L17] Big+18]
FLF19 ALP20 Le +21 KSS21 provide sinite-sample and asymptotic statistical guarantees for estimating the Wasserstein barycenter from samples.
However, computing these objects is challenging because of two fundamental obstacles. The first is that in general, barycenters and geometric medians can be complicated distributions which are much harder to represent (even approximately) than the input distributions. The second is that generically, these problems are computationally hard in high dimensions. For instance, Wasserstein barycenters and geometric medians of discrete distributions are NP-hard to compute (even approximately) in high dimension AB21b].
Algorithms for averaging on the Bures- Wasserstein manifold. Nevertheless, these computa- tional obstacles can be potentially averted in parametric settings. This paper as well as most of the literature [ÁIv+16 Bac+18 7P19 Che+20) on parametric settings focuses on the commonly arising setting where P is supported on Gaussian distributions/ As noted in [Ál++16], the Gaussian case also encompasses general location-scant scatter families.
There are two natural families of approaches for designing averaging algorithms this setting. Both exploit the fact that modulo simple re-centering of all distributions, the relevant space of probability distributions is isometric to the Bures- Wasserstein manifold, i.e. the cone of positive semidefinite matrices equipped with the Sares-Wasserstein metric (background is given in Section 2
The first approach is simply recognize the (regularized) Wasserstein barycenter problem as a convex optimization problem over the space of positive semidefinite matrices and apply off-the-shelf methods such as Euclidean GD or semidefinite programming solvers. However, these methods have received little prior attention for good reason: they suffer from severe scalability and parameter-tuning issues (see Section B.3 for numerics). Briefly, the underlying issue that these algorithms operate in the standard Euclidean geometry rather than the natural geometry of optimal transport. Moreover, this approach does not apply to the Wasserstein geometric median problem because even in one dimension, it is non-convex in the Fuclidean geometry.
much more effective approach in practice (see Section 3.3 for numerics) is to exploit the geometry of the Bures-Waas Wasserstein manifold via geodesic optimization. Prior work has extensively pursued this direction, investigating the effectiveness of (stochastic) Riemannian GD for computing Wasserstein Che+20
barycenters, see e.g.. [1v+16 Bac+18 ZP19
Challenges for geodesic optimization over the Bures- Wares-Wasserstenn manifold. Although geodesic optimization is natural for this problem, it comes with several important obstacles: the non-negative curvature of the Bures- Wasserstenn manifold necessitates new tools for analysis, and moreover both non-convex in the Bures- wasse-eassers geometry.
the barycenter and geometric median problems (These two issues are in fact intimately related, see Appendix AA.)) This prevents applying standard results in the geodesic optimization literature (see e.g., (ZS16) Bou20) since in general it is only possible to prove local convergence guarantees for non-convex problems.
For the Wasserstein barycenter problem, it is possible to interpret Riemannian GD (with step size one) as fixed-point iteration, and through this lens establish asymptotic convergence [Álv+16, Bac+18 ZP19]. Obtaining -symattipt rates of convergence is more challenging because requires developing quantitative proxies for the standard convexity inequalities needed to analyze GD. The first such result was achieved by [Che+20], showing that Riemannian GD converges to the Wasserstein barycenter at linear rate. Yet their convergence rate depends exponentially on the
In the setting of Gaussians, the Wasserstein barycenter was first studied in the 1990s OR93
KS94

---

Number of passes to convergencce
30
GD
GD SGD
1
20
01010320000011
200
1001
600
800
1000
dimneusion
Figure I: Passes until convergence error 10 to the barycenter, for r (3,5). This is dimension independent for Riemannian GD and LGD-ccosistenn with our main results. Details in Section
dimension d, and also their work does not extend to the Wasserstein geometric median or regularized Wasserstein barycenter.
1.1 Contributions
In this paper, we analyze first-order optimization algorithms on the Bures- Wasserstein manifold. We summarize main results here and overview techniques in the next section.
From exponential dimension dependence to dimension-free rates. In Section we show that for the Wasserstein barycenter problem, Riemannian GD) enjoys dimension-free convergence rates (Theorem 2). We make several comments to contextualize this result. First, our result eliminates the exponential dimension dependence of state-of-the of------tthrrrr convergence rates [Che+20), which aligns with the empirical performance this algorithm (see Figure D Second, our result stands in sharp contrast to the setting of discrete distributions in which there are computational complexity barriers to achieving even polynomial dimension dependence [AB21b]. Third, our result closes the gap between computation and statistical estimation for the Bures Wasserstein barycenter, since dimension-free sample complexity bounds were recently proven in [LFF99.
Moreover, in Theorem 3 we further refine this result by replacing the worst-case assumption of uniform bounds on the matrices' eigenvalues with significantly weaker average-case assumption.
Beyond barycenters. In Sections and 5 we show how our analysis techniques also enable proving fast convergence of Riemannian GD for computing regularized barycenters (Theorem and geometric medians (Theorem 5). To the best of our knowledge, these are the first guarantees for manifold beyond the barycenter.
Riemannian GD for notions averaging on the
1.2 Techniques
Here we briefly sketch the specific technical challenges we face and how we address them to analyze Riemannian GD for the three notions of Bures- Wasserstein average: barycenter, regularized barycenter, and geometric median. Although each analysis necessarily exploits particularities of its own objective, the common structure of our overarching analysis framework may be of interest for studying other geodesically non-convex optimization problems.
Overcoming non-convexity. As discuss in Appendix A.4 there a close connection between the second-order behavior of these objective functionals and the non-negative curvature of the Bures- Wasserstein manifold. In particular, while nog-niettive curvature used to prove smoothness properties for the three functionals, it also leads to them all being geodesically non-convex. To circumvent this issue, we establish gradient domination conditions known as inequalities [Po164 Loj63], which intuitively are quantitative proxies for strong convexity in non- convex settings (see e.g., [KNST6| Bol+17)) Proving such inequalities requires synthesizing general optimization principles with specialized arguments based optimal transport theory. We ultimately

---

show that these inequalities hold with constants depending on the conditioning of the iterates, i.e., the ratio between the maximum and minimum eigenvalues of the corresponding covariance matrices.
Overcoming ill-conditioned iterates. So long as smoothness and gradient domination inequalities hold at the current iterate, standard optimization results guarantee that the next iterate of GD makes progress. However, the amount of progress degrades if the iterates are poorly conditioned, since then our PL incquality degrades. Thus the second major obstacle is to control the regularity of the iterates. Here, the primary technical tool is shared across the analyses. Informally, it states that if the objective matrices, then the
is a sum of functions, each of whose gradients point towards well-conditionss GD iterates remain celliionnd. Formally, this is captured by the following geodesic convexity result, which may be of independent interest. Below, $ denotes the set of d x d positive definite matrices. See Appendix AA2 for review of the relevant geometric concepts, and see Appendix for the proof, discussion of tightness, and complementary results.
(a.2
Theorem 1. The functions -VAmin: Vmaax Sd +
are convex along generalized geodesics.
Using this theorem in conjunction with careful analysis of the objective functions, we establish global convergence guarantees for first-order geodesic optimization.
1.3 Other related work
Averages such barycenters and medians on general curved spaces have become popular due to far-ranging applications in domains such as machine learning, computer vision, analysis, radar signal processing [ABYI3), and brain-computer interfaces (YBL16 CBB17J. While their mathematical
properties such as existence and uniqueness are fairly well-understood [Afs]T]. their computation is [ac14] OP15].
an active area of research [Wei37 V200 Stu03
Yan BI13
For the Wasserstein barycenter problem in particular, there have been many approaches. These approaches vary significantly depending on if the setting is discrete continuous. In the discrete setting, the problem is NP-hard in high dimension [AB21b]. In low dimension (or more precisely fixed dimension), reasonable approximations can be obtained using fixed-support approximations which reduce the problem to a large linear program (CD14 Ben+15) CO015) Kro+19) 1in+19] Lin+20 Dvi21; Gum+21 [Ha+21], and it recently shown that eigh-precision (or even exact) solutions can be computed in polynomial time using computational geometry techniques (AB21a]
In the continuous setting, the problem is in general intractable since the optimal barycenter is intractable even to represent, let alone to compute. Nevertheless, in certain settings it has been empirically effective to parameterize and solve using neural networks [CAD20] FTC21 Kor+21]
stochastic gradient descent [i++20], or Riemannian optimization [Ál++16 Bac+ ZP19
Che+20
However, for continuous settings, the theory currently lags far behind the empirics. This is true even in the seemingly simple setting of Gaussians, in which case the barycenter has concise representation since it is also Gaussian. Riemmanian GD for this problem was first proposed and demonstrated to be empirically effective in [Álv+16), where was introduced convergence was proved in [6l++16], and then extended non-population and stochastic settings in [Bac+18]] [PP19]. aon-asymptotic convergence rates were first shown in [Che+20). but there is a large gap between these theoretical rates and what is observed in practice. It particular, previous rates depend exponentially on the dimension. The present paper improves this to dimension-free.
fixed-point algorithm. Asymptotic
For the other two problems we study, Bures- Wasserstein geometric medians and entropically- regularized barycenters, no convergence guarantees were previously known for the natural Ric- mannian GD algorithm.
Our work is in the midst of flurry of exciting recent developments about entropically regularized barycenters. Several recent works have, simultaneously with each other, extensively studied these objects in the particular setting of Gaussians, leading to the establishment of fundamental results such as the fact that the regularized barycenter of Gaussians is itself Gaussian [BL20; Jan+20; [GM21] Another recent and related line of work has established fundamental mathematical and statistical results for entropically regularized barycenters in the setting of general distributions, although with slightly different penalties than the KI. divergence studied here, typically the differential entropy b b [ro188 BCP19
CEK2 and sometimes even broader classes of regularizations (BCP19).

---

1.4 Organization
Section briefll recalls relevant preliminaries. We analyze Riemannian GD for computing Bures- Wasserstein barycenters, regularized barycenters, and geometric medians in Sections 3 4 and 5 respectively. We conclude in Section 66, We provide proofs as well as additional background and numerical results in the Appendix. Code reproducing all experiments in this paper can be found at
Preliminaries
Given probability measures and on E with finite second moment, the 2- 2-Wasserstein distance between y and v is defined as
/ /Ile u11² (x,
(2 := inf
(4)
xeII(a..)
where [1, u)) denotes the set of couplings of and v, i.e. the probability measures on R& Rd whose marginals are respectively l and v. and admit densities with respect to the Lebesgue measure on Rd. then the infimum attained, and the optimal coupling is supported ie., there exists map T' R - IR such that for π-a.c. (z, y) E - x Rd, it holds that y = T(x). The map I' is called the optimal transport map from an introduction to optimal transport, and to [Car92] for background on Riemannian geometry. The Riemannian structure of optimal transport introduced in the seminal work [Ou01] detailed treatments are in [AGS08] V for completeness we also provide a quick overview in AappendixAA
the graph of map,
We refer readers [Vi103] [Sn115] for
In this paper, mainly work with centered Gaussians, which can be identified with their covariance matrices. (Extensions the non-centered case are also discussed below.) We abuse notation via this identification: given E, E' E $ we write (2 E') for the 2-Wasserstein distance between centered Gaussians with covariance matrices Σ, E' respectively. Here, Sd denotes the space of symmetric d x d matrices, and $ denotes the open subset of Sd consisting of positive definite matrices. Throughout, all Gaussians are non-degenerate; that is, their covariances are non-singular.
The Wasserstein distance has closed-form expression for Gaussians:
(E, tr [ E' 22 23 1/2
(5)
Also, the optimal transport map from Σ Σ' the symmetric matrix
TΣEΣ - 1/2 (Σ ////////////// 1/2 1- GM(G-MM--... E.).
A (A 1/2) 1/2 A1/2 denotes the matrix geometric mean between two positive semidefinite matrices Bha07 Ch. 4]. The Wasserstein distance on S in fact arises from a Riemannian metric, which was first introduced by Bures in [Bur69]. Hence, the Riemannian manifold $4 endowed with this Wasserstein distance is referred to as the Bures- Wasserstein space. The geometry of this space is studied in detail in [Mod17; BJL19). For completeness, we provide additional background on the Bures- Wasserstein manifold in Appendix A
Above, GM(A,
Barycenters
In this section, we consider the Bures- Wasserstein barycenter
Ex E arg min W3.
Eesy
We refer to the introduction for discussion of the past work on the Bures- wasserasssrrttiin barycenter. We also remark that the case when is supported on possibly non-centered Gaussians is easily reduced to the centered case; the discussion in [Chc+20, $4].
5

---

3.1 Algorithms
We consider both Riemannian gradient descent (GD) and Riemannian stochastic gradient descent (SGD) algorithms for computing the Sures-Wasserstein barycenter, which are given as Algorithm [ and (
GD is useful for computing high-precision solutions due to its linear convergence (Theorem 2). and SGD is useful for large-scale or online settings because of its cheaper updates. We refer to [2P19 Che+20) for the derivation of the updates. Here, En is the initialization, which can be taken to be any matrix in the support of P. For SGD, also require sequence (me)==1 of step sizes and a sequence (K₂), of i.i.d. samples from P. Note that whereas SGD requires choosing step sizes, GD simply uses step size 1, justified in [ZP19].
Algorithm 1 GD for Barycenters
Algorithm 2 SGD for Barycenters
1: procedure BARY- RD P,T)
I: procedure BARY-SGID (mo)=D (oll (Ko) for 1 = 1, 1,..., do
2: for 1 = 1.... T do
2:
3:
St GM 2) dP()
(1 ( 11) Ne GM(ET11,
4:
4:
+ + 15
5:
return Σg
5: return Er
3.2
Convergence guarantees
Denote the barycenter functional by F(E) = / W2 dP, and denote the variance of P by var P 2F(E*). We assume that supported on matrices whose eigenvalues lie in the range [Amiu Amax]. and we let к := Amax/ Amin denote the condition number. Whereas the previous state-of-the-art convergence analysis for Algorithms and [Che+20] suffered a dependence of nd we show that the rates of convergence are in fact independent of the dimension d.
Theorem 2. Assume that P is supported covariance matrices whose eigenvalues lie in the range DAmin Amax] 0 < Amin I Amax < oC. Let K. == Amax/ Ami denote the condition number. Assume that we initialize at Eg supp
1. (GD) Let 24 denote the T-th iterate of GD (Algorithm) Ty. Then,
1 2/K
T
W
F(EE)) S F(E*) exp
(F(Eo) - F(E)))
4x3/2
2. (SGD) Let ISGD denote the T-th iterate of SGD (Algorithm from P. Then, with an appropriate choice of step sizes
where (Ki) are i.i.d.
1 2
48k var P " S EF SGD -F(E"")))F(()
T
In fact, using our new geodesic convexity results we can also relax conditioning assumption from requiring all matrices be uniformly well-conditioned, to being individually well-conditioned. This is a significant improvement when the eigenvalue ranges differ significantly between matrices.
Theorem 3.
Let K² == suppesapp(P) Amax
(2) /mmi (2). The conclusions of Theorem hold when
replacing &s with n² everywhere.
Actually, we deduce this from an even stronger statement: in Theorem 2 K can be replaced ev- erywhere by an average-case notion of conditioning, namely F = Amax/ Amin where Amin 1/2
Amin ( dP(E)) and Xmax -1/2 : >max (2) aP(E).
We give the proofs of these results in Appendix C.I
Namely. 5.3/2 (1
16 (2 (1++4))++1
suffices, where lo 32 1.
(t+10

---

3.3 Numerical experiments
There are two natural competitors of Riemannian GD when minimizing the barycenter functional: (i) solving an SDP (see Appendix C.3 for the SDP reformulation), and (ii) Euclidean GD (see Appendix C.2 for deseription and analysis of Euclidean GD).
High-precision algorithms for baryoenters
Largo-soale algorithms for barycenters
10a
DGD
BGD 1p=
10a
SGD 1==0,1/t
ESGD 1 x=01111
BGD
SQD Rx Rx1/66
ESGD m=1/t
Wgtx, X*)
W Ld
10-
SGD 4-10//t
ESGD 1=100/1
10-1
10-10 10
10-1
10 19
0.1
0.2
0.3
anermds
seconds
Figure 2: Riemannian vs. Euclidean GD.
Figure 3: Riemannian vs. Euclidean SGD.
In Figure 2 we compare Riemannian and Euclidean GD on dataset consisting of n = 50 covariance matrices of dimension d 50, each with condition number x 1000. Their eigenspaces are independent Haar distributed, and their eigenvalues are equally spaced in the interval [Amin Amax] [0.03, 30). Qualitatively similar results are observed for other inputs; see Appendix F We run 50 experiments and plot the average accuracy cut off at 10 10- X denotes the best iterate. We omit SDP solvers from the plot because their runtime is orders of magnitude slower: using the Splitting Cone Solver (SCS) 1ODo+16 16 1Do+19], the problem takes ~15 seconds to solve, and MOSEK (MOS2T] is even slower. For completeness, we also compare GD to the Riemannian algorithm WS17] in Appendix [P) and conclude that GD is superior. We observe that Euclidean GD's convergence rate is quite sensitive to its step size, which depends heavily on the conditioning of the problem. Riemannian GD was the clear winner in experiments, as its step size requires no tuning and it always performed no worse (in fact, often significantly better) than Euclidean GD.
Wolle
In Figure 3 we observe that Riemannian SGD typically outperforms Euclidean SGD, sometimes substantially. We average 300 x 300 covariance matrices drawn from distribution whose barycenter is known to be the identity, see Appendix [F for details. As Figure 3 shows, in practice it can be helpful tune the step sizes beyond the guidance given by our worst-case theoretical guarantees. In our experiments, Riemannian SGD was competitive on a wide range of problems with r) = 1/t.
We comment on Figure D which illustrates the dimension independence of the two Riemannian algorithms, a main result of this paper. It plots the number of passes until convergence (X,, S 10"" var P to the barycenter X*, for r f
E {3, 5). To compare algorithms on equal footing, the y-axis measures "passes" over the n 50 matrices: a pass constitutes one GD iteration, or n SGD iterations. The imput is generated as in Figure - Observe also the tradeoff between GD and SGD: SGD converges rapidly to low- low-precision solutions, but takes longer to converge to high-precision solutions.
4 Entropicaly-rreuuarizee barycenters
In this section, we consider the entropically--regulaizze barycenter bieg which minimizes
1 F, (b) : 2 dP
(.).
where KL denotes the Kullback-Leible divergence, and 0 is given regularization parameter. It suffices to consider the case when all measures are centered, see Remark 6 The following proposition justifies considering this problem on the Wasserrannn space; proof Appendix DD3
Proposition Suppose P is supported on centered Gaussians whose covariance matrices have eigenvalues lying in the range [L/ Vki, for some > 1. Then there exists unique minimizer Ding of F,, over P2 ) and this minimizer is a centered Gaussian distribution whose covariance matrix x also has eigenvalues in the range [L/VK, VR, vk..
7

---

As described in the introduction, prior work on the Wasserstein barycenter typically focuses on slightly different entropic penalty, the differential entropy f b In b. Note that the differential entropy penalty encourages be diffuse over all of Red (the minimizer blows up as 7 - oo). Here, we focus KL divergence penalty which has the advantage of interpolating between two well-studied problems: the Wasserstein barycenter problem (y = 0) and minimization of the KL divergence (7 oo). We take the standard Gaussian
canonical choice of reference distribution, and note that our method of analysis can be extended to other reference measures at the cost significant additional technical complexity. We thus choose exclusively focus on the standard Gaussian case.
4.1 Algorithm
Algorithm
iis Riemannian GD for minimizing F.. A derivation of the update rule is in Appendix D
Algorithm 3 GD for Regularized Barycenters
1: procedure RBARY-GD(Eo. P,T.7.1)
2: for 1 1.....7 do
3:
S: + GM dP(E) + (1- T)
4: Mi f 11
5: return Σx
4.2
Convergence guarantees
We provide two convergence guarantees for Algorithm[3 The first holds for all choices of the regu- larization parameter 7. However, this rate deteriorates with larger 7, and intuitively the optimization problem should become somewhat casier with larger regularization; hence, we prove a second rate of convergence to capture this behavior. We emphasize that as in $ our convergence rates are inmeesionninnepeeddee
The proof of each rate appears in Appendix D
Theorem 4. Fix 0 and suppose that P is supported covariance matrices with eigenvalues in [1/ VK. VK.. If Algorithm is initialized at a point in supp P and run with step size n = 1/(1 + 2//r) then the following hold.
1. For any choice of regularization parameter 7y >> and any T 1.
S exp(- Rd (1+ (F,(Eo) - FF())))
2. y 2 14x* is sufficiently large, then the following improved rate holds: for 2 1.
T
E,(ET) - < exp(
GVK (Ea) F,(E"))
For brevity, we omit guarantees in terms of the distance W2,,,, E+) to the minimizer.
4.3 Numerical experiments
Effect of regularization
In Figure we investigate the use of the regu- larization term 7 KLC (0, Ia)) to encode prior belief of isotropy. In Figure 4 we gener- n 100 i..dd 20 20 covariance matrices from a distribution whose barycenter is the iden- tity (see Appendix F). Then, for p E [0, 10] we compute the barycenter of perturbed dataset obtained by adding pere] each matrix for dif- ferent choices of 7. We see that for = 0 the barycenter quickly departs from isotropy, while for larger 7 the regularization yields averages which are more consistent with our prior belief.
100
0.0
2.5
5.0
7.5
10.0
perturbation size
Figure 4: Effect of regularization for varying 7.
8

---

Geometric medians
In this section, we consider the Wasserstein geometric median
/ W2(b,-)dl
b
median E arg min bEP2(RR)
See the introduction for a discussion of the literature on this problem. Observe that, contrast to the barycenter ( here we are minimizing the average unsquared Wasserstein distance.
The following basic result justifies the consideration of the geometric median problem on the Bures- Wasserstein space. It is proved in Appendix E.1
Suppose that P is supported on centered non-degenerate Gaussians whose covariance matrices have eigenvalues lying in the range [Amins Amax]. where 0 S Amin S Amax < co. Then, there exists a solution to 7 which is also a centered non-degenerate Gaussian distribution; moreover, its covariance matrix Σ median can be taken to have eigenvalues in [Aminn Amaz)-
Proposition 2.
Remark 1. Suppose now that P supported on non-degenerate Gaussian distributions which are not necessarily centered. Then, the proof of Proposition applies with minor modifications to show that the minimizer of the median functional still attained a Gaussian distribution. However, unlike the barycenter and entropically regularized barycenter, it is not the case that the mean of the Wasserstein geometric median is the Euclidean geometric median of the means, thus it not as straightforward to reduce to the centered case for this problem. Nevertheless, in Appendix E.2 describe reduction which allows the algorithm described in the next section to be applied in black box manner to the non-centered case, with corresponding convergence guarantees.
5.1 Algorithm
Since the Wasserstein distance W2(E,,) ) is neither geodesically convex nor geodesically smooth, nor Euclidean convex nor Euclidean smooth (see Remark (7). poses challenges for optimization. We therefore smooth the objective before optimization. Given desired target accuracy >> 0, let
W2.5 V W? +
(b) == (b, dP..
Algorithm provides pseudocode for running Riemannian GD on the smoothed functional Fg. See Appendix for a derivation the update rule.
Algorithm 4 Smoothed GD for Median
1: procedure MEDIAN-GD(Eo, e,
2: for t 1.... do
3:
S, + 17 + 4
1a) (2,-11111,, - dP())
4: Σ - S
5:
return Ex
5.2 Convergence guarantees
We show that Algorithm 4 finds approxiiit minimizer for the geometric median functional in O(kfe*) iterations. We emphasize that as in our other results, this convergence is dimension- independent. Below, let F : Fo denote the unregularized functional. Note that since F typically does not have a unique minimizer, we only guarantee a small suboptimality. The proof is in Appendix E.
Theorem 5. Assume that P is supported on covariance matrices whose eigenvalues lie in (Amins )max], 0 < Amia S Amax c0. Let k. == >mix/>min 0 < E < 1 denote a target accuracy. Assume that we initialize Algorithm A Do E supp P. Then, Algorithm 4 outputs Er satisfying F(ET)) F( S 38
Amin denote the condition number, and let
T > 32x F (Eo)

---

Median at different uccuracies
Robustness of median and baryoenter
100
leuyeeiis (30%)
1200
neelan (0%)
F(X) P((X))
10:6
0.00
X X2
1000
median (4356)
berveenter (60%)
(6G%)
naelisa
10-20
600
100
200
0
seconds
0
10
20]
30
petturbation tntitiplier
Figure 5: Evolution of median objective for vary-
Figure 6: Robustness of the Wasserstein median.
ing X denotes the best iterate.
5.3 Numerical experiments
In Figure we plot the suboptimality gap for the unregularized objective F [2(, Σ) dP as we optimize Fg using Algorithm [ for various ε. The regularization parameter & has natural trade-off: smaller results in better approximation to the (unregularized) geometric median, but slower convergence. The covariance matrices are generated as in Figure 2 with = d = 30 and (Amin Amax] = [0.01, 10). The promising empirics in Figure [ that our algorithm performs even better in practice than our worst-case theoretical results guarantee: few iterations may suffice for convergence, and also moderate regularization may suffice for high-precision approximations.
In Figure we illustrate the robustness of the Wasserstein geometric median up its breakdown point of 50% [FV]. We take random input matrices above, with n = d 20 and [Amin Amax 10), and compute their barycenter and approximate median (e = 1). We then perturb fraction (20%, 45%, and 55% for our figure) of the matrices by multiplying them by a constant greater than 1. The r-axis of the plot shows the size of the perturbation while the y-axis gives the distance of the original barycenter and median to the barycenter and median of this new, perturbed dataset.
We also implemented Euclidean GD for this geometric median problem; plots are omitted for brevity since the results are similar to those for the barycenter (c.f. Section 3.3) in that Euclidean GD depends much more heavily on parameter tuning. Note also that Euclidean GD does not come with convergence guarantees for this problem since it is non-convex in the Euclidean geometry.
Discussion
In this paper we revisited the problem of computing Bures- Wasserstein barycenters and explained the empirical efficacy of Riemannian (S)GD by proving convergence rates that improve from exponential An attractive feature of our analysis framework was that
dimension dependence to dimension-frn tools sufficiently general to prove similar dimension-fre guarantees for related problems of interest, namely Bures- Wasserstein geometric medians and aatropccall--rgularizze barycenters.
Our results suggest several interesting directions for future research. The focus of this paper was dimension-dependence, and while we also improved the dependence on other parameters along way, it is unclear if these other dependencies are optimal. Can the dependence stronger PL incqualities? Is the dependence on improvable via alternate methods of smoothing in the case of geometric medians, via fixed-pon acceleration schemes such as Anderson acceleration in the case of barycenters?
k be improved via
More broadly, conventional wisdom from the now-established field of geodesic optimization tells us that whenever possible, one should recast non-convex optimization problem as a convex one by changing the geometry. However, as demonstrated empirically in Section 3 for computing Bures- Wasserstein barycenters, optimal transport than in the convex geometry of Euclidean space. A full understanding of why and when non-convex geometry can be helpful in general optimization problems an intriguing direction with potentially significant implications for both the theory and practice of non-convex optimization.
significantly better to run GD in the non-convex geometry of
10

---

Acknowledgments and Disclosure of Funding
We are grateful to Victor-Emmanuel Brunel, Tyler Maunu, and Philippe Rigollet for stimulating conversations, to Pablo Parrilo for pointing out that Bures-Wassene Wasserstein barycenters have an SDP formulation (Appendix C.3), and to the anonymous reviewers for their thoughtful comments.
JA was supported by NSF Graduate Research Fellowship 1 122374 and a TwoSigma PhD fellowship. SC and AS were supported by the Department of Defense (DoD) through the National Defense Science Engineering Graduate Fellowship (NDSEG) Program.
References
[AB2Ia]
J. M. Altschuler and E. Boix-Adsera. "Wasserstein barycenters can be computed in polynomial time in fixed dimension". In: Journal of Machine Learning Research 22.44 (2021), pp- 1-19.
[AB21b]
J. M. Altschuler and E. Boix-Adsera. "Wasserstein barycenters are NP-hard to compute".
In: SIAM Journal on Mathematics of Data Science, to appear (2021).
[ABY13]
M. Arnaddon, Barbaresco, and L. Yang. "Riemannian medians and means with applications to radar signal processing". In: IEEE Journal of Selected Topics in Signal Processing 7.4 (2013), pp. 595-604.
[AC11]
M. Agueh and G. Carlier. "Baryeenters the Wasserstein space". In: SIAM Journal on Mathematical Analysis 43.2 (2011), pp. 904-924.
[AC17]
M. Agueh and G. Carlier. "Vers théorème de la limite centrale dans l'espace de Wasserstein?" ?" In: Comptes Rendus Mathématique. Académie des Sciences. Paris 355.7
(2017). pp. 812-818.
[Afs11]
Afsari. "Riemannian I, center of mass: existence, uniqueness, and convexity". In: Proceedings of the American Mathematical Society 139.2 (201 1), pp. 655-673.
[AGS08]
L. Ambrosio, N. Gigli, and G. Savaré. Gradient flows in metric spaces and in the space of probability measures. Second. Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, 2008, pp. x+334.
[ALP20]
A. Ahidar-Coutrix, T. Le Gouic, and Paris. "Convergence rates for empirical barycen- ters in metric spaces: curvature, convexity and extendable geodesics". In: Probability Theory and Related Fields 177.1-2 (2020), pp. 323-368.
[Ál++16]
P. C. Alvarez-Esteben et al. "A fixed-point approach to barycenters in Wasserstein space". In: Journal of Mathematical Analysis and Applications 441.2 (2016), pp. 744 762.
[Bac+18]
J. Backhoff-Veraguas al. "Bayesian learning with Wasserstein barycenters". In: arXiv e-prints, arXiv: 1eX.v:180..... 1888500033 (May 2018).
[Bac14]
M. Bacák. "Computing medians and means in Hadamard spaces". In: SIAM Journal on Optimization 24.3 (2014), pp. 1542-1566.
[BCP19]
J. Bigot, E. Cazelles, and Papadakis. "Penalization of barycenters in the Wasserstein space". In: SIAM Journal on Mathematical Analysis 51.3 (2019), pp. 2261-2285. J.-D. Benamou et al. "Iterative Bregman projections for regularized transportation problems". In: SIAM Journal on Scientific Computing 37.2 (2015), A1111-A1134
[Ben+15]
[Bha07]
R. Bhatia. Positive definite matrices. Princeton Series in Applied Mathematics. Princeton University Press, Princeton, NJ, 2007, pp. x+254.
[BI13]
D. A. Bini and B. lannazzo. "Computing the Karcher mean of symmetric positive definite matrices". In: Linear Algebra and its Applications 438.4 (2013), pp. 1700-1710. J. Bigot al. "Upper and lower risk bounds for estimating the Wasserstein barycenter of random measures on the real line". In: Electronic Journal of Statistics 12.2 (2018), pp. 2253-2289.
[Big+18]
[BJL19]
R. Bhatia, Jain, and Y. Lim. "On the Bures- Wasserstein distance between positive definite matrices". In: Expositiones Mathematicae 37.2 (2019), pp. 165-191. F. del Barrio and J.-M. Loubes. "The statistical effect of entropic regularization in optimal transportation". In: arXiv e-prints, 19X (2020).
[BL20]]

---

[Bol+17]
J. Bolte et al. "From error bounds to the complexity of first-order descent methods for convex functions". In: Mathematical Programming 165.2 (2017). pp. 471-507.
[Bou20]
N. Boumal. "An introduction to optimization on smooth manifolds", In: Available online, May (2020).
[BPC16]
N. Bonneel, G. Peyré, and M. Cuturi. "Wasserstein barycentric coordinates: histogram regression using optimal transport". In: ACM Transactions on Graphics 35.4 (2016).
[Bur69]
D. Bures. "An extension of Kakutani's theorem on infinite product measures to the tensor product of semifinite w -aleebrss" In: Transactions of the American Mathematical Society 135 (1969), pp. 199-212.
[CAD20]
S. Cohen, M. Arbel, and M. P. Deisenroth. "Estimating barycenters of measures in high dimensions". In: arXiv e-prints, vEXiv:2007.07100 (2020).
[Car92]
M. do Carmo. Riemannian geometry. Mathematics: Theory & Applications. Translated from the second Portuguese edition by Francis Flaherty. Birkhäuser Boston, Inc., Boston, MA, 1992 pp. xiv+300.
[CBB17]
M. Congedo, A. Barachant, and R. Bhatia. "Riemannian geometry for EEG-based brain- computer interfaces; a primer and a review". In: Brain-Computer Interfaces 4.3 (2017), pp. 155-174.
[CD14]
M. Cuturi and A. Doucet. "Fast computation of Wasserstein barycenters". In: Interna- tional Conference on Machine Learning. Vol. 32. 2014, pp. 685-693.
[CE10]
G. Carlier and I. Ekeland. "Matching for teams". In: Economic Theory 42.2 (2010),
pp. 397-418.
[CEK21]
G. Carlier, K. Hichinger, and A. Kroshnin. "Entropic-Wasserseinn characterization, regularity, and CLT". In: SIAM Journal on Mathematical Analysis 53.5 (2021), pp- 5880-5914.
Wasserstein barycenters: PDE
[Che+20]
Chewi al. "Gradient descent algorithms for Bures-Wasserse Wasserstein barycenters". In: Conference on Learning Theory. Vol. 125. 2020, pp. 1276-1304.
[C0015]
G. Carlier, A. Oberman, and E. Oudet. "Numerical methods for matching for teams and Wasserstein barycenters". In: ESAIM: Mathematical Modelling and Numerical Analysis 49.6 (2015), pp. 1621-1642.
[Dog+19]
P. L. Dognin al. "Wasserstein barycenter model ensembling". In: International Con-
ference on Learning Representation. 2019.
D. Dvinskikh. "Stochastic approximation versus sample average approximation for Wasserstein barycenters". In: Optimization Methods and Software 0.0 (2021), pp. 1-33. [Elv+20] Elvander et al. "Multi-margina"" optimal transport using partial information with applications in robust localization and sensor fusion". In: Signal Processing 171 (2020), p. 107474.
[Dvi21]
[FLF19]
R. Flamary, K. Lounici, and A. Ferrari. "Concentration bounds for linear Monge mapping estimation and optimal transport domain adaptation". In: arXiv preprint arXiv:1905. 10155 (2019).
[FTC21]
J. Fan, A. Taghvaei, and Y. Chen. "Scalable computations of Wasserstein barycenter via input convex neural networks". In: International Conference on Machine Learning. Vol. 139. 2021, pp. 1571-1581.
[FVJ09]
P. T. Fletcher, S. Venkatasubramanian, and S. Joshi. "The geometric median on Rieman- nian manifolds with application robust atlas estimation". In: Neurolmage 45.1 (2009), S143 S152.
[GPC15]
A. Gramfort, G. Peyré, and M. Cuturi. "Fast optimal transport averaging of neuroimaging data". In: International Conference on Information Processing in Medical Imaging. Springer. 2015, pp. 261-272.
[Gum+21] S. Guminov et al. "On a combination of alternating minimization and Nesterov's mo- mentum".. In: International Conference on Machine Learning. Vol. 139. 2021, pp. 3886- 3898.
[Haa+21]
I. Haasler et al. "Multimarginal optimal transport with a tree-structured cost and the Schrödinger bridge problem"". In: SIAM Journal on Control and Optimization 59,6 4 (2021). pp- 2428-2452
(Ho+17)
N. Ho et al. "Multilevel clustering via Wasserstein means". In: International Conference Machine Learning. PMLR. 2017, pp. 1501-1509.
12

---

[Jan+20]
H. Janati et al. "Entropic optimal transport between unbalanced Gaussian measures has closed form". In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle et al. Vol. 33. Curran Associates, Inc., 2020, pp- 10468-10479.
[KNS16]
H. Karimi, J. Nutini, and M. Schmidt. "Linear convergence of gradient and proximal- gradient methods under the
condition".. In: European Conference Machine Learning and Knowledge Discovery in Databases Volume 9851. ECML
PKDD 2016. Riva del Garda, Italy: Springer-Veriee 2016, pp. 795-811.
[Kor+21]
A. Korotin et al. "Continuous W barycenter estimation without minimax optimization". In: International Conference on Learning Representations. 2021.
[Kro+19]
A. Kroshmin et al. "On the complexity of approximating Wasserstein barycenters". In: International Conference on Machine Learning. Vol. 97. 2019, pp. 3530-35400
[Kro18]
A. Kroshnin. "Fréchet barycenters in the Monge-Kantorovich spaces". In: Journal of Convex Analysis 25.4 (2018), pp. 1371-1395.
[KS94]
M. Knott and C. S. Smith. "On generalization of cyclic monotonicity and distances among random vectors". In: Linear Algebra and its Applications 199 (1994), pp. 363- 371.
[KSS21]
A. Kroshmin, V. Spokoiny, and A. Suvorikova. "Statistical inference for Bures- Wasserstein barycenters". In: The Annals of Applied Probability 31.3 (2021), pp. 1264 1298.
+21]
T. Le Gouic al. "Fast convergence of empirical barycenters in Alexandrov spaces and the Wasserstein space". In: Journal of the European Math Society, to appear (2021).
[Li+20]
L. Li et al. "Continuous regularized Wasserstein barycenters". In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle et al. Vol. 33. Curran Associates, Inc., 2020, pp. 17755-17765.
[Lin+19]
T. Lin et al. "On the complexity of approximating multimarginal optimal transport". In:
arXiv preprint 1nXiv:1910.000 (2019).
[Lin+20]
T. Lin et al. "Fixed-support Wasserstein barycenters: computational hardness and fast al- gorithm".. In: Advances in Neural Information Processing Systems. Ed. by H. Larochelle et al. Vol. 33. Curran Associates, Inc., 2020, pp. 5368-533680
[LL17]
T. Le Gouic and J.-M. Loubes. "Existence and consistency of Wasserstein barycenters".
In: Probability Theory and Related Fields 168.3-4 (2017), pp. 901-917.
[LLR20]
T. Le Gouic, J.-M. Loubes, and P. Rigollet. "Projection to faimess in statistical learning".
In: arXiv e-prints, arXiv:2005.11720 (2020).
[Loj63]
Lojasiewicz. "A topological property of real analytic subsets (in French)". In: Coll. du CNRS, Les équations aux dérivées partielles 117.87-89 (1963), 2.
[MGM21]
A. Mallasto, A. Gerolin, and H. Q Minh. "Enntooy-reegularized 2-Wasserstein distance between Gaussian measures". In: Information Geometry (Aug. 2021).
[Mod17]
K. Modin. "Geometry of matrix decompositions seen through optimal transport and information geometry"' In: Journal of Geometric Mechanics 9.3 (2017), pp. 335-390.
[MOS21] MOSEK. MOSEK Optimizer API for C. 2021. URL: hww.::ttps::. :/tp:s://doce mosek. com/9
2/ /ac
[ODo+16]
B. O'Donoghue et al. "Conic optimization via operator splitting and homogeneous self-dual embedding".. In: Journal of Optimization Theory and Applications 169.3 (June 2016). pp. 1042-1068.
[ODo+19]
B. O'Donoghue et al. SCS: Splitting Conic Solver, version 2.1.3. https : //gthubb
com/ htmm/cvxgrp//se scs| Nov. 2019.
[OP15]
S.-i. Ohta and M. Pálfia. "Discrete-time gradient flows and law of large numbers in Alexandrov spaces". In: Calculus of Variations and Partial Differential Equations 54.2 (2015), pp. 1591-1610.
[OR93]
I. Olkin and S. T. Rachev. "Maximum submatrix traces for positive definite matrices".
In: SIAM Journal on Matrix Analysis and Applications 14.2 (1993), pp. 390-397.
[OnO1]
F. Otto. "The geometry of dissipative evolution equations: the porous medium equation".
In: Communications in Partial Differential Equations 26. 1-2 (2001), pp. 101-174.
[PC+19]
G. Peyré, M. Cuturi, et al. "Computational optimal transport: With applications to data science". In: Foundations and Trends in Machine Learning 1 11.5-6 (2019), pp. 307-607.
13

---

[Pol64]
B. T. Polyak. "Gradient methods for solving equations and inequalities (in Russian)". In:
USSR Computational Mathematics and Mathematical Physics 4.6 (1964), pp. 17-32.
[PZ16]
M. Panaretos and Y. Zemel. "Amplitude and phase variation point processes". In:
The Annals of Statisties 44.2 (2016), pp. 771-812.
[PZ19]
V. M. Panaretos and Y. Zemel. "Statistical aspects of Wasserstein distances". In: Annual Review of Statistics and its Application 6 (2019), pp. 405-431.
[Rab+11]
J. Rabin et al. "Wasserstein barycenter and its application to texture mixing". In: In- ternational Conference on Scale Space and Variational Methods in Computer Vision. Springer. 2011, pp. 435-446.
[RP15]
J. Rabin and N. Papadakis. "Convex color image segmentation with optimal transport distances". In: International Conference on Scale Space and Variational Methods in Computer Vision. Springer. 2015, pp. 256-269.
[RU02]
Rüschendorf and L. Uckelmann. "On the n-coupling problem". In: Journal of Multi- variate Analysis 81.2 (2002). pp. 242-258.
[San15]
F. Santambrogio. Optimal transport for applied mathematicians, Vol. 87. Progress in Nonlinear Differential Equations and their Applications. Calculus of variations, PDEs, and modeling. Birkhäuser/Springer, Cham, 2015, pp. xxvii+353.
[SLD18]
Srivastava, C. Li, and D. Dunson. "Scalable Bayes via barycenter in Wasserstein space". In: The Journal of Machine Learning Research 19.1 (2018), pp. 312-346. J. Solomon et al. "Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains". In: ACM Transactions on Graphics 34.4 (2015), pp. 1-11. K.-T. Sturm. "Probability measures on metric spaces of nonpositive curvature". In: Heat kernels and analysis on manifolds, graphs, and metric spaces (Paris, 2002). Vol. 338. Contemp. Math. Amer. Math. Soc., Providence, RI, 2003, pp. 357-390.
[5o1+15]
[Stu03]
[Vi103]
C. Villani. Topics in optimal transportation. Vol. 58. Graduate Studies in Mathematics. American Mathematical Society, Providence, RI, 2003, pp. xvi+370.
[Vi109]
C. Villani. Optimal transport. Vol. 338. Grundlehren der Mathematischen Wis- senschaften [Fundamental Principles of Mathematical Sciences). Old and new. Springer- Verlag, Berlin, 2009, pp. xxii+973.
[VZ00]
Y. Vardi and C.-H. Zhang. "The multivariate Lm-mddian and associated data depth". In: Proceedings of the National Academy of Sciences 97.4 (2000), pp. 1423-1426,
[Wei37]
E. Weiszfeld. "Sur le point pour lequel la somme des distances de n points donnés est minimum". In: Tohoku Mathematical Journal, First Series 43 (1937), pp. 355-386. M. Weber and S. Sra. "Riemannian optimization via Frank-Wolfe methods". In: arXiv e-prints, arXiv: 107.70077 (2017).
[WS17]
[Yan10]
L. Yang. "Riemannian median and its estimation". In: LMS Journal of Computation and Mathematics 13 (2010), pp. 461-479
[YBL16]
F. Yger, M. Berar, and F. Lotte. "Riemannian approaches in brain-computer interfaces: a review". In: IEEE Transactions on Neural Systems and Rehabilitation Engineering 25. 10 (2016), pp. 1753-1762.
[ZP19]
Y. Zemel and V. M. Panaretos. "Fréchet means and Procrustes analysis in Wasserstein space". In: Bernouli 25.2 (2019), pp. 932-976.
[ZS16]
H. Zhang and S. Sra. "First-order methods for geodesically convex optimization". In: Conference on Learning Theory. Vol. 49. 2016, pp. 1617-1638.
14


Reviews:
[{"review_id": "yE_xiN2GmJh", "reviewer": "Reviewer_t1br", "summary": "This work is with the recent active wave of work in numerical optimal transport to take ML and statistical estimation problems. More precisely, the authors study first-order optimization algorithms on Bures-Wasserstein manifold for computing the barycenter of Gaussian measures with respect to the 2-Wassertstein distance, the  regularized barycenter and geometric median.\n", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main contribution of the paper is to show new geodesic convexity results hence allowing to obtain a dimension-free convergence rate of Riemannian gradient descent. The techniques also enable to go beyond barycenters by tackling the problem the entropically-regularized barycenter and geometric median, hence providing first convergence guarantees for these problems. Although linear convergence of gradient descent on the Bures-Waasserstein manifold was established recently, the rate had exponential dependence on the dimension. Therefore, as far as I can say, the results of this paper are new and of interest to the literature. The paper is also fairly well-written. Having said this, I still have some comments:\n\n1) I found that the paper lacks discussion to prior work for each stated result (though a brief discussion is gathered in the introduction).\n2) P3, L106-107: the reference given there is not fair. Rather cite the original paper of Łojasiewicz, and those of J. Bolte and his co-authors who have the actual pioneers of this in optimization. \n3) P4, L151-153: please void self-citation for such a classical subject on which there are many standard monographs (including diCarmo).\n4) P4, L155: sentence in parentheses. You say a bit later that it is in [Che+20, §4]. Note also that for the regularized W_2, it is also in [Jan+20, Lemma 1].\n5) P5, Theorem 2: the notation of the iteration counter is t in the algorithms and n here. Unify notations.\n6) P5, Theorem 2: can you comment on optimality of this rate ? We know the optimal rate in the Euclidean (convex) case.\n7) P5, Theorem 2: the statement \"appropriate step-sizes\" (ans also the K_t's) has to be made precise here (not just in the appendix). It is well-known that such a choie crucial to get the 1/n rate. \nBTW, can the dependence in \\kappa^3 be improved again ?\nIs there any hope to have a variance reduced version for the case of finitely many distributions (i.e. finite sum) ?\n8) Figure 2: the difference between ESGD and SGD is quite sensitive to the constant in the stepsize. Can you comment on this ?\n9) P6, L208: n is now used for the number of Gaussians (see comment 5) above).\n10) Figure 3 and dimension-independence: an alternative way would be rather to run several examples of the experiment on Figure 1 for increasing dimension, estimate the rate from a semilog fit, and plot the estimated rate as a function of d.\n11) Proposition 1: a similar result is in Theorem 2 of [Jan+20] but with a different metric (debiased version of the entropically regularized 2-W). Discuss please.\n12) Theorem 5: why this form of smoothing ? There are other ways, probably with better guarantees, or even no smoothing with subgradient descent. the iteration complexity in \\varepsilon^{-4} is rather slow (and play be pessimistic). Any clue on optimality of this bound or if it is unavoidable ?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "None.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3h."}, "extra_text": {}}, {"review_id": "oWn5sHsj9p1", "reviewer": "Reviewer_vGkj", "summary": "\nThe subject of this paper is Riemannian (stochastic) gradient descent for Bures-Wasserstein barycenters of Gaussian measures. The authors first improve recent results by S. Chewi et al (2020) by removing the exponential dependence of the convergence rates w.r.t. the dimension d. This improvement is obtained by showing the geodesic convexity of $-\\sqrt{\\lambda_{\\min}}$ and $\\sqrt{\\lambda_\\max}$, which yields tighter bounds on the eigenvalues of the iterates of Bures-(S)GD than that of S. Chewi et al (2020). Next, the authors extend Bures-(S)GD to a KL-regularized barycenter problem that encodes an isotropy prior on the barycenter, and to a median version of the barycenter problem (where the sum of un-squared distances is minimized). Finally, in numerical experiments they compare Bures-Riemannian (S)GD to Euclidean (S)GD on the original barycenter problem, and investigate the effect of smoothness parameters in the regularized and median problems.\n", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I have greatly enjoyed reading this paper, which is well-written and makes important contributions, even though it relies on techniques developed by S. Chewi et al (2020). Here is a list of points I particularly appreciated:\n\n- Removing the exponential dependence in $d$ is a *tour de force* that closes the gap with the sample complexity bounds for Bures-Wasserstein distance estimation, which are also dimension-free (see e.g. R. Flamary et al. (2019))\n\n- The paper is very clearly-written, and mostly self-contained. I particularly appreciated the fact that all required geometric notions are recalled in appendix A, which does a great job at introducing them in an intuitive way.\n\n- The extensions to the regularized barycenter and geometric median are interesting, in that they show how the tools developed for the Riemannian descent on the BW barycenter problem can be applied to problems that seem less amenable to them at first sight.\n\n- Thm. 1 is of interest by itself, and could probably be used to optimize spectral functions of PSD matrices in the Bures geometry.\n\n\nA few comments:\n\n- While I appreciated the fact that they provide additional examples of how to apply proof techniques based on (generalized) geodesic convexity, I am not sure how significant the regularized barycenter and geometric problems are.\n\n- l.238: « Prior work focuses on a slightly different entropic penalty, the differential entropy  $\\int b \\ln b$. ». There are some works that study this problem with a broader class of penalties, see e.g. the work of J. Bigot et al. (2019) which seems quite relevant.\n\n\nReferences:\n\nFlamary, R., Lounici, K., & Ferrari, A. (2019). Concentration bounds for linear Monge mapping estimation and optimal transport domain adaptation. arXiv preprint arXiv:1905.10155.\n\nBigot, J., Cazelles, E., & Papadakis, N. (2019). Penalization of barycenters in the Wasserstein space. SIAM Journal on Mathematical Analysis, 51(3), 2261-2285.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "My only comment on that subject is that the fact that the algorithm for geometric medians computes an approximate solution that depends on the smoothness parameter $\\varepsilon$ could be mentioned earlier than in section 5. However, based on Figure 5 it seems that numerical precision is reached quite fast when taking $\\varepsilon$ to $0$.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "XJx_hRasFDf", "reviewer": "Reviewer_A7RV", "summary": "The authors study Riemannian gradient descent on the Bures-Wasserstein manifold and provide non-asymptotic convergence guarantees for three optimization tasks: The computation of Wasserstein barycenters, entropy-regularized Wasserstein barycenters, and geometric medians.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "After rebuttal: Thanks to the authors for their comments. I'm keeping my score of acceptance.\n\n-----------------------------------------\n\n*Assessment*\n\nNon-asymptotic convergence guarantees for regularized barycenters and geometric medians are to my knowledge novel and should be of great interest to the community. The theoretical results are clearly formulated and seem correct. The paper is very well written.\n\n*Additional Comments*\n\n- Related work: Wasserstein barycenters on the Bures manifold lie in a compact, geodesically convex set and can therefore be treated as a constrained problem. One can then get non-asymptotic convergence guarantees with constrained Riemannian methods, such as Frank-Wolfe. The PL inequality should ensure linear convergence for such methods too. See, e.g., [*].\n- I did not see any experimental results for a practical application. The current numerical evaluation on synthetic data is (to me) sufficient for acceptance, but an experiment illustrating the results in an application could strengthen the paper.\n\n[*] Weber, Sra: Projection-free nonconvex stochastic optimization on Riemannian manifolds \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "  "}, "extra_text": {}}, {"review_id": "Eyx6T72bINK", "reviewer": "Reviewer_Bkse", "summary": "The paper address the problem of computing the barycenter of a gaussian distributions using optimal transport metric. The paper propose two algorithm (section 3.1, section 4.1) to compute thw Entropically regularized barycenters and the Bures-Wasserstain barycenters and it studies its convergence guarantees.", "questions": "", "limitations": "", "rating": 3, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Even if the paper deals with an interesting and important topic, it is not well organized neither clear. For example, the general discussion is completely missing. In this current form, the paper can not be published. As a general comment,  I think it is more suitable for a journal, where there is more room for discussion of the literature and additional details. \n\nSome comments for the possible reorganization of the paper:\n1) Section 1 and Section 2 should be reorganized, separating the state of the art from the rest\n2) Results in section 3,4,5 should be well listed at the beginning, otherwise it is not clear the whole paper structure\n3) Discussion is a fundamental part of the paper, beacuse it allows the reader to understand the work and value your discovers\n4) The simulated data should be better explained, not left on the side of the figures as captions.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No, unfortunately the discussion is missing.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proves dimension-free convergence rates for the gradient descent algorithm for the Bures-Wasserstein Barycenter computation. It is a significantly novel result that provides a theoretical basis to something that is practically observed and as such should be of interest to many.

Decision: Accept (Spotlight)