Gradient descent algorithms for Bures-Wasserstein barycenters

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.

arXiv comment28 pages, 5 figures, to appear in COLT 2020

References in corpus (7)

• Fast Computation of Wasserstein Barycenters
• Amplitude and phase variation of point processes
• Stochastic Wasserstein Barycenters
• Averaging Stochastic Gradient Descent on Riemannian Manifolds
• First-order Methods for Geodesically Convex Optimization
• On the Complexity of Approximating Wasserstein Barycenter
• Convex Color Image Segmentation with Optimal Transport Distances

Cited by in corpus (15)

• Wasserstein barycenters are NP-hard to compute
• Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey
• Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks
• Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization
• Estimating Barycenters of Measures in High Dimensions
• Wasserstein barycenters can be computed in polynomial time in fixed dimension
• Asymptotics of smoothed Wasserstein distances
• Fast and Smooth Interpolation on Wasserstein Space
• Stochastic Saddle-Point Optimization for Wasserstein Barycenters
• Stochastic Approximation versus Sample Average Approximation for population Wasserstein barycenters
• Frank-Wolfe Methods in Probability Space
• Sampling From the Wasserstein Barycenter
• The Wasserstein Proximal Gradient Algorithm
• A novel notion of barycenter for probability distributions based on optimal weak mass transport
• Riemannian-geometric generalizations of quantum fidelities and Bures-Wasserstein distance