Gradient descent algorithms for Bures-Wasserstein barycenters

TL;DR

This paper develops a framework to analyze gradient descent methods for computing barycenters in the space of probability measures, providing the first global convergence rates even without geodesic convexity.

Contribution

The paper introduces a novel analysis using a PL inequality for Bures-Wasserstein barycenters, enabling global convergence rate results for first order algorithms.

Findings

Established a PL inequality for Gaussian measures on the Bures-Wasserstein manifold.

Derived the first global convergence rates for gradient descent in this setting.

Provided a framework applicable to non-convex barycenter functionals in optimal transport.

Abstract

We study first order methods to compute the barycenter of a probability distribution $P$ over the space of probability measures with finite second moment. We develop a framework to derive global rates of convergence for both gradient descent and stochastic gradient descent despite the fact that the barycenter functional is not geodesically convex. Our analysis overcomes this technical hurdle by employing a Polyak-Lojasiewicz (PL) inequality and relies on tools from optimal transport and metric geometry. In turn, we establish a PL inequality when $P$ is supported on the Bures-Wasserstein manifold of Gaussian probability measures. It leads to the first global rates of convergence for first order methods in this context.