[2604.01606] Random Coordinate Descent on the Wasserstein Space of Probability Measures

Statistics > Machine Learning arXiv:2604.01606 (stat) [Submitted on 2 Apr 2026] Title:Random Coordinate Descent on the Wasserstein Space of Probability Measures Authors:Yewei Xu, Qin Li View a PDF of the paper titled Random Coordinate Descent on the Wasserstein Space of Probability Measures, by Yewei Xu and 1 other authors View PDF HTML (experimental) 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-Ł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 adaptiv...