# Mirror and Preconditioned Gradient Descent in Wasserstein Space

arXiv: https://arxiv.org/abs/2406.08938 · HTML: https://mlsanity.com/paper/2406.08938/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 121 references that are in the map.
- First posted on arXiv 13 Jun 2024.

## Builds on (most cited first, 21 of 21)

- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Proximal Langevin Algorithm: Rapid Convergence Under Isoperimetry](https://mlsanity.com/paper/1911.01469/) (2019)
- [GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics](https://mlsanity.com/paper/2310.09254/) (2023)
- [Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifolds](https://mlsanity.com/paper/2403.06560/) (2024)
- [Information Newton's flow: second-order optimization method in probability space](https://mlsanity.com/paper/2001.04341/) (2020)
- [Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space](https://mlsanity.com/paper/2312.02849/) (2023)
- [First-order Conditions for Optimization in the Wasserstein Space](https://mlsanity.com/paper/2209.12197/) (2022)
- [Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces](https://mlsanity.com/paper/2402.04613/) (2024)
- [Fast sampling from constrained spaces using the Metropolis-adjusted Mirror Langevin algorithm](https://mlsanity.com/paper/2312.08823/) (2023)
- [Wasserstein Mirror Gradient Flow as the limit of the Sinkhorn Algorithm](https://mlsanity.com/paper/2307.16421/) (2023)
- [Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization](https://mlsanity.com/paper/2205.13098/) (2022)
- [Learning Elastic Costs to Shape Monge Displacements](https://mlsanity.com/paper/2306.11895/) (2023)
- [Ambiguity set and learning via Bregman and Wasserstein](https://mlsanity.com/paper/1705.08056/) (2017)
- [Gradient descent with a general cost](https://mlsanity.com/paper/2305.04917/) (2023)
- [Covariance-modulated optimal transport and gradient flows](https://mlsanity.com/paper/2302.07773/) (2023)
- [Accelerated gradient descent method for functionals of probability measures by new convexity and smoothness based on transport maps](https://mlsanity.com/paper/2305.05127/) (2023)
- [A Lagrangian approach to totally dissipative evolutions in Wasserstein spaces](https://mlsanity.com/paper/2305.05211/) (2023)
- [Some Convexity Criteria for Differentiable Functions on the 2-Wasserstein Space](https://mlsanity.com/paper/2306.09120/) (2023)
- [Sinkhorn Flow: A Continuous-Time Framework for Understanding and Generalizing the Sinkhorn Algorithm](https://mlsanity.com/paper/2311.16706/) (2023)
- [Mirror Duality in Convex Optimization](https://mlsanity.com/paper/2311.17296/) (2023)

## Built on by (most cited first, 12 of 12)

- [Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods](https://mlsanity.com/paper/2502.01384/) (2025)
- [Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces](https://mlsanity.com/paper/2412.20556/) (2024)
- [Worst-case generation via minimax optimization in Wasserstein space](https://mlsanity.com/paper/2512.08176/) (2025)
- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)
- [Convergence Rates for Distribution Matching with Sliced Optimal Transport](https://mlsanity.com/paper/2602.10691/) (2026)
- [Semi-Implicit Functional Gradient Flow for Efficient Sampling](https://mlsanity.com/paper/2410.17935/) (2024)
- [A Gradient Flow Perspective on Minimum MMD Estimation](https://mlsanity.com/paper/2607.03871/) (2026)
- [Subspace Langevin Monte Carlo](https://mlsanity.com/paper/2412.13928/) (2024)
- [Analyzing and Improving Diffusion Models for Time-Series Data Imputation: A Proximal Recursion Perspective](https://mlsanity.com/paper/2602.01182/) (2026)
- [Variational Inference via Entropic Transport Descent](https://mlsanity.com/paper/2606.25265/) (2026)
- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)
- [Trust-Region Optimization for Smooth Potential-Interaction Energies in Wasserstein Space](https://mlsanity.com/paper/2610.08883/) (2026)

## Abstract (arXiv)

As the problem of minimizing functionals on the Wasserstein space encompasses many applications in machine learning, different optimization algorithms on ℝd have received their counterpart analog on the Wasserstein space. We focus here on lifting two explicit algorithms: mirror descent and preconditioned gradient descent. These algorithms have been introduced to better capture the geometry of the function to minimize and are provably convergent under appropriate (namely relative) smoothness and convexity conditions. Adapting these notions to the Wasserstein space, we prove guarantees of convergence of some Wasserstein-gradient-based discrete-time schemes for new pairings of objective functionals and regularizers. The difficulty here is to carefully select along which curves the functionals should be smooth and convex. We illustrate the advantages of adapting the geometry induced by the regularizer on ill-conditioned optimization tasks, and showcase the improvement of choosing different discrepancies and geometries in a computational biology task of aligning single-cells.