# The Wasserstein Proximal Gradient Algorithm

Adil Salim
Visual Computing Center, KAUST
Anna Korba
Gatsby Computational Neuroscience Unit, UCL Giulia Luise
Computer Science Department, UCL

###### Abstract

Wasserstein gradient flows are continuous time dynamics that define curves of
steepest descent to minimize an objective function over the space of
probability measures (i.e., the Wasserstein space). This objective is
typically a divergence w.r.t. a fixed target distribution. In recent years,
these continuous time dynamics have been used to study the convergence of
machine learning algorithms aiming at approximating a probability
distribution. However, the discrete-time behavior of these algorithms might
differ from the continuous time dynamics. Besides, although discretized
gradient flows have been proposed in the literature, little is known about
their minimization power. In this work, we propose a Forward Backward (FB)
discretization scheme that can tackle the case where the objective function is
the sum of a smooth and a nonsmooth geodesically convex terms. Using
techniques from convex optimization and optimal transport, we analyze the FB
scheme as a minimization algorithm on the Wasserstein space. More precisely,
we show under mild assumptions that the FB scheme has convergence guarantees
similar to the proximal gradient algorithm in Euclidean spaces (resp. similar
to the associated Wasserstein gradient flow).

##  1 Introduction

The task of transporting an initial distribution $\mu_{0}$ to a target
distribution $\mu_{\star}$ is common in machine learning. This task can be
reformulated as the minimization of a cost functional defined over the set of
probability distributions. Wasserstein gradient flows [2] are suitable
continuous time dynamics to minimize such cost functionals. These flows have
found applications in various fields of machine learning such as reinforcement
learning [28, 36], sampling [34, 14, 5, 10] and neural networks optimization
[12, 23]. Indeed, Wasserstein gradient flows can be seen as the continuous
limit of several discrete time machine learning algorithms. The analysis of
continuous time dynamics is often easier than the analysis of their discrete
time counterparts. Therefore, many works focus solely on continuous time
analyses of machine learning algorithms such as variants of gradient descent
[30, 33, 12, 23, 13, 7]. However, as in optimization over Hilbert spaces, the
behavior of discrete time algorithms might differ from the behavior of their
continuous time limits.

In this paper, we focus on the resolution, by a discrete time algorithm, of a
minimization problem defined on the set ${\mathcal{P}}_{2}({\mathcal{X}})$ of
probability measures $\mu$ over ${\mathcal{X}}={\mathbb{R}}^{d}$ such that
$\int\|x\|^{2}d\mu(x)<\infty$. More precisely, $\mu_{\star}$ is defined as a
solution to

$\min_{\mu\in{\mathcal{P}}_{2}({\mathcal{X}})}{\mathcal{G}}(\mu):={\mathcal{E}}_{F}(\mu)+{\mathcal{H}}(\mu),$
(1)

where ${\mathcal{E}}_{F}$ is a potential energy ${\mathcal{E}}_{F}(\mu)=\int
F(x)d\mu(x)$ tied to a smooth convex function
$F:{\mathbb{R}}^{d}\to{\mathbb{R}}$, and ${\mathcal{H}}$ is a nonsmooth term
convex along the generalized geodesics defined by the Wasserstein distance.
The potential ${\mathcal{E}}_{F}$ plays the role of a data fitting term
whereas ${\mathcal{H}}$ can be seen as a regularizer.

##### Motivation for studying the template problem (1).

Many free energy minimization problems can be cast as Problem (1), see [1,
Proposition 7.7] or more generally [2, Section 9]. For instance,
${\mathcal{H}}$ can be an internal energy [2, Example 9.3.6]. In particular,
if ${\mathcal{H}}$ is the negative entropy, then ${\mathcal{G}}$ boils down to
the Kullback Leibler divergence (up to an additive constant) w.r.t. the Gibbs
measure $\mu_{\star}\propto\exp(-F)$. This remark has been used by several
authors to study sampling tasks as minimizing Problem (1) [34, 14, 18, 5, 10].
Another example where ${\mathcal{H}}$ is an internal energy is the case where
${\mathcal{H}}$ is a higher order entropy. In this case, $\mu_{\star}$ follows
a Barenblatt profile. Moreover, in the context of optimization of infinitely
wide two layers neural networks [23, 12], $\mu_{\star}$ denotes the optimal
distribution over the parameters of the network. In this context, $F=\int
k(x,y)d\mu_{\star}(y)$ is non-convex, with $k$ depending on the activation
functions of the network, and ${\mathcal{H}}$ is an interaction energy [2,
Example 9.3.4]. Moreover, ${\mathcal{G}}$ boils down to a Maximum Mean
Discrepancy w.r.t. $\mu_{\star}$ [3] under a well-posedness condition.
Alternatively, ${\mathcal{H}}$ can be a regularizer of the distribution on the
parameters of the network.

Related works. Wasserstein gradient flows are continuous time dynamics that
can minimize (1). Several time discretizations of such flows have been
considered [29, 2, 34]. However, these discretization schemes have been mainly
analyzed as numerical schemes to approximate the continuous gradient flow,
rather than optimization algorithms.

Numerous optimization algorithms to solve (1), relying on different time-
discretization schemes of the Wasserstein gradient flow of ${\mathcal{G}}$
have been proposed previously. For instance, [15, 36, 29] rely on the
implementation of the JKO scheme [17], which can be seen as a proximal scheme
(i.e., backward discretization) with the Wasserstein distance. In this case,
each step of the algorithm relies on evaluating the JKO operator of
${\mathcal{G}}$, exploiting efficient subroutines for this operator. When
${\mathcal{G}}$ is smooth, some gradient descent algorithms over the
${\mathcal{P}}_{2}({\mathcal{X}})$ (i.e., forward discretizations) have also
been proposed [11, 19, 18]. However, their analysis is notably challenging
without further assumptions. For example [11] is limited to the case of
Gaussian distributions, while a convergence rate for SVGD [19, 18] remains
unknown. Another time-discretization of Wasserstein gradient flows can be
found with the Langevin Monte Carlo (LMC) algorithm in the sampling
literature. LMC can be seen as a splitting algorithm involving a gradient step
for ${\mathcal{E}}_{F}$ and an exact solving of the flow of ${\mathcal{H}}$
(Forward-Flow discretization [34]) to solve (1) in the case where
${\mathcal{H}}$ is the negative entropy, since its gradient flow can be
computed exactly in discrete time. Several works provide non asymptotic
analyses of LMC in discrete time, see [14, 10, 5, 35, 31, 20] among others.
However, the convergence rate of the LMC algorithm does not match the
convergence rate of the associated continuous time dynamics and LMC is a
biased algorithm.

A natural approach to minimize the sum of a smooth and a nonsmooth convex
functions over a Hilbert space is to apply the proximal gradient algorithm
[4], which implements a gradient step for the smooth term and a proximal step
for the nonsmooth one. Indeed, many nonsmooth regularizers admit closed form
proximity operators111see www.proximity-operator.net. In this work we propose
and analyze a splitting algorithm to minimize the functional ${\mathcal{G}}$
over ${\mathcal{P}}_{2}({\mathcal{X}})$, that can be seen as an analogue to
the proximal gradient algorithm. More precisely, the proposed algorithm
implements a forward (gradient step) for the smooth term ${\mathcal{E}}_{F}$
and relies on the JKO (Jordan-Kinderlehrer-Otto) operator for the non-smooth
term ${\mathcal{H}}$ [17] only, which plays the role of a proximity operator.
This approach provides an algorithm with lower iteration complexity compared
to the "full" JKO scheme applied to ${\mathcal{E}}_{F}+{\mathcal{H}}$.

Contributions. In summary, the Wasserstein gradient flows to minimize (1) are
well understood and modelize many machine learning algorithms. However, little
is known about the minimization power of discretized gradient flows, that are
better modelizations of these machine learning algorithms. In this work, we
propose a natural optimization algorithm to solve Problem (1), which is a
Forward-Backward discretization of the Wasserstein gradient flow. This
algorithm is a generalization of a discretization scheme proposed in [34]. Our
main contribution is to prove non asymptotic rates for the proposed scheme,
under the assumptions that $F$ is smooth and convex and that ${\mathcal{H}}$
is convex along the generalized geodesics defined by the Wasserstein distance.
We show that the obtained rates fortunately match the ones of the proximal
gradient algorithm over Hilbert spaces (resp. the associated Wasserstein
gradient flow).

