# First-order Conditions for Optimization in the Wasserstein Space

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

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 25 Sep 2022.

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

- [Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling](https://mlsanity.com/paper/2504.10612/) (2025)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Mirror and Preconditioned Gradient Descent in Wasserstein Space](https://mlsanity.com/paper/2406.08938/) (2024)
- [Symmetric Mean-field Langevin Dynamics for Distributional Minimax Problems](https://mlsanity.com/paper/2312.01127/) (2023)
- [Distributionally Robust Optimization via Iterative Algorithms in Continuous Probability Spaces](https://mlsanity.com/paper/2412.20556/) (2024)
- [Non-geodesically-convex optimization in the Wasserstein space](https://mlsanity.com/paper/2406.00502/) (2024)
- [DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows](https://mlsanity.com/paper/2503.01140/) (2025)
- [Convergence Analysis of the Wasserstein Proximal Algorithm beyond Geodesic Convexity](https://mlsanity.com/paper/2501.14993/) (2025)
- [Hessian-guided Perturbed Wasserstein Gradient Flows for Escaping Saddle Points](https://mlsanity.com/paper/2509.16974/) (2025)
- [Convergence Rates for Distribution Matching with Sliced Optimal Transport](https://mlsanity.com/paper/2602.10691/) (2026)
- [Training Generalizable Collaborative Agents via Strategic Risk Aversion](https://mlsanity.com/paper/2602.21515/) (2026)
- [Wasserstein Gaussianization and Efficient Variational Bayes for Robust Bayesian Synthetic Likelihood](https://mlsanity.com/paper/2305.14746/) (2023)
- [Global Optimisation of Black-Box Functions with Generative Models in the Wasserstein Space](https://mlsanity.com/paper/2407.11917/) (2024)
- [Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation](https://mlsanity.com/paper/2603.23398/) (2026)
- [Static and Dynamic Approaches to Computing Barycenters of Probability Measures on Graphs](https://mlsanity.com/paper/2603.26940/) (2026)
- [Nonlinear Distributionally Robust Optimization](https://mlsanity.com/paper/2306.03202/) (2023)
- [A Particle-Flow Algorithm for Free-Support Wasserstein Barycenters](https://mlsanity.com/paper/2509.11435/) (2025)
- [Ensemble Distributionally Robust Bayesian Optimisation with Continuous Context](https://mlsanity.com/paper/2605.07565/) (2026)
- [Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization](https://mlsanity.com/paper/2606.27767/) (2026)

## Abstract (arXiv)

We study first-order optimality conditions for constrained optimization in the Wasserstein space, whereby one seeks to minimize a real-valued function over the space of probability measures endowed with the Wasserstein distance. Our analysis combines recent insights on the geometry and the differential structure of the Wasserstein space with more classical calculus of variations. We show that simple rationales such as "setting the derivative to zero" and "gradients are aligned at optimality" carry over to the Wasserstein space. We deploy our tools to study and solve optimization problems in the setting of distributionally robust optimization and statistical inference. The generality of our methodology allows us to naturally deal with functionals, such as mean-variance, Kullback-Leibler divergence, and Wasserstein distance, which are traditionally difficult to study in a unified framework.