# Variational Analysis in the Wasserstein Space

Nicolas Lanzetti∗
ETH Zürich
<EMAIL_ADDRESS>
&Antonio Terpin∗
ETH Zürich
<EMAIL_ADDRESS>&Florian Dörfler
ETH Zürich
<EMAIL_ADDRESS>
###### Abstract

We study optimization problems whereby the optimization variable is a
probability measure. Since the probability space is not a vector space, many
classical and powerful methods for optimization (e.g., gradients) are of
little help. Thus, one typically resorts to the abstract machinery of
infinite-dimensional analysis or other ad-hoc methodologies, not tailored to
the probability space, 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 probability spaces. 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 very general necessary
optimality conditions for optimality. Notably, our conditions (i) resemble the
rationales of Euclidean spaces, such as the Karush-Kuhn-Tucker and Lagrange
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 believe this framework lays the foundation for new
algorithmic and theoretical advancements in the study of optimization problems
in probability spaces, which we exemplify with numerous case studies and
applications to machine learning, drug discovery, and distributionally robust
optimization.

∗ Equal contribution. This research was supported by the Swiss National
Science Foundation under the NCCR Automation, grant agreement 51NF40_180545.

##  1 Main result

This work considers optimization problems of the form

$\inf_{\operatorname{\mu}\in\operatorname{\mathcal{C}}}\,\operatorname{\mathcal{J}}(\operatorname{\mu}),$
(1)

where $\operatorname{\mathcal{C}}\subseteq\mathcal{P}(\mathbb{R}^{d})$ is a
set of admissible probability measures and
$\operatorname{\mathcal{J}}:\mathcal{P}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$ is
a functional to minimize. This abstract problem setting stems from the
observation that numerous fields, including machine learning, robust
optimization, and biology, tackle their own version of (1), but with ad-hoc
methods that often cease to be effective as soon as the problem structure
deviates from idealized assumptions. Despite the recent efforts in the
literature [1], these problems still demand a comprehensive theory for the
optimization problem (1), which is the subject of this work. Specifically, we
derive novel necessary first-order optimality conditions for (1), for
arbitrary functionals and constraints. These formally resemble the rationales
of Euclidean spaces (e.g., Karush–Kuhn–Tucker and Lagrange conditions) and are
intuitive, informative, and easy to study. As a byproduct of our analysis, we
translate tools from variational analysis (e.g., generalized subgradients,
normal cones, tangent cones, etc.) to the Wasserstein space (i.e., the
probability space endowed with the Wasserstein distance). After practicing
these novel tools in numerous pedagogical examples, we tackle open problems
arising in machine learning, drug discovery, and distributionally robust
optimization (DRO), showcasing how our conditions result either in closed-form
solutions of (1) or computationally attractive algorithms.

Our main result are general first-order optimality conditions of (1):

[First-order optimality conditions] If
$\operatorname{\mu^{\ast}}\in\mathcal{P}(\mathbb{R}^{d})$ is an optimal
solution of (1) with finite second moment and provided that a constraint
qualification holds, then the “Wasserstein subgradients” are “aligned” with
the constraints at “optimality”, i.e.,

$\mathbf{0}_{\operatorname{\mu^{\ast}}}\in\partial\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}})+\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\mu^{\ast}}),$

where $\partial\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}})$ is the
“Wasserstein subgradient” of $\operatorname{\mathcal{J}}$ at
$\operatorname{\mu^{\ast}}$,
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\mu^{\ast}})$ is
the “Wasserstein normal cone” of $\operatorname{\mathcal{C}}$ at
$\operatorname{\mu^{\ast}}$ and $\mathbf{0}_{\operatorname{\mu^{\ast}}}$ is a
“null Wasserstein tangent vector” at $\operatorname{\mu^{\ast}}$.

As corollaries of our theorem, we obtain the “Wasserstein counterparts” of
Fermat’s rule in the unconstrained setting (i.e., the gradient vanishes at
optimality) and the Lagrange conditions for (in)equality-constrained settings
(i.e., the “gradients” of the objective and the constraint are “aligned” at
“optimality”, see Figure 1).

Before diving into variational analysis in the Wasserstein space, we
illustrate our optimality conditions by informally studying a simple and
accessible version of (1). For $\theta\neq 0$ and $\varepsilon>0$, consider
the problem

$\inf_{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})}\;\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\langle\theta,x\right\rangle\right]\qquad\text{subject
to}\qquad\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\|x\right\|^{2}\right]\leq\varepsilon^{2},$
(2)

depicted in Figure 1 for $\theta=\begin{bmatrix}1&1\end{bmatrix}^{\top}$. To
get some intuition, let us restrict to Dirac’s delta of the form $\delta_{x}$
for $x\in\mathbb{R}^{d}$. Accordingly, (2) reduces to
$\inf_{\left\|x\right\|^{2}\leq\varepsilon^{2}}\left\langle\theta,x\right\rangle.$
This optimization problem can be studied through standard first-order
optimality conditions in Euclidean spaces. Since the gradient of the objective
$\nabla_{x}{\left\langle\theta,x\right\rangle}=\theta$ never vanishes, the
optimal solution (if it exists) lies at the boundary. We thus seek the
Lagrange multiplier $\lambda>0$ such that

$0=\nabla_{x}{\left\langle\theta,x\right\rangle}+\lambda\nabla_{x}{\left\|x\right\|^{2}}=\theta+2\lambda
x\quad\text{and}\quad\varepsilon^{2}=\left\|x\right\|^{2},$ (3)

which yields $x=-\varepsilon\frac{\theta}{\left\|\theta\right\|}$ and
$\lambda=\frac{\left\|\theta\right\|}{2\varepsilon}$. Now back to (2): After
basic algebraic manipulations, our main result (stated informally above) tells
us that any solution
$\operatorname{\mu^{\ast}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ of (2) satisfies
the Lagrange-like condition

$0=2\lambda
x+\theta\quad\operatorname{\mu^{\ast}}-\text{a.e.}\quad\text{and}\quad\varepsilon^{2}=\mathbb{E}^{x\sim\operatorname{\mu^{\ast}}}\left[\left\|x\right\|^{2}\right]=\frac{\left\|\theta\right\|^{2}}{4\lambda^{2}},$
(4)

for some $\lambda\geq 0$ constant across the support of
$\operatorname{\mu^{\ast}}$; cf. Figure 1 for
$\theta=\begin{bmatrix}1&1\end{bmatrix}^{\top}$. We conclude that the mass of
any candidate solution is necessarily located at
$-\varepsilon\frac{\theta}{\left\|\theta\right\|}$. In particular, our
optimality condition in (4) mirrors its counterpart on $\mathbb{R}^{d}$ in
(3).

Figure 1: “Gradients” are “aligned” with the constraints at “optimality”. The
figure on the left depicts two empirical candidate solutions
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu}_{1}}$
and
${\color[rgb]{1,0.55078125,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.55078125,0}\operatorname{\mu}_{2}}$
for (1) with
$\operatorname{\mathcal{C}}=\\{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{2})\,|\,\mathbb{E}^{(x_{1},x_{2})\sim\operatorname{\mu}}\left[x_{1}^{2}+x_{2}^{2}\right]\leq\varepsilon^{2}\\}$
(i.e., bounded second moment) and
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\mathbb{E}^{(x_{1},x_{2})\sim{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu}}}\left[x_{1}+x_{2}\right]$,
of which we show the contours and the gradient vector field. The solid (resp.,
dashed) black arrows represent the gradient of the constraint function
$\mathbb{E}^{(x_{1},x_{2})\sim\operatorname{\mu}}\left[x_{1}^{2}+x_{2}^{2}\right]-\varepsilon^{2}$
at
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu}_{1}}$
(resp.,
${\color[rgb]{1,0.55078125,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.55078125,0}\operatorname{\mu}_{2}}$).
Here,
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu}_{1}}$
is indeed a candidate optimal solution: The gradient of the objective is
aligned with the gradient of the constraint. For
${\color[rgb]{1,0.55078125,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.55078125,0}\operatorname{\mu}_{2}}$,
instead, these two are not aligned. Thus,
${\color[rgb]{1,0.55078125,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0.55078125,0}\operatorname{\mu}_{2}}$
cannot be optimal. The figure on the right shows that
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu^{\ast}}}$
satisfies the optimality condition for (1) with
$\operatorname{\mathcal{C}}=\bar{\mathbb{B}}_{W_{2}}({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname{\hat{\mu}}};\varepsilon)=\\{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{2})\,|\,W_{2}\left(\operatorname{\mu},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname{\hat{\mu}}}\right)\leq\varepsilon\\}$
(i.e., Wasserstein ball centered at
${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname{\hat{\mu}}}$
of radius $\varepsilon$) and
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\mathbb{E}^{(x_{1},x_{2})\sim\operatorname{\mu}}\left[x_{1}+x_{2}\right]$,
of which the contours and the gradient vector field are shown. The black
arrows connecting
${\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname{\hat{\mu}}}$
and
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu^{\ast}}}$
represent the gradient of the constraint function
$W_{2}\left({\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu}},{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname{\hat{\mu}}}\right)^{2}-\varepsilon^{2}$.
Since
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu^{\ast}}}$
is optimal, the gradient of the objective and the constraint are aligned at
all the “particles” of
${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\operatorname{\mu^{\ast}}}$.

###  1.1 Introduction to the broader context

Optimization problems over the probability space are ubiquitous across a
variety of fields.

#####  DRO

DRO has emerged as a major paradigm for decision-making under uncertainty [2].
In DRO, the goal is to identify solutions that are robust against a range of
possible (adversary) probability measures, acknowledging the inherent
ambiguity in real-world data and other a-priori assumptions in uncertainty
modeling, where the true underlying probability measure is uncertain and
difficult to ascertain. Thus, DRO falls within the scope of (1) where
$\operatorname{\mathcal{J}}$ is a risk measure [3, 4, 5] and
$\operatorname{\mathcal{C}}$ is a so-called ambiguity set of probability
measures, often defined in terms of the Kullback-Leibler divergence [6, 7, 8]
or an optimal transport discrepancy [9, 10, 11, 12, 13].

##### Inverse problems

In inverse problems one seeks the state of a system given some noisy
observations. For instance, Bayesian inference amounts to solving

$\inf_{\operatorname{\mu}\in\mathcal{P}(\mathbb{R}^{d})}\mathbb{E}^{x\sim\operatorname{\mu}}\left[\sum_{i=1}^{N}-\log(p(z_{i}|x))\right]+\operatorname{KL}(\operatorname{\mu}||\operatorname{\hat{\mu}}),$
(5)

where $z_{i}$ for $i\in\\{1,\ldots,N\\}$ are the observations, $p(z_{i}|x)$ is
the probability of the observation $x_{i}$ given the state
$x\sim\operatorname{\mu}$, and $\operatorname{KL}$ is the Kullback-Leibler
divergence between a candidate posterior $\operatorname{\mu}$ and
$\operatorname{\hat{\mu}}\in\mathcal{P}(\mathbb{R}^{d})$, the prior [14].
Various inference problems (e.g., see [14, Table 1] and [15, 16]) result from
modifications of (5).

##### Reinforcement learning

The dual formulation [17] of the reinforcement learning (RL) problem seeks the
optimal stationary state-action distribution compatible with the dynamics
$\operatorname{\mu^{\ast}}\in\Delta(\mathcal{S}\times\mathcal{A})\subseteq\mathcal{P}(\mathcal{S}\times\mathcal{A})$,
where $\mathcal{S}$ is the state space and $\mathcal{A}$ the action space,
that maximizes the expected reward
$r:\mathcal{S}\times\mathcal{A}\to\bar{\mathbb{R}}$,

$\operatorname{\mu^{\ast}}\in\operatorname*{argmax}_{\operatorname{\mu}\in\Delta(\mathcal{S}\times\mathcal{A})}\mathbb{E}^{(s,a)\sim\operatorname{\mu}}\left[r(s,a)\right].$
(6)

Variations of (6) yield different problem settings [18, 19, 20].

##### Many others

A growing number of fields are tackling optimal decision problems formulated
in the probability space, including weather forecasting [21], single-cell
perturbation responses [22], control of dynamical systems [23, 24], neural
network training [25], mean-field control [26, 27], and finance [28, 2], among
others.

###  1.2 Related work

Our work studies (1) through the lens of optimal transport. The theory of
optimal transport, dating back to the seminal work of Monge [29] and
Kantorovich [30], defines a metric, the Wasserstein metric, on the space of
probability measures. While the probability space is not a vector space, which
makes most optimization tools in Banach space (e.g., [31, 32]) inapplicable,
the theory of optimal transport enables a notion of differentiability, called
Wasserstein differentiability, specifically tailored for the probability space
[33, 34, 35]. Wasserstein differentiability was exploited in [1] to derive
first-order optimality conditions for (1) for differentiable objective and
constraint sets being sublevel sets of differentiable functionals. In the
context of optimal control, [36, 37, 27] use Wasserstein differentiability to
derive optimality conditions for optimal control problems, whereas [24]
explores the properties of the dynamic programming algorithm in probability
spaces for discrete-time. The theory has also been applied to derive
algorithms to compute Wasserstein barycenters [38, 39, 40], to analyze over-
parametrized neural networks [41, 25], approximate inference [42], and
reinforcement learning [43, 19].

Unfortunately, many functionals over the probability space, including the
Wasserstein distance itself, fail to be differentiable. This, together with
the rigidity of so-far-considered feasible sets $\operatorname{\mathcal{C}}$,
effectively hinders the deployment of these tools in many practical instances.
From a technical standpoint, these limitations are intrinsic in the choice of
the variations in the probability space: Variations induced by the pushforward
of (sufficiently regular) _transport maps_ , as in [1], are not expressive
enough. For instance, the pushforward of an empirical probability measure is
always empirical (in particular, mass cannot be split). Thus, while attractive
(e.g., perturbations can be captured by well-behaved functions, which form a
vector space), a comprehensive theory of calculus requires a more general way
of perturbing probability measures.

Intuitively, we can approximate a Dirac’s delta by Gaussians with vanishing
variance. However, there is no transport map describing a variation from the
Dirac’s delta to the approximating Gaussians. In this work, we therefore adopt
a different approach and consider perturbations induced by _transport plans_.
This more general approach requires us to dive into the generalized notion of
Wasserstein subgradient proposed in [33, §10.3] and, importantly and contrary
to the literature, in the generalized tangent space first introduced in [33,
§12]. While these perturbations entail significant challenges (e.g., transport
plans do not form a vector space), we show in this work that, if judiciously
combined with traditional ideas from variational analysis [44, 45, 46, 47],
they result in general necessary optimality conditions for (1).

Our approach offers several advantages over alternative methods for
optimization in the probability space. First, our analysis is specifically
tailored to the probability space and, in particular, does not require the
introduction of non-negativity and normalization. For instance, if (1) is
unconstrained, the corresponding optimality condition simply predicates that,
at optimality, the Wasserstein gradient vanishes, just like in Euclidean
settings. Second, our optimality conditions hold in full generality, and, in
particular, do not rely on convexity-type assumptions or linearity
assumptions, which in some cases might allow one to study (1) through
infinite-dimensional linear programming or convex analysis. Third, analogously
to the traditional Karush-Kuhn-Tucker conditions, we demonstrate that our
optimality conditions can be used to both solve (1) in closed form, and devise
efficient numerical methods when a closed-form solution is not available.

###  1.3 More details on our contributions

