Random Coordinate Descent on the Wasserstein Space of Probability Measures

Yewei Xu; Qin Li

stat.MLApr 2, 2026

Random Coordinate Descent on the Wasserstein Space of Probability Measures

Authors: Yewei Xu, Qin Li

Organizations: Department of Mathematics, University of Wisconsin-Madison, Madison, WI 53706-1388, USA

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

Figures & tables

Method	Nonconvex		(Geo.) Convex		$m$ -PL / Strongly Convex
	Eucl.	Wass.	Eucl.	Wass.	Eucl.	Wass.
	$\\|\nabla f(x)\\|^{2}$	$\\|\nabla_{\mathbb{W}}E[\mu]\\|_{\mu}^{2}$	$f(x)-f_{*}$	$E[\mu]-E_{*}$	$f(x)-f_{*}$	$E[\mu]-E_{*}$
RCD/RWCD	$O\left(\frac{L_{\rm sum}}{\epsilon}\right)$	$O\left(\frac{L_{\rm sum}}{\epsilon}\right)$	$O\left(\frac{L_{\rm sum}}{\epsilon}\right)$	$O\left(\frac{L_{\rm sum}}{\epsilon}\right)$	$O\left(\frac{L_{\rm sum}}{m}\ln\frac{1}{\epsilon}\right)$	$O\left(\frac{L_{\rm sum}}{m}\ln\frac{1}{\epsilon}\right)$
RCP/RWCP	$O\left(\frac{\mathcal{C}}{\epsilon}\right)$	$O\left(\frac{\mathcal{C}}{\epsilon}\right)$	$O\left(\frac{\mathcal{C}}{\epsilon}\right)$	$O\left(\frac{\mathcal{C}}{\epsilon}\right)$	$O\left(\frac{\mathcal{C}}{m}\ln\frac{1}{\epsilon}\right)$	$O\left(\frac{\mathcal{C}}{m}\ln\frac{1}{\epsilon}\right)$

Table 1: Euclidean vs. Wasserstein convergence rates. Our Wasserstein coordinate methods recover the Euclidean complexity regimes with geometry-specific constants

L_{\rm sum}

and

\mathcal{C}

(cf. Theorems 12 and 14 ).

Figure 1: Example 1: RWCD vs. WGD on the 2D quadratic objective ( 28 ) using $N=2000$ particles. Horizontal axis: work measured in coordinate-gradient evaluations (one WGD step equals $d=2$ work units). Blue: RWCD median with $10$ – $90\%$ band across $50$ runs; orange: WGD. The three panels respectively show $E[\mu_{k}]$ , empirical mean of the first coordinate of $N$ samples, empirical mean of the second coordinate of $N$ samples.

Figure 2: Example 1: Particle cloud snapshots for the 2D quadratic experiment ( 28 ). Top row: RWCD (a median trial); bottom row: WGD. Snapshots are taken at work $=0,8,1200,2400$ , where one WGD iteration is counted as $d=2$ work units.

Figure 3: Example 2: RWCD vs. WGD for the quadratic potential and interaction objective with $N=2000$ and $d=50$ . The horizontal axis represents work measured in coordinate-gradient evaluations. The blue curve shows the RWCD median with a $10$ – $90\%$ confidence band across $50$ independent runs, while the orange curve represents WGD.

Figure 4: Example 3: RWCD vs. WGD for the MMD-type functional with an anisotropic Gaussian kernel in $d=50$ . The horizontal axis represents work measured in coordinate-gradient evaluations. The blue curve denotes the RWCD median with a $10$ – $90\%$ band across $50$ runs; the orange curve denotes WGD. Note that the RWCD band is visually indistinguishable from the median curve at this resolution.

Figure 5: Example 4: Coordinate-wise smoothness constants $H_{i}$ of the regularizer $\Psi_{\epsilon}$ . The $y$ -axis shows $H_{i}$ on a logarithmic scale, and the $x$ -axis indexes the coordinates after sorting by $H_{i}$ .