The remainder is organized as follows. In Section 2 we provide some background
knowledge in optimal transport and gradient flows. In Section 3, we introduce
the Forward Backward Euler discretization scheme. We study the FB scheme as an
optimization algorithm and present our main results, i.e., non-asymptotic
rates for the resolution of (1) in Section 4. In Section 5 we illustrate the
performance of our algorithm for a simple sampling task. We discuss the
broader impact of our work in Section 6. The convergence proofs and an
additional simulation are postponed to the appendix.

##  2 Preliminaries

In this section, we introduce the notations and recall fundamental definitions
and properties on optimal transport and gradient flows that will be used
throughout the paper.

###  2.1 Notations

In the sequel, ${\mathcal{P}}_{2}({\mathcal{X}})$ is the space of probability
measures $\mu$ on ${\mathcal{X}}$ with finite second order moment. Denote
${\mathcal{B}}({\mathcal{X}})$ the Borelian $\sigma$-field over
${\mathcal{X}}$. For any $\mu\in{\mathcal{P}}_{2}({\mathcal{X}})$,
$L^{2}(\mu)$ is the space of functions
$f:({\mathcal{X}},{\mathcal{B}}({\mathcal{X}}))\to({\mathcal{X}},{\mathcal{B}}({\mathcal{X}}))$
such that $\int\|f\|^{2}d\mu<\infty$. Note that the identity map $I$ is an
element of $L^{2}(\mu)$. For any $\mu\in{\mathcal{P}}_{2}({\mathcal{X}})$, we
denote by $\|\cdot\|_{\mu}$ and $\langle\cdot,\cdot\rangle_{\mu}$ respectively
the norm and the inner product of the space $L^{2}(\mu)$. For any measures
$\mu,\nu$, we write $\mu\ll\nu$ if $\mu$ is absolutely continuous with respect
to $\nu$, and we denote $Leb$ the Lebesgue measure over ${\mathcal{X}}$. The
set of regular distributions of the Wasserstein space is denoted
${\mathcal{P}}_{2}^{r}({\mathcal{X}}):=\\{\mu\in{\mathcal{P}}_{2}({\mathcal{X}}),\;\mu\ll
Leb\\}$. If $f,g:{\mathcal{X}}\to{\mathcal{X}}$, the composition $f\circ g$ of
$g$ by $f$ is sometimes denoted $f(g)$.

###  2.2 Optimal transport