More specifically, our main contribution consists of four key aspects. First,
we provide a solid foundation for a theory of variational analysis in the
Wasserstein space, introducing various concepts such as the generalized
subgradient, normal cone, and tangent cone in the Wasserstein space. To this
extent, we need to resort to very general perturbations induced by transport
plans. While invisible to the end user, these perturbations entail working
with a tangent space which is not a linear space and carefully dealing with
compactness issues. Second, we provide closed-form and easy-to-use expressions
for the subgradients of many functionals and for the normal cone of feasible
sets of practical interests. In particular, we show that the Wasserstein
distance is not regularly subdifferentiable and only admits a generalized
subgradient. This result, which we believe to be of independent interest, also
confirms that a general theory of optimality conditions is required already
for simple functions (in particular, the distance itself), and not only to
cover all corner cases. The key technique to characterize the general
subgradient of the Wasserstein distance involves approximation arguments and
its differentiability at regular measures, which is a promising technique to
explore in future work addressing the differentiability of other functionals
not covered in this work. Third, we derive general first-order optimality
conditions for (1). Inspired by classical variational analysis in Euclidean
spaces, we prove our result by reformulating (1) as an optimization problem
over the epigraph of $\operatorname{\mathcal{J}}$. This way, we can establish
our optimality conditions under very weak assumptions on the functionals – not
even continuity – and constraints. Notably, as we demonstrate with several
pedagogical examples, this complexity is hidden from the end user, and the
deployment of our optimality conditions is effectively analogous to what one
would do in Euclidean spaces. Fourth, we deploy our optimality to study a wide
variety of optimization problems of the form (1) arising in machine learning,
drug discovery, and DRO. Across all these settings, we show that our
optimality conditions both enable novel insights and, when the problem of
interest does not admit a closed-form solution, can be used to design
computational methods. We believe our tools enable the development of novel
algorithms and results in machine learning, robust optimization, and biology,
among others.

##  2 Subgradients and variational geometry

In this section, we translate various elements for variational analysis in the
Wasserstein space. After recalling preliminaries in measure theory and optimal
transport in Section 2.1, we present variations in the Wasserstein space in
Section 2.3. We then introduce Wasserstein subgradients in Section 2.4 and
study the variational geometry of the Wasserstein space in Section 2.5.

###  2.1 Preliminaries