Panel of Figure 6 — Figure 6: Example 4: RWCP vs. WPG on the composite objective $E=G+\Psi_{\epsilon}$ . Left: energy gap $E[\mu_{k}]-E_{*}$ ; Right: barycenter norm. Horizontal axis: work measured in coordinate-gradient evaluations. Blue: RWCP median with $10$ – $90\%$ band across $50$ runs; orange: WPG.

Figure 7: Example 5: Coordinate-wise smoothness constants $L_{i}$ of $E$ on $\mathcal{P}_{2}(\mathcal{B})$ . The $y$ -axis shows $L_{i}$ on a logarithmic scale, and the $x$ -axis indexes the coordinates after sorting by $L_{i}$ .

Panel of Figure 8 — Figure 8: Example 5: RWCD vs. WGD on the neural network objective ( 34 ). Left: energy gap $E[\mu_{k}]-E_{*}$ . Right: running best gradient norm $\min_{0\leq i<k}\|\nabla_{\mathbb{W}}E[\mu_{i}]\|_{\mu_{i}}^{2}$ . The horizontal axis measures work in coordinate-gradient evaluations. Blue: RWCD median with $10$ – $90\%$ band over $50$ runs. Orange: WGD.

Nonconvex Optimization

Nonconvex optimization develops methods for problems whose objectives or constraints can have multiple local optima, saddle points, or nonsmooth structure.

See all 145 papers

Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization
Jun 26, 2026Clément Bonet, Pierre-Cyril Aubin-Frankowski, Youssef MrouehMaximum Mean Discrepancy Nonconvex Optimization
CMAP, CNRS, Ecole Polytechnique, IP Paris · CERMICS, CNRS, ENPC, IP Paris · IBM Research
Stochastic Gradient Descent over P2
Sep 11, 2026Maria Oprea, Qin Li, Yunan YangStochastic Optimization Stochastic Gradient Descent
Center for Applied Mathematics, Cornell University · Department of Mathematics, University of Wisconsin–Madison, USA · Department of Mathematics, Cornell University, USA
Projected Riemannian Gradient Descent for the Bures-Wasserstein Barycenter: Dimension-Independent Linear Convergence at Unit Step Size
Sep 3, 2026A. AfhamRiemannian Optimization Optimization Convergence Analysis
Centre for Quantum Technologies, National University of Singapore, Singapore

Jun 26, 2026J/K move · Enter open · S save

Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

Clément Bonet, Pierre-Cyril Aubin-Frankowski, Youssef Mroueh

CMAP, CNRS, Ecole Polytechnique, IP Paris · CERMICS, CNRS, ENPC, IP Paris · IBM Research

Maximum Mean Discrepancy Nonconvex Optimization Optimal Transport

Open paper AskPDF arXiv

Optimizing functionals over the space of probability measures is now ubiquitous in machine learning. A widely used approach is to perform the optimization directly over the Wasserstein space, but many objective functionals of practical interest are non-convex along Wasserstein geodesics, making the analysis of standard first-order methods challenging. In this work, we study a class of objectives over the Wasserstein space that admit a difference-of-convex (DC) decomposition and we lift the classical convex-concave procedure (CCCP) to this setting. Under smoothness and strong convexity assumptions on the convex components of the decomposition, we prove almost stationarity along the iterates of the resulting algorithm. Our main focus is on the Maximum Mean Discrepancy (MMD) and the Energy Distance (ED) functionals, for which we develop explicit Wasserstein DC decompositions, and establish local convergence of the scheme under mild assumptions. Empirically, we show that well-chosen DC decompositions yield faster and more stable convergence than Wasserstein gradient descent on these MMD objectives.

Abstract

Nonconvex Optimization

Explore similar work

Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

Stochastic Gradient Descent over P2

Projected Riemannian Gradient Descent for the Bures-Wasserstein Barycenter: Dimension-Independent Linear Convergence at Unit Step Size