For every measurable map $T$ defined on
$({\mathcal{X}},{\mathcal{B}}({\mathcal{X}}))$ and for every
$\mu\in{\mathcal{P}}_{2}({\mathcal{X}})$, we denote $T_{\\#}\mu$ the
pushforward measure of $\mu$ by $T$ characterized by the transfer lemma:

$\int\phi(y)dT_{\\#}\mu(y)=\int\phi(T(x))d\mu(x)\quad\text{ for any measurable
and bounded function }\phi.$ (2)

Consider the $2$-Wasserstein distance defined for every
$\mu,\nu\in{\mathcal{P}}_{2}({\mathcal{X}})$ by

$W^{2}(\mu,\nu):=\inf_{\upsilon\in\Gamma(\mu,\nu)}\int\|x-y\|^{2}d\upsilon(x,y),$
(3)

where $\Gamma(\mu,\nu)$ is the set of couplings between $\mu$ and $\nu$ [32],
i.e. the set of nonnegative measures $\upsilon$ over
${\mathcal{X}}\times{\mathcal{X}}$ such that $P_{\\#}\upsilon=\mu$ (resp.
$Q_{\\#}\upsilon=\nu$) where $P:(x,y)\mapsto x$ (resp. $Q:(x,y)\mapsto y$).

We now recall the well-known Brenier theorem [8],[2, Section 6.2.3].

######  Theorem 1.

Let $\mu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$ and
$\nu\in{\mathcal{P}}_{2}({\mathcal{X}})$. Then,

  1. 1.

There exists an unique minimizer $\upsilon$ of (3). Besides, there exists an
uniquely determined $\mu$-almost everywhere (a.e.) map
$T_{\mu}^{\nu}:{\mathcal{X}}\to{\mathcal{X}}$ such that
$\upsilon=(I,T_{\mu}^{\nu})_{\\#}\mu$ where
$(I,T_{\mu}^{\nu}):(x,y)\mapsto(x,T_{\mu}^{\nu}(x))$. Finally, there exists a
convex function $f:{\mathcal{X}}\to{\mathbb{R}}$ such that
$T_{\mu}^{\nu}=\nabla f$ $\mu$-a.e.

  2. 2.

As a corollary,

$W^{2}(\mu,\nu)=\int\|x-\nabla
f(x)\|^{2}d\mu(x)=\inf_{T_{\\#}\mu=\nu}\int\|x-T(x)\|^{2}d\mu(x).$ (4)

  3. 3.

If $g:{\mathcal{X}}\to{\mathbb{R}}$ is convex, then $\nabla g$ is well defined
$\mu$-a.e. and if $\nu=\nabla g_{\\#}\mu$, then $T_{\mu}^{\nu}=\nabla g$
$\mu$-a.e.

  4. 4.

If $\nu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, then $T_{\nu}^{\mu}\circ
T_{\mu}^{\nu}=I$ $\mu$-a.e. and $T_{\mu}^{\nu}\circ T_{\nu}^{\mu}=I$
$\nu$-a.e.

Under the assumptions of Theorem 1, the map $T_{\mu}^{\nu}$ is called the
optimal transport (OT) map from $\mu$ to $\nu$. In this paper, as it is
commonly the case in the literature, we may refer to the space of probability
distributions ${\mathcal{P}}_{2}({\mathcal{X}})$ equipped with the
$2$-Wasserstein distance as the Wasserstein space.

###  2.3 Review of Gradient Flows and their discretizations

####  2.3.1 In an Euclidean space

Assume that ${\mathcal{X}}$ is an Euclidean space, consider a proper lower
semi-continuous function $G:{\mathcal{X}}\to(-\infty,+\infty]$ and denote
$D(G)=\\{x\in{\mathcal{X}},\;G(x)<\infty\\}$ its domain. We assume that $G$ is
convex, i.e., for every $x,z\in{\mathcal{X}}$ and for every
$\varepsilon\in[0,1]$, we have:

$G(\varepsilon z+(1-\varepsilon)x)\leq\varepsilon G(z)+(1-\varepsilon)G(x).$
(5)

Given $x\in{\mathcal{X}}$, $y\in{\mathcal{X}}$ is a subgradient of $G$ at $x$
if for every $z\in{\mathcal{X}}$,

$G(x)+\langle y,z-x\rangle\leq G(z).$

The (possibly empty) set of subgradients of $G$ at $x$ is denoted $\partial
G(x)$, and the map $x\mapsto\partial G(x)$ is called the subdifferential. If
$G$ is differentiable at $x$, then $\partial G(x)=\\{\nabla G(x)\\}$ where
$\nabla G(x)$ is the gradient of $G$ at $x$. The subdifferential of the convex
function $G$ allows to define the gradient flow of $G$: for every initial
condition ${\mathsf{x}}(0)=a$ such that $\partial G(a)\neq\emptyset$, there
exists an unique absolutely continuous function
${\mathsf{x}}:[0,+\infty)\to{\mathcal{X}}$ solution to the differential
inclusion [9, Th. 3.1], [26, Th. 2.7]:

${\mathsf{x}}^{\prime}(t)\in\partial G({\mathsf{x}}(t)).$ (6)

One can check that the gradient flow of $G$ is also characterized by the
following system of Evolution Variational Inequalities (EVI) :

$\forall z\in
D(G),\quad\frac{d}{dt}\|{\mathsf{x}}(t)-z\|^{2}\leq-2\left(G({\mathsf{x}}(t))-G(z)\right).$

In contrast to (6), the former characterization allows to define the gradient
flow without using the notion of subdifferential, a property that can be
practical in nonsmooth settings. Moreover, the non-asymptotic analysis of
discretized gradient flows in the optimization literature often relies on
discrete versions of the EVI.

The existence of Gradient Flows can be established as the limit of a proximal
scheme [26, Th. 2.14], [6, Th. 5.1] when the step-size $\gamma\to 0$. Defining
the proximity operator of $G$ as:

$\operatorname{prox}_{\gamma
G}(x):=\operatorname*{arg\,min}_{y\in{\mathcal{X}}}G(y)+\frac{1}{2\gamma}\|y-x\|^{2},$
(7)

the proximal scheme is written

$x_{n+1}=\operatorname{prox}_{\gamma G}(x_{n}),$ (8)

which corresponds to the proximal point algorithm to minimize the function
$G$, see [21]. The proximal scheme can be seen as a Backward Euler
discretization of the gradient flow. Indeed, writing the first order
conditions of (8), we have

$x_{n+1}\in x_{n}-\gamma\partial G(x_{n+1}),\quad\text{ or equivalently
}\quad\frac{x_{n+1}-x_{n}}{\gamma}\in-\partial G(x_{n+1}).$

Hence, each iteration of the proximal scheme requires solving an equation
which can be intractable in many cases. The Forward Euler scheme is a more
tractable integrator of the gradient flow of $G$, but is less stable and
requires the differentiability of $G$. Under this assumption, this scheme is
written

$\frac{x_{n+1}-x_{n}}{\gamma}=-\nabla G(x_{n})\quad\text{ or equivalently
}\quad x_{n+1}=x_{n}-\gamma\nabla G(x_{n}),$ (9)

which corresponds to the well-known gradient descent algorithm to minimize the
function $G$. Consider now the case where the function $G$ can be decomposed
as $G=F+H$, where $F$ is convex and smooth and $H$ is convex and nonsmooth. To
integrate the gradient flow of $G=F+H$, another approach is to use the Forward
and the Backward Euler schemes for the smooth term and nonsmooth term
respectively [26]. This approach is also motivated by the fact that in many
situations, the function $H$ is simple enough to implement its proximity
operator $\operatorname{prox}_{\gamma H}$. If $G=F+H$, the Forward Backward
Euler scheme is written

$\frac{x_{n+1}-x_{n}}{\gamma}\in-\nabla F(x_{n})-\partial H(x_{n+1}).$ (10)

Recalling the definition of the proximity operator, this scheme can be
rewritten as

$x_{n+1}=\operatorname{prox}_{\gamma H}(x_{n}-\gamma\nabla F(x_{n})),$ (11)

which corresponds to the proximal gradient algorithm to minimize the composite
function $G$.

####  2.3.2 In the Wasserstein space

Consider a proper lower semi continuous functional
${\mathcal{G}}:{\mathcal{P}}_{2}({\mathcal{X}})\to(-\infty,+\infty]$ and
denote
$D({\mathcal{G}})=\\{\mu\in{\mathcal{P}}_{2}({\mathcal{X}}),\;{\mathcal{G}}(\mu)<\infty\\}$
its domain. We assume that ${\mathcal{G}}$ is convex along generalized
geodesics defined by the 2-Wasserstein distance [2, Chap. 9], i.e. for every
$\mu,\pi\in{\mathcal{P}}_{2}({\mathcal{X}}),\nu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$
and for every $\varepsilon\in[0,1]$,

${\mathcal{G}}((\varepsilon
T_{\nu}^{\pi}+(1-\varepsilon)T_{\nu}^{\mu})_{\\#}\nu)\leq\varepsilon{\mathcal{G}}(\pi)+(1-\varepsilon){\mathcal{G}}(\mu).$
(12)

where $T_{\nu}^{\pi}$ and $T_{\nu}^{\mu}$ are the optimal transport maps from
$\nu$ to $\pi$ and from $\nu$ to $\mu$ respectively. Given
$\mu\in{\mathcal{P}}_{2}({\mathcal{X}})$, $\xi\in L^{2}(\mu)$ is a strong
Fréchet subgradient of ${\mathcal{G}}$ at $\mu$ [2, Chap. 10] if for every
$\phi\in L^{2}(\mu)$,

${\mathcal{G}}(\mu)+\varepsilon\langle\xi,\phi\rangle_{\mu}+o(\varepsilon)\leq{\mathcal{G}}((I+\varepsilon\phi)_{\\#}\mu).$

The (possibly empty) set of strong Fréchet subgradients of ${\mathcal{G}}$ at
$\mu$ is denoted $\partial{\mathcal{G}}(\mu)$. The map
$\mu\mapsto\partial{\mathcal{G}}(\mu)$ is called the strong Fréchet
subdifferential. Conveniently, the strong Fréchet subdifferential enables to
define the gradient flow of the functional ${\mathcal{G}}$ [2, Chap. 11].
However in the nonsmooth setting that will be considered in this paper, the
characterization of gradient flows through EVI will be more practical. The
gradient flow of ${\mathcal{G}}$ is the solution of the following system of
Evolution Variational Inequalities (EVI) [2, Th. 11.1.4]:

$\forall\pi\in
D({\mathcal{G}}),\quad\frac{d}{dt}W^{2}(\mu(t),\pi)\leq-2\left({\mathcal{G}}(\mu(t))-{\mathcal{G}}(\pi)\right).$

We shall perform a non-asymptotic analysis of a discretized gradient flow
scheme to minimize the functional ${\mathcal{G}}$. Our approach is to prove a
discrete EVI for this scheme.

The existence of gradient flows can be established as the limit of a
minimizing movement scheme [2, Th. 11.2.1], [17]. Defining the JKO operator of
${\mathcal{G}}$ as:

$\operatorname{JKO}_{\gamma{\mathcal{G}}}(\mu):=\operatorname*{arg\,min}_{\nu\in{\mathcal{P}}_{2}({\mathcal{X}})}{\mathcal{G}}(\nu)+\frac{1}{2\gamma}W^{2}(\nu,\mu),$
(13)

the JKO scheme is written

$\mu_{n+1}\in\operatorname{JKO}_{\gamma{\mathcal{G}}}(\mu_{n}).$

The JKO operator can be seen as a proximity operator by replacing the
Wasserstein distance by the Euclidean distance. Moreover, the JKO scheme can
be seen as a Backward Euler discretization of the gradient flow. More
precisely, under some assumptions on the functional ${\mathcal{G}}$, using [2,
Lemma 10.1.2]

$\frac{T_{\mu_{n+1}}^{\mu_{n}}-I}{\gamma}\in\partial{\mathcal{G}}(\mu_{n+1}).$

Using Brenier’s theorem, since $T_{\mu_{n+1}}^{\mu_{n}}\circ
T_{\mu_{n}}^{\mu_{n+1}}=I$ $\mu_{n}$-a.e., there exists a strong Fréchet
subgradient of ${\mathcal{G}}$ at $\mu_{n+1}$ denoted
$\nabla_{W}{\mathcal{G}}(\mu_{n+1})$ such that

$\mu_{n+1}=\left(I-\gamma\nabla_{W}{\mathcal{G}}(\mu_{n+1})\circ
T_{\mu_{n}}^{\mu_{n+1}}\right)_{\\#}\mu_{n}.$

Each iteration of the JKO scheme thus requires the minimization of a function
which can be intractable in many cases. As previously, the Forward Euler
scheme is more tractable and enjoys additionally a simpler geometrical
interpretation. Assume
$\partial{\mathcal{G}}(\mu)=\\{\nabla{\mathcal{G}}(\mu)\\}$ is a singleton for
any $\mu\in D({\mathcal{G}})$ (some examples are given [2, Sec. 10.4]). The
Forward Euler scheme for the gradient flow of ${\mathcal{G}}$ is written:

$\mu_{n+1}=(I-\gamma\nabla{\mathcal{G}}(\mu_{n}))_{\\#}\mu_{n},$ (14)

and corresponds to the iterations of the gradient descent algorithm over the
Wasserstein space to minimize ${\mathcal{G}}$. Although the Wasserstein space
is not a Riemannian manifold, it can still be equipped with a Riemannian
structure and interpretation [22, 25]. In particular, the Forward Euler scheme
can be seen as a Riemannian gradient descent where the exponential map at
$\mu$ is the map $\phi\mapsto(I+\phi)_{\\#}\mu$ defined on $L^{2}(\mu)$.

##  3 The Forward Backward Euler scheme

Recall that our goal is to minimize ${\mathcal{E}}_{F}+{\mathcal{H}}$, where
${\mathcal{E}}_{F}(\mu)=\int F(x)d\mu(x)$ for any
$\mu\in{\mathcal{P}}_{2}({\mathcal{X}})$, and ${\mathcal{H}}$ is a nonsmooth
functional. Throughout this paper, we assume the following on the potential
function $F$: there exists $L,\lambda\geq 0$ such that

  * •

A1. $F$ is $L$-smooth i.e. $F$ is differentiable and $\nabla F$ is
$L$-Lipschitz continuous; for all $(x,y)\in{\mathcal{X}}^{2}$:

$F(y)\leq F(x)+\langle\nabla F(x),y-x\rangle+\frac{L}{2}\|x-y\|^{2}.$ (15)

  * •

A2. $F$ is $\lambda$-strongly convex (we allow $\lambda=0$); for all
$(x,y)\in{\mathcal{X}}^{2}$:

$F(x)\leq F(y)-\langle\nabla F(x),y-x\rangle-\frac{\lambda}{2}\|x-y\|^{2}.$
(16)

Moreover, we assume the following on the function ${\mathcal{H}}$.

  * •

B1. ${\mathcal{H}}:{\mathcal{P}}_{2}({\mathcal{X}})\to(-\infty,+\infty]$ is
proper and lower semicontinuous. Moreover,
$D({\mathcal{H}})\subset{\mathcal{P}}_{2}^{r}({\mathcal{X}})$.

  * •

B2. There exists $\gamma_{0}>0$ such that $\forall\gamma\in(0,\gamma_{0})$,
$\operatorname{JKO}_{\gamma{\mathcal{H}}}(\mu)\neq\emptyset$ for every
$\mu\in{\mathcal{P}}_{2}({\mathcal{X}})$.

  * •

B3. ${\mathcal{H}}$ is convex along generalized geodesics.

Assumptions B1, B2, B3 are satisfied by examples given in the introduction,
such as internal or interaction energies see [2, Section 9 and 10]. Moreover,
${\mathcal{H}}$ is often nonsmooth, see e.g. [34, Section 3.1.1]. Therefore,
we use a Forward Backward Euler scheme to integrate the gradient flow of
${\mathcal{G}}$. Let $\gamma>0$ a step size. The proposed Forward Backward
Euler scheme is written, for $n\geq 0$:

$\displaystyle\nu_{n+1}$ $\displaystyle=(I-\gamma\nabla F)_{\\#}\mu_{n}$ (17)
$\displaystyle\mu_{n+1}$
$\displaystyle\in\operatorname{JKO}_{\gamma{\mathcal{H}}}(\nu_{n+1}).$ (18)

This scheme can be seen as a proximal gradient algorithm over the Wasserstein
space to minimize the composite function
${\mathcal{G}}={\mathcal{E}}_{F}+{\mathcal{H}}$.

######  Remark 1.

The FB scheme can be implemented since there exist subroutines to compute JKOs
as well documented in the optimal transport and PDE literature (see the review
of different strategies in [29]). The situation is similar to proximal
splitting algorithms in optimization, where algorithms rely on the computation
of the proximity operator. Moreover, many proximity operators admit a close
form222see www.proximity-operator.net or efficient subroutines. Therefore, for
some simple regularizer ${\mathcal{H}}$, one can hope to be able to compute
$\operatorname{JKO}_{\gamma{\mathcal{H}}}$. For instance, when ${\mathcal{H}}$
is the negative entropy (defined by
${\mathcal{H}}(\mu)=\int\log\left(\mu(x)\right)d\mu(x)$ if $\mu\ll Leb$ with
density $\mu$ and ${\mathcal{H}}(\mu):=+\infty$ else), we conjecture that a
technique similar to [16] can be applied. For this particular functional, the
JKO is known in closed form in the Gaussian case [34]. Other works of interest
for the proposed FB scheme have investigated efficient methods to implement
the JKO of a generic functional with respect to the entropy-regularized
Wasserstein distance [27].

In the next section, we study the non asymptotic properties of the FB scheme.

##  4 Non asymptotic analysis

We consider a fixed step size $\gamma<1/L$ and a probability distribution
$\pi\in{\mathcal{P}}_{2}({\mathcal{X}})$. Our main result (Proposition 8)
combines several ingredients: the identification of the optimal transport maps
between $\mu_{n},\nu_{n+1}$ and $\mu_{n+1}$ (see Equations (17) and (18)), the
proof of a generic lemma regarding generalized geodesic convexity (Lemma 4)
and a proof of a discrete EVI for our proximal gradient algorithm, relying on
the generalized geodesic convexity of ${\mathcal{H}}$.

###  4.1 Identification of optimal transport maps

Lemmas 2,3 identify the optimal transport maps from $\mu_{n}$ to $\nu_{n+1}$
and from $\nu_{n+1}$ to $\mu_{n+1}$ in the Forward Backward Euler scheme, as
soon as the step size is sufficiently small. In particular, Lemma 3 is a
consequence of [2, Lemma 10.1.2].

######  Lemma 2.

Assume A1, A2. Let $\mu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$ and
$\nu=(I-\gamma\nabla F)_{\\#}\mu$. Then if $\gamma<1/L$, the optimal transport
map from $\mu$ to $\nu$ corresponds to

$T_{\mu}^{\nu}=I-\gamma\nabla F.$

Moreover, $\nu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}}).$

######  Lemma 3.

Assume B1, B2. Let $\nu\in{\mathcal{P}}_{2}({\mathcal{X}})$. If
$\mu\in\operatorname{JKO}_{\gamma{\mathcal{H}}}(\nu)$, then $\mu\in
D({\mathcal{H}})\subset{\mathcal{P}}_{2}^{r}({\mathcal{X}})$ and the optimal
transport map from $\mu$ to $\nu$ satisfies $T_{\mu}^{\nu}\in
I+\gamma\partial{\mathcal{H}}(\mu)$. In other words, there exists a strong
Fréchet subgradient at $\mu$ denoted $\nabla_{W}{\mathcal{H}}(\mu)$ such that

$T_{\mu}^{\nu}=I+\gamma\nabla_{W}{\mathcal{H}}(\mu).$ (19)

Using Lemmas 2,3, if $\mu_{0}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, then
$\mu_{n},\nu_{n}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$ for every $n$ by
induction. This remark allows to consider optimal transport maps from
$\mu_{n}$ and $\nu_{n}$ to any $\pi\in{\mathcal{P}}_{2}({\mathcal{X}})$. The
next lemma extends [2, 10.1.1.B] for generalized geodesically convex
functionals ${\mathcal{H}}$.

######  Lemma 4.

Assume B1, B2, B3. Let $\nu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$,
$\mu,\pi\in{\mathcal{P}}_{2}({\mathcal{X}})$ and $T_{\nu}^{\mu},T_{\nu}^{\pi}$
the optimal transport maps from $\nu$ to $\mu$ and from $\nu$ to $\pi$
respectively. If $\xi\in\partial{\mathcal{H}}(\mu)$, then

$\langle\xi\circ
T_{\nu}^{\mu},T_{\nu}^{\pi}-T_{\nu}^{\mu}\rangle_{\nu}\leq{\mathcal{H}}(\pi)-{\mathcal{H}}(\mu).$
(20)

Lemma 4 is natural, holds for any functional convex along generalized
geodesics, and was not known to our knowledge. The following results rely on
this lemma.

###  4.2 A descent lemma

Without using any convexity assumption on $F$, we first obtain a descent
lemma. We denote $Y_{n+1}:=T_{\nu_{n+1}}^{\mu_{n+1}}$ the optimal transport
map between $\nu_{n+1}$ and $\mu_{n+1}$ in the Forward Backward Euler scheme
(17), (18), and $X_{n+1}:=Y_{n+1}\circ(I-\gamma\nabla F)$.

######  Theorem 5 (Descent).

Assume $\mu_{0}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, $\gamma<1/L$ and A1,
B1, B2, B3. Then for $n\geq 0$, there exists a strong Fréchet subgradient at
$\mu_{n+1}$ denoted $\nabla_{W}{\mathcal{H}}(\mu_{n+1})$ such that:

${\mathcal{G}}(\mu_{n+1})\leq{\mathcal{G}}(\mu_{n})-\gamma\left(1-\frac{L\gamma}{2}\right)\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2},$

where we use the notation $\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})$ to
denote $\nabla_{W}{\mathcal{H}}(\mu_{n+1})\circ X_{n+1}$.

Hence, the sequence $({\mathcal{G}}(\mu_{n}))_{n}$ is decreasing as soon as
the step-size is small enough.

###  4.3 Discrete EVI

To prove a discrete EVI and obtain convergence rates, we need the additional
convexity assumption A2 on the potential function $F$. We firstly prove the
two following lemmas.

######  Lemma 6.

Assume $\mu_{0}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, $\gamma<1/L$, and B1,
B2, B3. Then for $n\geq 0$ and $\pi\in{\mathcal{P}}_{2}({\mathcal{X}})$, there
exists a strong Fréchet subgradient at $\mu_{n+1}$ denoted
$\nabla_{W}{\mathcal{H}}(\mu_{n+1})$ such that:

$W^{2}(\mu_{n+1},\pi)\leq
W^{2}(\nu_{n+1},\pi)-2\gamma\left({\mathcal{H}}(\mu_{n+1})-{\mathcal{H}}(\pi)\right)-\gamma^{2}\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})\|_{\mu_{n+1}}^{2}.$

######  Lemma 7.

Assume $\mu_{0}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, $\gamma\leq 1/L$, and
A1, A2 with $\lambda\geq 0$. Then for $n\geq 0$, and
$\pi\in{\mathcal{P}}_{2}({\mathcal{X}})$

$W^{2}(\nu_{n+1},\pi)\leq(1-\gamma\lambda)W^{2}(\mu_{n},\pi)-2\gamma\left({\mathcal{E}}_{F}(\mu_{n})-{\mathcal{E}}_{F}(\pi)\right)+\gamma^{2}\|\nabla
F\|_{\mu_{n}}^{2}.$

We can now provide a discrete EVI for the functional
${\mathcal{G}}={\mathcal{E}}_{F}+{\mathcal{H}}$.

######  Proposition 8 (discrete EVI).

Assume $\mu_{0}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, $\gamma<1/L$, and
A1–B3 with $\lambda\geq 0$. Then for $n\geq 0$ and
$\pi\in{\mathcal{P}}_{2}({\mathcal{X}})$, there exists a strong Fréchet
subgradient at $\mu_{n+1}$ denoted $\nabla_{W}{\mathcal{H}}(\mu_{n+1})$ such
that the Forward Backward Euler scheme verifies:

$W^{2}(\mu_{n+1},\pi)\leq(1-\gamma\lambda)W^{2}(\mu_{n},\pi)-2\gamma\left({\mathcal{G}}(\mu_{n+1})-{\mathcal{G}}(\pi)\right).$
(21)

###  4.4 Convergence rates

When the potential function $F$ is convex, we easily get rates from the
discrete EVI inequality provided above. Theorem 9 is a direct consequence of
Proposition 8 by taking $\pi=\mu_{\star}$, and its corollaries provide rates
depending on the strong convexity parameter of $F$.

######  Theorem 9.

Assume $\mu_{0}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, $\gamma<1/L,$ and
A1–B3 with $\lambda\geq 0$. Then for every $n\geq 0$,

$W^{2}(\mu_{n+1},\mu_{\star})\leq(1-\gamma\lambda)W^{2}(\mu_{n},\mu_{\star})-2\gamma({\mathcal{G}}(\mu_{n+1})-{\mathcal{G}}(\mu_{\star})).$

######  Corollary 10 (Convex case rate).

Under the assumptions of Theorem 9, for $n\geq 0$:

${\mathcal{G}}(\mu_{n})-{\mathcal{G}}(\mu_{\star})\leq\frac{W^{2}(\mu_{0},\mu_{\star})}{2\gamma
n}.$

######  Corollary 11 (Strongly convex case rate).

Under the assumptions of Theorem 9, if $\lambda>0$, then for $n\geq 0$:

$W^{2}(\mu_{n},\mu_{\star})\leq(1-\gamma\lambda)^{n}W^{2}(\mu_{0},\mu_{\star}).$

Hence, as soon as $F$ is convex, we get sublinear rates in terms of the
objective function ${\mathcal{G}}$, while when $F$ is $\lambda$-strongly
convex with $\lambda>0$, we get linear rates in the Wasserstein distance for
the iterates of the Forward Backward Euler scheme. The rates match the rates
of the proximal gradient algorithm in Hilbert space in the convex and strongly
convex cases. [24]. Moreover, these rates are discrete time analogues of the
continuous time rates obtained in [2, Th. 11.2.1] for the gradient flow of
${\mathcal{G}}$.

##  5 Numerical experiments

We provide numerical experiments with a ground truth target distribution
$\mu_{\star}$ to illustrate the dynamical behavior of the FB scheme, similarly
to [30, Section 4.1]. We consider $F(x)=0.5|x|^{2}$, and ${\mathcal{H}}$ the
negative entropy. In this case, ${\mathcal{G}}(\mu)$ is (up to an additive
constant) the Kullback-Leibler divergence w.r.t. the standard Gaussian
distribution $\mu_{\star}$. We denote by $m_{\star}$ the mean and
$\sigma_{\star}$ the variance, and fix $m_{\star}=0$ and
$\sigma_{\star}^{2}=1.0$. We use the closed-form particle implementation of
the FB scheme [34, Section G.1]. This allows to show the dynamical behavior of
the FB scheme when $\gamma=0.1$, and $\mu_{0}$ is Gaussian with $m_{0}=10$ and
$\sigma_{0}=100$, in Figure 1. Note that $\lambda=1.0$.

$\begin{array}[]{cc}\includegraphics[width=173.44534pt]{Simulations/Histogram_t_0.png}&\includegraphics[width=173.44534pt]{Simulations/Particles}\\\
\includegraphics[width=173.44534pt]{Simulations/Histogram_t_10.png}&\includegraphics[width=173.44534pt]{Simulations/Histogram_t_100.png}\end{array}$
Figure 1: The particle implementation of the FB scheme illustrate the
convergence of $\mu_{n}$ to $\mu_{\star}$.

More precisely, the position of a set of particles initially distributed
according to $\mu_{0}$ is updated iteratively. The empirical distribution of
the particles, represented by histograms, approximate $\mu_{n}$. We see that
$\mu_{n}$ matches $\mu_{\star}$ after few iterations on Figure 1. The linear
convergence of $\mu_{n}$ to $\mu_{\star}$ in Wasserstein distance is
illustrated in a multidimensional case in Figure 2 (Supplementary material).

##  6 Broader impact

We proposed an unified analysis of a Forward-Backward discretization of
Wasserstein gradient flows considered in the machine learning litterature.
Although the proposed discretization still makes some approximations compared
to algorithms implemented in practice, discrete time dynamics are closer to
real world applications than continuous time dynamics usually considered.

Note that the implementation of the JKO operators is independent from the
analysis of the FB scheme. Presenting this paper at Neurips could motivate
research on the JKO of specific functionals ${\mathcal{H}}$ relevant in
machine learning problems. Such research might lead to interesting results
since many proximity operators admit closed form formulas in Euclidean spaces,
and many works in applied optimal transport deal with minimization problems
similar to JKOs. The results that we showed, together with efficient
implementations of some specific JKOs, would be very impactful for many
machine learning tasks to which we refer in the introduction.

## References

  * [1] Martial Agueh and Guillaume Carlier.  Barycenters in the wasserstein space.  SIAM Journal on Mathematical Analysis, 43(2):904–924, 2011.
  * [2] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré.  Gradient flows: in metric spaces and in the space of probability measures.  Springer Science & Business Media, 2008.
  * [3] Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton.  Maximum mean discrepancy gradient flow.  In Advances in Neural Information Processing Systems, pages 6481–6491, 2019.
  * [4] Heinz H Bauschke, Patrick L Combettes, et al.  Convex analysis and monotone operator theory in Hilbert spaces, volume 408.  Springer, 2011.
  * [5] Espen Bernton.  Langevin monte carlo and jko splitting.  arXiv preprint arXiv:1802.08671, 2018.
  * [6] Pascal Bianchi, Walid Hachem, and Adil Salim.  A constant step Forward-Backward algorithm involving random maximal monotone operators.  Journal of Convex Analysis, 2019.
  * [7] Adrien Blanchet and Jérôme Bolte.  A family of functional inequalities: Łojasiewicz inequalities and displacement convex functions.  Journal of Functional Analysis, 275(7):1650–1673, 2018.
  * [8] Yann Brenier.  Polar factorization and monotone rearrangement of vector-valued functions.  Communications on Pure and Applied Mathematics, 44(4):375–417, 1991\.
  * [9] Haïm Brézis.  Opérateurs maximaux monotones et semi-groupes de contractions dans les espaces de Hilbert.  North-Holland mathematics studies. Elsevier Science, Burlington, MA, 1973\.
  * [10] Xiang Cheng and Peter Bartlett.  Convergence of langevin mcmc in kl-divergence.  arXiv preprint arXiv:1705.09048, 2017.
  * [11] Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J Stromme.  Gradient descent algorithms for bures-wasserstein barycenters.  arXiv preprint arXiv:2001.01700, 2020.
  * [12] Lenaic Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  In Advances in Neural Information Processing Systems, pages 3036–3046, 2018.
  * [13] Andrew Duncan, Nikolas Nüsken, and Lukasz Szpruch.  On the geometry of stein variational gradient descent.  arXiv preprint arXiv:1912.00894, 2019.
  * [14] Alain Durmus, Szymon Majewski, and Blazej Miasojedow.  Analysis of langevin monte carlo via convex optimization.  Journal of Machine Learning Research, 20(73):1–46, 2019.
  * [15] Charlie Frogner and Tomaso Poggio.  Approximate inference with wasserstein gradient flows.  arXiv preprint arXiv:1806.04542, 2018.
  * [16] Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach.  Stochastic optimization for large-scale optimal transport.  In Advances in Neural Information Processing Systems, pages 3440–3448, 2016.
  * [17] Richard Jordan, David Kinderlehrer, and Felix Otto.  The variational formulation of the fokker–planck equation.  SIAM Journal on Mathematical Analysis, 29(1):1–17, 1998.
  * [18] Qiang Liu.  Stein variational gradient descent as gradient flow.  In Advances in Neural Information Processing Systems, pages 3115–3123, 2017.
  * [19] Qiang Liu and Dilin Wang.  Stein variational gradient descent: A general purpose bayesian inference algorithm.  In Advances in Neural Information Processing Systems, pages 2378–2386, 2016.
  * [20] Yi-An Ma, Niladri Chatterji, Xiang Cheng, Nicolas Flammarion, Peter Bartlett, and Michael I Jordan.  Is there an analog of nesterov acceleration for mcmc?  arXiv preprint arXiv:1902.00996, 2019.
  * [21] Bernard Martinet.  Brève communication. régularisation d’inéquations variationnelles par approximations successives.  Revue française d’informatique et de recherche opérationnelle. Série rouge, 4(R3):154–158, 1970.
  * [22] Robert J McCann.  Polar factorization of maps on riemannian manifolds.  Geometric & Functional Analysis GAFA, 11(3):589–608, 2001.
  * [23] Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit.  arXiv preprint arXiv:1902.06015, 2019.
  * [24] Yurii Nesterov.  Lectures on convex optimization, volume 137.  Springer, 2018.
  * [25] Felix Otto.  The geometry of dissipative evolution equations: the porous medium equation.  Communications in Partial Differential Equations, 26(1-2):101–174, 2001.
  * [26] Juan Peypouquet and Sylvain Sorin.  Evolution equations for maximal monotone operators: asymptotic analysis in continuous and discrete time.  arXiv preprint arXiv:0905.1270, 2009.
  * [27] Gabriel Peyré.  Entropic approximation of wasserstein gradient flows.  SIAM Journal on Imaging Sciences, 8(4):2323–2351, 2015.
  * [28] Pierre H Richemond and Brendan Maginnis.  On wasserstein reinforcement learning and the fokker-planck equation.  arXiv preprint arXiv:1712.07185, 2017.
  * [29] Filippo Santambrogio.  $\\{$Euclidean, metric, and Wasserstein$\\}$ gradient flows: an overview.  Bulletin of Mathematical Sciences, 7(1):87–154, 2017.
  * [30] Amirhossein Taghvaei and Prashant G Mehta.  Accelerated gradient flow for probability distributions.  In International Conference on Machine Learning, 2019.
  * [31] Santosh Vempala and Andre Wibisono.  Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices.  In Advances in Neural Information Processing Systems, pages 8092–8104, 2019.
  * [32] Cédric Villani.  Optimal transport: old and new, volume 338.  Springer Science & Business Media, 2008.
  * [33] Yifei Wang and Wuchen Li.  Accelerated information gradient flow.  arXiv preprint arXiv:1909.02102, 2019.
  * [34] Andre Wibisono.  Sampling as optimization in the space of measures: The langevin dynamics as a composite optimization problem.  arXiv preprint arXiv:1802.08089, 2018.
  * [35] Andre Wibisono.  Proximal langevin algorithm: Rapid convergence under isoperimetry.  arXiv preprint arXiv:1911.01469, 2019.
  * [36] Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin.  Policy optimization as wasserstein gradient flows.  arXiv preprint arXiv:1808.03030, 2018.

##  Appendix A Further numerical experiment

Minimization problems involving high dimensional Gaussian distributions have
several applications. We shall illustrate the linear convergence of the FB
scheme on such a problem, see also [11].

We consider a multidimensional extension of the simulation of Section 5. More
precisely, $F(x)=0.5\|x\|^{2}$ and, using the notations of Section 5, the
inital distribution is $\mu_{0}^{\otimes d}$ and the target distribution is
$\mu_{\star}^{\otimes d}$, where $\otimes$ denotes the product of measures.

The simulation can no longer be represented with histograms and particles,
however, we represent the linear convergence of $\mu_{n}$ to $\mu_{\star}$
predicted by Corollary 11 in Figure 2.

$\begin{array}[]{c}\includegraphics[width=433.62pt]{Simulations/Wass_d.pdf}\end{array}$
Figure 2: Linear convergence of $W^{2}(\mu_{n},\mu_{\star})$ to $0$ in
dimension $d=1000$.

##  Appendix B Proof of Lemma 2

The map $I-\gamma\nabla F$ is a pushforward from $\mu$ to $\nu$. Moreover,
denoting $u:x\mapsto\frac{1}{2}\|x\|^{2}-\gamma F(x)$, $\nabla
u=I-\gamma\nabla F$.

By elementary algebra, for any $(x,y)\in{\mathcal{X}}^{2}$ we have:

$\frac{1}{2}\|x\|^{2}=\frac{1}{2}\|y\|^{2}-\langle
x,y-x\rangle-\frac{1}{2}\|x-y\|^{2},$ (22)

and from the smoothness of $F$,

$F(y)\leq F(x)+\langle\nabla F(x),y-x\rangle+\frac{L}{2}\|x-y\|^{2}.$ (23)

Therefore, combining (22) and (23) multiplied by $\gamma$ gives:

$\frac{1}{2}\|x\|^{2}-\gamma F(x)\leq\frac{1}{2}\|y\|^{2}-\gamma F(y)-\langle
x-\gamma\nabla F(x),y-x\rangle-\frac{1}{2}(1-L\gamma)\|x-y\|^{2}.$ (24)

In other words,

$u(x)\leq u(y)-\langle\nabla
u(x),y-x\rangle-\frac{1}{2}(1-L\gamma)\|x-y\|^{2}.$ (25)

Therefore, if $\gamma\leq 1/L$, then $u$ is convex and $\nabla
u=T_{\mu}^{\nu}$ using Brenier’s theorem. Moreover, if $\gamma<1/L$ then $u$
is $(1-L\gamma)$-strongly convex. In consequence,

$(1-L\gamma)\|x-y\|^{2}\leq\langle x-y,\nabla u(x)-\nabla u(y)\rangle.$

Therefore, $\nabla u$ is injective. Furthermore, using the strong convexity of
$u$ and [2, Lemma 5.5.3] (see also [2, Th. 6.2.3, Th. 6.2.7]),
$\nu\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$.

##  Appendix C Proof of Lemma 3

Let $\mu\in\operatorname{JKO}_{\gamma{\mathcal{H}}}(\nu)$. Since
$D({\mathcal{H}})\subset{\mathcal{P}}_{2}^{r}({\mathcal{X}})$, [2, Lemma
10.1.2] implies $\mu\in D({\mathcal{H}})$ and
$\frac{1}{\gamma}(T_{\mu}^{\nu}-I)\in\partial{\mathcal{H}}(\nu)$ is a strong
subgradient of ${\mathcal{H}}$ at $\nu$.

##  Appendix D Proof of Lemma 4

Since $\xi\in\partial{\mathcal{H}}(\mu)$, for every $\phi\in L^{2}(\mu)$,

${\mathcal{H}}((I+\varepsilon\phi)_{\\#}\mu)\geq{\mathcal{H}}(\mu)+\varepsilon\langle\xi,\phi\rangle_{\mu}+o(\varepsilon).$

Applying the last inequality to $\phi=T_{\nu}^{\pi}\circ T_{\mu}^{\nu}-I$ and
using the transfer lemma ($\mu={T_{\nu}^{\mu}}_{\\#}\nu$) we have

${\mathcal{H}}((T_{\nu}^{\mu}+\varepsilon(T_{\nu}^{\pi}-T_{\nu}^{\mu}))_{\\#}\nu)={\mathcal{H}}((I+\varepsilon\phi)_{\\#}\mu),$
$\langle\xi,\phi\rangle_{\mu}=\langle\xi(T_{\nu}^{\mu}),T_{\nu}^{\pi}-T_{\nu}^{\mu}\rangle_{\nu},$

and

$\frac{{\mathcal{H}}((T_{\nu}^{\mu}+\varepsilon(T_{\nu}^{\pi}-T_{\nu}^{\mu}))_{\\#}\nu)-{\mathcal{H}}(\mu)}{\varepsilon}\geq\langle\xi(T_{\nu}^{\mu}),T_{\nu}^{\pi}-T_{\nu}^{\mu}\rangle_{\nu}+o(1).$
(26)

Using the generalized geodesic convexity of ${\mathcal{H}}$,

${\mathcal{H}}((T_{\nu}^{\mu}+\varepsilon(T_{\nu}^{\pi}-T_{\nu}^{\mu}))_{\\#}\nu)\leq\varepsilon{\mathcal{H}}(\pi)+(1-\varepsilon){\mathcal{H}}(\mu).$

Plugging the last inequality into (26),

${\mathcal{H}}(\pi)-{\mathcal{H}}(\mu)\geq\langle\xi(T_{\nu}^{\mu}),T_{\nu}^{\pi}-T_{\nu}^{\mu}\rangle_{\nu}+o(1).$
(27)

We get the conclusion by letting $\varepsilon\to 0$.

##  Appendix E Proof of Theorem 5

Denote $Y_{n+1}:=T_{\nu_{n+1}}^{\mu_{n+1}}$ the optimal transport map between
$\nu_{n+1}$ and $\mu_{n+1}$ and $\nabla_{W}{\mathcal{H}}(\mu_{n+1})$ the
strong Fréchet subgradient of ${\mathcal{H}}$ evaluated at $\mu_{n+1}$ defined
by Lemma 3:
$T_{\mu_{n+1}}^{\nu_{n+1}}=I+\gamma\nabla_{W}{\mathcal{H}}(\mu_{n+1})$. Since
$\mu_{n+1},\nu_{n+1}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$,
$(I+\gamma\nabla_{W}{\mathcal{H}}(\mu_{n+1}))\circ Y_{n+1}=I$ using Brenier’s
theorem. We thus have $\nu_{n+1}$-a.e.:

$Y_{n+1}=I-\gamma\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1}).$ (28)

We firstly bound the ${\mathcal{H}}$ term. By taking
$\mu=\mu_{n+1},\pi=\mu_{n}$ and $\nu=\nu_{n+1}$ in Lemma 4, we have:

${\mathcal{H}}(\mu_{n+1})\leq{\mathcal{H}}(\mu_{n})-\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(T_{\nu_{n+1}}^{\mu_{n+1}}),T_{\nu_{n+1}}^{\mu_{n}}-T_{\nu_{n+1}}^{\mu_{n+1}}\rangle_{\nu_{n+1}}.$
(29)

We now identify $T_{\nu_{n+1}}^{\mu_{n}}$ and $T_{\nu_{n+1}}^{\mu_{n+1}}$.
Recall that $Y_{n+1}=T_{\nu_{n+1}}^{\mu_{n+1}}$. Moreover, using Brenier’s
theorem and Lemma 2, $\nu_{n+1}\in{\mathcal{P}}_{2}^{r}({\mathcal{X}})$ and
$T_{\nu_{n+1}}^{\mu_{n}}=(I-\gamma\nabla F)^{-1}$. Therefore,

${\mathcal{H}}(\mu_{n+1})\leq{\mathcal{H}}(\mu_{n})-\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1}),(I-\gamma\nabla
F)^{-1}-Y_{n+1}\rangle_{\nu_{n+1}}.$

Using the transfer lemma, with $Y_{n+1}=X_{n+1}\circ(I-\gamma\nabla F)^{-1}$,
the last inequality is equivalent to

${\mathcal{H}}(\mu_{n+1})\leq{\mathcal{H}}(\mu_{n})-\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1}),I-X_{n+1}\rangle_{\mu_{n}}.$
(30)

Then, we can bound the potential term. Using Equation (28), and
$X_{n+1}=Y_{n+1}\circ(I-\gamma\nabla F)$, we have

$X_{n+1}=I-\gamma\nabla F-\gamma\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1}).$
(31)

Since $F$ is $L$-smooth, we have [4],

$F(z)\leq F(x)+\langle\nabla
F(x),z-x\rangle+\frac{L}{2}\|x-z\|^{2},\qquad\forall\,\,x,z\in{\mathcal{X}}.$
(32)

Replacing $z$ by $X_{n+1}(x)$, we obtain

$\displaystyle F(X_{n+1}(x))\leq F(x)-$ $\displaystyle\gamma\langle\nabla
F(x),\nabla F(x)+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1}(x))\rangle$ (33)
$\displaystyle+$ $\displaystyle\frac{L\gamma^{2}}{2}\|\nabla
F(x)+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1}(x))\|^{2}.$ (34)

Integrating w.r.t. $\mu_{n}$,

${\mathcal{E}}_{F}(\mu_{n+1})\leq{\mathcal{E}}_{F}(\mu_{n})-\gamma\langle\nabla
F,\nabla F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\rangle_{\mu_{n}}\\\
+\frac{L\gamma^{2}}{2}\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}.$ (35)

Then, recalling (31) and summing equations (30) and (35), we get

$\displaystyle{\mathcal{H}}(\mu_{n+1})+{\mathcal{E}}_{F}(\mu_{n+1})\leq$
$\displaystyle{\mathcal{H}}(\mu_{n})+{\mathcal{E}}_{F}(\mu_{n})$
$\displaystyle-\gamma\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1}),\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\rangle_{\mu_{n}}$
$\displaystyle-\gamma\langle\nabla F,\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\rangle_{\mu_{n}}+\frac{L\gamma^{2}}{2}\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}$
$\displaystyle\leq$
$\displaystyle{\mathcal{H}}(\mu_{n})+{\mathcal{E}}_{F}(\mu_{n})-\gamma\left(1-\frac{L\gamma}{2}\right)\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}.$

##  Appendix F Proof of Lemma 6

Recall (28),

$Y_{n+1}=I-\gamma\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1})=T_{\nu_{n+1}}^{\mu_{n+1}}.$

