Random Coordinate Descent on the Wasserstein Space of Probability Measures

TL;DR

This paper introduces a randomized coordinate descent framework tailored for optimization on the Wasserstein space of probability measures, offering improved efficiency and convergence guarantees in high-dimensional and ill-conditioned scenarios.

Contribution

The authors develop novel Wasserstein coordinate descent algorithms with rigorous convergence analysis, extending coordinate descent techniques to the space of probability measures.

Findings

Methods achieve faster convergence than full-gradient approaches on ill-conditioned problems.

Theoretical guarantees hold under non-convex, Polyak-Łojasiewicz, and geodesically convex conditions.

Numerical experiments demonstrate significant speedups in practical applications.

Abstract

Optimization over the space of probability measures endowed with the Wasserstein-2 geometry is central to modern machine learning and mean-field modeling. However, traditional methods relying on full Wasserstein gradients often suffer from high computational overhead in high-dimensional or ill-conditioned settings. We propose a randomized coordinate descent framework specifically designed for the Wasserstein manifold, introducing both Random Wasserstein Coordinate Descent (RWCD) and Random Wasserstein Coordinate Proximal{-Gradient} (RWCP) for composite objectives. By exploiting coordinate-wise structures, our methods adapt to anisotropic objective landscapes where full-gradient approaches typically struggle. We provide a rigorous convergence analysis across various landscape geometries, establishing guarantees under non-convex, Polyak-{\L}ojasiewicz, and geodesically convex conditions. Our theoretical results mirror the classic convergence properties found in Euclidean space, revealing a compelling symmetry between coordinate descent on vectors and on probability measures. The developed techniques are inherently adaptive to the Wasserstein geometry and offer a robust analytical template that can be extended to other optimization solvers within the space of measures. Numerical experiments on ill-conditioned energies demonstrate that our framework offers significant speedups over conventional full-gradient methods.