All the maps considered in this work are tacitly assumed to be Borel, i.e.,
measurable w.r.t. the Borel topology. The set of Borel probability measures on
$\mathbb{R}^{d}$ is $\mathcal{P}(\mathbb{R}^{d})$, and we denote the set of
finite second moment probability distributions by
$\mathcal{P}_{2}(\mathbb{R}^{d})$. We write
$\operatorname{\mu}\ll\operatorname{\bar{\mu}}$ to indicate that
$\operatorname{\mu}$ is absolutely continuous w.r.t.
$\operatorname{\bar{\mu}}$, and by
$\mathcal{P}_{2,\mathrm{abs}}(\mathbb{R}^{d})$ the set of absolutely
continuous probability measures with finite second moment. The support
$\operatorname{supp}(\operatorname{\mu})\subseteq\mathbb{R}^{d}$ of a
probability measure $\operatorname{\mu}\in\mathcal{P}(\mathbb{R}^{d})$ is the
closed set $\left\\{x\in\mathbb{R}^{d}\,|\,\operatorname{\mu}(U)>0\text{ for
each neighborhood $U$ of $x$ }\right\\}.$ The identity map on $\mathbb{R}^{d}$
is $\mathrm{Id}_{\mathbb{R}^{d}}(x)=x$ and when clear from the context, we
simply write $\mathrm{Id}$. The gradient of a function
$h:\mathbb{R}^{d}\to\mathbb{R}$ at $x\in\mathbb{R}^{d}$ is $\nabla{h}(x)$, and
the partial derivatives of a function
$c:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ at
$(x,y)\in\mathbb{R}^{d}$ are denoted $\nabla_{x}{c}(x,y)$ and
$\nabla_{y}{c}(x,y)$, respectively. For two expressions $f(x),g(x)$, we write
$f(x)=o\left(g(x)\right)$ if, for $g(x)\to 0$, we have $\frac{f(x)}{g(x)}\to
0$. The _pushforward_ of a probability measure
$\operatorname{\mu}\in\mathcal{P}(\mathbb{R}^{n})$ through a map
$T:\mathbb{R}^{d}\to\mathbb{R}^{m}$ (cf. [48, Definition 1.2.2]), denoted by
${T}_{\\#}{\mu}\in\mathcal{P}(\mathbb{R}^{m})$, is defined by
$({T}_{\\#}{\operatorname{\mu}})(B)=\mu(T^{-1}(B))$ for all Borel sets
$B\subseteq\mathbb{R}^{m}$, and it is a probability measure; see [48, Lemma
1.2.3]. Then, for any ${T}_{\\#}{\operatorname{\mu}}$-integrable
$\phi:\mathbb{R}^{m}\to\mathbb{R}$, it holds
$\int_{\mathbb{R}^{m}}\phi(y)\mathrm{d}({T}_{\\#}{\operatorname{\mu}})(y)=\int_{\mathbb{R}^{n}}\phi(T(x))\mathrm{d}\operatorname{\mu}(x)$;
see [48, Corollary 1.2.6]. Finally (cf. [48, Lemma 1.2.7]), for any
$T:\mathbb{R}^{n}\to\mathbb{R}^{m}$ and $S:\mathbb{R}^{m}\to\mathbb{R}^{t}$
measurable, ${(T\circ
S)}_{\\#}{\operatorname{\mu}}={S}_{\\#}{({T}_{\\#}{\operatorname{\mu}})}$.

###  2.2 Optimal transport

Given $\operatorname{\mu}_{1}\in\mathcal{P}_{2}(\mathbb{R}^{n})$ and
$\operatorname{\mu}_{2}\in\mathcal{P}_{2}(\mathbb{R}^{m})$, their product
measure is $\operatorname{\mu}_{1}\times\operatorname{\mu}_{2}$. We say that
$T:\mathbb{R}^{n}\to\mathbb{R}^{m}$ is a _transport map_ from
$\operatorname{\mu}_{1}$ to $\operatorname{\mu}_{2}$ if
${T}_{\\#}{\mu_{1}}=\mu_{2}$ or, equivalently,
$\int_{\mathbb{R}^{m}}\phi(y)\mathrm{d}\operatorname{\mu}_{2}(y)=\int_{\mathbb{R}^{m}}\phi(y)\mathrm{d}({T}_{\\#}{\operatorname{\mu}_{1}})(y)$
for all $\phi\in C^{0}_{b}(\mathbb{R}^{m})$ [48, Lemma 1.2.5], where
$C^{0}_{b}(\mathbb{R}^{m})$ denotes the space of real-valued bounded
continuous functions on $\mathbb{R}^{m}$. For some $i,m\in\mathbb{N}$, $1\leq
i\leq m$, we denote by $\pi_{i}:(\mathbb{R}^{d})^{m}\to\mathbb{R}^{d}$ the
projection map on the $i^{\mathrm{th}}$ component, i.e.
$\pi_{i}(x_{1},x_{2},\ldots,x_{m})=x_{i}$, which we use to _marginalize_ a
probability measure via the pushforward, ${\pi_{i}}_{\\#}{\mu_{1}}$. A
_transport plan_ between $\operatorname{\mu}_{1}$ and $\operatorname{\mu}_{2}$
is a probability measure
$\operatorname{\gamma}\in\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d})$
so that ${\pi_{1}}_{\\#}{\operatorname{\gamma}}=\operatorname{\mu}_{1}$ and
${\pi_{2}}_{\\#}{\operatorname{\gamma}}=\operatorname{\mu}_{2}$. A transport
map may not exist between $\operatorname{\mu}_{1}$ and
$\operatorname{\mu}_{2}$ (for instance, when
$\operatorname{\mu}_{1}=\frac{1}{N}\sum_{i=1}^{N}\delta_{x_{i}}$ and
$\operatorname{\mu}_{2}=\frac{1}{M}\sum_{j=1}^{M}\delta_{y_{j}}$ with $M>N$),
but a transport plan always does (e.g., the product measure
$\operatorname{\gamma}=\operatorname{\mu}_{1}\times\operatorname{\mu}_{2}$).
We collect them in the set (of _couplings_)
$\Gamma(\operatorname{\mu}_{1},\operatorname{\mu}_{2})$. Given a lower semi-
continuous function
$c:\mathbb{R}^{n}\times\mathbb{R}^{m}\to\bar{\mathbb{R}}\coloneqq\mathbb{R}\cup\\{\pm\infty\\}$,
the optimal transport problem reads:

$W_{c}\left(\operatorname{\mu}_{1},\operatorname{\mu}_{2}\right)\coloneqq\min_{\operatorname{\gamma}\in\Gamma(\operatorname{\mu}_{1},\operatorname{\mu}_{2})}\int_{\mathbb{R}^{n}\times\mathbb{R}^{m}}c(x,y)\mathrm{d}\operatorname{\gamma}(x,y).$
(7)

The celebrated Wasserstein distance is a special case of (7) (cf. [48,
Definition 3.1.3]):

$W_{2}\left(\operatorname{\mu}_{1},\operatorname{\mu}_{2}\right)\coloneqq\left(\min_{\operatorname{\gamma}\in\Gamma(\operatorname{\mu}_{1},\operatorname{\mu}_{2})}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\|x-y\right\|^{2}\mathrm{d}\operatorname{\gamma}(x,y)\right)^{\frac{1}{2}},$
(8)

where $\left\|\cdot\right\|$ is the standard Euclidean norm in
$\mathbb{R}^{d}$. We write
$\Gamma_{c}(\operatorname{\mu}_{1},\operatorname{\mu}_{1})$ and
$\Gamma_{o}(\operatorname{\mu}_{1},\operatorname{\mu}_{1})$ for the set of
minimizers of (7) and (8), respectively. The Wasserstein distance is a
distance on $\mathcal{P}_{2}(\mathbb{R}^{d})$ [49, §6]. We define the
Wasserstein ball of radius $\varepsilon$ as
$\mathbb{B}_{W_{2}}(\varepsilon;\operatorname{\bar{\mu}})\coloneqq\left\\{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})\,|\,W_{2}\left(\operatorname{\mu},\operatorname{\bar{\mu}}\right)<\varepsilon\right\\}.$

Throughout the work, we use two notions of convergence for probability
measures. A sequence
$(\operatorname{\mu}_{n})_{n\in\mathbb{N}}\subseteq\mathcal{P}(\mathbb{R}^{d})$
(i) narrowly converges to $\operatorname{\mu}\in\mathcal{P}(\mathbb{R}^{d})$,
denoted by $\operatorname{\mu}_{n}\rightharpoonup\operatorname{\mu}$, if for
all $\phi\in C^{0}_{b}(\mathbb{R}^{d})$ we have
$\lim_{n\to\infty}\int_{\mathbb{R}^{d}}\phi(x)\mathrm{d}\operatorname{\mu}_{n}(x)\to\int_{\mathbb{R}^{d}}\phi(x)\mathrm{d}\operatorname{\mu}(x)$
(cf. [48, Definition 2.1.5]) and (ii) converges in the Wasserstein topology to
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$,
$\operatorname{\mu}_{n}\to\operatorname{\mu}$, if
$\lim_{n\to\infty}W_{2}\left(\operatorname{\mu},\operatorname{\bar{\mu}}\right)=0$.
The narrow topology is weaker than the Wasserstein topology on
$\mathcal{P}_{2}(\mathbb{R}^{d})$. Indeed, by [33, Proposition 7.1.5],
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$ if and only if
$\operatorname{\mu}_{n}\rightharpoonup\operatorname{\bar{\mu}}$ and
$\int_{\mathbb{R}^{d}}\left\|x\right\|^{2}\mathrm{d}\operatorname{\mu}_{n}(x)\to\int_{\mathbb{R}^{d}}\left\|x\right\|^{2}\mathrm{d}\operatorname{\bar{\mu}}(x)$
or, equivalently,
$\int_{\mathbb{R}^{d}}\phi(x)\mathrm{d}\operatorname{\mu}_{n}(x)\to\int_{\mathbb{R}^{d}}\phi(x)\mathrm{d}\operatorname{\bar{\mu}}(x)$
for all real-valued continuous $\phi$ with $\absolutevalue{\phi(x)}\leq
C(1+\left\|x\right\|^{2})$.

###  2.3 Variations in the Wasserstein space

In the Euclidean space $\mathbb{R}^{d}$, a variation at $x\in\mathbb{R}^{d}$
can be interpreted as an “arrow” $\operatorname{\upsilon}\in\mathbb{R}^{d}$
rooted at $x$. In the same spirit, in the probability space, a variation at
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ is a “(weighted)
collection of arrows” for each point in the support of
$\operatorname{\bar{\mu}}$ [33, Chapter 12]. Formally, a _variation at
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$_ is a probability
measure
$\operatorname{\xi}\in\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d})$
whose first marginal is $\operatorname{\bar{\mu}}$ (i.e.,
${\pi_{1}}_{\\#}{\operatorname{\xi}}=\operatorname{\bar{\mu}}$) (see Figure 2,
left). We can disintegrate [33, Theorem 5.3.1] $\operatorname{\xi}$ to obtain
a collection
$\\{\operatorname{\xi}_{x}\\}_{x\in\mathbb{R}^{d}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$
where each $\operatorname{\xi}_{x}$ denotes the probability measure over the
“tangent” vectors (i.e., the “(weighted) collection of arrows”) at each
$x\in\operatorname{supp}\operatorname{\bar{\mu}}$.

The distance between two variations
$\operatorname{\xi}_{1},\operatorname{\xi}_{2}\in\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d})$
at $\operatorname{\bar{\mu}}$ can be defined in terms of the distance between
their (weighted) arrows. To do so, we have to account for the fact that
different arrows starting from the same point can be coupled in different ways
(cf. Figure 2, middle and right). We express this coupling as a transport plan
$\operatorname{\alpha}\in\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{d})$
so that ${\pi_{12}}_{\\#}{\operatorname{\alpha}}=\operatorname{\xi}_{1}$ and
${\pi_{13}}_{\\#}{\operatorname{\alpha}}=\operatorname{\xi}_{2}$, where
$(\bar{x},\operatorname{\upsilon}_{1},\operatorname{\upsilon}_{2})\in\operatorname{supp}(\operatorname{\alpha})$
if and only if the arrows $\operatorname{\upsilon}_{1}\in\mathbb{R}^{d}$ and
$\operatorname{\upsilon}_{2}\in\mathbb{R}^{d}$ are anchored at $\bar{x}$, and
$\operatorname{\upsilon}_{1}$ and $\operatorname{\upsilon}_{2}$ are coupled.
Note that
${\pi_{1}}_{\\#}{\operatorname{\gamma}}={\pi_{1}}_{\\#}{\operatorname{\xi}_{1}}={\pi_{1}}_{\\#}{\operatorname{\xi}_{2}}=\operatorname{\bar{\mu}}$,
since $\operatorname{\xi}_{1},\operatorname{\xi}_{2}$ are variations at
$\operatorname{\bar{\mu}}$. Then, the distance between
$\operatorname{\xi}_{1}$ and $\operatorname{\xi}_{2}$ with common anchor
$\operatorname{\bar{\mu}}$ amounts to the minimum distance that can be
obtained among all such couplings (cf. Figure 2, middle):

$\displaystyle
W_{\operatorname{\bar{\mu}}}\left(\operatorname{\xi}_{1},\operatorname{\xi}_{2}\right)$
$\displaystyle\coloneqq\left(\min_{\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\|\operatorname{\upsilon}_{1}-\operatorname{\upsilon}_{2}\right\|^{2}\mathrm{d}\operatorname{\alpha}(\bar{x},\operatorname{\upsilon}_{1},\operatorname{\upsilon}_{2})\right)^{\frac{1}{2}}.$
(9)

where $\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})$ is the set
of all couplings, as defined above, between $\operatorname{\xi}_{1}$ and
$\operatorname{\xi}_{2}$:

$\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})\coloneqq\left\\{\operatorname{\alpha}\in\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{d})\,|\,{\pi_{12}}_{\\#}{\operatorname{\alpha}}=\operatorname{\xi}_{1},{\pi_{13}}_{\\#}{\operatorname{\alpha}}=\operatorname{\xi}_{2}\right\\}.$

In view of [33, Proposition 12.4.6],
$W_{\operatorname{\bar{\mu}}}\left(\operatorname{\xi}_{1},\operatorname{\xi}_{2}\right)=0\Rightarrow
W_{2}\left(\operatorname{\xi}_{1},\operatorname{\xi}_{2}\right)=0$. Similarly,
we can also define the inner product and norm:

$\displaystyle\left\langle\operatorname{\xi}_{1},\operatorname{\xi}_{2}\right\rangle_{\operatorname{\bar{\mu}}}$
$\displaystyle\coloneqq\max_{\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\langle\operatorname{\upsilon}_{1},\operatorname{\upsilon}_{2}\right\rangle\mathrm{d}\operatorname{\alpha}(\bar{x},\operatorname{\upsilon}_{1},\operatorname{\upsilon}_{2}),$
(10)
$\displaystyle\left\|\operatorname{\xi}_{i}\right\|_{\operatorname{\bar{\mu}}}$
$\displaystyle\coloneqq\left(\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\|\operatorname{\upsilon}_{i}\right\|^{2}\mathrm{d}\operatorname{\xi}_{i}(\bar{x},\operatorname{\upsilon}_{i})\right)^{\frac{1}{2}},\quad
i\in\\{1,2\\}.$ (11)

With these definitions, the Euclidean-like identity
$W_{\operatorname{\bar{\mu}}}\left(\operatorname{\xi}_{1},\operatorname{\xi}_{2}\right)^{2}=\left\|\operatorname{\xi}_{1}\right\|_{\operatorname{\bar{\mu}}}^{2}-2\left\langle\operatorname{\xi}_{1},\operatorname{\xi}_{2}\right\rangle_{\operatorname{\bar{\mu}}}+\left\|\operatorname{\xi}_{1}\right\|_{\operatorname{\bar{\mu}}}^{2}$,
which we exploit in our proofs, holds. The identity uncovers the motivation
behind the $\min$ and $\max$ in (9) and (10): They are the only two terms
involving couplings and they appear with opposite signs. Intuitively, to
minimize the distance between two variations we need to “maximally align” the
arrows at each particle; cf. Figure 2.

In Euclidean spaces, given a variation $\operatorname{\upsilon}$ at a point
$x$, we can construct a new point by altering $x$ according to $v$: $y=x+v$.
Analogously, given a variation $\operatorname{\xi}$ at a probability measure
$\operatorname{\bar{\mu}}$, we obtain a new probability measure summing to all
$x\in\operatorname{supp}(\operatorname{\bar{\mu}})$ the variations
$\operatorname{\upsilon}$ allocating mass according to the weight of
$(x,\operatorname{\upsilon})\in\operatorname{supp}(\operatorname{\xi})$:
$\operatorname{\mu}={(\pi_{1}+\pi_{2})}_{\\#}{\operatorname{\xi}}$. However,
differently from the Euclidean counterpart where the variation displacing $x$
to $y$ is $\operatorname{\upsilon}=y-x$, in the probability spaces we have
multiple ways of connecting $\operatorname{\bar{\mu}}$ and
$\operatorname{\mu}$, described by the transport plans
$\Gamma(\operatorname{\bar{\mu}},\operatorname{\mu})$. Each of these
$\operatorname{\gamma}\in\Gamma(\operatorname{\bar{\mu}},\operatorname{\mu})$
induces a variation
$\operatorname{\xi}={(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{\operatorname{\gamma}}$.
Another difference is that while $\operatorname{\upsilon}$ describes a
geodesic in the Euclidean space, i.e.,
$\left\|\operatorname{\upsilon}\right\|=\left\|y-x\right\|$ and
$\left\|\varepsilon\operatorname{\upsilon}\right\|=\varepsilon\left\|y-x\right\|$
for $\varepsilon\in\mathbb{R}$ this is not the case for the generic variation
$\operatorname{\xi}$:
$\left\|\operatorname{\xi}\right\|_{\operatorname{\bar{\mu}}}\geq
W_{2}\left(\operatorname{\bar{\mu}},\operatorname{\mu}\right)$ and
$\left\|{(\pi_{1},\varepsilon(\pi_{2}-\pi_{1}))}_{\\#}{\operatorname{\gamma}}\right\|_{\operatorname{\bar{\mu}}}\geq\varepsilon
W_{2}\left(\operatorname{\bar{\mu}},\operatorname{\mu}\right)$. In light of
this, to define a meaningful _tangent space_ , we consider only the variations
that if “scaled enough” (i.e., for sufficiently small $\bar{\varepsilon}$
below) are optimal and, thus, describe geodesics in the Wasserstein space: The
_tangent space_
$\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
of $\mathcal{P}_{2}(\mathbb{R}^{d})$ at
$\operatorname{\bar{\mu}}\in\mathcal{P}(\mathbb{R}^{d})$ is the closure with
respect to $W_{\operatorname{\bar{\mu}}}$ of

$\displaystyle\biggl{\\{}\operatorname{\xi}\in\mathcal{P}_{2}(\mathbb{R}^{d}\times\mathbb{R}^{d})\,|$
$\displaystyle\exists\bar{\varepsilon}>0,\forall\varepsilon\in(-\bar{\varepsilon},\bar{\varepsilon})\;$
(12)
$\displaystyle{\pi_{1}}_{\\#}{\operatorname{\xi}}=\operatorname{\bar{\mu}},$
$\displaystyle{(\pi_{1},\pi_{1}+\varepsilon\pi_{2})}_{\\#}{\operatorname{\xi}}\in\Gamma_{o}(\operatorname{\bar{\mu}},{(\pi_{1}+\varepsilon\pi_{2})}_{\\#}{\operatorname{\xi}})\biggr{\\}}.$

Figure 2: Wasserstein geometry. On the left, we consider example variations
for
${\color[rgb]{0.55859375,0.40234375,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.55859375,0.40234375,0.07421875}\operatorname{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$
(denoted by a brown square),
${\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\xi}_{1}}$
in dotted red and
${\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\xi}_{2}}$
in solid blue. They are scaled version of each other:
${\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\xi}_{1}}=2{\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\xi}_{2}}$.
The other two images show two possible sums (in dotted green) of the
variations
${\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\xi}_{1}}$
(in dotted red) and
${\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\xi}_{2}}$
(in solid blue) at
${\color[rgb]{0.55859375,0.40234375,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.55859375,0.40234375,0.07421875}\operatorname{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$:
Each sum results from a different coupling
$\operatorname{\alpha}\in\Gamma^{1}({\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\xi}_{1}},{\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\xi}_{2}})$
of
${\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\xi}_{1}}$
and
${\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\xi}_{2}}$
(i.e., different coupling of the arrows starting from the same square). The
coupling resulting in the sum on the left is the one maximizing the inner
product (cf. (10)) of the two variations (and, equivalently, the one attaining
the minimum in the distance between the two variations; cf. (9)).

Critically, the tangent space (12) is not a vector space. However, we can
equip it with an “almost” linear structure (see Figure 2 for an intuition):

######  Definition 2.1 (An “almost linear” structure of the tangent space).

Given $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$,
$\operatorname{\xi}_{1},\operatorname{\xi}_{2}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
and $\tau\in\mathbb{R}$, we define:

  1. (i)

_A “local” zero:_
$\mathbf{0}_{\operatorname{\bar{\mu}}}\coloneqq\operatorname{\bar{\mu}}\times\delta_{0}.$

  2. (ii)

_A scalar multiplication:_
$\tau\operatorname{\xi}_{1}\coloneqq{(\pi_{1},\tau\pi_{2})}_{\\#}{\operatorname{\xi}_{1}}.$

  3. (iii)

_An addition:_ For
$\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})$,
$\operatorname{\xi}_{1}\,+_{\operatorname{\alpha}}\,\operatorname{\xi}_{2}\coloneqq{(\pi_{1},\pi_{2}+\pi_{3})}_{\\#}{\operatorname{\alpha}}.$

In the appendix (Propositions A.6 and A.8), we show that Definition
2.1(ii)-(iii) are well-posed (i.e., $\tau\operatorname{\xi}$ and
$\operatorname{\xi}_{1}\,+_{\operatorname{\alpha}}\,\operatorname{\xi}_{2}$
belong to the tangent space) as well as additional properties of this “almost
linear” structure of the tangent space.

######  Remark 2.2.

The sum of $n$ elements $\operatorname{\xi}_{1},\ldots,\operatorname{\xi}_{n}$
results from iteratively applying Definition 2.1(iii), and we write
$\operatorname{\xi}_{1}\,+_{\operatorname{\alpha}}\,\ldots\,+_{\operatorname{\alpha}}\,\operatorname{\xi}_{n}={(\pi_{1},\pi_{2}+\ldots+\pi_{n+1})}_{\\#}{\operatorname{\alpha}}$
for some
$\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi}_{1},\ldots,\operatorname{\xi}_{n})$.

##### Comparison with the literature

Most of the literature uses the tangent space of optimal transport maps,

$\overline{\left\\{\nabla{\varphi}\,|\,\varphi\in
C_{c}^{\infty}(\mathbb{R}^{d})\right\\}}^{L^{2}(\mathbb{R}^{d};\operatorname{\bar{\mu}})}\\\
=\overline{\left\\{\varepsilon(T-\mathrm{Id})\,|\,{(\mathrm{Id},T)}_{\\#}{\operatorname{\bar{\mu}}}\in\Gamma_{o}(\operatorname{\bar{\mu}},{T}_{\\#}{\operatorname{\bar{\mu}}}),\exists\bar{\varepsilon}>0,\forall\varepsilon\in(-\bar{\varepsilon},\bar{\varepsilon})\right\\}}^{L^{2}(\mathbb{R}^{d};\operatorname{\bar{\mu}})},$
(13)

where $C_{c}^{\infty}(\mathbb{R}^{d})$ is the set of cylindrical and
infinitely differentiable functions. This (simplified) tangent space was first
introduced in [33, Chapter 8] (see e.g. [1, 27, 36, 37] and [33, Theorem
8.5.1] for the proof of the equality in (13)) and can be embedded in (12)
since for $\nabla{\varphi}$ with $\varphi\in C_{c}^{\infty}(\mathbb{R}^{d})$,
the tangent vector
$\operatorname{\xi}={(\mathrm{Id}_{,}\nabla{\varphi})}_{\\#}{\operatorname{\bar{\mu}}}$
belongs of
$\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
of $\mathcal{P}_{2}(\mathbb{R}^{d})$. In fact, (i) its first marginal is
$\operatorname{\bar{\mu}}$ and (ii)
${(\pi_{1},\pi_{1}+\varepsilon\pi_{2})}_{\\#}{\operatorname{\xi}}={(\mathrm{Id},\mathrm{Id}+\varepsilon\nabla{\varphi})}_{\\#}{\operatorname{\bar{\mu}}}$
is, for $\varepsilon$ sufficiently small, a transport plan. The latter follows
from $\mathrm{Id}+\varepsilon\nabla{\varphi}$ being the gradient of a convex
function and, thus, an optimal transport map [1, Proposition 2.3]. While
attractive for its simplicity (it is a vector space), the tangent space (13)
limits the perturbations to transport maps. For absolutely continuous
measures, this comes with no loss of generality in virtue of Brenier’s theorem
[50]. However, for empirical probability measures (such as the ones in data-
driven applications, where one has only access to a dataset consisting of a
finite number of observations), the restriction to transport maps effectively
limits the possible perturbations (e.g., a Dirac’s delta can only be
transported to another Dirac’s delta with a transport map). Finally, the
tangent space (12), with $\varepsilon$ restricted to be positive, was first
defined in [33, Chapter 12]. Instead, we drop the non-negativity of
$\varepsilon$, which simplifies various technical results.

###  2.4 Wasserstein subgradients

We now define the _regular_ and _general_ subgradient, along the lines of [44,
Definition 8.3] and [33, §10], for the Wasserstein space. The definition is
inspired from the Euclidean setting, where $\bar{\operatorname{\upsilon}}$ is
a subgradient of $f:\mathbb{R}^{d}\to\bar{\mathbb{R}}$ at $x\in\mathbb{R}^{d}$
if
$f(x)-f(\bar{x})\geq\left\langle\bar{\operatorname{\upsilon}},x-\bar{x}\right\rangle+o\left(\left\|x-\bar{x}\right\|\right)$
for all $x\in\mathbb{R}^{d}$. That is, if $\bar{\operatorname{\upsilon}}$
linearly lower-bounds $f$ at $\bar{x}$. General subgradients are then defined
as the limits of regular subgradients. In the probability space, we can
proceed analogously. In particular, we use the (local) inner product (10) and
replace the displacement “$x-\bar{x}$” with
${(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{\operatorname{\gamma}}$ where
$\operatorname{\gamma}\in\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\mu})$
is an optimal transport plan between $\operatorname{\bar{\mu}}$ and
$\operatorname{\mu}$ (which reduces to
${(\mathrm{Id},T-\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}}$ when
$\operatorname{\gamma}$ is induced by a transport map $T$):

######  Definition 2.3 (Wasserstein subgradients).

Consider a functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
and a probability measure
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ with
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\in\mathbb{R}$. For
$\operatorname{\bar{\operatorname{\xi}}}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$,
we say that

  1. (i)

$\operatorname{\bar{\operatorname{\xi}}}$ is a _regular subgradient of
$\operatorname{\mathcal{J}}$ at $\operatorname{\bar{\mu}}$_, written
$\operatorname{\bar{\operatorname{\xi}}}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$,
if for all $\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ and all
$\operatorname{\xi}\in{(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\mu})}$,
it holds that
$\operatorname{\mathcal{J}}(\operatorname{\mu})-\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\geq\left\langle\operatorname{\bar{\operatorname{\xi}}},\operatorname{\xi}\right\rangle_{\operatorname{\bar{\mu}}}+o\left(W_{2}\left(\operatorname{\bar{\mu}},\operatorname{\mu}\right)\right).$
Then, $\operatorname{\mathcal{J}}$ is regularly subdifferentiable at
$\operatorname{\bar{\mu}}$.

  2. (ii)

$\operatorname{\bar{\operatorname{\xi}}}$ is a _(general) subgradient of
$\operatorname{\mathcal{J}}$ at $\operatorname{\bar{\mu}}$_, written
$\operatorname{\bar{\operatorname{\xi}}}\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
if there are sequences $\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$
with $\operatorname{\mathcal{J}}(\mu_{n})\to\operatorname{\mathcal{J}}(\mu)$
and
$\operatorname{\xi}_{n}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\mu}_{n})$
with $\operatorname{\xi}_{n}\to\operatorname{\bar{\operatorname{\xi}}}$. Then,
$\operatorname{\mathcal{J}}$ is subdifferentiable at
$\operatorname{\bar{\mu}}$.

With this definition, we can then define regular supergradients as the
elements of
$-\hat{\partial}(-\operatorname{\mathcal{J}})(\operatorname{\bar{\mu}})$, and
supergradients as elements of
$-\partial(-\operatorname{\mathcal{J}})(\operatorname{\bar{\mu}})$). We
consider the Wasserstein convergence for the general subgradient since the
first marginal may differ (we consider sequences
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$). For this reason, the
general subgradient may not be in the tangent space. The regular subgradient,
instead, is in the tangent space by definition. Similarly, we can then define
differentiability:

######  Definition 2.4 (Differentiable functional).

A functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
is _differentiable_ at
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ if it admits both
a regular subgradient and supergradient at $\operatorname{\bar{\mu}}$; i.e.,
$-\hat{\partial}(-\operatorname{\mathcal{J}})(\operatorname{\bar{\mu}})\cap\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\neq\emptyset.$

Wasserstein subgradients enjoy similar properties to their Euclidean
counterpart:

######  Proposition 2.5 (Characterization of the (sub)gradients).

Consider a functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
and a probability measure
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$, where
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\in\mathbb{R}$. Then,

  1. (i)

Either $\operatorname{\mathcal{J}}$ is differentiable at
$\operatorname{\bar{\mu}}$ or at least one between
$\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$ and
$\hat{\partial}(-\operatorname{\mathcal{J}})(\operatorname{\bar{\mu}})$ is
empty.

  2. (ii)

If $G:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$ is a lower bound on
$\operatorname{\mathcal{J}}$, i.e.,
$G(\operatorname{\mu})\leq\operatorname{\mathcal{J}}(\operatorname{\mu})$ for
all $\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ and
$G(\operatorname{\bar{\mu}})=\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$,
then
$\hat{\partial}G(\operatorname{\bar{\mu}})\subseteq\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$.

If $\operatorname{\mathcal{J}}$ is differentiable at
$\operatorname{\bar{\mu}}$, the following statements hold:

  1. (iii)

For all
$\operatorname{\xi}_{1},\operatorname{\xi}_{2}\in-\hat{\partial}(-\operatorname{\mathcal{J}})(\operatorname{\bar{\mu}})\cap\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$,
$\operatorname{\xi}_{1}=\operatorname{\xi}_{2}$. Thus, we call this unique
element the gradient of $\operatorname{\mathcal{J}}$ at
$\operatorname{\bar{\mu}}$, written
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})$. In particular,
it satisfies
$\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})=\\{\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})\\}\subseteq\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$.

  2. (iv)

A tangent element
$\operatorname{\xi}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
is the gradient of $\operatorname{\mathcal{J}}$ at $\operatorname{\bar{\mu}}$
if and only if for all $\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})$,
$\operatorname{\xi}^{\prime}\in{(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\mu})}$,
$\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi},\operatorname{\xi}^{\prime})$
it holds:

$\operatorname{\mathcal{J}}(\operatorname{\mu})-\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})=\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\langle
x_{2},x_{3}-x_{1}\right\rangle\mathrm{d}\operatorname{\alpha}(x_{1},x_{2},x_{3})+o\left(W_{2}\left(\operatorname{\bar{\mu}},\operatorname{\mu}\right)\right).$
(14)

In particular, Proposition 2.5(iii) allows us to talk about _the_ gradient.
The (general) subgradient might not be a singleton, even for a differentiable
functional; see [44, §8.B] for an example in $\mathbb{R}^{d}$. Proposition
2.5(iv), instead, is the analogous of
$f(y)-f(x)=\nabla{f}(x)^{\top}(y-x)+o\left(\left\|x-y\right\|\right)$: Namely,
Wasserstein gradients provide a “linear approximation” from below of
$\operatorname{\mathcal{J}}$, with errors that are small in terms of the
Wasserstein distance.

##### Comparison with the literature

Our definition of regular subgradient ((i) in Definition 2.4) coincides with
the definition of extended Fréchet subdifferential, introduced in [33,
Definition 10.3.1] and employed, for instance, in [27]. It is a weaker notion
of subdifferentiability compared to the ones in [51, 1], whereby subgradients
are of the form ${(\mathrm{Id},\phi)}_{\\#}{\operatorname{\mu}}{}$ for some
Borel map $\phi:\mathbb{R}^{d}\to\mathbb{R}^{d}$. From an optimization
perspective, these definitions are however not satisfying since, as we will
discuss shortly, the Wasserstein distance itself fails to be subdifferentiable
at measures that are not absolutely continuous. Inspired by classical
variational analysis in Euclidean spaces [44], we therefore introduce the
definition of general Wasserstein subgradients ((ii) in Definition 2.3). To
the best of our knowledge, this definition is novel.

####  2.4.1 Why do we need such machinery?

We argue that our level of sophistication is required for at least two
reasons. First, already in Euclidean spaces, variational analysis is needed to
deal with non-differentiable and non-convex settings; the (general)
subgradient usually appears in the necessary conditions, while the horizon
subgradient (which we shall discuss in Section 2.5.1) appears in the
constraint qualification (e.g., see [44, Theorem 8.15] or [46, §5.1]). We do
not expect things to simplify in the more general probability space. Second,
the Wasserstein distance itself fails to be regularly subdifferentiable. The
existing theory of gradient flows [33, 1] bypasses this complexity resorting
to absolutely continuous measures. However, this introduces a substantial
practical and theoretical limitation since it excludes any data-driven
setting. We characterize below the general subgradient which we show instead
to be always non-empty.

##### The Wasserstein distance is not regularly subdifferentiable

To start, we study the regular differentiability properties of the Wasserstein
distance:

######  Proposition 2.6 (Wasserstein subgradients and optimal plans).

Consider the functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq\frac{1}{2}W_{2}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)^{2}$.
Then, for any $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$:

  1. (i)

$\operatorname{\mathcal{J}}$ is regularly superdifferentiable at
$\operatorname{\bar{\mu}}$ and all
${(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{\operatorname{\gamma}}$ with
$\operatorname{\gamma}\in\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$
are supergradients of $\operatorname{\mathcal{J}}$ at
$\operatorname{\bar{\mu}}$.

  2. (ii)

If $\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$ is a
singleton and induced by an optimal transport map (e.g., when
$\operatorname{\bar{\mu}}$ is absolutely continuous), then
$\operatorname{\mathcal{J}}$ is regularly sub-differentiable, and thus
differentiable, at $\operatorname{\bar{\mu}}$.

  3. (iii)

If
$\operatorname{\xi}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$,
then
${(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{\operatorname{\xi}}\in\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$.

  4. (iv)

If $\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$ is not a
singleton, then $\operatorname{\mathcal{J}}$ is not regularly
subdifferentiable at $\operatorname{\mu}$.

When the set of optimal transport plans consists of a unique optimal transport
map $T:\mathbb{R}^{d}\to\mathbb{R}^{d}$, then
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})={(\mathrm{Id},\mathrm{Id}-T)}_{\\#}{\operatorname{\bar{\mu}}}$,
which is formally analogous the gradient of
$x\mapsto\frac{1}{2}\left\|x-\hat{x}\right\|^{2}$ being $x-\hat{x}$ in
Euclidean spaces (with “$\mathrm{Id}$ as $x$” and “$T$ as $\hat{x}$”).
However, Proposition 2.6 also indicates, perhaps surprisingly, that the
squared Wasserstein distance from a reference measure
$\operatorname{\hat{\mu}}$ is not subdifferentiable at all measures
$\operatorname{\bar{\mu}}$ for which there exist multiple optimal transport
plans between $\operatorname{\bar{\mu}}$ and $\operatorname{\hat{\mu}}$.
Instead, it is regularly differentiable, and thus differentiable, whenever
there is a unique optimal transport plan between $\operatorname{\bar{\mu}}$
and $\operatorname{\hat{\mu}}$ and this plan is induced by an optimal
transport map, a result first established in [33, Theorem 10.2.6]. By
Brenier’s theorem [52], if $\operatorname{\bar{\mu}}$ is absolutely
continuous, then there is a unique optimal transport plan induced by an
optimal transport map between $\operatorname{\bar{\mu}}$ and
$\operatorname{\hat{\mu}}$, and so the squared Wasserstein distance is
differentiable at absolutely continuous measures. This result is sharp: Even
when a single optimal transport plan exists, the Wasserstein distance may lack
regular subgradients, as we illustrate in the next example.

######  Example 2.7 (Non regular-subdifferentiability of the squared
Wasserstein distance).

Let $\operatorname{\bar{\mu}}\coloneqq\delta_{(0,0)}$ (blue in the picture)
and
$\operatorname{\hat{\mu}}\coloneqq\frac{1}{2}\delta_{(-1,0)}+\frac{1}{2}\delta_{(1,0)}$
(red in the picture). Applying the definition of Wasserstein distance (cf.
(8)),
$\operatorname{\mathcal{J}}({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}})=\frac{1}{2}W_{2}\left({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}},{\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\hat{\mu}}}\right)^{2}=\frac{1}{2}$,
and there is a unique optimal transport plan between
$\operatorname{\bar{\mu}}$ and $\operatorname{\hat{\mu}}$; in particular,
$\Gamma_{o}({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}},{\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\hat{\mu}}})=\\{{\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}}\times{\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\hat{\mu}}}\\}$.
We now show that
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\frac{1}{2}W_{2}\left(\operatorname{\mu},{\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\hat{\mu}}}\right)^{2}$
is not regularly subdifferentiable at
${\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}}$.
By Proposition 2.6(iii), it suffices to prove that
$\operatorname{\xi}={(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}}\times{\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\hat{\mu}}})}=\frac{1}{2}\delta_{(0,-1)}+\frac{1}{2}\delta_{(0,1)}$
is not a regular subgradient of $\operatorname{\mathcal{J}}$.

Consider
${\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}=\frac{1}{2}\delta_{(-\frac{1}{n},-\frac{1}{n})}+\frac{1}{2}\delta_{(\frac{1}{n},\frac{1}{n})}}$.
Then, there is a unique optimal transport plan
$\Gamma_{o}({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}},{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}})=\\{{\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}}\times{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}}\\}$,
and
$\operatorname{\mathcal{J}}({\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}})=\frac{1}{2}W_{2}\left({\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}},{\color[rgb]{0.71875,0.20703125,0.17578125}\definecolor[named]{pgfstrokecolor}{rgb}{0.71875,0.20703125,0.17578125}\operatorname{\hat{\mu}}}\right)^{2}=\frac{1}{2}((1-\frac{1}{n})^{2}+\frac{1}{n^{2}})=\frac{1}{2}-\frac{1}{n}+\frac{1}{n^{2}}$,
and
$W_{2}\left({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}},{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}}\right)=(\frac{1}{n^{2}}+\frac{1}{n^{2}})^{1/2}=\frac{\sqrt{2}}{n}$.
We now show that such a sequence contradicts Definition 2.3(i) by proving
that, along this “direction”, $\operatorname{\xi}$ does not lower bound $J$.
For this, we consider the variations
$\operatorname{\xi}_{n}\in{(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{\Gamma_{o}({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}},{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}})}$.
Using the expressions of
${\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}}$
and
${\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}}$,
we conclude that the only variation between $\operatorname{\bar{\mu}}$ and
$\operatorname{\mu}_{n}$ is
$\operatorname{\xi}_{n}={(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}}\times{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}})}=\frac{1}{2}\delta_{(0,-\frac{1}{n})}+\frac{1}{2}\delta_{(0,\frac{1}{n})}$.

With this, we compute the inner product between $\operatorname{\xi}$ and
$\operatorname{\xi}_{n}$:

$\left\langle\operatorname{\xi},\operatorname{\xi}_{n}\right\rangle_{\operatorname{\bar{\mu}}}=\max_{\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi},\operatorname{\xi}_{n})}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}^{d}}\left\langle\operatorname{\upsilon},\operatorname{\upsilon}_{n}\right\rangle\mathrm{d}\operatorname{\alpha}(x,\operatorname{\upsilon},\operatorname{\upsilon}_{n}).$

Since $\operatorname{\alpha}$ is a coupling between
$\frac{1}{2}\delta_{(0,-1)}+\frac{1}{2}\delta_{(0,1)}$ and
$\frac{1}{2}\delta_{(0,-\frac{1}{n})}+\frac{1}{2}\delta_{(0,\frac{1}{n})}$, it
is fully parametrized by a single coefficient $\tau\in[0,1]$, and the max
reduces to a line search:

$\displaystyle\left\langle\operatorname{\xi},\operatorname{\xi}_{n}\right\rangle_{\operatorname{\bar{\mu}}}$
$\displaystyle=\max_{\tau\in[0,1]}\frac{\tau}{2}\left(\left\langle(1,0),\left(\frac{1}{n},\frac{1}{n}\right)\right\rangle+\left\langle(-1,0),\left(-\frac{1}{n},-\frac{1}{n}\right)\right\rangle\right)$
$\displaystyle\qquad\qquad+\frac{(1-\tau)}{2}\left(\left\langle(1,0),\left(-\frac{1}{n},-\frac{1}{n}\right)\right\rangle+\left\langle(-1,0),\left(\frac{1}{n},\frac{1}{n}\right)\right\rangle\right)=\frac{1}{n}.$

Thus,

$\displaystyle\liminf_{n\to\infty}\frac{\operatorname{\mathcal{J}}({\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}})-\operatorname{\mathcal{J}}({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}})-\left\langle\operatorname{\xi},\operatorname{\xi}_{n}\right\rangle_{\operatorname{\bar{\mu}}}}{W_{2}\left(\operatorname{\bar{\mu}},{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}}\right)}=\liminf_{n\to\infty}\frac{-\frac{2}{n}+\frac{1}{n^{2}}}{\frac{\sqrt{2}}{n}}=-\sqrt{2}<0$

and so
$\operatorname{\mathcal{J}}({\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}})-\operatorname{\mathcal{J}}({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}})<\left\langle\operatorname{\xi},\operatorname{\xi}_{n}\right\rangle_{\operatorname{\bar{\mu}}}+o\left(W_{2}\left({\color[rgb]{0.12890625,0.359375,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.12890625,0.359375,0.6875}\operatorname{\bar{\mu}}},{\color[rgb]{0.3828125,0.44921875,0.07421875}\definecolor[named]{pgfstrokecolor}{rgb}{0.3828125,0.44921875,0.07421875}\operatorname{\mu}_{n}}\right)\right)$
That is, $\operatorname{\mathcal{J}}$ is not regularly subdifferentiable at
$\operatorname{\bar{\mu}}$ and, in particular, uniqueness of the optimal plan
does not imply regular subdifferentiability.

The absence of regular subgradients prompts us to study the general
subgradient of the Wasserstein distance. We will do so in the next section.

##### The (general) subgradient of the Wasserstein distance

The subdifferentiability of the Wasserstein distance at absolutely continuous
probability measures (cf. Proposition 2.6(ii)) fuels the hope that the general
subgradient can be characterized by approximating each probability measure via
absolutely continuous ones. We illustrate the intuition by approximating a
Dirac’s delta at 0 with Gaussian probability measures with vanishing variance:

######  Example 2.8 (Subgradients via Gaussians).