Since $(Y_{n+1},T_{\nu_{n+1}}^{\pi})_{\\#}\nu_{n+1}$ is a coupling between
$\mu_{n+1}$ and $\pi$, we can upper bound the Wasserstein distance between
$\mu_{n+1}$ and $\pi$ as:

$\displaystyle W^{2}($
$\displaystyle\mu_{n+1},\pi)\leq\|Y_{n+1}-T_{\nu_{n+1}}^{\pi}\|_{\nu_{n+1}}^{2}$
$\displaystyle=$
$\displaystyle\|I-T_{\nu_{n+1}}^{\pi}\|_{\nu_{n+1}}^{2}-2\gamma\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1}),I-T_{\nu_{n+1}}^{\pi}\rangle_{\nu_{n+1}}+\gamma^{2}\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1})\|_{\nu_{n+1}}^{2}$
$\displaystyle=$
$\displaystyle\|I-T_{\nu_{n+1}}^{\pi}\|_{\nu_{n+1}}^{2}-2\gamma\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1}),Y_{n+1}-T_{\nu_{n+1}}^{\pi}\rangle_{\nu_{n+1}}-\gamma^{2}\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1})\|_{\nu_{n+1}}^{2}.$
(36)

where $\|I-T_{\nu_{n+1}}^{\pi}\|_{\nu_{n+1}}^{2}=W^{2}(\nu_{n+1},\pi)$.
Moreover, using Lemma 4 with $\mu=\mu_{n+1}$ and $\nu=\nu_{n+1}$,

