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

arXiv: https://arxiv.org/abs/2106.08502 · HTML: https://mlsanity.com/paper/2106.08502/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 116 references that are in the map.
- First posted on arXiv 16 Jun 2021.

## Builds on (most cited first, 30 of 33)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Convex Optimization: Algorithms and Complexity](https://mlsanity.com/paper/1405.4980/) (2014)
- [Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Lojasiewicz Condition](https://mlsanity.com/paper/1608.04636/) (2016)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [First-order Methods for Geodesically Convex Optimization](https://mlsanity.com/paper/1602.06053/) (2016)
- [Efficient constrained sampling via the mirror-Langevin algorithm](https://mlsanity.com/paper/2010.16212/) (2020)
- [Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization](https://mlsanity.com/paper/2102.01752/) (2021)
- [Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks](https://mlsanity.com/paper/2007.04462/) (2020)
- [Multilevel Clustering via Wasserstein Means](https://mlsanity.com/paper/1706.03883/) (2017)
- [Scalable Bayes via Barycenter in Wasserstein Space](https://mlsanity.com/paper/1508.05880/) (2015)
- [From error bounds to the complexity of first-order descent methods for convex functions](https://mlsanity.com/paper/1510.08234/) (2015)
- [Entropy-Regularized 2-Wasserstein Distance between Gaussian Measures](https://mlsanity.com/paper/2006.03416/) (2020)
- [Continuous Regularized Wasserstein Barycenters](https://mlsanity.com/paper/2008.12534/) (2020)
- [Fixed-Support Wasserstein Barycenters: Computational Hardness and Fast Algorithm](https://mlsanity.com/paper/2002.04783/) (2020)
- [On the Complexity of Approximating Multimarginal Optimal Transport](https://mlsanity.com/paper/1910.00152/) (2019)
- [Wasserstein barycenters are NP-hard to compute](https://mlsanity.com/paper/2101.01100/) (2021)
- [Fast Optimal Transport Averaging of Neuroimaging Data](https://mlsanity.com/paper/1503.08596/) (2015)
- [Wasserstein barycenters can be computed in polynomial time in fixed dimension](https://mlsanity.com/paper/2006.08012/) (2020)
- [Riemannian Optimization via Frank-Wolfe Methods](https://mlsanity.com/paper/1710.10770/) (2017)
- [Fr'echet Means and Procrustes Analysis in Wasserstein Space](https://mlsanity.com/paper/1701.06876/) (2017)
- [On a Combination of Alternating Minimization and Nesterov's Momentum](https://mlsanity.com/paper/1906.03622/) (2019)
- [Estimating Barycenters of Measures in High Dimensions](https://mlsanity.com/paper/2007.07105/) (2020)
- [Concentration bounds for linear Monge mapping estimation and optimal transport domain adaptation](https://mlsanity.com/paper/1905.10155/) (2019)
- [The statistical effect of entropic regularization in optimal transportation](https://mlsanity.com/paper/2006.05199/) (2020)
- [Bayesian Learning with Wasserstein Barycenters](https://mlsanity.com/paper/1805.10833/) (2018)
- [Wasserstein Barycenter Model Ensembling](https://mlsanity.com/paper/1902.04999/) (2019)
- [A proof of the Caffarelli contraction theorem via entropic regularization](https://mlsanity.com/paper/1904.06053/) (2019)

## Built on by (most cited first, 30 of 32)

- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [The Schr"odinger Bridge between Gaussian Measures has a Closed Form](https://mlsanity.com/paper/2202.05722/) (2022)
- [Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space](https://mlsanity.com/paper/2304.05398/) (2023)
- [Learning Gaussian Mixtures Using the Wasserstein-Fisher-Rao Gradient Flow](https://mlsanity.com/paper/2301.01766/) (2023)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Wasserstein Iterative Networks for Barycenter Estimation](https://mlsanity.com/paper/2201.12245/) (2022)
- [Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space](https://mlsanity.com/paper/2312.02849/) (2023)
- [Sobolev Transport: A Scalable Metric for Probability Measures with Graph Metrics](https://mlsanity.com/paper/2202.10723/) (2022)
- [Provable Acceleration of Heavy Ball beyond Quadratics for a Class of Polyak-Lojasiewicz Functions when the Non-Convexity is Averaged-Out](https://mlsanity.com/paper/2206.11872/) (2022)
- [Scalable Unbalanced Sobolev Transport for Measures on a Graph](https://mlsanity.com/paper/2302.12498/) (2023)
- [Measure Estimation in the Barycentric Coding Model](https://mlsanity.com/paper/2201.12195/) (2022)
- [Generalized Sobolev Transport for Probability Measures on a Graph](https://mlsanity.com/paper/2402.04516/) (2024)
- [Bridging the Gap Between Variational Inference and Wasserstein Gradient Flows](https://mlsanity.com/paper/2310.20090/) (2023)
- [Optimal Transport for Measures with Noisy Tree Metric](https://mlsanity.com/paper/2310.13653/) (2023)
- [Continuized Acceleration for Quasar Convex Functions in Non-Convex Optimization](https://mlsanity.com/paper/2302.07851/) (2023)
- [Personalised Federated Learning On Heterogeneous Feature Spaces](https://mlsanity.com/paper/2301.11447/) (2023)
- [Wasserstein PAC-Bayes Learning: Exploiting Optimisation Guarantees to Explain Generalisation](https://mlsanity.com/paper/2304.07048/) (2023)
- [An Optimal Transport-driven Approach for Cultivating Latent Space in Online Incremental Learning](https://mlsanity.com/paper/2211.16780/) (2022)
- [Stochastic variance-reduced Gaussian variational inference on the Bures-Wasserstein manifold](https://mlsanity.com/paper/2410.02490/) (2024)
- [Variational Online Mirror Descent for Robust Learning in Schr"odinger Bridge](https://mlsanity.com/paper/2504.02618/) (2025)
- [Convergence Rates for Distribution Matching with Sliced Optimal Transport](https://mlsanity.com/paper/2602.10691/) (2026)
- [Rieoptax: Riemannian Optimization in JAX](https://mlsanity.com/paper/2210.04840/) (2022)
- [Finite sample bounds for barycenter estimation in geodesic spaces](https://mlsanity.com/paper/2502.14069/) (2025)
- [Metric Graph Kernels via the Tropical Torelli Map](https://mlsanity.com/paper/2505.12129/) (2025)
- [Busemann Functions in the Wasserstein Space: Existence, Closed-Forms, and Applications to Slicing](https://mlsanity.com/paper/2510.04579/) (2025)
- [Meta-learning Pathologies from Radiology Reports using Variance Aware Prototypical Networks](https://mlsanity.com/paper/2210.13979/) (2022)
- [On Barycenter Computation: Semi-Unbalanced Optimal Transport-based Method on Gaussians](https://mlsanity.com/paper/2410.08117/) (2024)
- [Gaussian mixture layers for neural networks](https://mlsanity.com/paper/2508.04883/) (2025)
- [Exact Solutions to the Quantum Schr"odinger Bridge Problem](https://mlsanity.com/paper/2509.25980/) (2025)
- [Stochastic Gradient Variational Inference with Price's Gradient Estimator from Bures-Wasserstein to Parameter Space](https://mlsanity.com/paper/2602.18718/) (2026)

## Abstract (arXiv)

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 GD empirically converges rapidly, in fact faster than off-the-shelf methods such as Euclidean GD and SDP solvers. This stands in stark contrast to the best-known theoretical results for Riemannian GD, 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 Riemannian GD for these problems.