# Random Coordinate Descent on the Wasserstein Space of Probability Measures

arXiv: https://arxiv.org/abs/2604.01606 · HTML: https://mlsanity.com/paper/2604.01606/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 1 of its 46 references that are in the map.
- First posted on arXiv 2 Apr 2026.

## Builds on (most cited first, 1 of 1)

- [Convergence of coordinate ascent variational inference for log-concave measures via optimal transport](https://mlsanity.com/paper/2404.08792/) (2024)

## Built on by (most cited first, 1 of 1)

- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)

## Abstract (arXiv)

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