$-2\gamma\langle\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1}),Y_{n+1}-T_{\nu_{n+1}}^{\pi}\rangle_{\nu_{n+1}}\leq-2\gamma\left({\mathcal{H}}(\mu_{n+1})-{\mathcal{H}}(\pi)\right).$

Plugging the latter inequality into (F), we get the result.

##  Appendix G Proof of Lemma 7

Since $(I-\gamma\nabla F,T_{\mu_{n}}^{\pi})_{\\#}\mu_{n}$ is a coupling
between $\nu_{n+1}$ and $\pi$, we can upper bound the Wasserstein distance
between $\nu_{n+1}$ and $\pi$ as:

$\displaystyle W^{2}(\nu_{n+1},\pi)$ $\displaystyle\leq\|(I-\gamma\nabla
F)-T_{\mu_{n}}^{\pi}\|_{\mu_{n}}^{2}$
$\displaystyle=\|I-T_{\mu_{n}}^{\pi}\|_{\mu_{n}}^{2}-2\gamma\langle\nabla
F,I-T_{\mu_{n}}^{\pi}\rangle_{\mu_{n}}+\gamma^{2}\|\nabla F\|_{\mu_{n}}^{2}.$
(37)

where $\|I-T_{\mu_{n}}^{\pi}\|_{\mu_{n}}^{2}=W^{2}(\mu_{n},\pi)$. Moreover,
since $F$ is $\lambda$-strongly convex, we have:

$F(x)\leq F(z)+\langle\nabla
F(x),x-z\rangle-\frac{\lambda}{2}\|x-z\|^{2},\qquad\forall\,\,x,z\in{\mathcal{X}}.$
(38)

Replacing $z$ by $T_{\mu_{n}}^{\pi}(x)$ and multiplying by $2\gamma$, we
obtain

$-2\gamma\langle\nabla
F(x),x-T_{\mu_{n}}^{\pi}(x)\rangle\leq-2\gamma\left(F(x)-F\circ
T_{\mu_{n}}^{\pi}(x)\right)-\gamma\lambda\|x-T_{\mu_{n}}^{\pi}(x)\|^{2}.$

Integrating w.r.t. $\mu_{n}$ results in

$-2\gamma\langle\nabla
F,I-T_{\mu_{n}}^{\pi}\rangle_{\mu_{n}}\leq-2\gamma\left({\mathcal{E}}_{F}(\mu_{n})-{\mathcal{E}}_{F}(\pi)\right)-\gamma\lambda
W^{2}(\mu_{n},\pi).$

Plugging the latter inequality into (G) gives the result.

##  Appendix H Proof of Proposition 8

Recall that $Y_{n+1}=T_{\nu_{n+1}}^{\mu_{n+1}}$. Combining Lemma 6 and Lemma
7, we firstly get

$\displaystyle W^{2}(\mu_{n+1},\pi)\leq$
$\displaystyle(1-\gamma\lambda)W^{2}(\mu_{n},\pi)-2\gamma\left({\mathcal{E}}_{F}(\mu_{n})+{\mathcal{H}}(\mu_{n+1})-{\mathcal{E}}_{F}(\pi)-{\mathcal{H}}(\pi)\right)$
$\displaystyle\quad+\gamma^{2}\|\nabla
F\|_{\mu_{n}}^{2}-\gamma^{2}\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1})\|_{\nu_{n+1}}^{2}.$
(39)

