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

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 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.

arXiv commentv3: Fixed error in Theorem 3, see footnote 2

References in corpus (6)

• Statistical Aspects of Wasserstein Distances
• Amplitude and phase variation of point processes
• Entropic Optimal Transport between Unbalanced Gaussian Measures has a Closed Form
• On the Complexity of Approximating Multimarginal Optimal Transport
• Continuous Regularized Wasserstein Barycenters
• Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization

Cited by in corpus (3)

• Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey
• On the Wasserstein median of probability measures
• Riemannian-geometric generalizations of quantum fidelities and Bures-Wasserstein distance