Let
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\frac{1}{2}W_{2}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)^{2}$
and fix $\operatorname{\bar{\mu}}=\delta_{0}$ and
$\operatorname{\hat{\mu}}=\mathcal{N}\left(0,I_{d}\right)$ being the Gaussian
centered in $0$ with covariance matrix $I_{d}$, the identity matrix in
$\mathbb{R}^{d\times d}$. In this case, there is a unique plan
$\operatorname{\bar{\mu}}\times\operatorname{\hat{\mu}}$ from
$\operatorname{\bar{\mu}}$ to $\operatorname{\hat{\mu}}$, but it is not
induced by a transport map (indeed, the mass at 0 needs to be split, see the
picture below for an intuition). Here, we show that
$\operatorname{\xi}={(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{(\operatorname{\bar{\mu}}\times\operatorname{\hat{\mu}})}$
belongs to the general subgradient
$\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$. To do so, we
approximate $\delta_{0}$ with Gaussians with vanishing variance. Specifically,
consider the sequence
$\operatorname{\mu}_{n}=\mathcal{N}\left(0,\frac{1}{n^{2}}I_{d}\right)$ so
that
$W_{2}\left(\operatorname{\bar{\mu}},\operatorname{\mu}_{n}\right)=\frac{1}{n}\searrow
0$ and $\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$ and
$\operatorname{\mathcal{J}}(\operatorname{\mu}_{n})\to\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$.
Since both $\operatorname{\bar{\mu}}$ and $\operatorname{\mu}_{n}$ are
Gaussians, the optimal transport plans is a singleton
$\Gamma_{o}(\operatorname{\mu}_{n},\operatorname{\hat{\mu}})=\\{{(\mathrm{Id},T_{n})}_{\\#}{\operatorname{\mu}_{n}}\\}$,
where the optimal transport plan is induced by the optimal transport map
$T_{n}(x)=nx$. By Proposition 2.6, $\operatorname{\mathcal{J}}$ is
differentiable at $\operatorname{\mu}$ and, in particular,
$\operatorname{\xi}_{n}\coloneqq{(\mathrm{Id},\mathrm{Id}-T_{n})}_{\\#}{\operatorname{\mu}_{n}}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\mu}_{n}).$
We now show that $\operatorname{\xi}_{n}\to\operatorname{\xi}$ and so
$\operatorname{\xi}$ is a (general) subgradient of
$\operatorname{\mathcal{J}}$. Convergence (in the Wasserstein topology) of the
first marginal follows directly from
${(\pi_{1})}_{\\#}{\operatorname{\xi}_{n}}=\operatorname{\mu}_{n}$,
${(\pi_{1})}_{\\#}{\operatorname{\xi}}=\operatorname{\bar{\mu}}$, and
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$. For narrow convergence of
the second marginal, consider $\phi\in C^{0}_{b}(\mathbb{R}^{d})$. Then,

$\displaystyle\lim_{n\to\infty}\int_{\mathbb{R}^{d}}\phi(y)\mathrm{d}\operatorname{\xi}_{n}(x,y)$
$\displaystyle=\lim_{n\to\infty}\int_{\mathbb{R}^{d}}\phi(x-T_{n}(x))\mathrm{d}\operatorname{\mu}_{n}(x)$
$\displaystyle=\lim_{n\to\infty}\int_{\mathbb{R}^{d}}\phi(x-T_{n}(x))\mathrm{d}({(T_{n}^{-1})}_{\\#}{\operatorname{\hat{\mu}}})(x)$
$\displaystyle=\int_{\mathbb{R}^{d}}\lim_{n\to\infty}\phi\left(\frac{1}{n}x-x\right)\mathrm{d}\operatorname{\hat{\mu}}(x)$
$\displaystyle=\int_{\mathbb{R}^{d}}\phi(0-x)\mathrm{d}\operatorname{\hat{\mu}}(x)$
$\displaystyle=\int_{\mathbb{R}^{d}}\phi(y)\mathrm{d}\operatorname{\xi}(x,y),$

where we used dominated convergence to exchange limit and integral and
continuity of $\phi$. Thus,
${(\pi_{2})}_{\\#}{\operatorname{\xi}_{n}}\rightharpoonup{(\pi_{2})}_{\\#}{\operatorname{\xi}}$,
and so
$\operatorname{\xi}\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$,
and the Wasserstein distance is sub-differentiable at
$\operatorname{\bar{\mu}}$.

Since $\mathcal{P}_{2,\mathrm{abs}}(\mathbb{R}^{d})$ is dense in
$\mathcal{P}_{2}(\mathbb{R}^{d})$ [40, Theorem 2.2.7], we can always construct
a sequence
$(\operatorname{\mu}_{n})_{n\in\mathbb{N}}\subseteq\mathcal{P}_{2,\mathrm{abs}}(\mathbb{R}^{d})$
with $\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$, compute
$\operatorname{\xi}_{n}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\mu}_{n})$,
which exists by absolute continuity of $\operatorname{\mu}_{n}$, and study the
limit of $\operatorname{\xi}_{n}$. One way to construct such
$\operatorname{\mu}_{n}$ is via convolution with a Gaussian kernel; see [33,
Lemma 7.1.10]. Overall, this procedure yields a sufficient characterization
(for the sake of necessary optimality conditions) of the (general) subgradient
of the Wasserstein distance:

######  Proposition 2.9 (General subgradient of the squared Wasserstein
distance).

For
$\operatorname{\bar{\mu}},\operatorname{\hat{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$,
$\operatorname{\mathcal{J}}:\mathbb{R}^{d}\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\frac{1}{2}W_{2}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)^{2}$,
we have
$\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\neq\emptyset$
and
$\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\subseteq{(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})}$.

In particular, Proposition 2.9 establishes subdifferentiability of the
Wasserstein distance at _all_ probability measures
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$.

####  2.4.2 Examples of subgradients

Next, we characterize the subgradients of several functionals of interest. We
start with general optimal transport discrepancies, defined in (7), which
arise, e.g., in uncertainty propagation in dynamical systems [53, 24] and DRO
[10]:

######  Proposition 2.10 (Optimal transport discrepancy).

Let $\operatorname{\hat{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ and let
$c:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}_{\geq 0}$ be twice
continuously differentiable with bounded Hessian. Suppose that $c$ satisfies
the twist condition (i.e., $y\mapsto\nabla_{x}{c}(x,y)$ is injective for all
$x\in\mathbb{R}^{d}$). Consider the functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq
W_{c}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)$ as defined in
(7). Then,

  1. (i)

$\operatorname{\mathcal{J}}$ is proper, non-negative, and lower semi-
continuous.

  2. (ii)

At any $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$,
$\operatorname{\mathcal{J}}$ is subdifferentiable with
$\emptyset\neq\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\subseteq{(\pi_{1},\nabla_{x}{c})}_{\\#}{\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})}$,
where $\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$ is the
set of transport plans between $\operatorname{\bar{\mu}}$ and
$\operatorname{\hat{\mu}}$ that yields the minimum for
$W_{c}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)$. Moreover, if
$\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$ is a singleton
and induced by an optimal transport map (e.g., $\operatorname{\bar{\mu}}$ is
absolutely continuous), then $\operatorname{\mathcal{J}}$ is differentiable
with
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})={(\pi_{1},\nabla_{x}{c})}_{\\#}{\operatorname{\xi}},$
where $\operatorname{\xi}$ is the unique element in
$\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$.

As already observed for the Wasserstein distance, optimal transport
discrepancies are subdifferentiable (but not regularly subdifferentiable), and
their subgradient is characterized in terms of the gradient of the
transportation cost and the set of optimal transport plans. As a notable
special case, consider $\operatorname{\bar{\mu}}=\delta_{\bar{x}}$ and
$\operatorname{\hat{\mu}}=\delta_{\hat{x}}$. Then,
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})=W_{c}\left(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}}\right)=c(x,\hat{x})$
and
$\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}}))=\\{\delta_{(\bar{x},\hat{x})}\\}$.
Thus, $\operatorname{\mathcal{J}}$ is Wasserstein differentiable with gradient
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})={(\pi_{1},\nabla_{x}{c})}_{\\#}{\delta_{(\bar{x},\hat{x})}}=\delta_{\bar{x}}\times\delta_{\nabla_{x}{c}(\bar{x},\hat{x})}$,
which is formally analogous to the standard Euclidean gradient of $x\mapsto
c(x,\hat{x})$. We now exemplify our results in the case of strictly convex
transportation costs:

######  Example 2.11 (Strictly convex transportation cost).

Consider $c(x,y)=h(x-y)$ for $h:\mathbb{R}^{d}\to\mathbb{R}_{\geq 0}$ strictly
convex and twice continuously differentiable with bounded Hessian. By strict
convexity of $h$, $y\mapsto\nabla_{x}{c}(x,y)=\nabla{h}(x-y)$ is injective.
Thus, $c$ satisfies the assumptions of Proposition 2.10 and the subgradient of
$\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq
W_{c}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)$ at
$\operatorname{\bar{\mu}}$ is
$\emptyset\neq\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\subseteq{(\pi_{1},\nabla{h}\circ(\pi_{1}-\pi_{2}))}_{\\#}{\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})}$.
Moreover, $\operatorname{\mathcal{J}}$ is differentiable whenever
$\Gamma_{c}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$ is a singleton
and induced by an optimal transport map and, in particular, at all absolutely
continuous measures.

Next, we recall and extend, by characterizing their general subgradients, the
existing subdifferentiability results for expected values:

######  Proposition 2.12 (Expected value).

Let $V:\mathbb{R}^{d}\to(-\infty,+\infty]$ be a proper, $\lambda$-convex,
lower semicontinuous functional whose negative part has $2$-growth, i.e.
$V(x)\geq-A-B\left\|x\right\|^{2}$ for all $x\in\mathbb{R}^{d}$ and for some
$A,B\in\mathbb{R}$. Consider the functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq\mathbb{E}^{x\sim\operatorname{\mu}}\left[V(x)\right].$
Then,

  1. (i)

$\operatorname{\mathcal{J}}$ is proper and lower semi-continuous.

  2. (ii)

At every $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ at which
$\operatorname{\mathcal{J}}$ is finite,
$\operatorname{\xi}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
is in $\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
(resp., $\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$) if and
only if for all $(x,y)\in\operatorname{supp}{(\operatorname{\xi})}$,
$y\in\hat{\partial}V(x)$ (resp., $\partial V(x)$).

  3. (iii)

Since the Euclidean regular and general subgradients of $V$ coincides (cf.
Proposition A.1), $\hat{\partial}V(x)=\partial V(x)$, we have that
$\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})=\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$.

  4. (iv)

If $V$ is differentiable and $V(x)\leq A+B\left\|x\right\|^{2}$,
$\operatorname{\mathcal{J}}$ is differentiable at any
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ and
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})={(\mathrm{Id},\nabla{V})}_{\\#}{\operatorname{\bar{\mu}}}.$

Any twice continuously differentiable function $V$ with bounded Hessian is
$\lambda$-convex and satisfies $\absolutevalue{V(x)}\leq
A+B\left\|x\right\|^{2}$ for some $A,B\in\mathbb{R}$. Thus, the corresponding
functional $\operatorname{\mathcal{J}}$ is differentiable and, in particular,
$(x,y)\in\operatorname{supp}(\nabla{\operatorname{\mathcal{J}}})$ if and only
if $y=\nabla{V}(x)$. If, instead, $V$ does not have negative quadratic growth,
then $\operatorname{\mathcal{J}}$ is not proper and might evaluate to
$-\infty$ (e.g., $V(x)=-e^{x}$ and
$\operatorname{\mu}=\mathcal{N}\left(0,1\right)$). Similarly, if $V$ grows
more than quadratically, then $\operatorname{\mathcal{J}}$ will attain
$+\infty$ (e.g., $V(x)=e^{x}$ and
$\operatorname{\mu}=\mathcal{N}\left(0,1\right)$) and cannot be differentiable
there.

######  Example 2.13 (Expected values of linear functions).

Consider the same functional $\operatorname{\mathcal{J}}$ as in Section 1,
namely
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\langle\theta,x\right\rangle\right]$
and $V(x)=\left\langle\theta,x\right\rangle$ in Proposition 2.12. Since $V$ is
a linear function, it satisfies all the assumptions of Proposition 2.12, and
in particular $\operatorname{\mathcal{J}}$ is Wasserstein differentiable with
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu})={(\mathrm{Id},\theta)}_{\\#}{\operatorname{\mu}}$.

Finally, we extend the existing subdifferentiability results for interaction-
type functionals which model the energy associated with the interaction
between particles in a distribution; e.g., see [54] for an introduction, [55]
for an example in physics, and [56, §3.3] for one in economics.

######  Proposition 2.14 (Interaction energy).

Let $U:\mathbb{R}^{d}\to\mathbb{R}$ be twice continuously differentiable with
bounded Hessian and consider the functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq\frac{1}{2}\mathbb{E}^{(x,y)\sim\operatorname{\mu}\times\operatorname{\mu}}\left[U(x-y)\right].$
Then,

  1. (i)

$\operatorname{\mathcal{J}}$ is proper and continuous.

  2. (ii)

$\operatorname{\mathcal{J}}$ is Wasserstein differentiable and its Wasserstein
gradient at $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ is
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})={(\mathrm{Id},(\nabla{U}\ast\operatorname{\bar{\mu}}))}_{\\#}{\operatorname{\bar{\mu}}},$
where $\ast$ denotes the convolution operation,

$(\nabla{U}\ast\operatorname{\bar{\mu}})(x)\coloneqq\int_{\mathbb{R}^{d}}U(x-y)\mathrm{d}\operatorname{\bar{\mu}}(y).$

We exemplify the result in the case of quadratic functions:

######  Example 2.15 (Quadratic interaction energies).

Consider $U(z)=\frac{1}{2}\left\langle Qz,z\right\rangle$ for
$Q\in\mathbb{R}^{d\times d}$. Then, $U$ is twice continuosly differentiable
with $\nabla{U}(z)=Qz$ and constant (and, thus, bounded) Hessian $Q$. In
virtue of Proposition 2.14, the functional
$\operatorname{\mathcal{J}}(\operatorname{\mu})=\frac{1}{2}\mathbb{E}^{(x,y)\sim\operatorname{\mu}\times\operatorname{\mu}}\left[U(x-y)\right]$
is Wasserstein differentiable with
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu})={(\mathrm{Id},Q\mathbb{E}^{x\sim\operatorname{\mu}}\left[x\right])}_{\\#}{\operatorname{\mu}}.$

####  2.4.3 Subdifferential calculus

We now show some useful calculus rules in the Wasserstein space. More can be
studied, and we expect this to be a source of interesting future work.

######  Proposition 2.16 (Subdifferential calculus).

Consider the proper functionals
$\operatorname{\mathcal{J}}_{1},\operatorname{\mathcal{J}}_{2}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
a monotone map $g:\bar{\mathbb{R}}\to\bar{\mathbb{R}}$ with
$g(\infty)=\infty$, and
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ so that
$\operatorname{\mathcal{J}}_{i}(\operatorname{\bar{\mu}})\in\mathbb{R}$, $g$
is continuously differentiable at
$\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})\in\mathbb{R}$ with
$g(\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}}))\in\mathbb{R}$,
and $\operatorname{\mathcal{J}}_{1}$ continuous in the Wasserstein topology at
$\operatorname{\bar{\mu}}$. Then, the following rules hold:

  1. (i)

Sum rule:

$\displaystyle\hat{\partial}(\operatorname{\mathcal{J}}_{1}+\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})$
$\displaystyle\supseteq\hat{\partial}(\operatorname{\mathcal{J}}_{1})(\operatorname{\bar{\mu}})+\hat{\partial}(\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})$
$\displaystyle\coloneqq\left\\{\operatorname{\xi}_{1}\,+_{\operatorname{\alpha}}\,\operatorname{\xi}_{2}\,|\,\operatorname{\xi}_{i}\in\hat{\partial}\operatorname{\mathcal{J}}_{i}(\operatorname{\bar{\mu}}),\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})\right\\}.$

In particular, if $\operatorname{\mathcal{J}}_{2}$ is differentiable at
$\operatorname{\bar{\mu}}$, then

$\displaystyle\hat{\partial}(\operatorname{\mathcal{J}}_{1}+\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})$
$\displaystyle=\hat{\partial}(\operatorname{\mathcal{J}}_{1})(\operatorname{\bar{\mu}})+\\{\nabla{\operatorname{\mathcal{J}}}_{2}(\operatorname{\bar{\mu}})\\}$

and, thus,

$\displaystyle\partial(\operatorname{\mathcal{J}}_{1}+\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})$
$\displaystyle=\partial\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})+\left\\{\nabla{\operatorname{\mathcal{J}}_{2}}(\operatorname{\bar{\mu}})\right\\}.$

  2. (ii)

Chain rule: If
$g^{\prime}(\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}}))>0$, then