Multiplying (35) by $2\gamma$,

$\displaystyle-2\gamma{\mathcal{E}}_{F}(\mu_{n})\leq$
$\displaystyle-2\gamma{\mathcal{E}}_{F}(\mu_{n+1})$
$\displaystyle-2\gamma^{2}\langle\nabla
F,\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\rangle_{\mu_{n}}-2\gamma^{2}\|\nabla
F\|_{\mu_{n}}+L\gamma^{3}\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}.$

Moreover, using the transfer lemma,
$\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}=\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1})\|_{\nu_{n+1}}^{2}$.
Therefore,

$\displaystyle-2\gamma$
$\displaystyle{\mathcal{E}}_{F}(\mu_{n})+\gamma^{2}\|\nabla
F\|_{\mu_{n}}^{2}-\gamma^{2}\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(Y_{n+1})\|_{\nu_{n+1}}^{2}$
$\displaystyle\leq$
$\displaystyle-2\gamma{\mathcal{E}}_{F}(\mu_{n+1})-\gamma^{2}\|\nabla
F\|_{\mu_{n}}^{2}-\gamma^{2}\|\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}-2\gamma^{2}\langle\nabla
F,\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\rangle_{\mu_{n}}$
$\displaystyle+L\gamma^{3}\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}$
$\displaystyle\leq$
$\displaystyle-2\gamma{\mathcal{E}}_{F}(\mu_{n+1})-\gamma^{2}(1-L\gamma)\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}.$

