First-order Conditions for Optimization in the Wasserstein Space

TL;DR

This paper develops first-order optimality conditions for constrained optimization in the Wasserstein space, integrating geometric and variational methods to handle complex functionals like KL divergence and mean-variance.

Contribution

It introduces a unified framework for first-order optimality conditions in Wasserstein space, applicable to various complex functionals and optimization problems.

Findings

Derived conditions analogous to classical derivatives in Wasserstein space

Applied framework to distributionally robust optimization and statistical inference

Unified treatment of diverse functionals like mean-variance and KL divergence

Abstract

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.