$\displaystyle\hat{\partial}(g\circ\operatorname{\mathcal{J}}_{1})(\operatorname{\bar{\mu}})$
$\displaystyle=g^{\prime}(\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})){\hat{\partial}\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})}$
$\displaystyle\coloneqq\\{g^{\prime}(\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}}))\operatorname{\xi}\,|\,\operatorname{\xi}\in\hat{\partial}\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})\\}\quad\text{and}$
$\displaystyle\partial(g\circ\operatorname{\mathcal{J}}_{1})(\operatorname{\bar{\mu}})$
$\displaystyle=g^{\prime}(\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})){\partial\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})}.$

With $g_{i}(x)=\rho_{i}x$ for $\rho_{1},\rho_{2}\in\mathbb{R}_{\geq 0}$,
Proposition 2.16 readily characterizes the subgradients of the linear
combination of functionals:

######  Corollary 2.17 (Linear combination rule).

Consider proper functionals
$\operatorname{\mathcal{J}}_{1},\operatorname{\mathcal{J}}_{2}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\rho_{1},\rho_{2}\in\mathbb{R}_{\geq 0}$, and
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$. Then,

$\hat{\partial}(\rho_{1}\operatorname{\mathcal{J}}_{1}+\rho_{2}\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})\supseteq\rho_{1}\hat{\partial}\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})+\rho_{2}\hat{\partial}\operatorname{\mathcal{J}}_{2}(\operatorname{\bar{\mu}}).$

If additionally $\operatorname{\mathcal{J}}_{2}$ is differentiable,

$\displaystyle\hat{\partial}(\rho_{1}\operatorname{\mathcal{J}}_{1}+\rho_{2}\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})$
$\displaystyle=\rho_{1}\hat{\partial}\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})+\rho_{2}\nabla{\operatorname{\mathcal{J}}_{2}}(\operatorname{\bar{\mu}})\quad\text{and}$
$\displaystyle\partial(\rho_{1}\operatorname{\mathcal{J}}_{1}+\rho_{2}\operatorname{\mathcal{J}}_{2})(\operatorname{\bar{\mu}})$
$\displaystyle=\rho_{1}\partial\operatorname{\mathcal{J}}_{1}(\operatorname{\bar{\mu}})+\rho_{2}\nabla{\operatorname{\mathcal{J}}_{2}}(\operatorname{\bar{\mu}}).$

Next, we deploy Proposition 2.16 to study the variance via Proposition 2.12:

######  Corollary 2.18 (Variance).

Consider the functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
$\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq\mathrm{Var}^{x\sim\operatorname{\mu}}\left[V(x)\right].$
If both $V$ and $x\mapsto V(x)^{2}$ satisfy the assumptions of Proposition
2.12, at every $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ at
which $\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\in\mathbb{R}$,

$\displaystyle\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
$\displaystyle=\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})={(\pi_{1},2(\pi_{1}-\mathbb{E}^{x\sim\operatorname{\bar{\mu}}}\left[V(x)\right])\pi_{2})}_{\\#}{\biggl{\\{}\operatorname{\xi}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})\,|\,\operatorname{supp}{(\operatorname{\xi})}\subseteq\hat{\partial}V\biggr{\\}}}$

When $V$ is continuously differentiable and $V(x)^{2},V(x)^{4}\leq
A+B\left\|x\right\|^{2}$ for some $A,B>0$, $\operatorname{\mathcal{J}}$ is
differentiable and its gradient is
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}})={(\mathrm{Id},2(V-\mathbb{E}^{x\sim\operatorname{\mu}}\left[V(x)\right])\nabla{V})}_{\\#}{\operatorname{\bar{\mu}}}.$

###  2.5 Variational geometry

In this section, we study the geometry of the constraint set
$\operatorname{\mathcal{C}}$, which we will combine in Section 3 with the
subdifferentability of the objective function to derive our necessary
optimality conditions. Similarly to Euclidean settings, all the necessary
information is encoded in tangent and normal cones (see Figure 3 for an
intuition in $\mathbb{R}^{d}$). We start with the tangent cone:

######  Definition 2.19 (Tangent cone).

The _tangent cone_ of
$\operatorname{\mathcal{C}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$ at
$\operatorname{\bar{\mu}}\in\operatorname{\mathcal{C}}$, denoted by
$\operatorname{T}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$, is
characterized by all the
$\operatorname{\bar{\operatorname{\xi}}}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
such that for some
$(\operatorname{\mu}_{n})_{n\in\mathbb{N}}\subseteq\operatorname{\mathcal{C}}$,
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$,
$(\operatorname{\xi}_{n})_{n\in\mathbb{N}}\subseteq\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$,
$\operatorname{\xi}_{n}\in{(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\mu}_{n})}$
and $\tau_{n}\searrow 0$, one has
$W_{\operatorname{\bar{\mu}}}\left(\frac{1}{\tau_{n}}\operatorname{\xi}_{n},\operatorname{\bar{\operatorname{\xi}}}\right)\to
0$.

We highlight the formal similarity with the tangent cone in Euclidean spaces,
defined by all the $\bar{\operatorname{\upsilon}}$ for which there is
$(x_{n})_{n\in\mathbb{N}}\subset\operatorname{\mathcal{C}}$,
$x_{n}\to\bar{x}$, and $\tau_{n}\searrow 0$ so that
$\frac{1}{\tau_{n}}(x_{n}-\bar{x})\to\bar{\operatorname{\upsilon}}$. In
particular, if $\operatorname{\mu}_{n}=\delta_{x_{n}}$ and
$\operatorname{\bar{\mu}}=\delta_{\bar{x}}$, then
$\Gamma_{o}(\operatorname{\mu}_{n},\operatorname{\bar{\mu}})=\\{\delta_{(x_{n},\bar{x})}\\}$
and $\xi_{n}=(\bar{x},x_{n}-\bar{x})$. Namely, the definitions in the
probability space and in the Euclidean space agree. Intuitively, the tangent
cone is a restriction of the tangent space, and it describes the variations
admissible within the set $\operatorname{\mathcal{C}}$. Dual to tangent cones
are the normal cones:

######  Definition 2.20 (Normal cone).

The _regular normal cone_ of
$\operatorname{\mathcal{C}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$ at
$\operatorname{\bar{\mu}}\in\operatorname{\mathcal{C}}$, denoted by
$\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$,
contains all the
$\operatorname{\bar{\operatorname{\xi}}}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
such that for all $\operatorname{\mu}\in\operatorname{\mathcal{C}}$, one has
$\left\langle\operatorname{\bar{\operatorname{\xi}}},\operatorname{\xi}\right\rangle_{\operatorname{\bar{\mu}}}\leq
o\left(\left\|\operatorname{\xi}\right\|_{\operatorname{\bar{\mu}}}\right)$
for
$\operatorname{\xi}={(\pi_{1},\pi_{2}-\pi_{1})}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\nu})}$.
The _normal cone_ of $\operatorname{\mathcal{C}}$ at
$\operatorname{\bar{\mu}}$, denoted by
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$, is
defined as the limit of regular normals; i.e.,
$\operatorname{\bar{\operatorname{\xi}}}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})$
belongs to
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$ if
there exists
$(\operatorname{\mu}_{n})_{n\in\mathbb{N}}\subset\operatorname{\mathcal{C}}$,
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$,
$\operatorname{\xi}_{n}\in\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\mu}_{n})$
so that $\operatorname{\xi}_{n}\to\operatorname{\bar{\operatorname{\xi}}}$.

Figure 3: Variational Geometry in $\mathbb{R}^{d}$. On the left, the arrows
depict some examples of regular normals (e.g., at the cusp). The interior
points have a trivial normal cone, indicated by a single blue dot. At the same
location $x\in\operatorname{\mathcal{C}}$ they may have different scales and
directions. There also may not be any regular normal. In the middle, the
solid, blue, arrows depict some examples of normals. They are the limit points
of sequences of regular normals. The sequence is colored with increasing
opacity as the regular normals approach the general normal. On the right, in
dotted red there are sequences $\tau_{n}(x_{n}-x)\to
w\in\operatorname{T}_{\mathbb{R}^{2}}(x)$, with the tangent vector $w$ being
depicted in solid red. The inner product between tangent vectors and regular
normals is non-positive, whereas the same might not hold with general normals.

Similarly to the subgradient definition Definition 2.3, for general normals we
consider the Wasserstein convergence of regular normals. Thus, general normals
may not be in the tangent space, and need not be. Analogously to the tangent
cone, if we specify the definition of normal cone to Dirac’s deltas, we
recover the notion of normal cone in Euclidean spaces; see also Proposition
2.22.

Another similarity to Euclidean spaces (cf. Figure 3) is the dual relationship
between normal and tangent cones:

######  Proposition 2.21 (Normal and tangent cones relationships).

At any
$\operatorname{\bar{\mu}}\in\operatorname{\mathcal{C}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$,
the sets
$\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$
and $\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$
are closed cones. Additionally, if
$\operatorname{\bar{\operatorname{\xi}}}\in\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$
then
$\left\langle\operatorname{\bar{\operatorname{\xi}}},\operatorname{\xi}\right\rangle_{\operatorname{\bar{\mu}}}\leq
0$ for all
$\operatorname{\xi}\in\operatorname{T}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}}).$

As a first example, we study the normal cones of trivial constraint sets:

######  Proposition 2.22 (Trivial normal cones).

Let $\operatorname{\mathcal{C}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$ and
$\operatorname{\bar{\mu}}\in\operatorname{\mathcal{C}}$. Then,

  1. (i)

if $\operatorname{\mathcal{C}}=\mathcal{P}_{2}(\mathbb{R}^{d})$, then
$\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\operatorname{N}_{\mathcal{P}_{2}(\mathbb{R}^{d})}(\operatorname{\bar{\mu}})=\\{\mathbf{0}_{\operatorname{\bar{\mu}}}\\}$;

  2. (ii)

if
$\operatorname{\bar{\mu}}\in\mathrm{int}_{W_{2}}{\operatorname{\mathcal{C}}}$,
then
$\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\\{\mathbf{0}_{\operatorname{\bar{\mu}}}\\}$;

  3. (iii)

for $\Xi\subseteq\mathbb{R}^{d}$, if
$\operatorname{\mathcal{C}}\subseteq\\{\delta_{x}\in\mathcal{P}_{2}(\mathbb{R}^{d})\,|\,x\in\Xi\\}$,
then at any
$\operatorname{\bar{\mu}}=\delta_{\bar{x}}\in\operatorname{\mathcal{C}}$

$\displaystyle\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$
$\displaystyle=\\{\delta_{\bar{x}}\times\nu\,|\,\mathbb{E}^{x\sim\nu}\left[x\right]\in\hat{\operatorname{N}}_{\Xi}(\bar{x})\\}\quad\text{and}$
$\displaystyle\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})$
$\displaystyle=\\{\delta_{\bar{x}}\times\nu\,|\,\mathbb{E}^{x\sim\nu}\left[x\right]\in\operatorname{N}_{\Xi}(\bar{x})\\},$

where $\hat{\operatorname{N}}_{\Xi}(\bar{x})\\}$ and
$\operatorname{N}_{\Xi}(\bar{x})\\}$ are the regular and general normal cones
in the Euclidean space, respectively.

In particular, analogously to the Euclidean setting, the normal cone at
$\operatorname{\bar{\mu}}$ trivializes if the feasible set covers the whole
space (Proposition 2.22(i)) or if $\operatorname{\bar{\mu}}$ lies in its
interior (Proposition 2.22(ii)); cf. Figure 3. Moreover, when we restrict
ourselves to the space of Dirac’s deltas, the normal cone is consistent with
its Euclidean counterpart (Proposition 2.22(iii)); namely, a “tangent vector”
$\nu$ lies in the normal cone if and only if its “average” tangent vector lies
in the Euclidean normal cone.

Next, we characterize the normal cone of sublevel sets of functionals:

######  Proposition 2.23 (Level sets).

Suppose
$\operatorname{\mathcal{C}}\coloneqq\left\\{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})\,|\,\operatorname{\mathcal{J}}(\operatorname{\mu})\leq
0\right\\}$ for a proper, continuous functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$.
Then, at any $\operatorname{\bar{\mu}}\in\operatorname{\mathcal{C}}$ so that
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})=0$, we have

$\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})\supseteq\left\\{\lambda\operatorname{\xi}\,|\,\lambda\geq
0,\operatorname{\xi}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\right\\}$
(15)

and, if
$\mathbf{0}_{\operatorname{\bar{\mu}}}\not\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
and
$\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\neq\emptyset$,
then

$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})\subseteq\left\\{\lambda\operatorname{\xi}\,|\,\lambda\geq
0,\operatorname{\xi}\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\right\\}\cup\partial^{\infty}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}}).$
(16)

If $\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})<0$, then
$\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\\{\mathbf{0}_{\operatorname{\bar{\mu}}}\\}$.

Proposition 2.23 intimately relates to Lagrange multipliers in the Wasserstein
space: The normal cone of a sublevel set of a functional
$\operatorname{\mathcal{J}}{}$ consists of any scaling (by a multiplier) of
the subgradients of $\operatorname{\mathcal{J}}{}$. In particular, it includes
the case where the constraint is not active (i.e., the multiplier is zero and
the normal cone trivializes) and the case where the constraint is active
(i.e., the multiplier is non-zero). The condition
$\mathbf{0}_{\operatorname{\bar{\mu}}}\not\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
resembles the same “constraint qualification” required in Euclidean settings
(cf. [44, Proposition 10.3]), and in practice it is often not restrictive. As
an example, we consider the normal cone of (closed) Wasserstein balls:

######  Example 2.24 (Normal cone of closed Wasserstein balls).

For some $\varepsilon>0$, consider
$\operatorname{\mathcal{C}}=\bar{\mathbb{B}}_{W_{2}}(\varepsilon;\operatorname{\hat{\mu}})=\\{\mu\in\mathcal{P}_{2}(\mathbb{R}^{d})\,|\,W_{2}\left(\operatorname{\mu},\operatorname{\hat{\mu}}\right)^{2}\leq\varepsilon^{2}\\}$.
Since at any interior point the normal cone of $\operatorname{\mathcal{C}}$
trivializes, we only need to consider $\operatorname{\bar{\mu}}$ on the
boundary. By Proposition 2.9, the subgradient of the Wasserstein distance is
not empty and is contained in111Here, the factor 2 arises since there is no
$\frac{1}{2}$ in front of the Wasserstein distance; cf. Corollary 2.17.
${(\pi_{1},2\pi_{1}-2\pi_{2})}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})}$.
Since $\operatorname{\bar{\mu}}$ lies at the boundary, we have
$\operatorname{\bar{\mu}}\neq\operatorname{\hat{\mu}}$ and so
$\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})$ does not
contain the “identity plan”
${(\mathrm{Id},\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}}$. Thus,
$\mathbf{0}_{\operatorname{\bar{\mu}}}$ does not belong to the subgradient of
the Wasserstein distance, and the constraint qualification in Proposition 2.22
holds true. With this, we conclude that the normal cone satisfies

$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})\subseteq\\{\operatorname{\xi}\,|\,\operatorname{\xi}\in{(\pi_{1},2\lambda(\pi_{1}-\pi_{2}))}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})},\lambda\geq
0\\}.$