Plugging the last inequality into (H),

$\displaystyle W^{2}(\mu_{n+1},\pi)\leq$
$\displaystyle(1-\gamma\lambda)W^{2}(\mu_{n},\pi)-2\gamma\left({\mathcal{E}}_{F}(\mu_{n+1})+{\mathcal{H}}(\mu_{n+1})-{\mathcal{E}}_{F}(\pi)-{\mathcal{H}}(\pi)\right)$
$\displaystyle-\gamma^{2}(1-\gamma L)\|\nabla
F+\nabla_{W}{\mathcal{H}}(\mu_{n+1})(X_{n+1})\|_{\mu_{n}}^{2}.$

##  Appendix I Proof of Corollary 10

Let ${{\mathcal{L}}}_{n}:=2\gamma
n({\mathcal{G}}(\mu_{n})-{\mathcal{G}}(\mu_{\star}))+W^{2}(\mu_{n},\mu_{\star})$.
From Theorem 5,
${\mathcal{G}}(\mu_{n+1})-{\mathcal{G}}(\mu_{\star})\leq{\mathcal{G}}(\mu_{n})-{\mathcal{G}}(\mu_{\star})$
if $\gamma<1/L$. Therefore,

$2\gamma
n({\mathcal{G}}(\mu_{n+1})-{\mathcal{G}}(\mu_{\star}))+2\gamma({\mathcal{G}}(\mu_{n+1})-{\mathcal{G}}(\mu_{\star}))+W^{2}(\mu_{n+1},\mu_{\star})\leq
2\gamma
n({\mathcal{G}}(\mu_{n})-{\mathcal{G}}(\mu_{\star}))+W^{2}(\mu_{n},\mu_{\star}),$

where we used Theorem 9 with $\lambda=0$ (recall that $\lambda\geq 0$). In
other words, ${{\mathcal{L}}}_{n+1}\leq{{\mathcal{L}}}_{n}$. Finally,

$2\gamma
n({\mathcal{G}}(\mu_{n})-{\mathcal{G}}(\mu_{\star}))\leq{{\mathcal{L}}}_{n}\leq{{\mathcal{L}}}_{0}=W^{2}(\mu_{0},\mu_{\star}).$

##  Appendix J Proof of Corollary 11

Since the ${\mathcal{G}}(\mu_{n+1})-{\mathcal{G}}(\mu_{\star})$ is
nonnegative, from Theorem 9,

$W^{2}(\mu_{n+1},\mu_{\star})\leq(1-\gamma\lambda)W^{2}(\mu_{n},\mu_{\star}).$

We get the result by iterating.