Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent

1 Introduction . Averaging multiple data sources is among the most classical and fundamental subroutines in data science . However , a modern challenge is that data is often more complicated than points in Rd . In this paper , we study the task of averaging probability distributions on Rd , a setting that commonly arises in machine learning and statistics [ CD14 ; Ho+17 ; SLD18 ; Dog+19 ] , computer vision and graphics [ Rab+11 ; Sol+15 ] , probability theory [ KS94 ; RU02 ] , and signal processing [ Elv+20 ] ; see also the surveys [ PC+19 ; PZ19 ] and the references within . The Wasserstein barycenter [ AC11 ; Rab+11 ] has emerged as a particularly canonical notion of average . Formally , let P2 ( Rd ) denote the space of probability measures on Rd with finite second moment , let P be a probability measure over P2 ( Rd ) , and let W2 denote the 2-Wasserstein distance ( i.e . the standard optimal transport distance ) . Then , the Wasserstein barycenter of P is a solution of minimize b∈P2 ( Rd ) ∫ W 22 ( b , · ) dP . ( 1 ) A related notion of average is the entropically-regularized Wasserstein barycenter of P , which is defined to be a solution of minimize b∈P2 ( Rd ) ∫ W 22 ( b , · ) dP + ent ( b ) , ( 2 ) where ent is an entropic penalty which allows for incorporating prior knowledge into the average . Lastly , a third related notion of average with better robustness properties ( e.g. , with a breakdown point of 50 % [ FVJ09 ] ) is the Wasserstein geometric median of P , which is defined to be a solution of minimize b∈P2 ( Rd ) ∫ W2 ( b , · ) dP . ( 3 ) 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Importantly , while these three notions of average can be defined using other metrics in lieu of W2 , the Wasserstein distance is critical for many applications since it enables capturing geometric features of the distributions [ CD14 ] . The many applications of Wasserstein barycenters and geometric medians ( see e.g. , [ CE10 ; Rab+11 ; CD14 ; GPC15 ; RP15 ; Sol+15 ; BPC16 ; SLD18 ; LLR20 ] ) have inspired significant research into their mathematical and statistical properties since their introduction roughly a decade ago [ AC11 ; 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 [ CE10 ; AC11 ; COO15 ] . And on the statistical side , [ PZ16 ; AC17 ; LL17 ; Big+18 ; FLF19 ; ALP20 ; Le +21 ; KSS21 ] provide finite-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 computational obstacles can be potentially averted in parametric settings . This paper as well as most of the literature [ Álv+16 ; Bac+18 ; ZP19 ; Che+20 ] on parametric settings focuses on the commonly arising setting where P is supported on Gaussian distributions.1 As noted in [ Álv+16 ] , the Gaussian case also encompasses general location-scatter families . There are two natural families of approaches for designing averaging algorithms in this setting . Both exploit the fact that modulo a 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 Bures-Wasserstein metric ( background is given in Section 2 ) . The first approach is simply to 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 3.3 for numerics ) . Briefly , the underlying issue is 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 Euclidean geometry . A much more effective approach in practice ( see Section 3.3 for numerics ) is to exploit the geometry of the Bures-Wasserstein manifold via geodesic optimization . Prior work has extensively pursued this direction , investigating the effectiveness of ( stochastic ) Riemannian GD for computing Wasserstein barycenters , see e.g. , [ Álv+16 ; Bac+18 ; ZP19 ; Che+20 ] . Challenges for geodesic optimization over the Bures-Wasserstein manifold . Although geodesic optimization is natural for this problem , it comes with several important obstacles : the non-negative curvature of the Bures-Wasserstein manifold necessitates new tools for analysis , and moreover both the barycenter and geometric median problems are non-convex in the Bures-Wasserstein geometry . ( These two issues are in fact intimately related , see Appendix A.4 . ) 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 a fixed-point iteration , and through this lens establish asymptotic convergence [ Álv+16 ; Bac+18 ; ZP19 ] . Obtaining non-asymptotic rates of convergence is more challenging because it 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 a linear rate . Yet their convergence rate depends exponentially on the 1In the setting of Gaussians , the Wasserstein barycenter was first studied in the 1990s [ OR93 ; KS94 ] . 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 our main results here and overview our techniques in the next section . From exponential dimension dependence to dimension-free rates . In Section 3 , 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-art convergence rates [ Che+20 ] , which aligns with the empirical performance of this algorithm ( see Figure 1 ) . 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 [ FLF19 ] . Moreover , in Theorem 3 , we further refine this result by replacing the worst-case assumption of uniform bounds on the matrices ’ eigenvalues with a significantly weaker average-case assumption . Beyond barycenters . In Sections 4 and 5 , we show how our analysis techniques also enable proving fast convergence of Riemannian GD for computing regularized barycenters ( Theorem 4 ) and geometric medians ( Theorem 5 ) . To the best of our knowledge , these are the first guarantees for Riemannian GD for notions of averaging on the Bures-Wasserstein manifold beyond the barycenter . 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 we discuss in Appendix A.4 , there is a close connection between the second-order behavior of these objective functionals and the non-negative curvature of the BuresWasserstein manifold . In particular , while non-negative curvature is 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 Polyak-Łojasiewicz inequalities [ Pol64 ; Loj63 ] , which intuitively are quantitative proxies for strong convexity in nonconvex settings ( see e.g. , [ KNS16 ; Bol+17 ] ) . Proving such inequalities requires synthesizing general optimization principles with specialized arguments based on 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 inequality 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 is a sum of functions , each of whose gradients point towards well-conditioned matrices , then the GD iterates remain well-conditioned . Formally , this is captured by the following geodesic convexity result , which may be of independent interest . Below , Sd++ denotes the set of d× d positive definite matrices . See Appendix A.2 for a review of the relevant geometric concepts , and see Appendix B for the proof , discussion of tightness , and complementary results . Theorem 1 . The functions − √ λmin , √ λmax : Sd++ → R 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 as 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 [ ABY13 ] , and brain-computer interfaces [ YBL16 ; CBB17 ] . While their mathematical properties such as existence and uniqueness are fairly well-understood [ Afs11 ] , their computation is an active area of research [ Wei37 ; VZ00 ; Stu03 ; Yan10 ; BI13 ; Bac14 ; OP15 ] . For the Wasserstein barycenter problem in particular , there have been many approaches . These approaches vary significantly depending on if the setting is discrete or 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 ; COO15 ; Kro+19 ; Lin+19 ; Lin+20 ; Dvi21 ; Gum+21 ; Haa+21 ] , and it was recently shown that high-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 [ Li+20 ] , or Riemannian optimization [ Álv+16 ; Bac+18 ; 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 a 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 it was introduced as a fixed-point algorithm . Asymptotic convergence was proved in [ Álv+16 ] , and then extended to non-population and stochastic settings in [ Bac+18 ; ZP19 ] . Non-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 . For the other two problems we study , Bures-Wasserstein geometric medians and entropicallyregularized barycenters , no convergence guarantees were previously known for the natural Riemannian GD algorithm . Our work is in the midst of a 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 ; MGM21 ] . 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 KL divergence studied here , typically the differential entropy∫ b ln b [ Kro18 ; BCP19 ; CEK21 ] and sometimes even broader classes of regularizations [ BCP19 ] .

Summary:  The 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.