Random Coordinate Descent on the Wasserstein Space of Probability Measures

• URL: http://arxiv.org/abs/2604.01606v1
• Date: Thu, 02 Apr 2026 04:36:56 GMT
• Title: Random Coordinate Descent on the Wasserstein Space of Probability Measures
• Abstract summary: Optimization over the space of probability measures endowed with the Wasserstein-2 geometry is central to modern machine learning and mean-field modeling.<n>We propose a coordinate descent framework specifically for the Wasserstein manifold, introducing both Random Wasserstein Coordinate Descent (RWCD) and Random Wasserstein Coordinate Descent (RWCP)<n> Numerical experiments on ill-conditioned energies demonstrate that our framework offers significant speedups over conventional full-gradient methods.
• Score: 1.5888343554854376
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization [13.544454495712037]
  This paper presents a class of objectives over the Wasserstein space.<n>Under convex assumptions, we prove almost station along the gradient of DC decompositions.<n>Our main is on the Maximum Discreprepancy (MMD) and the Energy (ED) scheme.
  arXiv  Detail & Related papers  (2026-06-26T06:53:03Z)
• Fitting Unknown Number of Hyperplanes with Manifold Optimization [57.48093263119306]
  Fitting unknown number of linearplanes to data is a fundamental yet challenging problem in machine learning.<n>Existing approaches often struggle with optimum optima or lack geometric consistency.
  arXiv  Detail & Related papers  (2026-05-27T14:02:20Z)
• $L^2$ over Wasserstein: Statistical Analysis for Optimal Transport [0.0]
  We introduce the $L2$ over Wasserstein space by characterising distances and geodesic geometry.<n>We ensemble statistical convergence results of the optimal transport machinery using the empirical measure within the $L2$ over Wasserstein framework.<n>We refine Schwartz's consistency theorem to the Wasserstein topology and deduce posterior convergence of the same machinery in the $L2$ over Wasserstein space.
  arXiv  Detail & Related papers  (2026-05-20T16:29:01Z)
• Fast Estimation of Wasserstein Distances via Regression on Sliced Wasserstein Distances [70.94157767200342]
  We propose a fast estimation method based on regressing Wasserstein distance on sliced Wasserstein distances.<n>We show that accurate models can be learned from a small number of distribution pairs.<n>Our method consistently provides a better approximation of Wasserstein distance than the state-of-the-art Wasserstein embedding model, Wasserstein Wormhole.
  arXiv  Detail & Related papers  (2025-09-24T19:30:53Z)
• A Particle-Flow Algorithm for Free-Support Wasserstein Barycenters [0.31727619150610836]
  We develop a free-support algorithm for computing Wasserstein barycenters that avoids entropic regularization.<n>In our approach, barycenter atoms evolve as particles advected by averaged optimal-transport displacements.<n>This yields a geometry-aware particle-flow update that preserves sharp features of the Wasserstein barycenter while remaining computationally tractable.
  arXiv  Detail & Related papers  (2025-09-14T21:05:04Z)
• Mirror and Preconditioned Gradient Descent in Wasserstein Space [10.92458731910518]
  We focus on lifting two explicit algorithms: mirror descent and preconditioned gradient descent. We prove guarantees of convergence of some Wasserstein-gradient-based discrete-time schemes for new pairings of objective functionals and regularizers.
  arXiv  Detail & Related papers  (2024-06-13T09:07:22Z)
• Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space [21.12668895845275]
  We extend the family of continuous optimization methods in the Wasserstein space by extending the gradient on flow into the gradient descent.<n>By leveraging the property of Wasserstein space, we construct differential equations (SDEs) to approximate the corresponding discrete Euclidean dynamics.<n>Finally, we establish convergence rates of the proposed flows, which align with those known in the Euclidean setting.
  arXiv  Detail & Related papers  (2024-01-24T15:35:44Z)
• Tangential Wasserstein Projections [0.4297070083645048]
  We develop a notion of projections between sets of probability measures using the geometric properties of the 2-Wasserstein space. The idea is to work on regular tangent cones of the Wasserstein space using generalized geodesics.
  arXiv  Detail & Related papers  (2022-07-29T14:59:58Z)
• Learning High Dimensional Wasserstein Geodesics [55.086626708837635]
  We propose a new formulation and learning strategy for computing the Wasserstein geodesic between two probability distributions in high dimensions. By applying the method of Lagrange multipliers to the dynamic formulation of the optimal transport (OT) problem, we derive a minimax problem whose saddle point is the Wasserstein geodesic. We then parametrize the functions by deep neural networks and design a sample based bidirectional learning algorithm for training.
  arXiv  Detail & Related papers  (2021-02-05T04:25:28Z)
• Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization [94.18714844247766]
  Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport. We present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures.
  arXiv  Detail & Related papers  (2021-02-02T21:01:13Z)
• Continuous Regularized Wasserstein Barycenters [51.620781112674024]
  We introduce a new dual formulation for the regularized Wasserstein barycenter problem. We establish strong duality and use the corresponding primal-dual relationship to parametrize the barycenter implicitly using the dual potentials of regularized transport problems.
  arXiv  Detail & Related papers  (2020-08-28T08:28:06Z)
• On Projection Robust Optimal Transport: Sample Complexity and Model Misspecification [101.0377583883137]
  Projection robust (PR) OT seeks to maximize the OT cost between two measures by choosing a $k$-dimensional subspace onto which they can be projected. Our first contribution is to establish several fundamental statistical properties of PR Wasserstein distances. Next, we propose the integral PR Wasserstein (IPRW) distance as an alternative to the PRW distance, by averaging rather than optimizing on subspaces.
  arXiv  Detail & Related papers  (2020-06-22T14:35:33Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.