# Variational Analysis in the Wasserstein Space

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 15 Jun 2024.

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

- [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)
- [Training Generalizable Collaborative Agents via Strategic Risk Aversion](https://mlsanity.com/paper/2602.21515/) (2026)
- [Flows and Diffusions on the Neural Manifold](https://mlsanity.com/paper/2507.10623/) (2025)
- [Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation](https://mlsanity.com/paper/2603.23398/) (2026)
- [Trust-Region Optimization for Smooth Potential-Interaction Energies in Wasserstein Space](https://mlsanity.com/paper/2610.08883/) (2026)

## Abstract (arXiv)

We study optimization problems whereby the optimization variable is a probability measure. Since the space of probability measures is not a vector space, many classical methods for optimization (e.g., gradients) do not directly apply. Thus, one typically resorts to the abstract machinery of infinite-dimensional analysis or other ad-hoc methodologies, not tailored to the space of probability measures, which however involve projections or rely on convexity-type assumptions. We believe instead that these problems call for a comprehensive methodological framework for calculus in the space of probability measures. In this work, we combine ideas from optimal transport, variational analysis, and Wasserstein gradient flows to equip the Wasserstein space (i.e., the space of probability measures endowed with the Wasserstein distance) with a variational structure, both by combining and extending existing results and introducing novel tools. Our theoretical analysis culminates in general necessary optimality conditions. These conditions (i) resemble the optimality conditions in Euclidean spaces, such as the KKT conditions, (ii) are intuitive, informative, and easy to study, and (iii) yield closed-form solutions or can be used to design computationally attractive algorithms. We accompany our theoretical results with numerous examples and present applications to machine learning, drug discovery, and distributionally robust optimization.