In particular, if $\operatorname{\bar{\mu}}$ is absolutely continuous, then
$\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})=\\{{(\mathrm{Id},T)}_{\\#}{\operatorname{\bar{\mu}}}\\}$,
where $T$ is the unique optimal transport map $T$, and
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})\subseteq\\{{(\mathrm{Id},2\lambda(\mathrm{Id}-T))}_{\\#}{\operatorname{\bar{\mu}}}\,|\,\lambda\geq
0\\}.$

We also study the feasible set of the optimization problem (2), presented in
Section 1, namely the second moment constraint:

######  Example 2.25 (Second moment constraint).

For $\varepsilon>0$, consider
$\operatorname{\mathcal{C}}=\\{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})\,|\,\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\|x\right\|^{2}\right]\leq\varepsilon^{2}\\}$.
Then,
$\operatorname{\mathcal{C}}=\bar{\mathbb{B}}_{W_{2}}(\varepsilon;\delta_{0})$
and, thus, for any
$\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ in the interior
of $\operatorname{\mathcal{C}}$ the normal cone trivializes, whereas for all
the ones with
$\mathbb{E}^{x\sim\operatorname{\bar{\mu}}}\left[\left\|x\right\|^{2}\right]=\varepsilon^{2}$
we have
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})\subseteq\left\\{{(\mathrm{Id},2\lambda\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}}\,|\,\lambda\geq
0\right\\}.$

To prove our first-order optimality conditions in Theorem 3.2, we will require
a technical property for the constraint set:

######  Definition 2.26 (Sequential normal compactness).

A set $\operatorname{\mathcal{C}}$ is sequentially normally compact (SNC) at
$\operatorname{\bar{\mu}}$ if for any sequence
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}}$ and
$\operatorname{\xi}_{n}\in\hat{\operatorname{N}}_{\operatorname{\mathcal{C}}}(\operatorname{\mu}_{n})$
so that $\operatorname{\xi}_{n}\rightharpoonup\operatorname{\xi}$, we have
$\operatorname{\xi}\to\operatorname{\xi}$.

In particular, the normal cones of the sets in Examples 2.24 and 2.25 are SNC:

######  Proposition 2.27 (The Wasserstein ball is SNC).

The closed Wasserstein ball around $\operatorname{\hat{\mu}}$ of radius
$\varepsilon>0$,
$\bar{\mathbb{B}}_{W_{2}}(\varepsilon;\operatorname{\hat{\mu}})$, is SNC at
all
$\operatorname{\bar{\mu}}\in\bar{\mathbb{B}}_{W_{2}}(\varepsilon;\operatorname{\hat{\mu}})$.

####  2.5.1 Variational geometry and epigraphs

In Section 3, we provide the first-order optimality conditions for general,
constrained, optimization problems and present differentiable functionals as a
special case. From a technical standpoint, however, we _first_ establish
optimality conditions for differentiable functions, and _then_ study the non-
differentiable case via an epigraphical argument. This will result in a
“constraint qualification” which is typical in Euclidean spaces and we find in
the Wasserstein space as well. More formally, recall that the epigraph of a
proper functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
is the set

$\widetilde{\mathrm{epi}}\operatorname{\mathcal{J}}\coloneqq\left\\{(\operatorname{\mu},\beta)\,|\,\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d}),\operatorname{\mathcal{J}}(\operatorname{\mu})\leq\beta\right\\}\subset\mathcal{P}_{2}(\mathbb{R}^{d})\times\mathbb{R}.$

Then, we immediately observe that

$\displaystyle\inf_{\operatorname{\mu}\in\operatorname{\mathcal{C}}}\operatorname{\mathcal{J}}(\operatorname{\mu})=\inf_{(\operatorname{\mu},\beta)\in\widetilde{\mathrm{epi}}(\operatorname{\mathcal{J}})\cap(\operatorname{\mathcal{C}}\times\mathbb{R})}\beta=\inf_{(\operatorname{\mu},\beta)\in\widetilde{\mathrm{epi}}(\operatorname{\mathcal{J}})\cap(\operatorname{\mathcal{C}}\times\mathbb{R})}\tilde{l}(\operatorname{\mu},\beta),$

with $\tilde{l}:\mathcal{P}_{2}(\mathbb{R}^{d})\times\mathbb{R}\to\mathbb{R}$
being the linear function $\tilde{l}(\operatorname{\mu},\beta)=\beta$. Being
linear, $\tilde{l}$ is supposedly easy to optimize. Unfortunately, this
argument does not directly carry over to the Wasserstein space: While on
$\mathbb{R}^{d}$ the epigraph is a subset of $\mathbb{R}^{d+1}$, on
$\mathcal{P}_{2}(\mathbb{R}^{d})$ is a subset of
$\mathcal{P}_{2}(\mathbb{R}^{d})\times\mathbb{R}$ and not
$\mathcal{P}_{2}(\mathbb{R}^{d+1})$. Thus, one has to study the differential
structure and variational geometry of
$\mathcal{P}_{2}(\mathbb{R}^{d})\times\mathbb{R}$. To avoid this burden, we
take a different angle. Since each $x\in\mathbb{R}^{d}$ can be “embedded” in
the probability space as the Dirac’s delta
$\delta_{x}\in\mathcal{P}_{2}(\mathbb{R}^{d})$, we define the epigraph as

$\operatorname{epi}\operatorname{\mathcal{J}}\coloneqq\left\\{\operatorname{\mu}\times\delta_{\beta}\,|\,\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d}),\operatorname{\mathcal{J}}(\operatorname{\mu})\leq\beta\right\\}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d+1}).$
(17)

In particular, contrary to
$\widetilde{\mathrm{epi}}\operatorname{\mathcal{J}}$,
$\operatorname{epi}\operatorname{\mathcal{J}}$ is now a subset of
$\mathcal{P}_{2}(\mathbb{R}^{d+1})$. Then, the optimization of
$\operatorname{\mathcal{J}}$ over $\operatorname{\mathcal{C}}$ is equivalent
to the optimization of $l:\mathcal{P}_{2}(\mathbb{R}^{d+1})\to\mathbb{R}$,
$l(\bm{\operatorname{\mu}})=\mathbb{E}^{\bm{x}\sim\bm{\operatorname{\mu}}}\left[\left\langle(0,\ldots,0,1),\bm{x}\right\rangle\right]$,
over
$\operatorname{epi}\operatorname{\mathcal{J}}\cap\operatorname{\mathcal{C}}_{e}$,
with
$\operatorname{\mathcal{C}}_{e}\coloneqq\\{\operatorname{\mu}\times\delta_{\beta}\,|\,\operatorname{\mu}\in\operatorname{\mathcal{C}},\beta\in\mathbb{R}\\}.$
Note that we use bold symbols when working in $\mathbb{R}^{d+1}$ and
$\mathcal{P}_{2}(\mathbb{R}^{d+1})$ (e.g., $\bm{x}=(x,\beta)$ with
$x\in\mathbb{R}^{d}$ and $\beta\in\mathbb{R}$). Given $\bm{\mu}$, we can
“extract” $\operatorname{\mu}$ and $\delta_{\beta}$ via
${(\pi_{1},\ldots,\pi_{d})}_{\\#}{\bm{\operatorname{\mu}}}=\operatorname{\mu}$
(i.e., we forget the last component) and
${(\pi_{d+1})}_{\\#}{\bm{\operatorname{\mu}}}=\delta_{\beta}$ (i.e., we keep
only the last component). Similarly, given a variation
$\bm{\operatorname{\xi}}\in\operatorname{T}_{\mathcal{P}_{2}(\mathbb{R}^{d+1})}(\operatorname{\mu})$,
we can extract the variation corresponding to $\operatorname{\mu}$ via
${T}_{\\#}{\bm{\operatorname{\xi}}}$, where

$T\coloneqq(\pi_{1},\ldots,\pi_{d},\pi_{d+2},\ldots,\pi_{2d+1})$ (18)

(i.e., we forget the last component of the first and the second marginal), and
the variation “corresponding” to $\delta_{\beta}$ via
${(\pi_{d+1},\pi_{d+2})}_{\\#}{\bm{\operatorname{\xi}}}$ (i.e., we only keep
the last component of the first and the second marginal). With this formalism,
the gradient of $l$ follows directly from Proposition 2.12 and it reads
$\bm{\operatorname{\mu}}\times\bm{\delta_{(0,\ldots,0,1)}}$. Thus, the
complexity of our main result moves to prove first-order optimality conditions
in the differentiable case subject to epigraphical constraints, for which, in
particular, we need to study the variational geometry of
$\operatorname{epi}\operatorname{\mathcal{J}}$, which we do next. To start,
observe that epigraphical normals “point downwards in expectation”:

######  Proposition 2.28 (Characterization of the epigraphical normals).

For
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$,
consider $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ at which
$\operatorname{\mathcal{J}}$ is finite. For all
$\bm{\operatorname{\xi}}\in\operatorname{N}_{\operatorname{epi}\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}}\times\delta_{\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})})$
(i.e., at the boundary of $\operatorname{epi}\operatorname{\mathcal{J}}$) we
have
$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\bm{\operatorname{\upsilon}}\right\rangle\right]\leq
0$.

Proposition 2.28 is reminiscent of the well-known fact that, in Euclidean
settings, epigraphical normals never point upwards. We depict this idea below
on the right.

For the (discontinuous) function $f$ with graph depicted in black, the
epigraph is the region in light green above the graph. The regular
subgradients at the discontinuity point (black-filled point) correspond to the
thick blue arrows, which touch the plane at $y=-1$. These arrows are also
related to the normals of the epigraph, which contains their $\lambda$-scaled
version as well. These are depicted in dashed blue, and they intersect the
$-\lambda$ plane.

We rigorously show that the relationship between the non-horizontal normals
(in expectation) to $\operatorname{epi}\operatorname{\mathcal{J}}$ and the
subgradients of $\operatorname{\mathcal{J}}$ translates to the probability
space:

######  Proposition 2.29 (Subgradients and epigraphical normals).

Let
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
and any $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ at which
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\in\mathbb{R}$ and let
$\bm{\operatorname{\xi}}\in\operatorname{N}_{\operatorname{epi}\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}}\times\delta_{\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})})$
with
$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\bm{v}\right\rangle\right]=-1$.
Then,
${T}_{\\#}{\bm{\operatorname{\xi}}}\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
with $T$ as in (18).

Note that Proposition 2.29 characterizes _all_ non-horizontal normals. Indeed,
if
$\bm{\operatorname{\xi}}\in\operatorname{N}_{\operatorname{epi}\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}}\times\delta_{\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})})$
satisfies
$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\bm{v}\right\rangle\right]=-\lambda<0$,
then, since the normal cone is a cone (cf. Proposition 2.21),
$\frac{1}{\lambda}\bm{\operatorname{\xi}}$ also belongs to
$\operatorname{N}_{\operatorname{epi}\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}}\times\delta_{\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})})$
and it satifies
$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\frac{1}{\lambda}\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\bm{v}\right\rangle\right]=\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\frac{1}{\lambda}\bm{v}\right\rangle\right]=-1$.
Thus,
${(\pi_{1},\ldots,\pi_{d})}_{\\#}{\frac{1}{\lambda}\bm{\operatorname{\xi}}}\in\partial\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$.
In Euclidean spaces, these normals are referred to as horizon subgradients (we
depict one in red in the figure above). Thus, we introduce the natural
counterpart in the probability space:

######  Definition 2.30 (Wasserstein horizon subgradients).

For
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
and any $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ at which
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\in\mathbb{R}$, we say
that $\operatorname{\xi}$ is a _horizon subgradient of
$\operatorname{\mathcal{J}}$ at $\operatorname{\bar{\mu}}$_,
$\operatorname{\xi}\in\partial^{\infty}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$,
if there exists
$\bm{\operatorname{\xi}}\in\operatorname{N}_{\operatorname{epi}\operatorname{\mathcal{J}}}(\operatorname{\bar{\mu}}\times\delta_{\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})})$
such that ${T}_{\\#}{\bm{\operatorname{\xi}}}=\operatorname{\xi}$ and
$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})}\left[\left\langle(0,\ldots,0,1),\bm{\operatorname{\upsilon}}\right\rangle\right]=0$,
with $T$ as in (18).

The horizon subgradients relate to the constraint qualifications and, as we
shall see in Section 3, are required for a comprehensive theory of
optimization. Similar ideas are necessary already in Euclidean spaces and,
perhaps surprisingly, they naturally carry over to the probability space.
Since they are defined starting from the general normals to the epigraph, they
are not guaranteed to be in the tangent space. Nonetheless, for sufficiently
well-behaved functionals they are trivial:

######  Proposition 2.31 (Horizon subgradients for strictly continuous
functions).

Let
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
and any $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$. If
$\operatorname{\mathcal{J}}$ is _strictly continuous_ at
$\operatorname{\bar{\mu}}$, i.e.,

$\limsup_{\operatorname{\mu}\to\operatorname{\bar{\mu}}}\frac{\left\|\operatorname{\mathcal{J}}(\operatorname{\mu})-\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\right\|}{W_{2}\left(\operatorname{\mu},\operatorname{\bar{\mu}}\right)}<\infty,$

then the horizon subgradient of $\operatorname{\mathcal{J}}$ at
$\operatorname{\bar{\mu}}$ is trivial; i.e.,
$\partial^{\infty}\operatorname{\mathcal{J}}{\operatorname{\bar{\mu}}}=\\{\mathbf{0}_{\operatorname{\bar{\mu}}}\\}$.

For instance, all locally Lipschitz functionals are strictly continuous and
therefore have trivial horizon subgradients. This is also the case for many of
the functionals discussed so far:

######  Proposition 2.32 (Some horizon subgradients).

At any $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$, the
functionals defined in Propositions 2.12, 2.14, 2.10, 2.9 and 2.18 have
trivial horizon subgradient
$\partial^{\infty}\operatorname{\mathcal{J}}=\\{\mathbf{0}_{\operatorname{\bar{\mu}}}\\}$.

With horizon subgradients, we can finally provide a full characterization of
epigraphical normals:

######  Proposition 2.33 (Subgradients and epigraphical normals).

Let
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
and $\operatorname{\bar{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ at which
$\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})\in\mathbb{R}$. With $T$
as in (18), we have that
$\bm{\operatorname{\xi}}\in\operatorname{N}_{\operatorname{epi}\operatorname{\mathcal{J}}}(\operatorname{\mu}\times\delta_{\operatorname{\mathcal{J}}(\operatorname{\mu})})$
if and only if either

  1. (i)

$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\bm{\operatorname{\upsilon}}\right\rangle\right]=-\lambda,\lambda>0$
and
$\frac{1}{\lambda}{T}_{\\#}{\bm{\operatorname{\xi}}}\in\partial\operatorname{\mathcal{J}}(\operatorname{\mu})$;
or

  2. (ii)

$\mathbb{E}^{(\bm{x},\bm{\operatorname{\upsilon}})\sim\bm{\operatorname{\xi}}}\left[\left\langle(0,\ldots,0,1),\bm{\operatorname{\upsilon}}\right\rangle\right]=0$
and
${T}_{\\#}{\bm{\operatorname{\xi}}}\in\partial^{\infty}\operatorname{\mathcal{J}}(\operatorname{\mu})$.

##### Discussion

The definition of horizon subgradients introduced in Definition 2.30 follows
the geometrical ideas used for Euclidean spaces in [47, Definition 1.18]. It
departs from the alternative definition of the horizon subgradients as the
scaled limit of regular subgradients [44, Definition 8.3], according to which
$\operatorname{\xi}\in\partial^{\infty}\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
if there are sequences
$\operatorname{\mu}_{n}\to\operatorname{\bar{\mu}},\operatorname{\mathcal{J}}(\operatorname{\mu}_{n})\to\operatorname{\mathcal{J}}(\operatorname{\bar{\mu}})$
and
$\operatorname{\xi}_{n}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\mu}_{n})$
so that there exists $(\tau_{n})_{n\in\mathbb{N}}\subseteq\mathbb{R}_{\geq
0},\tau_{n}\searrow 0$ and
$\tau_{n}\operatorname{\xi}_{n}\to\operatorname{\bar{\operatorname{\xi}}}$.
Nonetheless, it is well known that, at least in Euclidean settings, the two
definitions are identical, cf. [44, Theorem 8.9]. Finally, we could have
defined the subgradients also as the non-horizontal epigraphical normals, as
in [47, Definition 1.18]. Nonetheless, to ease the presentation, we opted for
the more standard Definition 2.3, as in [44, Definition 8.3].

