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

TL;DR

This paper demonstrates that Riemannian gradient descent for Gaussian barycenters under the Bures-Wasserstein metric converges at a rate independent of dimension, supported by new convexity results and extending to related averaging problems.

Contribution

It provides the first dimension-free convergence guarantees for Riemannian GD on Gaussian barycenters and related averaging problems, with new geodesic convexity insights.

Findings

Riemannian GD converges faster than Euclidean methods in practice.

New geodesic convexity results enable dimension-free convergence analysis.

First guarantees for Riemannian GD on entropically-regularized barycenter and geometric median.

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.