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

翻译：在配备Wasserstein-2几何结构的概率测度空间上进行优化，是现代机器学习与平均场建模的核心问题。然而，传统依赖全Wasserstein梯度的方法在高维或病态条件下常面临高计算开销。我们针对Wasserstein流形提出了一种随机化坐标下降框架，分别引入随机Wasserstein坐标下降法（RWCD）和随机Wasserstein坐标近端-梯度法（RWCP）用于处理复合目标函数。通过利用坐标方向结构，我们的方法能够适应全梯度方法通常难以处理的各向异性目标地形。我们针对不同地形几何结构提供了严格的收敛性分析，在非凸、Polyak-Łojasiewicz以及测地凸条件下建立了理论保证。理论结果复现了欧氏空间中经典的收敛性质，揭示了向量坐标下降与概率测度坐标下降之间引人注目的对称性。所发展的技术天然适应Wasserstein几何结构，并提供了一个稳健的分析模板，可推广至测度空间内的其他优化求解器。针对病态能量的数值实验表明，我们的框架相比传统全梯度方法能显著加速。