Random Coordinate Descent on the Wasserstein Space of Probability Measures

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.

• The paper introduces Random Wasserstein Coordinate Descent (RWCD) and Proximal-Gradient methods that reduce per-iteration cost using coordinate-wise updates.
• The methodology adapts pushforward operations on particle ensembles to define coordinate updates, enabling rigorous convergence analysis in nonconvex and geodesically convex settings.
• Numerical experiments validate that RWCD outperforms full Wasserstein gradient descent, especially in high-dimensional, anisotropic, and ill-conditioned scenarios.

Random Coordinate Descent on the Wasserstein Space of Probability Measures

Motivation and Context

This work addresses a core computational bottleneck in measure-valued optimization: the high cost of full Wasserstein gradient computations in high-dimensional or ill-conditioned settings. The classical approach of optimization in the Wasserstein-2 ($\mathbb{W}_2$) space—central, for example, to Bayesian inference, optimal transport, sampling, and mean-field models—relies on gradient-based schemes analogous to Euclidean gradient descent (GD). However, the per-iteration cost of full gradient calculations grows prohibitively with dimension, particularly when objective landscapes are highly anisotropic.

The paper proposes to translate the coordinate descent paradigm from finite-dimensional spaces to the Wasserstein manifold, thereby introducing Random Wasserstein Coordinate Descent (RWCD) and Random Wasserstein Coordinate Proximal-Gradient (RWCP) methods. These solvers enable inexpensive per-iteration cost by exploiting coordinate-wise smoothness, and adapt step sizes according to coordinate geometry, closely paralleling the efficiency gains in Euclidean RCD/RCP.

Methodology

Analytical Framework

The central methodological innovation involves formulating coordinate updates in Wasserstein space. Since probability measures do not afford canonical coordinates as $\mathbb{R}^d$ vectors do, the approach is to represent measures by particle ensembles and introduce updates as pushforwards along single-coordinate directions. Given a measure $\mu$, a coordinate-restricted pushforward is defined via

$\mu \mapsto (\mathrm{Id} + U_{i} T)_{\#}\mu$

where $U_i$ is the projection matrix for the $i$-th coordinate and $T$ is a transport direction. The update modifies only the $i$-th marginal of $\mu$, mirroring coordinate updates in Euclidean space.

RWCD for a functional $E[\mu]$ is then

$\mathbb{R}^d$0

with randomized coordinate sampling weighted by the local smoothness $\mathbb{R}^d$1. The RWCP method generalizes this to the composite setting, leveraging a coordinate-proximal subproblem in each iteration. Optimal sampling and step size schedules are deduced by symmetry to classical coordinate descent analyses, with each iteration requiring only a partial Wasserstein gradient.

Theoretical Complexity

A rigorous convergence theory is developed, including nonconvex, geodesically convex, and Polyak-Łojasiewicz (PL) regimes. Critical technical challenges stem from the nonlinearity and lack of linear structure in $\mathbb{R}^d$2, requiring adapted notions of coordinate-wise smoothness, geodesic convexity, and PL conditions in Wasserstein space. The coordinate descent lemma in this context relies on integrating the Wasserstein gradient along axis-restricted transport interpolants, rather than standard linear increments.

Notably, RWCD and RWCP achieve the following rates—matching their Euclidean counterparts up to geometry-dependent constants:

• Nonconvex: $\mathbb{R}^d$3 iterations to $\mathbb{R}^d$4-stationarity
• Geodesically convex: $\mathbb{R}^d$5 for an $\mathbb{R}^d$6-optimality gap
• $\mathbb{R}^d$7-PL or strongly convex: $\mathbb{R}^d$8 for an $\mathbb{R}^d$9-gap

where $\mu$0 is the aggregate coordinate-wise smoothness. These rates are formally identical to classical RCD/RCP (modulo the replacement of the global $\mu$1 with $\mu$2 and detailed Wasserstein geometric structure).

Numerical Experiments

A diverse suite of high-dimensional experiments substantiates the theoretical claims.

2D Quadratic Objective: Extreme Anisotropy

The first experiment addresses a 2D objective with vastly disparate smoothness along coordinates ($\mu$3). RWCD achieves a much accelerated reduction in energy relative to Wasserstein gradient descent (WGD). The energy landscape's anisotropy causes WGD to take very small steps globally, while RWCD progresses quickly along less stiff coordinates. This is visible in differential convergence rates for each marginal.

Figure 1: RWCD and WGD convergence on a $\mu$4D quadratic objective sharply illustrate the benefit of aggressive coordinate-wise step sizes.

Figure 2: Particle clouds show rapid collapse to minimizer under RWCD compared to WGD.

High-dimensional Quadratic and Interaction Functionals

In $\mu$5, with log-spaced eigenvalues generating ill-conditioning, RWCD substantially outperforms WGD in terms of work-normalized convergence. The effect is robust across repeated trials.

Figure 3: In high dimensions with large anisotropy, RWCD achieves much faster decay of the energy than WGD, measured per coordinate-gradient evaluation.

MMD-type Functionals

RWCD is benchmarked on a nonconvex MMD functional with anisotropic Gaussian kernels. Performance gains are consistent, and the variance across repeated trials is negligible, demonstrating the scheme's robustness to both nonconvexity and inhomogeneous landscapes.

Figure 4: RWCD matches or exceeds WGD even on nonconvex MMD-type functionals; variance in outcomes is minimal.

Composite Regularized Objectives

A more complex scenario is constructed using a smooth but inseparable regularizer with widely varying coordinate-wise smoothness constants $\mu$6.

Figure 5: The diverse $\mu$7 profile for a non-separable regularizer highlights coordinate anisotropy, essential for RWCP's effectiveness.

RWCP, by adapting step sizes and sampling, dominates full Wasserstein proximal-gradient (WPG) in reducing both the energy gap and the barycenter error, across many orders of magnitude.

Figure 6: RWCP demonstrates rapid decay of energy and barycentric error vis-à-vis WPG on non-separable composite energies.

Mean-field Neural Network Optimization

Finally, a mean-field two-layer NN functional in $\mu$8 parameter dimensions is optimized. The smoothness spectrum for parameter groups illustrates the significant anisotropy present.

Figure 7: The highly disparate $\mu$9 values for neural network mean-field functional parameters.

RWCD yields consistent acceleration over WGD both for energy reduction and in driving down the (running best) Wasserstein gradient norm.

Figure 8: RWCD versus WGD on mean-field neural network objective: coordinate-adaptive steps produce more rapid optimization progress.

Implications and Future Directions

The coordinate descent dictionary developed here closes a major gap in measure-valued optimization. By porting RCD and RCP to Wasserstein space and analyzing their convergence with explicit geometric constants, the methodology enables scalable algorithms in high-dimensional Bayesian inference, inverse problems, mean-field learning, and scalable sampling.

The approach is extensible: analogues to block-coordinate descent (BCD) and specialized partitioning could be integrated, and the framework is compatible with particle, kernel, or grid-based representations. Further, these results suggest persistent benefits for landscape anisotropy exploitation in large-scale probabilistic modeling, reinforcing the need for geometric adaptivity in infinite-dimensional optimization solvers.

Conclusion

This work successfully extends coordinate descent and proximal-gradient methods to the Wasserstein space, attaining computational complexity in line with the best Euclidean methods. The practical upshot is a considerable speedup for ill-conditioned, high-dimensional measure optimization, confirmed rigorously and numerically. The cross-pollination of finite- and infinite-dimensional coordinate solvers presented here is poised for significant impact in scalable Bayesian inference, mean-field learning, and distributional robustness.