##  3 Optimality conditions

In this section, we prove our main result. To start, we need a notion of local
optimality:

######  Definition 3.1 (Local optimality in the Wasserstein space).

A functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
attains a local minimum over a constraint set
$\operatorname{\mathcal{C}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$ at
$\operatorname{\mu^{\ast}}\in\operatorname{\mathcal{C}}$ if there exists
$\varepsilon>0$ such that for all
$\operatorname{\mu}\in\mathbb{B}_{W_{2}}(\varepsilon;\operatorname{\mu^{\ast}})\cap\operatorname{\mathcal{C}}$
we have
$\operatorname{\mathcal{J}}(\operatorname{\mu})\geq\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}}).$

Next is the main result of this work, which encompasses the celebrated Karush-
Kuhn-Tucker and Lagrange conditions: first-order necessary optimality
conditions in the Wasserstein space.

######  Theorem 3.2 (First-order optimality conditions).

If a proper functional $\operatorname{\mathcal{J}}$ attains a local minimum at
$\operatorname{\mu^{\ast}}\in\operatorname{\mathcal{C}}\subseteq\mathcal{P}_{2}(\mathbb{R}^{d})$,
with $\operatorname{\mathcal{C}}$ SNC, and the constraint qualification
$\partial^{\infty}\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}})\cap\left(-\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\mu^{\ast}})\right)\subseteq\\{\mathbf{0}_{\operatorname{\mu^{\ast}}}\\}$
(19) holds, then $\displaystyle\mathbf{0}_{\operatorname{\mu^{\ast}}}$
$\displaystyle\in\partial\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}})+\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\mu^{\ast}})$
$\displaystyle\coloneqq\left\\{\operatorname{\xi}_{1}\,+_{\operatorname{\alpha}}\,\operatorname{\xi}_{2}\,|\,\operatorname{\xi}_{1}\in\partial\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}}),\operatorname{\xi}_{2}\in\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\mu^{\ast}}),\operatorname{\alpha}\in\Gamma^{1}(\operatorname{\xi}_{1},\operatorname{\xi}_{2})\right\\}.$
When $\operatorname{\mathcal{J}}$ is differentiable, (19) is automatically
satisfied and the optimality condition reads
$-\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}})\in\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\mu^{\ast}})$
even if $\operatorname{\mathcal{C}}$ is not SNC. Geometrically, the constraint
qualification (19) requires that the feasible set $\operatorname{\mathcal{C}}$
is not “tangent” to the epigraph at $\operatorname{\mu^{\ast}}$. Similarly to
the Euclidean setting, at optimality at least one element of the subgradient
of the cost functional (direction of improvement) must align and have the same
magnitude (but opposite direction) with one element of the normal cone (the
“blocked” directions). We refer to Section 3.1 for further discussion. In the
unconstrained case, the normal cone trivializes (cf. Proposition 2.22) and we
recover Fermat’s rule:

######  Theorem 3.3 (Fermat’s rule in the Wasserstein space).

If a proper functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\bar{\mathbb{R}}$
attains a local minimum at $\operatorname{\mu^{\ast}}$, then
$\mathbf{0}_{\operatorname{\mu^{\ast}}}\in\hat{\partial}\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}}),$
which implies
$\mathbf{0}_{\operatorname{\mu^{\ast}}}\in\partial\operatorname{\mathcal{J}}(\operatorname{\mu^{\ast}})$.
When $J$ is differentiable, this condition reduces to
$\mathbf{0}_{\operatorname{\mu^{\ast}}}=\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}}).$

###  3.1 Discussion

Before showcasing our first-order necessary optimality conditions in various
examples, we discuss several related aspects.

##### Relation to classical variational analysis

The conditions in Theorems 3.2 and 3.3 very much resemble their counterpart in
Euclidean spaces and Asplund spaces: the negative subgradient of the objective
function must lie in the normal cone of the feasible set, provided that a
constraint qualification (i.e., the intersection of the horizon subgradient
and “negative” the normal cone is trivial) and regularity condition on the
feasible set hold (i.e., the set $\operatorname{\mathcal{C}}$ is SNC); e.g.,
see [44, 47] for the corresponding statement in Euclidean spaces (where the
SNC condition always holds and thus does not appear) and [46, §5.1] for the
infinite-dimensional version. Note that these results in Euclidean spaces
generalize, among others, the celebrated Karush-Kuhn-Tucker and Lagrange
conditions.

##### Comparison with the literature

General necessary optimality conditions in the Wasserstein space are studied
in [1], in the smooth (or non-smooth but convex) setting where the constraint
is the level set of a real-valued functional, and in [36, 37, 27] in the
different setting of optimal control problems in the Wasserstein space. With
this work, we provide a general theory of variational analysis in the
Wasserstein space and significantly generalize existing necessary conditions
for static optimization.

##### Computational aspects

Our optimality conditions transform the problem of minimizing a function over
the infinite-dimensional probability space into an inclusion problem at the
level of transport plans which, in turn, amounts to a set of conditions over
$\mathbb{R}^{d}\times\mathbb{R}^{d}$. Depending on the application, they can
be solved in (quasi) closed-form or addressed via (nonlinear) function
approximators [57, 58], as we show in our examples. For this reason, we argue
that our necessary conditions are significantly more tractable than the
initial infinite-dimensional optimization problem. Nonetheless, more rigorous
statements of the computational aspects of necessary conditions for optimality
depend on the specifics of the problem at hand, as is already the case in
Euclidean settings.

##### Sufficient conditions

As in Euclidean settings, our conditions are not sufficient for optimality. We
expect sufficient conditions for optimality to be intimately related to
geodesic convexity [33, §7] and second-order calculus in the Wasserstein space
[59]; see [1] for preliminary results. We leave this topic to future research.

###  3.2 Pedagogical examples

In the remainder of this section, we illustrate our necessary conditions
across several examples. Despite its pedagogical approach, this section
culminates with Example 3.6, which already generalizes existing insights in
DRO. We start by solving rigorously the example in Section 1:

######  Example 3.4 (Expected value subject to second moment constraints).

For $\theta\neq 0$ and $\varepsilon>0$, consider the problem

$\inf_{\operatorname{\mu}\in\mathcal{P}_{2}(\mathbb{R}^{d})}\;\operatorname{\mathcal{J}}(\operatorname{\mu})\coloneqq\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\langle\theta,x\right\rangle\right]\qquad\text{subject
to}\qquad\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\|x\right\|^{2}\right]\leq\varepsilon^{2}.$

By Example 2.13, $\operatorname{\mathcal{J}}$ is differentiable with gradient
$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu})={(\mathrm{Id},\theta)}_{\\#}{\operatorname{\mu}}$.
The normal cone of the second moment constraint was studied in Example 2.25,
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\\{\mathbf{0}_{\operatorname{\bar{\mu}}}\\}$
if
$\mathbb{E}^{x\sim\operatorname{\bar{\mu}}}\left[\left\|x\right\|^{2}\right]<\varepsilon^{2}$
and
$\operatorname{N}_{\operatorname{\mathcal{C}}}(\operatorname{\bar{\mu}})=\\{{(\mathrm{Id},2\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}}\\}$
if
$\mathbb{E}^{x\sim\operatorname{\bar{\mu}}}\left[\left\|x\right\|^{2}\right]=\varepsilon^{2}$.
Any candidate optimal solution
$\operatorname{\mu^{\ast}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ must satisfy the
optimality conditions in Theorem 3.2. We study two cases. If
$\mathbb{E}^{x\sim\operatorname{\mu^{\ast}}}\left[\left\|x\right\|^{2}\right]<\varepsilon^{2}$,
it holds

$\displaystyle\mathbf{0}_{\operatorname{\mu^{\ast}}}=\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}})\Leftrightarrow
0$
$\displaystyle=\left\|\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}})\right\|_{\operatorname{\mu^{\ast}}}^{2}=\int_{\mathbb{R}^{d}}\left\|\theta\right\|^{2}\mathrm{d}\operatorname{\mu^{\ast}}(x)>0,$

which is absurd. Therefore, the optimal solution (if exists) lies at the
boundary, and a $\lambda>0$ and a “sum coupling”
$\operatorname{\alpha}\in\Gamma^{1}(\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}}),\lambda{(\mathrm{Id},2\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}})$
must exist so that

$\displaystyle\mathbf{0}_{\operatorname{\mu^{\ast}}}=\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}})\,+_{\operatorname{\alpha}}\,\lambda{(\mathrm{Id},2\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}}\Leftrightarrow
0$
$\displaystyle=\left\|\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}})\,+_{\operatorname{\alpha}}\,{(\mathrm{Id},2\mathrm{Id})}_{\\#}{\operatorname{\bar{\mu}}}\right\|_{\operatorname{\mu^{\ast}}}^{2}$
$\displaystyle=\int_{\mathbb{R}^{d}}\left\|\theta+2\lambda
x\right\|^{2}\mathrm{d}\operatorname{\mu^{\ast}}(x)>0.$

The above is by all means a “Lagrange’s condition” in the probability space,
which yields $\operatorname{\mu^{\ast}}=\delta_{-\frac{\theta}{2\lambda}}$ and
$\lambda=\frac{\left\|\theta\right\|}{2\varepsilon}$, since
$\mathbb{E}^{x\sim\operatorname{\mu^{\ast}}}\left[\left\|x\right\|^{2}\right]=\varepsilon^{2}$.

We now show that necessary conditions for optimality can be used to produce
certificates that an optimal solution does not exist:

######  Example 3.5 (A certificate of unfeasibility).

Consider the problem of finding the worst-case mean-variance risk over the
full probability space. In this case, we seek for the minimum of the
functional

$\operatorname{\mathcal{J}}(\operatorname{\mu})=-\left(\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\langle\theta,x\right\rangle\right]+\frac{\rho}{2}\mathrm{Var}^{x\sim\operatorname{\mu}}\left[\left\langle\theta,x\right\rangle\right]\right).$

By Propositions 2.12, 2.16 and 2.18, $\operatorname{\mathcal{J}}$ is
differentiable with gradient

$\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu})=-{(\mathrm{Id},(1+\rho(\left\langle\theta,\mathrm{Id}\right\rangle-\mathbb{E}^{x\sim\operatorname{\mu}}\left[\left\langle\theta,x\right\rangle\right]))\theta)}_{\\#}{\operatorname{\mu}}.$

To find the unconstrained minimizer of $\operatorname{\mathcal{J}}$, we can
deploy Theorem 3.3:

$\displaystyle\mathbf{0}_{\operatorname{\mu}}=\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu})\Leftrightarrow
0$
$\displaystyle=\left\|\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu})\right\|^{2}_{\operatorname{\mu}}$
$\displaystyle=\int_{\mathbb{R}^{d}}\left\|(1+\rho(\left\langle\theta,x\right\rangle-\mathbb{E}^{y\sim\operatorname{\mu}}\left[\left\langle\theta,y\right\rangle\right]))\theta\right\|^{2}\mathrm{d}\operatorname{\mu}(x)$

Thus, at any optimal
$\operatorname{\mu^{\ast}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$, it holds

$0=1+\rho(\left\langle\theta,x\right\rangle-\mathbb{E}^{y\sim\operatorname{\mu^{\ast}}}\left[\left\langle\theta,y\right\rangle\right])\quad\operatorname{\mu^{\ast}}-\mathrm{a.e.},$
(20)

but taking the expected value w.r.t. $\operatorname{\mu^{\ast}}$ of both sides
in (20) yields (it holds $\operatorname{\mu^{\ast}}-\mathrm{a.e.}$)

$0=\mathbb{E}^{x\sim\operatorname{\mu^{\ast}}}\left[0\right]=\mathbb{E}^{x\sim\operatorname{\mu^{\ast}}}\left[1+\rho(\left\langle\theta,x\right\rangle-\mathbb{E}^{y\sim\operatorname{\mu^{\ast}}}\left[\left\langle\theta,y\right\rangle\right])\right]=1,$

a contradiction. Thus, the infimum of $\operatorname{\mathcal{J}}$ is not
attained.

While not a novel result, the above example shows that with only simple
algebraic manipulations we can use the optimality conditions in Theorem 3.2 to
prove non-existence of a solution, in a way that mirrors arguments in
Euclidean spaces.

We now consider an example in the setting of DRO, a ubiquitous framework for
decision-making under uncertainty. We seek to evaluate the worst-case mean-
variance of a linear portfolio with allocation $\theta$ but now constrained to
a closed Wasserstein ball of a given radius $\varepsilon$ and center
$\operatorname{\bar{\mu}}$. Formally, this example is similar to Example 3.5,
but now the worst-case is taken over a Wasserstein ball instead of the entire
probability space. As we shall see below, our necessary conditions lead to a
closed-form solution for the worst-case probability measure, generalizing [1,
Section 4] to non-absolutely continuous measures (which, among others, allows
us to study the data-driven setting where $\operatorname{\hat{\mu}}$ is
empirical) and [60, Appendix E], which does not provide a closed-form solution
and assume positive variance of $\operatorname{\bar{\mu}}$.

######  Example 3.6 (A constrained problem: mean-variance DRO).

Consider the mean-variance functional
$\operatorname{\mathcal{J}}:\mathcal{P}_{2}(\mathbb{R}^{d})\to\mathbb{R}^{d}$
defined in Example 3.5 and the constraint
$\operatorname{\mu}\in\bar{\mathbb{B}}_{W_{2}}(\operatorname{\hat{\mu}};\varepsilon)$
for some $\operatorname{\hat{\mu}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ and
$\varepsilon>0$. We computed the gradient of $\operatorname{\mathcal{J}}$ in
Example 3.5, and we recall from Example 2.24 that the normal cone of
$\bar{\mathbb{B}}_{W_{2}}(\operatorname{\hat{\mu}};\varepsilon)$ at any
$\operatorname{\bar{\mu}}\in\bar{\mathbb{B}}_{W_{2}}(\operatorname{\hat{\mu}};\varepsilon)$
satisfies
$\operatorname{N}_{\bar{\mathbb{B}}_{W_{2}}(\operatorname{\hat{\mu}};\varepsilon)}(\operatorname{\bar{\mu}})\subseteq\\{\operatorname{\xi}\,|\,\operatorname{\xi}\in{(\pi_{1},2\lambda(\pi_{1}-\pi_{2}))}_{\\#}{\Gamma_{o}(\operatorname{\bar{\mu}},\operatorname{\hat{\mu}})},\lambda\geq
0\\}$. Then, by Theorem 3.2, at any optimal
$\operatorname{\mu^{\ast}}\in\mathcal{P}_{2}(\mathbb{R}^{d})$ an optimal
transport plan
$\operatorname{\gamma}\in\Gamma_{o}(\operatorname{\mu^{\ast}},\operatorname{\hat{\mu}})$,
a Lagrange multiplier $\lambda\geq 0$, and a “sum coupling”
$\operatorname{\alpha}\in\Gamma^{1}(\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu}),{(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{\operatorname{\gamma}})$
exist so that
$\mathbf{0}_{\operatorname{\mu^{\ast}}}=\nabla{\operatorname{\mathcal{J}}}(\operatorname{\mu^{\ast}})\,+_{\operatorname{\alpha}}\,2\lambda{(\pi_{1},\pi_{1}-\pi_{2})}_{\\#}{\operatorname{\gamma}}$.
Equivalently,

$\displaystyle 0$