# Understanding the training of infinitely deep and wide ResNets
with Conditional Optimal Transport

Raphaël Barboni
ENS – PSL Université
<EMAIL_ADDRESS>Gabriel Peyré
CNRS and ENS – PSL Univ.
<EMAIL_ADDRESS>François-Xavier Vialard
LIGM, Univ. Gustave Eiffel, CNRS
<EMAIL_ADDRESS>
###### Abstract

We study the convergence of gradient flow for the training of deep neural
networks. If Residual Neural Networks are a popular example of very deep
architectures, their training constitutes a challenging optimization problem
due notably to the non-convexity and the non-coercivity of the objective. Yet,
in applications, those tasks are successfully solved by simple optimization
algorithms such as gradient descent. To better understand this phenomenon, we
focus here on a “mean-field” model of infinitely deep and arbitrarily wide
ResNet, parameterized by probability measures over the product set of layers
and parameters and with constant marginal on the set of layers. Indeed, in the
case of shallow neural networks, mean field models have proven to benefit from
simplified loss-landscapes and good theoretical guarantees when trained with
gradient flow for the Wasserstein metric on the set of probability measures.
Motivated by this approach, we propose to train our model with gradient flow
w.r.t. the conditional Optimal Transport distance: a restriction of the
classical Wasserstein distance which enforces our marginal condition. Relying
on the theory of gradient flows in metric spaces we first show the well-
posedness of the gradient flow equation and its consistency with the training
of ResNets at finite width. Performing a local Polyak-Łojasiewicz analysis, we
then show convergence of the gradient flow for well-chosen initializations: if
the number of features is finite but sufficiently large and the risk is
sufficiently small at initialization, the gradient flow converges towards a
global minimizer. This is the first result of this type for infinitely deep
and arbitrarily wide ResNets. In addition, this work is the occasion to study
in more detail the conditional Optimal Transport metric, particularly its
dynamic formulation. Some of our results in this direction might be
interesting on their own.

##  1 Introduction

Understanding the training dynamic of neural networks is an important problem
in Machine Learning as it brings the hope of understanding the good
performances of these models. This training is however an involved
optimization problem, usually solved by performing (stochastic) gradient
descent for the training risk, an optimization procedure which, though simple,
often manages to find a global minimum of the risk despite its non-convexity.
This phenomenon is now correctly understood in some simple cases such as the
one of linear networks [HM16, BHL18, ZLG19, Bah+22]. In the more realistic
case of non-linear architectures, most works have focused on _Multi-Layer
Perceptrons (MLP)_ [LY17, Du+19, ALS19, Zou+20, Lee+19, Che+20, Ngu21] and
convergence towards a minimizer of the risk can be obtained with great
probability over a random initialization provided that the network is
sufficiently wide, a regime referred to as “overparameterization”. Taking the
limit of infinite width, many works have also studied the convergence gradient
descent for the training of neural networks in the limit of an infinite number
of parameters [CB18, MMN18, JMM20, Woj20, NP23]. In those works, the neural
network is trained by modeling the parameters as a probability measure over
the parameter space and performing a Wasserstein gradient flow over the set of
probability measures. Notably, [CB18] establish a result of optimality at
convergence: if the gradient flow converges then its limit is a global
minimizer of the training risk.

We will focus in this work on the case of the _Residual Neural Network
(ResNet)_ architecture. ResNets were first introduced by [He+16] for
applications in computer vision but the architecture has since distinguished
itself by obtaining state-of-the-art results in several other machine-learning
applications. A key feature of ResNets is the extensive use of _skip
connections_ : each layer consists of the addition of a perturbation (called
_residual_) to the output of the previous layer. The presence of skip
connections has indeed been identified to ease the training of deeper neural
networks [RVL12, Sze+17] by mitigating the _vanishing / exploding gradient_
phenomena, a common problem encountered when training deep neural networks
[BSF94, GB10]. The ResNet architecture has thus permitted the training of
neural networks of almost arbitrary depth [He+16a]. Considering the limit
where the depth tends to infinity [Che+18] introduced the _Neural Ordinary
Differential Equation ( $\mathrm{NODE}$)_ architecture: ResNets being seen as
the _explicit Euler_ discretization scheme of an ODE, passing to the limit of
infinite depth leads to a model performing the integration of an ODE with a
parametric velocity field. An important contribution of $\mathrm{NODEs}$ is to
provide a theoretical framework upon which many other works have been based to
study very deep neural network architectures. [Che+18] proposed a method based
on _adjoint sensitivity analysis_ to compute the gradient of $\mathrm{NODEs}$
efficiently without automatic differentiation. [San+21] proposed a new
architecture based on a second-order ODE which can be trained with reduced
computational complexity. Inspired by methods from medical imaging and shape
analysis, [Via+20] proposed a new algorithm for the training of deep ResNets.
[EHL19, EMW21] studied the training and generalization properties of deep
ResNets borrowing tools from the mathematical theory of _Optimal Control_.

##### Notations

For a metric space $X$, $\mathcal{P}(X)$ is the set of Borel probability
measures over $X$. This set is endowed with the _narrow topology_ , which is
the topology of convergence against the set $\mathcal{C}_{b}(X)$ of bounded
continuous functions. For $x\in X$, we note $\delta_{x}\in\mathcal{P}(X)$ the
Dirac measure at $x$. For $p\geq 1$, $\mathcal{P}_{p}(X)$ is the subset of
$\mathcal{P}(X)$ of probability measures with finite $p$-order moment, endowed
with the _Wasserstein distance_ $\mathcal{W}_{p}$ [Vil09, San15]. When $X$ is
a Hilbert space we define on $\mathcal{P}_{p}(X)$ the $p$-Energy
$\mathcal{E}_{p}(\mu)\coloneqq\int_{X}|x|^{p}\mathrm{d}\mu(x)$. If
$\mu\in\mathcal{P}(X)$ and $f:X\mapsto Y$ is a measurable map between
topological spaces we denote by $f_{\\#}\mu\in\mathcal{P}(Y)$ the pushforward
of $\mu$ by $f$. If $\\{f_{i}:X_{i}\to Y_{i}\\}_{1\leq i\leq n}$ is a family
of mappings then $(f_{1},...,f_{n})$ designates the product map
$(f_{1},...,f_{n}):(x_{1},...,x_{n})\mapsto(f_{1}(x_{1}),...,f_{n}(x_{n}))$
and if $X=X^{1}\times...\times X^{n}$ is a product space we designate by
$\pi^{i}$ the projection $\pi^{i}:(x^{1},...,x^{n})\in X\mapsto x^{i}\in
X^{i}$.

###  1.1 Mean-field model of neural network

We consider in this paper neural networks parameterized by probability
measures. Such “mean-field models” can be thought of as neural networks of
arbitrary width and were studied before by [CB18, MMN18, RV18, Woj20, NP23].
Provided with the _parameter space_ $\Omega\subset\mathbb{R}^{p}$, the input
space $\mathbb{R}^{d}$ and with a Borel map
$\phi:\Omega\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ (the _feature map_), we
consider mappings $F_{\mu}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ parameterized by
measures $\mu\in\mathcal{P}(\Omega)$ and defined by:

$\forall\mu\in\mathcal{P}(\Omega),\quad
F_{\mu}:x\in\mathbb{R}^{d}\mapsto\int_{\Omega}\phi(\omega,x)\mathrm{d}\mu(\omega).$
(1)

##### Single hidden layer perceptron

The above definition encompasses as a particular case standard neural network
architectures. For example, for
$\Omega=\mathbb{R}^{d}\times\mathbb{R}^{d}\times\mathbb{R}$ and
$\phi:((u,w,b),x)\mapsto u\sigma(w^{\top}x+b)$ with a real-valued function
$\sigma:\mathbb{R}\to\mathbb{R}$ (called _activation_), considering the atomic
measure $\mu=\frac{1}{m}\sum_{i=1}^{m}\delta_{(u_{i},w_{i},b_{i})}$ one
recovers the classical model of a _Single Hidden Layer (SHL) Perceptron_ of
width $m\geq 0$:

$F_{\mu}:x\mapsto\frac{1}{m}\sum_{i=1}^{m}u_{i}\sigma(w_{i}^{\top}x+b_{i}).$

##### Convolutional Layer

Closer to applications, eq. 1 also encompasses the convolutional architecture
originally used in ResNet for image recognition [He+16]. Consider integers
$n,c,m,k\geq 0$ and $\Omega=\mathbb{R}^{c\times 1\times k\times
k}\times\mathbb{R}^{1\times c\times k\times k}\times\mathbb{R}^{1\times
n\times n}$, the set of parameters of the form $(u,w,b)$ where $u$ and $w$ are
convolutional kernels of size $c\times 1\times k\times k$ and $1\times c\times
k\times k$ respectively and $b$ is a bias term of size $1\times n\times n$.
Then for an image input $x\in\mathbb{R}^{c\times n\times n}$ of size $n\times
n$ with $c$ channels, and $(u,w,b)\in\Omega$ consider the feature map
$\varphi:((u,w,b),x)\mapsto u\star\sigma(w\star x+b)\in\mathbb{R}^{c\times
n\times n}$ where $\star$ denotes the discrete convolution operator and the
activation $\sigma$ is applied component-wise. Then for an empirical measure
$\mu=\frac{1}{m}\sum_{i=1}^{m}\delta_{(u_{i},w_{i},b_{i})}$, eq. 1 gives on
input $x$:

$\displaystyle F_{\mu}(x)=\frac{1}{m}\sum_{i=1}^{m}u_{i}\star\sigma(w_{i}\star
x+b_{i}),$

which is the output of a ResNet residual with $m$ intermediary channels
(without normalization) [He+16]. However, the definition in eq. 1 does not
model some popular architectures such as normalization layers which play an
important role in the success of ResNets.

##### Attention layers

Finally, eq. 1 also models _attention layers_ at the heart of _Transformers_
architectures [Vas+17]. Consider as parameter space
$\Omega=\mathbb{R}^{c\times d}\times\mathbb{R}^{c\times
d}\times\mathbb{R}^{d\times d}$, the set of triplets $(K,Q,V)$ where
$K\in\mathbb{R}^{c\times d}$ is the _key_ matrix, $Q\in\mathbb{R}^{c\times d}$
is the _query_ matrix and $V\in\mathbb{R}^{d\times d}$ is the _value_ matrix.
One can define the feature map $\phi$ as an _attention_ layer with parameter
$(K,Q,V)\in\Omega$. For an input sequence of _tokens_
$\mathbf{x}=(x^{i})_{1\leq i\leq N}\in(\mathbb{R}^{d})^{N}$ of length $N\geq
0$:

$\displaystyle\phi((K,Q,V),\mathbf{x})\coloneqq\mathrm{Attention}((K,Q,V),\mathbf{x})=\left(\sum_{j=1}^{N}\frac{e^{\langle
Kx^{i},Qx^{j}\rangle}}{\sum_{j=1}^{N}e^{\langle
Kx^{i},Qx^{j}\rangle}}Vx^{j}\right)_{1\leq i\leq N}\in(\mathbb{R}^{d})^{N}\,.$

Then for a measure an empirical measure
$\mu=\frac{1}{m}\sum_{k=1}^{m}\delta_{(K_{k},Q_{k},V_{k})}$, eq. 1 defines a
_multi-head attention_ layer with $m$ heads:

$\displaystyle
F_{\mu}(\mathbf{x})=\frac{1}{m}\sum_{k=1}^{m}\mathrm{Attention}((K_{k},Q_{k},V_{k}),\mathbf{x})\,.$

Note that with this definition, we are only able to describe Transformer
architectures taking as input sequences of tokens of fixed length $N$.
However, our setting could be adapted to model Transformer architecture taking
as inputs sequences of tokens of various finite lengths by considering
different feature maps depending on the length of the input sequence.

##### Functional properties

[WW22] study the class of mappings defined by eq. 1 from a functional analysis
viewpoint. More precisely, they study the “Barron space”

$\mathcal{B}\coloneqq\\{f:x\mapsto\int
u\sigma(w^{\top}x+b)\mathrm{d}\mu(u,w,b),\,\mu\in\mathcal{P}_{2}(\mathbb{R}\times\mathbb{R}^{d}\times\mathbb{R})\\}$

in the case $\sigma=\mathrm{ReLU}$ and show $\mathcal{B}$ can be endowed with
a norm

$\displaystyle\forall
f\in\mathcal{B},\quad\|f\|_{\mathcal{B}}\coloneqq\inf\\{\int|u|(\|w\|+|b|)\mathrm{d}\mu,\,\mu\in\mathcal{P}_{2}(\mathbb{R}\times\mathbb{R}^{d}\times\mathbb{R}),\,f=F_{\mu}\\},$

for which it is a Banach space, continuously embedded in the space of
Lipschitz functions. In fact, $\mathcal{B}$ is the smallest Banach space
containing all single hidden layer neural networks of finite width [WW22,
Thm.3.7 and Thm.3.8] and it is thus dense in the set of continuous functions
provided with the compact-open topology [Cyb89]. Note however that the case of
convolutional or attention layers is different as the maps $F_{\mu}$ defined
by eq. 1 are then equivariant by translation of the input images or
permutation of the tokens respectively.

In our case, due to Assumption 1 on $\phi$, we have that if $f=F_{\mu}$ is
given by eq. 1 then $\|f\|_{\mathcal{B}}\leq\mathcal{E}_{2}(\mu)$ and the
second-order moments $\mathcal{E}_{2}(\mu)$ will be a key ingredient in
controlling the training dynamic of our model (_e.g._ by eq. 4). Finally, note
that eq. 1 is not the only way to represent neural networks of arbitrary width
and [WW22] propose various equivalent representations of the functional space
$(\mathcal{B},\|.\|_{\mathcal{B}})$.

###  1.2 Mean-field $\mathrm{NODEs}$

We proceed then to the definition of _Neural ODEs ( $\mathrm{NODEs}$)_
modeling ResNets whose depth tends to infinity with a proper rescaling of the
residual layers [Che+18]. Our $\mathrm{NODE}$ model is an ODE whose velocity
field (or _residual_) belongs to the class of mappings parameterized by
measure defined in eq. 1. Similar models of “mean-field $\mathrm{NODEs}$” or
“mean-field limit of ResNets” were studied by [Lu+20, Din+22, Iso23].

######  Definition 1.1 ($\mathrm{NODE}$).

For a family of probability measures
$\mu=\\{\mu(.|s)\\}_{s\in[0,1]}\in\mathcal{P}(\Omega)^{[0,1]}$ and input
$x\in\mathbb{R}^{d}$ we define the output of the $\mathrm{NODE}$ model as
$\mathrm{NODE}_{\mu}(x)\coloneqq x_{\mu}(1)$ where $(x_{\mu}(s))_{s\in[0,1]}$
satisfies the _Forward ODE_ :

$\frac{\mathrm{d}}{\mathrm{d}s}x_{\mu}(s)=F_{\mu(.|s)}(x_{\mu}(s)),\quad
x_{\mu}(0)=x.$ (2)

When there is no ambiguity, we simply note $x(s)$.

##### The parameter set $\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$

To justify the well-posedness of eq. 2 it is first necessary to define the
adequate set of parameters we will consider. Given a topological space $Z$, we
define $\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times Z)$ as the set of
probability measures $\mu\in\mathcal{P}_{2}([0,1]\times Z)$ whose marginal
w.r.t. $[0,1]$ is the Lebesgue measure $\mathrm{Leb}([0,1])$:

$\displaystyle\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times
Z)\coloneqq\left\\{\mu\in\mathcal{P}_{2}([0,1]\times
Z)\;:\;\pi^{1}_{\\#}\mu=\mathrm{Leb}([0,1])\right\\}.$

Given $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega))$, using a
result of disintegration of measures [ABM14, Thm.4.2.4], there exists a
$\mathrm{d}s$-a.e. uniquely determined family of probability measure
$\mu(.|s)\in\mathcal{P}_{2}(Z)$ such that for every measurable $f:[0,1]\times
Z\to\mathbb{R}$:

$s\in[0,1]\mapsto\int_{Z}f(s,z)\mathrm{d}\mu(z|s)$ is measurable and
$\int_{[0,1]\times
Z}f(s,z)\mathrm{d}\mu(s,z)=\int_{0}^{1}\int_{Z}f(s,z)\mathrm{d}\mu(z|s)\mathrm{d}s$.

In the following, we will consider as parameters probability measures
$\mu\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$. Therefore, every
parameter $\mu\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ is
naturally associated with a (almost everywhere uniquely defined) family of
probability measures $\\{\mu(.|s)\\}_{s\in[0,1]}$. We will provide this set of
parameters with a modification of the Wasserstein-$2$ distance [Vil09, San15]
that takes into account the marginal constraint by considering a restriction
of Kantorovitch’s original optimal coupling problem to the set of couplings
that are the identity on the first variable $s\in[0,1]$. The solution of this
new optimization problem induces the _Conditional Optimal Transport_ (COT)
distance on the parameter set [HHT23].

##### Well-posedness of $\mathrm{NODEs}$

The following assumption on $\phi$ will be sufficient to show the well-
posedness of eq. 2 for any parameter
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. This is the content
of Proposition 1.1.

######  Assumption 1.

Assume $\phi:\Omega\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ is measurable and

  1. 1.

(quadratic growth) grows at most quadratically w.r.t. $\omega$ and linearly
w.r.t. $x$: there exists a constant $C$ s.t.

$\displaystyle\forall
x\in\mathbb{R}^{d},\omega\in\Omega,\quad\|\phi(\omega,x)\|\leq
C(1+\|x\|)(1+\|\omega\|^{2})\,.$

  2. 2.

(local Lipschitzness) is locally Lipschitz with respect to $x$ with a
Lipschitz constant that grows at most quadratically with $\omega$: for every
$R\geq 0$, there exists a constant $C(R)$ s.t.

$\displaystyle\forall x,x^{\prime}\in
B(0,R),\,\forall\omega\in\Omega,\quad\|\phi(\omega,x)-\phi(\omega,x^{\prime})\|\leq
C(R)(1+\|\omega\|^{2})\|x-x^{\prime}\|\,.$

######  Proposition 1.1 (Well-posedness of the flow).

Assume $\mu\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ and $\phi$
satisfies Assumption 1. Then for every $x\in\mathbb{R}^{d}$ there exists a
unique weak solution to eq. 2, that is an absolutely continuous path
$(x(s))_{s\in[0,1]}$ such that for every $s\in[0,1]$:

$\displaystyle x(s)=x+\int_{0}^{s}F_{\mu(.|r)}(x(r))\mathrm{d}r.$ (3)

###### Proof.

The result follows Caratheodory’s theorem for the existence and uniqueness of
absolutely continuous solutions [Hal69, Sec.I.5]. Indeed, given
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ the map
$(s,x)\mapsto F_{\mu(.|s)}(x)$ is measurable w.r.t. $s$ and, thanks to
Assumption 1 (local Lipschitzness), locally Lipschitz w.r.t. $x$ with a
Lipschitz constant that is integrable w.r.t. $s$. Moreover, the solutions of
eq. 3 are defined up to time $s=1$ thanks to the growth assumption in
Assumption 1, and if $C$ is the growth constant we get the following bound on
the solution:

$\displaystyle\forall
s\in[0,1],\quad\|x(s)\|\leq\exp(C(1+\mathcal{E}_{2}(\mu)))(\|x(0)\|+C(1+\mathcal{E}_{2}(\mu))).$
(4)

∎

###  1.3 Supervised learning problem

We consider training our $\mathrm{NODE}$ model for a supervised learning task.
That is, given a data distribution
$\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}\ni(x,y)\sim\mathcal{D}$ and loss
$\ell:\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}\to\mathbb{R}_{+}$, we
associate to a parameterization
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ the risk:

$\displaystyle L(\mu)\coloneqq\
\mathbb{E}_{x,y}\ell(\mathrm{NODE}_{\mu}(x),y)=\mathbb{E}_{x,y}\ell(x_{\mu}(1),y).$
(5)

In the following, we assume the data distribution $\mathcal{D}$ has compact
support and $\ell$ is a smooth loss. In particular, remark that in the case of
an empirical data distribution
$\mathcal{D}=\frac{1}{N}\sum_{i=1}^{N}\delta_{(x_{i},y_{i})}$, we recover the
empirical risk

$\displaystyle
L(\mu)=\frac{1}{N}\sum_{i=1}^{N}\ell(\mathrm{NODE}_{\mu}(x_{i}),y_{i}).$

Training the $\mathrm{NODE}$ model then amounts to find a parameterization
$\mu$ which minimizes the risk defined in eq. 5 _i.e._ to solve the following
_risk minimization_ task:

$\text{Find}\quad\mu^{*}\in\operatorname*{arg\,min}_{\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)}L(\mu)\,.$

In this work, we consider solving the above risk minimization task by
performing gradient flow over the parameter $\mu$. We are interested in
providing sufficient conditions to ensure that such a gradient flow can
“efficiently train” our $\mathrm{NODE}$ _i.e._ finds a minimizer $\mu^{*}$ of
the risk.

###  1.4 Related works and contributions

Due to the popularity and performance of the ResNet architecture, many works
have studied its training dynamic and the convergence of this training dynamic
toward a global minimizer of the risk.

##### ResNets and $\mathrm{NODEs}$ of finite width

[ALS19, Du+19, LZB20] give convergence results for the training of deep neural
networks with gradient descent and their results can be applied to ResNets.
However, their condition for convergence depends on the ResNet’s depth and
their result can therefore not be applied to infinitely deep models of
$\mathrm{NODEs}$. [Mar+23] give local convergence results for the training of
$\mathrm{NODEs}$ based on a local Polyak-Łojasiewicz (P-Ł) condition. They
assume a model of finite width and their result therefore does not hold in the
mean field limit where residuals are of the form eq. 1. They also consider
parameter initializations that are Lipschitz w.r.t. depth which is not
consistent with applications where the parameters are initialized at random,
independently at each layer. As a comparison eq. 1 models residual of both
finite and infinite width and we only assume the family
$\\{\mu(.|s)\\}_{s\in[0,1]}$ is measurable w.r.t. $s\in[0,1]$

#####  $\mathrm{NODEs}$ with linear parameterization of the residuals

Relying on a similar local P-Ł analysis, [BPV22] show a local convergence
result for $\mathrm{NODEs}$ where the residuals are linear functions on a
Hilbert space of arbitrary dimension. For SHL residuals this amounts to only
training the outer weights and keeping the inner weights unchanged. While
preserving the expressivity of the model, one can see that this assumption
does not model the way ResNets are trained in practice. It also simplifies the
analysis by naturally providing the set of residuals with a Hilbert space
structure. In contrast, our model allows for residuals that are not
necessarily linear w.r.t. their parameters and encompasses (among other) the
case of SHL residuals where both the outer and inner weights are trained. This
is done at the cost of considering the parameter set
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ which has no Hilbert space
structure and is a positively curved metric space when provided with the
Conditional OT distance [AGS08, Sec.7.3]. Another notable difference is that
in [BPV22] the risk is shown to admit no spurious critical points (saddles,
local minima, …) whereas in this work, the risk may admit saddles whenever
$\mathrm{D}_{\omega}\phi$ is not surjective. This is for example happening for
the SHL architecture defined by eq. 32 when the feature distribution is not
sufficiently spread.

#####  $\mathrm{NODEs}$ in the mean-field limit

Some works have proposed models for ResNets of infinite depth similar to
Definition 1.1. [EMW21] study properties of the functional space induced by
considering the flow of functions of the form eq. 1 and define a notion of
norm which they use to provide bounds on the Rademacher complexity of this
class of function. [Che+23] also provide bounds on this Rademacher complexity
which they use to prove an upper bound on the generalization error of trained
ResNets.

Closer to our work are the works of [Lu+20, Din+21, Din+22] studying gradient
flow dynamics for the minimization of the risk $L$ for the ResNet model of
Definition 1.1. [Lu+20] consider gradient flows w.r.t. the true Wasserstein
distance on the space of measures. While this point of view motivates a new
training strategy, it is not consistent with the way ResNets are trained in
practice, that is with a layer-wise-$L^{2}$ metric. [Din+21, Din+22] show
existence of gradient flows similar to Definition 3.2. A regularization of the
risk is first assumed for the result to hold [Din+21], an assumption which is
then removed in [Din+22]. Those three works give an optimality result at
convergence: if the parameter converges then its limit is a global minimizer
of the risk, they do not however provide proofs of convergence. This
convergence assumption seems hard to justify _a priori_ as the loss-landscape
of ResNets can have non-compact subsets and cases where the gradient flow
fails to converge have been identified for simple architectures [BHL18].

As a comparison, a key contribution of our work is to provide the parameter
set with the appropriate metric structure allowing us to identify the gradient
flow equation, derived formally by _adjoint sensitivity analysis_ with a
_curve of maximal slope_ of the risk. In this direction, closest to our work
is the recent paper of [Iso23] considering a metric similar to the one of
Proposition 2.1 on the space of $\mathcal{P}_{2}(\Omega)$-valued
$L^{2}$-function and studying the training of the $\mathrm{NODE}$ defined in
Definition 1.1 via gradient flow for this metric. Similarly to our work, the
author shows a (generalization of the) Polyak-Łojasiewicz inequality is
satisfied by the risk and concludes to convergence towards a global minimizer.
Though the work of [Iso23] might be considered concurrent to ours, there are
nonetheless major differences. To obtain a convergence result, [Iso23] makes
strong regularity assumption on the residuals, namely analyticity of the
feature map $\phi$, as well as on the parameter initialization $\mu_{0}$
([Iso23, Assumption 1.9]). Also, the author relies on a notion of _generalized
semi-convexity_ of the risk to define its differential in the space of
parameters and to define the gradient-flow equation. This semi-convexity is
obtained by adding a positive regularization term to the risk whereas we show
convergence for the _unregularized_ risk (eq. 5) which is commonly used in
practice. For this reason, the two convergence results are fundamentally
different: [Iso23] considers a risk for which gradient flow curves stay in a
strongly compact set [Iso23, Prop.5.4] and for which a P-Ł inequality is
satisfied around every first-order critical point [Iso23, Thm.6.1], implying
those are necessarily global minima. In contrast, we consider the minimization
of a risk that may admit non-compact sub-level sets as well as spurious first-
order critical points.

##### ResNets as a discretization of $\mathrm{NODEs}$

While it is not addressed in the present work, an interesting question is the
one of the consistency of the $\mathrm{NODE}$ model with ResNets of finite
depth. [Mar+23] shows the convergence of ResNets of finite width towards
$\mathrm{NODEs}$, at initialization and during training, when the depth tends
to infinity. This convergence is uniform over finite training time intervals
but can be made uniform over the whole training dynamic under a convergence
condition. For ResNets of arbitrary width, with layers of the form eq. 1,
[Din+21, Din+22] give a result of uniform convergence over finite training
time intervals. Adding a regularization term, [TG23] show the
$\Gamma$-convergence of the risk associated with ResNets to the one associated
with $\mathrm{NODEs}$.

##### Conditional Optimal transport

While the literature on Conditional OT distances is relatively sparse, a
detailed study was recently provided by the work of [HHT23]. Most notably,
they show the existence of Monge maps in the case where the measures are
absolutely continuous. Some of our results in section 2 are already present in
their work but we still wrote the proofs for completeness and the arguments
might differ. However, an original contribution of our work is to study the
dynamical properties of the Conditional OT distance. Analogously to the
Wasserstein case [AGS08], we show that absolutely continuous curves are
solutions to certain continuity equations (Proposition 2.4).

##### Contributions

Our main contribution is to propose a model for ResNets of infinite depth and
arbitrary width together with a metric space structure that is consistent with
the layer-wise-$L^{2}$-metric used in practice when training ResNets with
gradient descent and automatic differentiation. Our model thus allows a
rigorous analysis of the training of ResNets at infinite depth and arbitrary
width.

In detail, the ResNet model of Definition 1.1 is parameterized over the set
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ of probability measures
whose first marginal is the Lebesgue measure on $[0,1]$, which we provide in
section 2 with the metric structure of a $L^{2}$-Wasserstein (or _Conditional
Optimal Transport_) distance $d$ (Proposition 2.1). In section 3 we leverage
results from the theory of gradient flows in metric spaces [AGS08, San17] to
define the gradient flow of the risk $L$. This gradient flow equation
corresponds to both notions of _curve of maximal slope_ of the risk and the
usual gradient flow of ResNets obtained by _adjoint sensitivity analysis_
[Che+18]. We conclude this part by showing well-posedness results for the
gradient flow equation, that is existence in arbitrary time (Theorem 3.2),
uniqueness (Theorem 3.3) and stability w.r.t. initialization (Theorem 3.4).
Finally, we study in section 4 the asymptotic behavior of gradient flow
curves. We show that the risk $L$ satisfies a Polyak-Łojasiewicz (P-Ł)
property around well-chosen initializations. The risk has no saddles in these
regions and decreases at a constant rate along the gradient flow. Based on
previous works on the convergence of curves of maximal slope under the P-Ł
assumption [SMP23, HM19], we can then formulate a convergence result: for
initializations with a sufficiently large but finite number of features and
sufficiently low risk the gradient flow converges towards a global minimizer
(Theorem 4.2). Our results are to be compared with the ones of [Lu+20,
Din+22]. Both works give a result of optimality under a convergence assumption
but do not give conditions guaranteeing convergence of the gradient flow.
Moreover, their results hold under the assumption of an infinite number of
features whereas our convergence conditions can be obtained with a finite
number of features.

In addition to this, we studied in section 2 theoretical properties of the
space $\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ provided with the
Conditional Optimal Transport distance $d$. The literature on this subject
still being sparse, some of our results might be of their own interest. In
particular, we provide in Proposition 2.4 a characterization of absolutely
continuous curves analogous to the one in the Wasserstein space [AGS08,
Thm.8.3.1].

##  2 Metric structure of the parameter set
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$

We define here a notion of distance $d$ over the parameter set
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and study its properties.
Importantly, the characterization of absolutely continuous curves in the
metric space $(\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega),d)$ will be
used in section 3 to define the notion of gradient flow for the risk $L$.

In the rest of this work, we will assume for simplicity that $\Omega$ is the
Euclidean space $\mathbb{R}^{p}$ for some $p\geq 1$. However, the presented
results could probably be adapted to the case where $\Omega$ is a smooth
manifold embedded in $\mathbb{R}^{p}$. In particular, we will extensively use
the fact that $\Omega$ is a complete, separable metric space. We recall the
definition of the _Wasserstein_ $\mathcal{W}_{2}$ distance over the space
$\mathcal{P}_{2}(\mathbb{R}^{p})$, obtained as the solution to Kantorovitch’s
optimal transport problem:

$\displaystyle\forall\mu,\mu^{\prime}\in\mathcal{P}_{2}(\mathbb{R}^{p}),\quad\mathcal{W}_{2}(\mu,\mu^{\prime})\coloneqq\min_{\gamma\in\Gamma(\mu,\mu^{\prime})}\left(\int_{\mathbb{R}^{p}\times\mathbb{R}^{p}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma(\omega,\omega^{\prime})\right)^{1/2},$
(6)

where $\Gamma(\mu,\mu^{\prime})$ is the set of _couplings_ between $\mu$ and
$\mu^{\prime}$, that is the set of probability measures
$\gamma\in\mathcal{P}(\mathbb{R}^{p}\times\mathbb{R}^{p})$ respecting the
marginal conditions $\pi^{1}_{\\#}\gamma=\mu$ and
$\pi^{2}_{\\#}\gamma=\mu^{\prime}$. We note
$\Gamma_{o}(\mu,\mu^{\prime})\subset\Gamma(\mu,\mu^{\prime})$ the subset of
_optimal couplings_ achieving the equality in eq. 6. We refer to the books of
[Vil09, San15] for further properties of the Wasserstein distance.

###  2.1 Conditional Optimal Transport distance

The Conditional Optimal Transport distance $d$ is a modification of the
Wasserstein distance $\mathcal{W}_{2}$ with the supplementary constraint that
the transport plan should preserve the marginal over $[0,1]$. This constraint
is introduced to closely model the training dynamic of ResNets where the
gradients are computed over the weights of each layer independently. For this
purpose, it is natural to define a “layer-wise-$L^{2}$” Wasserstein distance,
that is a $L^{2}$-distance over the set of families of probability measures in
$\mathcal{P}_{2}(\Omega)$ index over $s\in[0,1]$.

######  Proposition 2.1 (Conditional OT distance).

Consider the function
$d:\left(\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)\right)^{2}\to\mathbb{R}_{+}$
defined by:

$\displaystyle
d(\mu,\mu^{\prime})\coloneqq\left(\int_{0}^{1}\mathcal{W}_{2}\left(\mu(.|s),\mu^{\prime}(.|s)\right)^{2}\mathrm{d}s\right)^{1/2}.$

Then, $d$ defines a metric on
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$.

###### Proof.

One essentially needs to justify the existence of the integral in the
definition of $d$. That $d$ is a metric then follows from the properties of
the Wasserstein and $L^{2}$ metrics respectively.

For a complete separable metric space $\Omega$ and Borel probability measures
$\mu,\nu\in\mathcal{P}(\Omega)$ it is known [Vil09, Thm.5.10] that the Monge-
Kantorovitch problem admits the dual formulation:

$\displaystyle\mathcal{W}_{2}(\mu,\nu)^{2}=\sup\\{\int_{\Omega}\varphi\mathrm{d}\mu+\int_{\Omega}\psi\mathrm{d}\nu\\}$

where the supremum is taken over all pairs
$(\varphi,\psi)\in\mathcal{C}_{b}(\Omega)\times\mathcal{C}_{b}(\Omega)$ such
that $\varphi(x)+\psi(y)\leq\|x-y\|^{2}$. We also have the alternate
formulation:

$\displaystyle\mathcal{W}_{2}(\mu,\nu)^{2}=\sup_{\varphi\in\mathcal{C}_{b}(\Omega)}\\{\int_{\Omega}\varphi\mathrm{d}\mu+\int_{\Omega}\varphi^{c}\mathrm{d}\nu\\},$

where for $\varphi:\Omega\to\mathbb{R}$ the $c$-transform $\varphi^{c}$ of
$\varphi$ is defined as [San15, Def.1.10]:

$\displaystyle\forall\omega\in\Omega,\quad\varphi^{c}(\omega)\coloneqq\inf_{\omega^{\prime}\in\Omega}\|\omega^{\prime}-\omega\|^{2}-\varphi(\omega^{\prime}).$

Consider $(\varphi_{n})_{n\geq 0}$ a sequence of functions in
$\mathcal{C}_{b}(\Omega)$ such that for any
$\varphi\in\mathcal{C}_{b}(\Omega)$ we can find a subsequence $m(n)$ with
$\varphi_{m(n)}\to\varphi$ for the compact-open topology (uniform convergence
on compact subsets) and $\|\varphi_{m(n)}\|_{\infty}$ is uniformly bounded.
Then we also have $\varphi_{m(n)}^{c}\to\varphi^{c}$ uniformly on compact
subsets with $\|\varphi^{c}_{m(n)}\|_{\infty}\leq\|\varphi_{m(n)}\|_{\infty}$
uniformly bounded, whence:

$\displaystyle\mathcal{W}_{2}(\mu,\nu)^{2}=\sup_{n\in\mathbb{N}}\\{\int_{\Omega}\varphi_{n}\mathrm{d}\mu+\int_{\Omega}\varphi_{n}^{c}\mathrm{d}\nu\\}.$

Thus, for
$\mu,\mu^{\prime}\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$, the
application
$s\mapsto\mathcal{W}_{2}\left(\mu(.|s),\mu^{\prime}(.|s)\right)^{2}$ is
measurable as it can be expressed as the supremum of countably many measurable
functions. ∎

Alternatively, the distance $d$ can be viewed as an optimal transport distance
with the additional constraint that the transport plans should be the identity
on the first marginals. This new formulation will show itself convenient for
calculations and will in particular allow to estimate from above the distance
$d$ easily. Given
$\mu,\mu^{\prime}\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ it is
convenient to work with the two following sets of “couplings” between $\mu$
and $\mu^{\prime}$:

$\displaystyle\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})$
$\displaystyle\coloneqq\\{\gamma\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega^{2}),\
\gamma(.|s)\in\Gamma(\mu(.|s),\mu^{\prime}(.|s))\ \text{for $\mathrm{d}s$-a.e.
$s\in[0,1]$}\\},$ $\displaystyle\Gamma^{\mathrm{diag}}(\mu,\mu^{\prime})$
$\displaystyle\coloneqq\\{\gamma\in\Gamma(\mu,\mu^{\prime}),\ \int
f(s,s^{\prime})\mathrm{d}\gamma(s,\omega,s^{\prime},\omega^{\prime})=\int_{0}^{1}f(s,s)\mathrm{d}s\
\forall f\in\mathcal{C}([0,1]^{2})\\}.$

Note that these to sets are closely related as if
$\gamma\in\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})$ then
$\tilde{\gamma}\coloneqq(\pi^{1},\pi^{2},\pi^{1},\pi^{3})_{\\#}\gamma\in\Gamma^{\mathrm{diag}}(\mu,\mu^{\prime})$
and conversely, if $\tilde{\gamma}\in\Gamma^{\mathrm{diag}}(\mu,\mu^{\prime})$
then
$\gamma\coloneqq(\pi^{1},\pi^{2},\pi^{4})_{\\#}\tilde{\gamma}\in\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})$.
In both cases, we have for any measurable $f:\Omega^{2}\to\mathbb{R}$:

$\displaystyle\int_{[0,1]\times\Omega^{2}}f(\omega,\omega^{\prime})\mathrm{d}\gamma(s,\omega,\omega^{\prime})=\int_{([0,1]\times\Omega)^{2}}f(\omega,\omega^{\prime})\mathrm{d}\tilde{\gamma}(s,\omega,s^{\prime},\omega^{\prime}).$
(7)

In the same way the Wasserstein distance $\mathcal{W}_{2}(\mu,\mu^{\prime})$
can be obtained of the solution of a minimization problem over the set
$\Gamma(\mu,\mu^{\prime})$ (eq. 6), the Conditional OT distance $d$ can be
obtained as the solution of minimization problems over the sets
$\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})$ and
$\Gamma^{\mathrm{diag}}(\mu,\mu^{\prime})$.

######  Proposition 2.2.

Let $\mu,\mu^{\prime}\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$
then:

$\displaystyle d(\mu,\mu^{\prime})^{2}$
$\displaystyle=\min_{\gamma\in\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})}\int_{[0,1]\times\Omega^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma(s,\omega,\omega^{\prime})$
$\displaystyle=\min_{\gamma\in\Gamma^{\mathrm{diag}}(\mu,\mu^{\prime})}\int_{([0,1]\times\Omega)^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma(s,\omega,s^{\prime},\omega^{\prime}).$

We note respectively $\Pi^{\mathrm{Leb}}_{o}(\mu,\mu^{\prime})$ and
$\Gamma^{\mathrm{diag}}_{o}(\mu,\mu^{\prime})$ the set of optimal couplings in
both minimization problems. Then for
$\gamma\in\Pi^{\mathrm{Leb}}_{o}(\mu,\mu^{\prime})$ we have for
$\mathrm{d}s$-a.e. $s\in[0,1]$:

$\displaystyle\gamma(.|s)\in\Gamma_{o}(\mu(.|s),\mu^{\prime}(.|s))\quad\text{i.e.}\quad\int_{\Omega^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma(\omega,\omega^{\prime}|s)=\mathcal{W}_{2}(\mu(.|s),\mu^{\prime}(.|s))^{2}.$

###### Proof.

Our proof technique is similar to the one of [HHT23, Prop.3.3] and relies on
the possibility of choosing an optimal transport plan
$\gamma(.|s)\in\Gamma_{o}(\mu(.|s),\mu^{\prime}(.|s))$ for every $s\in[0,1]$
in a measurable way.

We show equality with the first minimization problem on
$\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})$, equality between the two minimization
problems then comes from eq. 7. Assume there exists a Borel map
$\gamma:s\mapsto\gamma(.|s)\in\mathcal{P}(\Omega^{2})$ (where
$\mathcal{P}(\Omega^{2})$ is provided with the narrow topology) such that
$\gamma(.|s)\in\Gamma_{o}(\mu(.|s),\mu^{\prime}(.|s))$ for every $s\in[0,1]$.
With such a map one can define a Borel probability measure that we also note
$\gamma$ over $[0,1]\times\Omega$ which is the measure whose disintegration
w.r.t. the Lebesgue measure on $[0,1]$ is $\\{\gamma(.|s)\\}_{s\in[0,1]}$. In
other words the measure $\gamma$ is defined as:

$\displaystyle\int_{[0,1]\times\Omega^{2}}f(s,\omega,\omega^{\prime})\mathrm{d}\gamma(s,\omega,\omega^{\prime})\coloneqq\int_{0}^{1}\int_{\Omega^{2}}f(s,\omega,\omega^{\prime})\mathrm{d}\gamma(\omega,\omega^{\prime}|s)\mathrm{d}s,\quad\forall
f\in\mathcal{C}_{b}([0,1]\times\Omega^{2}).$

Such $\gamma$ will be a solution to our first optimization problem as we have:

$\displaystyle\int_{[0,1]\times\Omega^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma(s,\omega,\omega^{\prime})=d(\mu,\mu^{\prime})^{2}\leq\inf_{\gamma\in\Pi^{\mathrm{Leb}}(\mu,\mu^{\prime})}\int_{[0,1]\times\Omega^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma(s,\omega,\omega^{\prime}).$

To show the existence of such $\gamma$ we use a measurable selection result,
that is considering the set-valued mapping
$s\in[0,1]\mapsto\Gamma_{o}(\mu(.|s),\mu^{\prime}(.|s))\subset\mathcal{P}(\Omega^{2})$
we show it admits a measurable section. Consider the set:

$\displaystyle\mathcal{G}^{*}\coloneqq\\{(s,\gamma),\,\gamma\in\Gamma_{o}(\mu(.|s),\mu^{\prime}(.|s))\\}\subset[0,1]\times\mathcal{P}_{2}(\Omega^{2}).$

Using [Bog07, Thm.6.9.6], as for every $s\in[0,1]$ the set
$\Gamma_{o}(\mu(.|s),\mu^{\prime}(.|s))$ is narrowly compact, it is sufficient
to show that
$\mathcal{G}^{*}\in\mathcal{B}([0,1]\times\mathcal{P}_{2}(\Omega^{2}))$. Let
$\\{f_{n}\\}_{n\in\mathbb{N}}$ be dense in $\mathcal{C}_{b}(\Omega)$ for the
compact-open topology. Then, for every $n\in\mathbb{N}$ as the mapping
$(s,\gamma)\mapsto\int_{\Omega^{2}}f_{n}\mathrm{d}\gamma-\int_{\Omega}f_{n}\mathrm{d}\mu(.|s)$
is measurable, so are the sets:

$\displaystyle\mathcal{G}_{n}$
$\displaystyle\coloneqq\\{(s,\gamma),\,\int_{\Omega^{2}}f_{n}(\omega)\mathrm{d}\gamma(\omega,\omega^{\prime})=\int_{\Omega}f_{n}(\omega)\mathrm{d}\mu(\omega|s)\\}$
$\displaystyle\mathcal{G}_{n}^{\prime}$
$\displaystyle\coloneqq\\{(s,\gamma),\,\int_{\Omega^{2}}f_{n}(\omega^{\prime})\mathrm{d}\gamma(\omega,\omega^{\prime})=\int_{\Omega}f_{n}(\omega^{\prime})\mathrm{d}\mu^{\prime}(\omega^{\prime}|s)\\}.$

Also, as the mapping
$(s,\gamma)\mapsto\mathcal{W}_{2}(\mu(.|s),\mu^{\prime}(.|s))^{2}-\int_{\Omega^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma$
is measurable by Proposition 2.1, so is the set:

$\displaystyle\mathcal{G}_{o}\coloneqq\\{(s,\gamma),\,\int_{\Omega^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma=\mathcal{W}_{2}(\mu(.|s),\mu^{\prime}(.|s))^{2}\\}\in\mathcal{B}([0,1]\times\mathcal{P}_{2}(\Omega^{2})).$

Finally we have that
$\mathcal{G}^{*}=\mathcal{G}_{o}\cap\left(\bigcap_{n\in\mathbb{N}}\mathcal{G}_{n}\cap\mathcal{G}_{n}^{\prime}\right)$
is a Borel set, which completes the proof. ∎

######  Remark 2.1 (Comparison with $\mathcal{W}_{2}$).

Note that, for
$\mu,\mu^{\prime}\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$, we
have that
$\Gamma^{\mathrm{diag}}(\mu,\mu^{\prime})\subset\Gamma(\mu,\mu^{\prime})$.
Hence from the previous result, it follows:

$\displaystyle\mathcal{W}_{2}(\mu,\mu^{\prime})\leq d(\mu,\mu^{\prime})$

and the topology induced by $d$ on
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ is stronger than the
Wasserstein topology. It is in fact strictly stronger and, for example, the
sequence $\mu_{n}=\int_{0}^{1}\delta_{(-1)^{\lfloor 2ns\rfloor}}\mathrm{d}s$
and the measure
$\mu=\frac{1}{2}\int_{0}^{1}(\delta_{1}+\delta_{-1})\mathrm{d}s$ in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\mathbb{R})$ are such that
$\mathcal{W}_{2}(\mu_{n},\mu)\to 0$ but $d(\mu_{n},\mu)\geq 1$.

The following result states that the metric space
$(\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega),d)$ is complete.

######  Proposition 2.3 (Completeness).

$(\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega),d)$ is a complete metric
space.

###### Proof.

The proof is analogous to the proof of completeness of the Wasserstein space
$\mathcal{P}_{2}([0,1]\times\Omega)$ (see [Vil09, Thm.6.18])

Let $(\mu_{n})_{n\geq 0}$ be a Cauchy sequence in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega))$. Then for any
$\nu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$

$\displaystyle d(\nu,\mu_{n})\leq d(\nu,\mu_{0})+\sum_{1\leq i\leq
n}d(\mu_{i},\mu_{i-1})\leq d(\nu,\mu_{0})+\sum_{i\geq
1}d(\mu_{i},\mu_{i-1})<\infty,$

implying that the sequence $(\mu_{n})$ is bounded. Hence it is tight and by
Prokhorov’s theorem, it admits a subsequence $(\mu_{n_{k}})_{k\geq 0}$
converging narrowly to some
$\mu_{\infty}\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. Then by
narrow-lower-semi-continuity of $d$ (Lemma 2.1) we have for every $l\geq 0$:

$\displaystyle
d(\mu_{\infty},\mu_{n_{l}})\leq\liminf_{k\to\infty}d(\mu_{n_{k}},\mu_{n_{l}})$

and by taking the $\limsup$ w.r.t. $l$:

$\displaystyle\limsup_{l\to\infty}d(\mu_{\infty},\mu_{n_{l}})\leq\limsup_{\begin{subarray}{c}k\to\infty\\\
l\to\infty\end{subarray}}d(\mu_{n_{k}},\mu_{n_{l}})=0.$

Hence $(\mu_{n_{k}})$ $d$-converges to $\mu_{\infty}$, implying that the whole
sequence $(\mu_{n})$ $d$-converges to $\mu_{\infty}$. ∎

######  Lemma 2.1 (narrow-lower-semi-continuity of $d$).

Let $(\mu_{n})_{n\geq 0}$ and $(\nu_{n})_{n\geq 0}$ be sequences in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ such that
$(\mu_{n},\nu_{n})\xrightarrow[n\to\infty]{}(\mu,\nu)$ narrowly for some
$\mu,\nu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. Then:

$\displaystyle d(\mu,\nu)\leq\liminf_{n\to\infty}d(\mu_{n},\nu_{n}).$

###### Proof.

Up to extraction of a subsequence one can consider
$d(\mu_{n},\nu_{n})\to\liminf d(\mu_{n},\nu_{n})$. Then for every $n\geq 0$
consider some $\gamma_{n}\in\Gamma_{o}^{\mathrm{diag}}(\mu_{n},\nu_{n})$. In
particular $\gamma_{n}\in\Gamma(\mu_{n},\nu_{n})$ and by [Vil09, Lem.4.4] the
sequence $(\gamma_{n})$ is tight. Hence it admits a subsequence
$(\gamma_{n_{k}})_{k\geq 0}$ narrowly converging to some $\gamma$ which is in
$\Gamma^{\mathrm{diag}}(\mu,\nu)$ by the properties of narrow convergence.
Thus applying [Vil09, Lem.4.3] and using the characterization of $d$ in
Proposition 2.2:

$\displaystyle
d(\mu,\nu)^{2}\leq\int_{([0,1]\times\Omega)^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma\leq\liminf_{k\to\infty}\int_{([0,1]\times\Omega)^{2}}\|\omega-\omega^{\prime}\|^{2}\mathrm{d}\gamma_{n_{k}}=\liminf_{n\to\infty}d(\mu_{n},\nu_{n})^{2},$

from which the result follows. ∎

###  2.2 Dynamical formulation of Conditional Optimal Transport

We analyze here the properties of absolutely continuous curves in the metric
space $\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ provided with the
Conditional OT metric. Similarly to the (unconstrained) Wasserstein metric, we
show that absolutely continuous curves obey a certain continuity equation.
This characterization will be crucial to define the gradient flow equation
used in the training of our $\mathrm{NODE}$ model.

##### Absolutely continuous curves for the Wasserstein distance
$\mathcal{W}_{2}$

For $T>0$, consider $I=(0,T)$ an open interval and $(\mu_{t})_{t\in I}$ a
family of probability measures over the Euclidean space $\mathbb{R}^{p}$.
Given a Borel velocity field $v:(t,x)\in I\times\mathbb{R}^{p}\mapsto
v_{t}(x)\in\mathbb{R}^{p}$ such that
$\int_{I}\int_{\mathbb{R}^{p}}\|v_{t}\|\mathrm{d}\mu_{t}\mathrm{d}t<\infty$,
we say that $(\mu_{t})_{t\in I}$ satisfies the continuity equation
$\partial_{t}\mu_{t}+\mathrm{div}(v_{t}\mu_{t})$ in the weak sense if:

$\displaystyle\int_{I}\int_{\mathbb{R}^{p}}\left(\partial_{t}\varphi(t,x)+\langle\nabla\varphi(t,x),v_{t}(x)\rangle\right)\mathrm{d}\mu_{t}(x)\mathrm{d}t=0,\quad\forall\varphi\in\mathcal{C}^{1}_{c}(I\times\mathbb{R}^{p}).$
(8)

Equivalently ([San15, Prop.4.2]), when the mapping $t\mapsto\mu_{t}$ is
narrowly continuous, this amounts to have that for every
$\varphi\in\mathcal{C}^{1}_{c}(\mathbb{R}^{p})$ the map
$t\mapsto\mu_{t}(\varphi)\coloneqq\int\varphi\mathrm{d}\mu_{t}$ is absolutely
continuous and verifies:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}(\varphi)=\int\langle\nabla\varphi,v_{t}\rangle\mathrm{d}\mu_{t},\quad\text{for
$\mathrm{d}t$-a.e. $t\in I$.}$

An important property of the Wasserstein space
$\mathcal{P}_{2}(\mathbb{R}^{p})$ endowed with the distance $\mathcal{W}_{2}$
is the characterization of absolutely continuous curves: a narrowly continuous
curve $(\mu_{t})_{t\in I}$ is absolutely continuous in
$\mathcal{P}_{2}(\mathbb{R}^{p})$ if and only if it is solution to the
continuity equation eq. 8 for some velocity field $v$ with
$\int_{I}\|v_{t}\|_{L^{2}(\mu_{t})}\mathrm{d}t<\infty$ [AGS08, Thm.8.3.1]. We
refer to the book by [AGS08] for a detailed study of absolutely continuous
curves in $(\mathcal{P}_{2}(\mathbb{R}^{p}),\mathcal{W}_{2})$.

##### Absolutely continuous curves for the Conditional OT distance

Similarly to the characterization of absolutely continuous curves in the
Wasserstein space $\mathcal{P}_{2}(\mathbb{R}^{p})$, an adaptation of [AGS08,
Thm.8.3.1] provides an analogous characterization of absolutely continuous
curves in $\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ provided with
the Conditional OT distance $d$. This characterization allows us to (formally)
provide the metric space
$(\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega),d)$ with a kind of “
differential structure” by seeing tangent vectors as velocity fields. This
identification will be crucial for defining the gradient flow equation which
will take the form of a continuity equation with an appropriate velocity field
(Definition 3.2).

######  Proposition 2.4 (adapted from [AGS08], Thm.8.3.1).

Assume $\Omega=\mathbb{R}^{p}$. Let $I=(0,T)$ for some $T>0$ and
$(\mu_{t})_{t\in I}$ an absolutely continuous curve in
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$. Then there exists a
unique Borel velocity field $v:(t,s,\omega)\in I\times[0,1]\times\Omega\mapsto
v_{t}(s,\omega)\in\Omega$ such that for a.e. $t\in I$:

$\displaystyle v_{t}\in
L^{2}(\mu_{t}),\quad\|v_{t}\|_{L^{2}(\mu_{t})}\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|,$

and $\mu$ is a weak solution of the continuity equation:

$\displaystyle\partial_{t}\mu_{t}+\mathrm{div}((0,v_{t})\mu_{t})=0\quad\text{over
$I\times[0,1]\times\Omega$.}$ (9)

We will refer to such $v_{t}$ as the _tangent velocity field_ of the curve
$(\mu_{t})_{t\in I}$. Conversely, if $(\mu_{t})_{t\in I}$ is a narrowly
continuous curve satisfying eq. 9 for some Borel velocity field $v_{t}$ with
$\|v_{t}\|_{L^{2}(\mu_{t})}\in L^{1}(I)$, then $(\mu_{t})_{t\in I}$ is
absolutely continuous and
$|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}|\leq\|v_{t}\|_{L^{2}(\mu_{t})}$ for
a.e. $t\in I$.

###### Proof.

_Part 1: AC curve $\Rightarrow$ Continuity equation._

Note that this part is the easiest as $d$-absolute continuity implies
$\mathcal{W}_{2}$-absolute continuity for which the result is well-known,
originally proven in [AGS08, Thm.8.3.1]. Therefore we here only adapt this
proof to our specific setting.

Up to reparameterization, one can assume without loss of generality that
$\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|\in L^{\infty}(I)$.

First we show that for
$\varphi\in\mathcal{C}^{1}_{c}(I\times[0,1]\times\Omega)$ the map
$t\mapsto\mu_{t}(\varphi)\coloneqq\int_{0}^{1}\int_{\Omega}\varphi\mathrm{d}\mu_{t}$
is absolutely continuous. Indeed, for $t,u\in I$, introducing a coupling
$\gamma_{t,u}\in\Pi^{\mathrm{Leb}}_{o}(\mu_{t},\mu_{u})$ we have:

$\displaystyle\left|\mu_{t}(\varphi)-\mu_{u}(\varphi)\right|\leq\left|\int_{0}^{1}\int_{\Omega^{2}}(\varphi(s,\omega)-\varphi(s,\omega^{\prime}))\mathrm{d}\gamma_{t,u}(s,\omega,\omega^{\prime})\right|\leq\|\nabla_{\omega}\varphi\|_{\infty}d(\mu_{t},\mu_{u}),$

from which absolute continuity follows. Then considering the map:

$\displaystyle
H(s,\omega,\omega^{\prime})\coloneqq\left\\{\begin{array}[]{cc}\|\nabla_{\omega}\varphi(s,\omega)\|&\text{if
$\omega=\omega^{\prime}$},\\\ &\\\
\frac{|\varphi(s,\omega)-\varphi(s,\omega^{\prime})|}{\|\omega-\omega^{\prime}\|}&\text{else},\end{array}\right.$

we have for every $t,u\in I$:

$\displaystyle\frac{|\mu_{t}(\varphi)-\mu_{u}(\varphi)|}{|t-u|}\leq\frac{1}{|t-u|}\int_{0}^{1}\int_{\Omega^{2}}\|\omega-\omega^{\prime}\|H(s,\omega,\omega^{\prime})\mathrm{d}\gamma_{t,u}(s,\omega,\omega^{\prime})\leq\frac{d(\mu_{t},\mu_{u})}{|t-u|}\|H\|_{L^{2}(\gamma_{t,u})}$

As $u\to t$ we have $d(\mu_{u},\mu_{t})\to 0$ and by the properties of $L^{2}$
spaces [[, _e.g._]Prop.3.11]cannarsa2015introduction we can take a sequence
$u_{n}\to t$ such that $\mathcal{W}_{2}(\mu_{u_{n}}(.|s),\mu_{t}(.|s))\to 0$
for $\mathrm{d}s$-a.e. $s\in[0,1]$. This implies for those $s\in[0,1]$ that
$\mu_{u_{n}}(.|s)\to\mu_{t}(.|s)$ narrowly and that
$\gamma_{t,u_{n}}(.|s)\to\gamma(.|s)\in\Gamma_{o}(\mu_{t}(.|s),\mu_{t}(.|s))$,
i.e. the trivial transport plan
$\gamma(.|s)=(\mathrm{Id},\mathrm{Id})_{\\#}\mu_{t}(.|s)$. Thus we have that
$\gamma_{t,u_{n}}\to(\pi^{1},\pi^{2},\pi^{2})_{\\#}\mu_{t}$ narrowly since, by
Lebesgue’s theorem, given a bounded continuous function
$f\in\mathcal{C}_{b}([0,1]\times\Omega\times\Omega)$:

$\displaystyle\int
f\mathrm{d}\gamma_{t,u_{n}}=\int_{0}^{1}\left(\int_{\Omega^{2}}f(s,\omega,\omega^{\prime})\mathrm{d}\gamma_{t,u_{n}}(\omega,\omega^{\prime}|s)\right)\mathrm{d}s\xrightarrow[n\to\infty]{}\int_{0}^{1}\left(\int_{\Omega}f(s,\omega,\omega)\mathrm{d}\mu_{t}(\omega|s)\right)\mathrm{d}s.$

Hence, at a point where $t\mapsto\mu_{t}$ is metrically differentiable:

$\displaystyle\limsup_{u\to
t}\frac{|\mu_{t}(\varphi)-\mu_{u}(\varphi)|}{|t-u|}\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|\|\nabla_{\omega}\varphi\|_{L^{2}(\mu_{t})}.$

Consider
$\bm{\mu}=\int_{I}\mu_{t}\mathrm{d}t\in\mathcal{P}(I\times[0,1]\times\Omega)$
the measure whose disintegration w.r.t. Lebesgue’s measure on $I$ is
$(\mu_{t})_{t\in I}$. Then for
$\varphi\in\mathcal{C}^{1}_{c}(I\times[0,1]\times\Omega)$ we have:

$\displaystyle\int_{I}\int_{[0,1]\times\Omega}\partial_{t}\varphi(t,s,\omega)\mathrm{d}\mu_{t}(s,\omega)\mathrm{d}t$
$\displaystyle=\lim_{h\to
0}\int_{I}\int_{[0,1]\times\Omega}\frac{\varphi(t,s,\omega)-\varphi(t-h,s,\omega)}{h}\mathrm{d}\mu_{t}(s,\omega)\mathrm{d}t$
$\displaystyle=\lim_{h\to
0}\int_{I}\frac{1}{h}\left(\int_{[0,1]\times\Omega}\varphi(t,s,\omega)\mathrm{d}\mu_{t}(s,\omega)-\int_{[0,1]\times\Omega}\varphi(t,s,\omega)\mathrm{d}\mu_{t-h}(s,\omega)\right)\mathrm{d}t.$

Thus by the previous inequality and Fatou’s lemma:

$\displaystyle\left|\int_{I}\int_{[0,1]\times\Omega}\partial_{t}\varphi(t,s,\omega)\mathrm{d}\mu_{t}(s,\omega)\mathrm{d}t\right|\leq\left(\int_{I}\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|^{2}\mathrm{d}t\right)^{1/2}\left(\int_{I\times[0,1]\times\Omega}\|\nabla_{\omega}\varphi(t,s,\omega)\|^{2}\mathrm{d}\bm{\mu}(t,s,\omega)\right)^{1/2}.$

Consider the subspace
$V\coloneqq\\{\nabla_{\omega}\varphi,\,\varphi\in\mathcal{C}^{1}_{c}(I\times[0,1]\times\Omega)\\}$
and $\mathcal{V}$ its closure in $L^{2}(I\times[0,1]\times\Omega,\bm{\mu})$.
Then by the previous inequality the linear functional
$\mathcal{A}:V\to\mathbb{R}$ defined by:

$\displaystyle\mathcal{A}(\nabla_{\omega}\varphi)\coloneqq-\int_{I\times[0,1]\times\Omega}\partial_{t}\varphi(t,s,\omega)\mathrm{d}\bm{\mu}(t,s,\omega)$

is continuous on $V$ and thus, by Hahn-Banach’s theorem, can be extended to a
unique continuous linear functional on $\mathcal{V}$. Therefore, by Lax-
Milgram’s theorem, the minimization problem:

$\displaystyle\min\left\\{\frac{1}{2}\int_{I\times\mathbb{R}^{p+1}}\|w(t,s,\omega)\|^{2}\mathrm{d}\bm{\mu}(t,s,\omega)-\mathcal{A}(w),\,w\in\mathcal{V}\right\\}$

admits a unique solution $v\in\mathcal{V}$ which is characterized by the
property that:

$\displaystyle\int_{I\times\mathbb{R}^{p+1}}\langle
v(t,s,\omega),\nabla_{\omega}\varphi(t,s,\omega)\rangle\mathrm{d}\bm{\mu}(t,s,\omega)=\mathcal{A}(\nabla_{\omega}\varphi),\quad\forall\varphi\in\mathcal{C}^{1}_{c}(I\times[0,1]\times\Omega),$

which is the desired continuity equation by definition of $\mathcal{A}$.

Finally, let $(\nabla_{\omega}\varphi_{n})\subset V$ be a sequence converging
to $v\in L^{2}(\bm{\mu})$. Considering an interval $J\subset I$ and some
$\eta\in\mathcal{C}^{1}_{c}(J)$ with $0\leq\eta\leq 1$ we have by the previous
arguments:

$\displaystyle\int_{I\times[0,1]\times\Omega}\eta(t)\|v(t,s,\omega)\|^{2}\mathrm{d}\bm{\mu}(t,s,\omega)$
$\displaystyle=\lim_{n\to\infty}\int_{I\times[0,1]\times\Omega}\eta\langle
v,\nabla_{\omega}\varphi_{n}\rangle\mathrm{d}\bm{\mu}$
$\displaystyle=\lim_{n\to\infty}\mathcal{A}(\nabla_{\omega}(\eta\varphi_{n}))$
$\displaystyle\leq\left(\int_{J}\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|^{2}\mathrm{d}t\right)^{1/2}\lim_{n\to\infty}\left(\int_{J\times[0,1]\times\Omega}\|\nabla_{\omega}\varphi_{n}\|^{2}\mathrm{d}\bm{\mu}\right)^{1/2}$
$\displaystyle=\left(\int_{J}\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|^{2}\mathrm{d}t\right)^{1/2}\left(\int_{J\times[0,1]\times\Omega}\|v\|^{2}\mathrm{d}\bm{\mu}\right)^{1/2}.$

Hence approximating the characteristic function of $J$ with such an $\eta$ we
get:

$\displaystyle\int_{J}\int_{[0,1]\times\Omega}\|v_{t}\|^{2}\mathrm{d}\mu_{t}\mathrm{d}t\leq\int_{J}\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|^{2}\mathrm{d}t,$

implying
$\|v_{t}\|_{L^{2}(\mu_{t})}\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|$
for a.e. $t\in I$.

_Part 2: Continuity equation $\Rightarrow$ AC curve._

This part of the proof is new as, according to [AGS08, Thm.8.3.1], the
continuity equation only ensures $\mathcal{W}_{2}$-absolute continuity as
explained in Remark 2.1. We show here that the specific form of the velocity
field ensures $d$-absolute continuity.

For $(t,s)\in I\times[0,1]$, we note $v_{t,s}$ the Borel vector field
$v_{t,s}:\omega\in\Omega\mapsto v_{t}(s,\omega)$. Note that by Jensen’s
inequality:

$\displaystyle\int_{I}\int_{0}^{1}\|v_{t,s}\|_{L^{2}(\mu_{t}(.|s))}\mathrm{d}s\mathrm{d}t\leq\int_{I}\|v_{t}\|_{L^{2}(\mu_{t})}\mathrm{d}t<+\infty,$

and we have that for $\mathrm{d}s$-a.e. $s\in[0,1]$,
$t\mapsto\|v_{t,s}\|_{L^{2}(\mu_{t}(.|s))}\in L^{1}(I)$. Also if
$\varphi\in\mathcal{C}_{c}^{1}(I\times\Omega)$ and
$\chi\in\mathcal{C}_{c}^{1}([0,1])$ then by definition of the continuity
equation:

$\displaystyle\int_{I}\int_{0}^{1}\int_{\Omega}\left(\partial_{t}\varphi+\langle\nabla_{\omega}\varphi,v_{t,s}\rangle\right)\chi(s)\mathrm{d}\mu_{t}(.|s)\mathrm{d}s\mathrm{d}t=0.$

Hence if $J\subset[0,1]$ is an interval, approaching the characteristic
function of $J$ with $\chi$ we get:

$\displaystyle\int_{I}\int_{J}\int_{\Omega}\left(\partial_{t}\varphi+\langle\nabla_{\omega}\varphi,v_{t,s}\rangle\right)\mathrm{d}\mu_{t}(.|s)\mathrm{d}s\mathrm{d}t=0,$

and hence for $\mathrm{d}s$-a.e. $s\in[0,1]$:

$\displaystyle\int_{I}\int_{\Omega}\left(\partial_{t}\varphi+\langle\nabla_{\omega}\varphi,v_{t,s}\rangle\right)\mathrm{d}\mu_{t}(.|s)\mathrm{d}t=0.$

Now if we consider $(\varphi_{n})$ a countable dense sequence in
$\mathcal{C}^{1}_{c}(I\times\Omega)$ endowed with the usual topology then we
can find a set $\Lambda\subset[0,1]$ of full Lebesgue’s measure such that for
every $s\in\Lambda$ the above equation holds for every test function
$\varphi\in\mathcal{C}_{c}^{1}(I\times\Omega)$. In other words we have shown
that, for $\mathrm{d}s$-a.e. $s\in[0,1]$, $\mu_{t}(.|s)$ solves the continuity
equation:

$\displaystyle\partial_{t}\mu_{t}(.|s)+\mathrm{div}(v_{t,s}\mu_{t}(.|s))=0\quad\text{over
$I\times\Omega$.}$

Note that, without loss of generality, we can consider the curve
$(\mu_{t}(.|s))_{t\in I}$ to be narrowly continuous. Indeed, as it is a
solution of the continuity equation we know that the curve
$(\mu_{t}(.|s))_{t\in I}$ admits a narrowly continuous representative
$\tilde{\mu}_{t}(.|s)$ [AGS08, Lem.8.1.2] and that this representative is
characterized by that for every $\varphi\in\mathcal{C}^{1}_{c}(\Omega)$ and
every $t\in I$:

$\displaystyle\tilde{\mu}_{t}(.|s)(\varphi)=\int_{0}^{t}\int_{\Omega}\left(\chi^{\prime}(u)\varphi+\chi(u)\langle\nabla_{\omega}\varphi,v_{u,s}\rangle\right)\mathrm{d}\mu_{u}(.|s)\mathrm{d}u,$

where $\chi\in\mathcal{C}^{1}(I)$ is any function such that $\chi=0$ on a
neighbourhood of $0$ and $\chi(u)=1$ for $u\geq t$ (the definition does not
depend on $\chi$ by definition of the continuity equation). Then it follows
that for any $t\in I$ and any $f\in\mathcal{C}_{b}([0,1]\times\Omega)$ the map
$s\mapsto\tilde{\mu}_{t}(.|s)(f)$ is measurable and integrating w.r.t. $s$ we
get that $\tilde{\mu}_{t}\coloneqq\int_{0}^{1}\tilde{\mu}_{t}(.|s)\mathrm{d}s$
defines a probability measure over $[0,1]\times\Omega$ whose disintegration is
$\\{\tilde{\mu}_{t}(.|s)\\}_{s\in[0,1]}$. Moreover for any
$\varphi\in\mathcal{C}_{c}^{1}([0,1]\times\Omega)$ we have:

$\displaystyle\tilde{\mu}_{t}(\varphi)=\int_{0}^{t}\int_{0}^{1}\int_{\Omega}\left(\chi^{\prime}(u)\varphi+\chi(u)\langle\nabla_{\omega}\varphi,v_{u,s}\rangle\right)\mathrm{d}\mu_{u}(.|s)\mathrm{d}s\mathrm{d}u=\mu_{t}(\varphi)$

and hence in fact the equality $\tilde{\mu}_{t}=\mu_{t}$.

Then using [AGS08, Thm.8.3.1] with the assumption that $t\mapsto\mu_{t}(.|s)$
is narrowly continuous we hence have that for $\mathrm{d}s$-a.e. $s\in[0,1]$
the curve $t\in I\mapsto\mu_{t}(.|s)$ is absolutely continuous and

$\displaystyle\mathcal{W}_{2}(\mu_{t_{1}}(.|s),\mu_{t_{2}}(.|s))^{2}\leq(t_{2}-t_{1})\int_{t_{1}}^{t_{2}}\int_{\Omega}\|v_{t,s}\|^{2}\mathrm{d}\mu_{t}(.|s)\mathrm{d}t,\quad\forall
t_{1}<t_{2}\in I.$

Hence integrating w.r.t. $s\in[0,1]$ gives:

$\displaystyle
d(\mu_{t_{1}},\mu_{t_{2}})^{2}\leq(t_{2}-t_{1})\int_{t_{1}}^{t_{2}}\int_{[0,1]\times\Omega}\|v_{t}\|^{2}\mathrm{d}\mu_{t}\mathrm{d}t,\quad\forall
t_{1}<t_{2}\in I,$

which shows that $(\mu_{t})_{t\in I}$ is $d$-absolutely continuous with
$\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|\leq\|v_{t}\|_{L^{2}(\mu_{t})}$
for a.e. $t\in I$. ∎

######  Remark 2.2.

Note that to study absolutely continuous curves, we introduce the
supplementary time variable $t\geq 0$. This time variable will model the
optimization time in the Definition 3.2 of the gradient flow equation. It is
not to be interverted with the $\mathrm{NODE}$ flow time $s\in[0,1]$.

As a consequence of Proposition 2.4 we recover two useful results about
absolutely continuous curves in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. Those are stated in the
following Lemmas 2.2 and 2.3. The first result is a result of approximation
along absolutely continuous curves. It states that the tangent velocity field
$(v_{t})_{t\in I}$ defined in Proposition 2.4 indeed furnishes a first-order
approximation of the curve $(\mu_{t})_{t\in I}$ at every time $t\in I$. It
will be particularly useful to differentiate quantities related to $\mu_{t}$
(Corollaries 3.2 and 3.3). The second result is an application and gives the
differential of the square-distance $d(\mu_{t},\mu^{\prime})^{2}$ along an
absolutely continuous curve $(\mu_{t})_{t\in I}$.

######  Lemma 2.2 (Adapted from [AGS08, Prop.8.4.6]).

Let $(\mu_{t})_{t\in I}$ be an absolutely continuous curve in
$(\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega),d)$ and let
$v:I\times\mathbb{R}^{p+1}\to\mathbb{R}^{p}$ be the unique velocity field
satisfying the conclusions of Proposition 2.4. Then for $\mathrm{d}t$-a.e.
$t\in I$ it holds that for any choice of
$\gamma^{h}_{t}\in\Pi^{\mathrm{Leb}}_{o}(\mu_{t+h},\mu_{t})$:

$\displaystyle\lim_{h\to
0}\left(\pi^{1},\pi^{2},\frac{1}{h}(\pi^{3}-\pi^{2})\right)_{\\#}\gamma^{h}_{t}=\left(\pi^{1},\pi^{2},v_{t})\right)_{\\#}\mu_{t}\quad\text{in
$\mathcal{W}_{2}([0,1]\times\Omega\times\Omega)$}$

and

$\displaystyle\lim_{h\to
0}\frac{d(\mu_{t+h},(\mathrm{Id}+h(0,v_{t}))_{\\#}\mu_{t})}{|h|}=0.$

###### Proof.

The proof only needs to be slightly adapted from the one of [AGS08, Lem.8.4.6]
but we rewrite it here for completeness.

Let $(\varphi_{n})_{n\geq 0}$ be a countable dense sequence in
$\mathcal{C}^{1}_{c}([0,1]\times\Omega)$. Then for $\mathrm{d}t$-a.e. $t\in I$
we have $\lim_{h\to
0}\frac{1}{|h|}d(\mu_{t+h},\mu_{t})=\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|$
and for every $n\geq 0$:

$\displaystyle\lim_{h\to
0}\frac{\mu_{t+h}(\varphi_{n})-\mu_{t}(\varphi_{n})}{h}=\int\langle\nabla_{\omega}\varphi_{n},v_{t}\rangle\mathrm{d}\mu_{t}.$

Introducing some $\gamma^{h}_{t}\in\Pi^{\mathrm{Leb}}_{o}(\mu_{t+h},\mu_{t})$
we consider:

$\displaystyle\nu^{h}\coloneqq\left(\pi^{1},\pi^{2},\frac{1}{h}(\pi^{3}-\pi^{2})\right)_{\\#}\gamma^{h}_{t}.$

Then for any sequence $(h_{n})$ converging to $0$ the sequence $(\nu^{h_{n}})$
is tight in $\mathcal{P}([0,1]\times\Omega\times\Omega)$ and we can consider a
narrow limit point $\nu^{0}$. The marginal of $\nu^{h}$, and hence of
$\nu^{0}$, on $[0,1]\times\Omega$ is $\mu_{t}$ which allows to write by
disintegration $\nu^{0}=\int\nu^{0}_{s,\omega}\mathrm{d}\mu_{t}(s,\omega)$.
Then we have for every $n\geq 0$:

$\displaystyle\frac{\mu_{t+h}(\varphi_{n})-\mu_{t}(\varphi_{n})}{h}$
$\displaystyle=\frac{1}{h}\int\left(\varphi_{n}(s,\omega^{\prime})-\varphi_{n}(s,\omega)\right)\mathrm{d}\gamma^{h}_{t}(s,\omega,\omega^{\prime})$
$\displaystyle=\frac{1}{h}\int\left(\varphi_{n}(s,\omega+hz)-\varphi_{n}(s,\omega)\right)\mathrm{d}\nu^{h}(s,\omega,z)$

and taking the limit $h\to 0$, by Lebesgue’s theorem:

$\displaystyle\int\langle\nabla_{\omega}\varphi_{n},v_{t}\rangle\mathrm{d}\mu_{t}=\int_{[0,1]\times\Omega}\int_{\Omega}\langle
z,\nabla_{\omega}\varphi_{n}(s,\omega)\rangle\mathrm{d}\nu^{0}_{s,\omega}(z)\mathrm{d}\mu_{t}(s,\omega).$

Let us note, for $(s,\omega)\in[0,1]\times\Omega$,
$\tilde{v}_{t}(s,\omega)\coloneqq\int_{\Omega}z\mathrm{d}\nu^{0}_{s,\omega}(z)$
the first moment of $\nu^{0}_{s,\omega}$. Then from the last equality and by a
density argument it follows:

$\displaystyle\mathrm{div}((0,\tilde{v}_{t}-v_{t})\mu_{t})=0,$

and in particular the continuity equation eq. 9 is satisfied with the vector
field $(0,\tilde{v}_{t})$. Let us now show:

$\displaystyle\int_{[0,1]\times\Omega}\int_{\Omega}\|z\|^{2}\mathrm{d}\nu^{0}_{s,\omega}(z)\mathrm{d}\mu_{t}(s,\omega)\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|^{2}.$

Indeed we have:

$\displaystyle\int_{[0,1]\times\Omega}\int_{\Omega}\|z\|^{2}\mathrm{d}\nu^{0}_{s,\omega}(z)\mathrm{d}\mu_{t}(s,\omega)$
$\displaystyle\leq\liminf_{h\to
0}\int_{[0,1]\times\Omega\times\Omega}\|z\|^{2}\mathrm{d}\nu^{h}(s,\omega,z)$
$\displaystyle=\liminf_{h\to
0}\int_{[0,1]\times\Omega\times\Omega}\frac{1}{h^{2}}\|\omega^{\prime}-\omega\|^{2}\mathrm{d}\gamma^{h}_{t}(s,\omega,\omega^{\prime})$
$\displaystyle=\liminf_{h\to
0}\frac{d(\mu_{t+h},\mu_{t})^{2}}{h^{2}}=\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|^{2}.$

Whence by definition of $\tilde{v}_{t}$ and Jensen’s inequality:

$\displaystyle\|\tilde{v}_{t}\|_{L^{2}(\mu_{t})}\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|=\|v_{t}\|_{L^{2}(\mu_{t})}$

from which it follows that $\tilde{v}_{t}=v_{t}$ in $L^{2}(\mu_{t})$ because
of the minimality of $\|v_{t}\|_{L^{2}(\mu_{t})}$ and the strict convexity of
the $L^{2}$-norm. Moreover the above inequality is strict whenever
$\nu^{0}_{s,\omega}$ is not a Dirac mass in a set of $\mu_{t}$ positive
measure. This implies that $\nu^{0}_{s,\omega}$ is a Dirac mass for
$\mathrm{d}\mu_{t}$-a.e. $(s,\omega)\in[0,1]\times\Omega$ and that
$\nu^{0}=(\pi^{1},\pi^{2},v_{t})_{\\#}\mu_{t}$. This proves the narrow
convergence of $\nu^{h}$ towards the desired measure and together with the
convergence of the second moments we have $\mathcal{W}_{2}$ convergence.

Let us now estimate the distance between $\mu_{t+h}$ and
$(\pi^{1},\pi^{2}+h(0,v_{t}))_{\\#}\mu_{t}$ with the coupling
$\gamma\coloneqq(\pi^{1},\pi^{2}+h(0,v_{t}),\pi^{3})_{\\#}\gamma^{h}_{t}\in\Pi^{\mathrm{Leb}}((\pi^{1},\pi^{2}+h(0,v_{t}))_{\\#}\mu_{t},\mu_{t+h})$.
By the preceding result:

$\displaystyle\frac{d((\pi^{1},\pi^{2}+hv_{t})_{\\#}\mu_{t},\mu_{t+h})^{2}}{h^{2}}$
$\displaystyle\leq\int_{[0,1]\times\Omega\times\Omega}\frac{1}{h^{2}}\|\omega+hv_{t}(s,\omega)-\omega^{\prime}\|^{2}\mathrm{d}\gamma^{h}_{t}(s,\omega,\omega^{\prime})$
$\displaystyle=\int_{[0,1]\times\Omega\times\Omega}\|v_{t}(s,\omega)-z\|^{2}\mathrm{d}\nu^{h}(s,\omega,z)\xrightarrow[h\to
0]{}0.$

∎

######  Lemma 2.3 (Adapted from [AGS08, Thm.8.4.7]).

Let $(\mu_{t})_{t\in I}$ be an absolutely continuous curve in
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$, let
$v:I\times[0,1]\times\omega\to\Omega$ be its _tangent vector field_ and let
$\mu^{\prime}\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$. Then for
$\mathrm{d}t$-a.e. $t\in I$:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}d(\mu_{t},\mu_{t}^{\prime})^{2}=2\int_{0}^{1}\int_{\Omega}\langle\omega-\omega^{\prime},v_{t}(s,\omega)\rangle\mathrm{d}\gamma(s,\omega,\omega^{\prime}),\quad\forall\gamma\in\Pi^{\mathrm{Leb}}_{o}(\mu_{t},\mu_{t}^{\prime}).$

###### Proof.

Having shown Lemma 2.2, the proof is the same as the one of [AGS08,
Thm.8.4.7]. ∎

##  3 Gradient flow dynamics

To train the $\mathrm{NODE}$ model of Definition 1.1 we consider performing
_Gradient Flow_ on the parameter $\mu$ for the risk $L$ and for the
Conditional OT metric described in the previous section. However the parameter
set $\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ provided with the
distance $d$ lacks a proper differential structure. We will thus in this
section give a sense to the notion of gradient flow of $L$. First, motivated
by formal computations we will introduce a definition of gradient flow that is
consistent with the one proposed by [Che+18] for the training of
$\mathrm{NODEs}$ of finite width. Then, we will show this definition to be
equivalent to the notion of _curve of maximal slope_ from the theory of
gradient flow in metric spaces [AGS08, San17]. Finally, this equivalence will
allow us to show well-posedness results for the gradient flow equation.

###  3.1 Backward equation and adjoint variables

The computation of the gradient will make use of a new ODE linked to eq. 3.
This ODE should be understood as running backward over the time variable
$s\in[0,1]$ with the initial condition at $s=1$. In the same way eq. 3 models
the processing of the data by a ResNet of infinite depth, the _adjoint
variables_ $p$ solutions to this _backward ODE_ should be considered as
modeling the quantities calculated when performing back-propagation over a
deep ResNet.

######  Definition 3.1 (Adjoint variable).

Let $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and
$(x,y)\in\mathbb{R}^{d+d^{\prime}}$. Let $(x_{\mu}(s))_{s\in[0,1]}$ be the
solution to eq. 2 with parameter $\mu$ and $x_{\mu}(0)=x$. Then we call
_adjoint variable associated to $\mu$, $x$ and $y$_ the solution
$(p_{\mu,x,y}(s))_{s\in[0,1]}$ to the _backward ODE_ :

$\forall s\in[0,1],\quad
p_{\mu,x,y}(s)=\nabla_{x}\ell(x_{\mu}(1),y)+\int_{s}^{1}\mathrm{D}_{x}F_{\mu(.|r)}(x_{\mu}(r))^{\top}p_{\mu,x,y}(r)\mathrm{d}r.$
(10)

When no ambiguity, the dependence of $p$ w.r.t. $\mu$, $x$ and $y$ is omitted
and we simply write $p(s)$.

The following proposition states the well-posedness of the backward equation
under suitable assumptions on the feature map $\phi$ and gives a useful
representation of the adjoint variables.

######  Proposition 3.1.

Let $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and
$(x,y)\in\mathbb{R}^{d+d^{\prime}}$. Assume $\phi$ satisfies Assumptions 1, 2
and 3. Then there exists a unique solution to eq. 10 which is given by:

$\forall s\in[0,1],\quad
p_{\mu,x,y}(s)=\Phi_{\mu,x}(s)^{-\top}\Phi_{\mu,x}(1)^{\top}\nabla_{x}\ell(x_{\mu}(1),y).$
(11)

where we define $\Phi_{\mu,x}$ to be the (matrix) solution of the linear ODE:

$\forall
s\in[0,1],\quad\Phi_{\mu,x}(s)=\mathrm{Id}+\int_{0}^{s}\mathrm{D}_{x}F_{\mu(.|r)}(x_{\mu}(r))\Phi_{\mu,x}(r)\mathrm{d}r.$
(12)

When no ambiguity we simply note $\Phi_{\mu}(s)$ or even $\Phi(s)$.

###### Proof.

Note that eq. 10 is non-autonomous linear ODE w.r.t. the variable $p$. Thus,
the representation eq. 11 follows from the existence and uniqueness of $\Phi$
and to prove the result it suffices to show the map
$s\mapsto\mathrm{D}_{x}F_{\mu(.|s)}(x(s))$ is integrable.

First, as $\phi$ is continuously differentiable w.r.t. with integrable
differential for almost every fixed $s\in[0,1]$ the map $x\mapsto
F_{\mu(.|s)}(x)$ is continuously differentiable with differential given by:

$\displaystyle\mathrm{D}_{x}F_{\mu(.|s)}(x)=\int_{\Omega}\mathrm{D}_{x}\phi_{\omega}(x)\mathrm{d}\mu(\omega|s).$

Moreover, by continuity of $s\mapsto x(s)$ the integrand
$\mathrm{D}_{x}\phi_{\omega}(x(t))$ is measurable and so is the map
$s\mapsto\mathrm{D}_{x}F_{\mu(.|s)}(x(s))$. Finally integrability follows as
$\mathrm{D}_{x}\phi_{\omega}(x(s))$ has 2-growth w.r.t. $\omega$ and
$\int_{\Omega}\|\omega\|^{2}\mathrm{d}\mu(\omega|s)$ is integrable on $[0,1]$.
∎

The following result gives an alternate point of view on the adjoint variable
$p$. Geometrically, it follows from eq. 11 that $p$ lives in the co-tangent
space of the flow $x$. Therefore in the case of a general (not necessarily
with finite support) data distribution
$\mathcal{D}\in\mathcal{P}(\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}})$ it is
convenient to see $p$ as the gradient of a potential $\psi$ over the variables
$(x,y)\in\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}$.

######  Lemma 3.1.

Let $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. Then for every
$(x,y)\in\mathbb{R}^{d+d^{\prime}}$ the associated adjoint variable $p$ can be
expressed for every $s\in[0,1]$ as:

$p_{\mu,x,y}(s)=\nabla_{x}\psi_{\mu}(s,x_{\mu}(s),y),$ (13)

where $\psi_{\mu}$ is the unique solution to the transport equation:

$\partial_{s}\psi_{\mu}+\langle\nabla_{x}\psi_{\mu},F_{\mu(.|s)}\rangle=0,\quad\psi_{\mu}(1,x,y)=\ell(x,y),\,\forall(x,y)\in\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}.$
(14)

###### Proof.

The solution to the transport equation can be given in the characteristic
form:

$\forall
s\in[0,1],x\in\mathbb{R}^{d},\quad\psi_{\mu}(s,x_{\mu}(s),y)=\psi_{\mu}(1,x_{\mu}(1),y)=\ell(x_{\mu}(1),y).$

One can then check that the r.h.s. of eq. 13 is indeed a solution of eq. 10. ∎

###  3.2 The gradient flow equation

We motivate here by formal computations a definition of a gradient flow
equation for the risk $L$. This _adjoint sensitivity analysis_ consists in
using a Lagrangian form of the risk minimization problem to obtain an
expression of the gradient w.r.t. the parameter $\mu$.

As the inputs are processed by our model through eq. 2 one can consider for a
parameter $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and every
time $s\in[0,1]$ the distribution $\rho_{\mu}(.|s)\coloneqq(x\mapsto
x_{\mu}(s),\mathrm{Id})_{\\#}\mathcal{D}$ of the data at time $s$. Then
$\\{\rho_{\mu}(.|s)\\}_{s\in[0,1]}$ is a narrowly continuous solution to the
continuity equation:

$\partial_{s}\rho^{*}_{\mu}(.|s)+\mathrm{div}_{x}(F_{\mu(.|s)}\rho^{*}_{\mu}(.|s))=0$
(15)

and the risk associated to $\mu$ is
$L(\mu)=\int_{\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}}\ell(x,y)\mathrm{d}\rho_{\mu}(x,y|1)=\rho_{\mu}(.|1)(\ell)$.
We introduce a Lagrange multiplier $\psi$ to penalize the above continuity
equation. For a parameter
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$, a measurable family
$\rho=\\{\rho(.|s)\\}_{s\in[0,1]}$ of probability measures over
$\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}$ and a smooth test function
$\psi:[0,1]\times\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}\to\mathbb{R}$,
consider the lagrangian $\mathcal{L}$ defined as:

$\mathcal{L}(\mu,\rho,\psi)\coloneqq\rho(.|1)(\ell)-\rho(.|1)(\psi(1))-\rho(.|0)(\psi(0))+\int_{0}^{1}\int_{\mathbb{R}^{d+d^{\prime}}}\left(\partial_{s}\psi+\langle\nabla_{x}\psi,F_{\mu(.|s)}\rangle\right)\mathrm{d}\rho(.|s)\mathrm{d}s.$
(16)

Using the definition of $F$ and inverting integrals, the variation of
$\mathcal{L}$ w.r.t. $\mu$ is given for every $\rho$ and $\psi$ by:

$\displaystyle\frac{\delta\mathcal{L}}{\delta\mu}(\mu,\rho,\psi):(s,\omega)\mapsto\int_{\mathbb{R}^{d+d^{\prime}}}\langle\nabla_{x}\psi(s,x,y),\phi(\omega,x)\rangle\mathrm{d}\rho(x,y|s).$

Also, if $\rho=\rho_{\mu}$ is the solution of eq. 15 for the parameter $\mu$,
we have the relation $\mathcal{L}(\mu,\rho_{\mu},\psi)=L(\mu)$ for any test
function $\psi$. Hence the variation of $L$ w.r.t. $\mu$ is:

$\displaystyle\frac{\delta
L}{\delta\mu}(\mu)=\frac{\delta\mathcal{L}}{\delta\mu}(\mu,\rho_{\mu},\psi)+\frac{\delta\mathcal{L}}{\delta\rho}(\mu,\rho_{\mu},\psi)\frac{\delta\rho_{\mu}}{\delta\mu}(\mu),$

where the _Lagrange multiplier_ $\psi$ can be chosen arbitrarily. Also the
variation of $\mathcal{L}$ w.r.t. the family of probability measures $\rho$,
seen as the probability measure whose disintegration on $[0,1]$ is
$\\{\rho_{s}\\}_{s\in[0,1]}$ (with the fixed initial condition
$\rho(.|0)=\mathcal{D}$), can be formally given by:

$\displaystyle\frac{\delta\mathcal{L}}{\delta\rho}(\mu,\rho,\psi)=\left(\ell-\psi(1)\right)\delta_{s=1}+\partial_{s}\psi+\langle\nabla_{x}\psi,F_{\mu(.|s)}\rangle.$

We see that taking $\psi=\psi_{\mu}$ to be a solution of eq. 14 cancels
$\frac{\delta\mathcal{L}}{\delta\rho}$ for every $\rho$ and hence:

$\displaystyle\frac{\delta L}{\delta\mu}(\mu)$
$\displaystyle=\frac{\delta\mathcal{L}}{\delta\mu}(\mu,\rho_{\mu},\psi_{\mu}).$

By Proposition 2.4 we know that for every absolutely continuous curve
$(\mu_{t})$ passing through $\mu$, its variation at $\mu$ is given by
$\partial_{t}\mu_{t}=-\mathrm{div}((0,v)\mu)$ for some $v\in L^{2}(\mu)$. A
notion of gradient of $L$ (for the “differential” structure of
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$) at $\mu$ could thus be
defined as the unique solution to the variational problem:

$\displaystyle\nabla L(\mu)\in\,\operatorname*{arg\,min}_{v\in
L^{2}(\mu)}\frac{1}{2}\|v\|^{2}_{L^{2}(\mu)}-\langle\nabla_{\omega}\frac{\delta
L}{\delta\mu}(\mu),v\rangle_{L^{2}(\mu)}.$

This problem admits a unique solution $v^{*}\in L^{2}(\mu)$, provided that
$\nabla_{\omega}\frac{\delta L}{\delta\mu}(\mu)\in L^{2}(\mu)$, and using the
relation eq. 13 between the adjoint variable $p_{\mu}$ and the potential
$\psi_{\mu}$ we have:

$\displaystyle v^{*}$ $\displaystyle=\nabla_{\omega}\frac{\delta
L}{\delta\mu}(\mu)=\nabla_{\omega}\frac{\delta\mathcal{L}}{\delta\mu}(\mu,\rho_{\mu},\psi_{\mu})$
$\displaystyle\text{that is}\quad v^{*}$
$\displaystyle:(s,\omega)\mapsto\int_{\mathbb{R}^{d+d^{\prime}}}\mathrm{D}_{\omega}\phi(\omega,x)^{\top}\nabla_{x}\psi_{\mu}(x,y)\mathrm{d}\rho_{\mu}(x,y|s)=\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{\mu}(s))^{\top}p_{\mu,x,y}(s).$

If the above calculations are purely formal they motivate the following
definition of gradient flow for $L$. In particular, this definition will be
shown in the next section to be equivalent to the appropriate notion of
gradient flow in metric spaces.

######  Definition 3.2 (Gradient flow).

Let $I\subset\mathbb{R}$ be an interval. For
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ let us define:

$\displaystyle\nabla
L[\mu]:(s,\omega)\mapsto\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{\mu}(s))^{\top}p_{\mu,x,y}(s).$
(17)

We say that a curve $t\in
I\mapsto\mu_{t}\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ is a
_gradient flow_ for $L$ if it is locally absolutely continuous curve
satisfying the continuity equation eq. 9 with $v_{t}=-\nabla L[\mu_{t}]$ for
$\mathrm{d}t$-a.e. $t\in I$.

The following result is a useful representation formula for the gradient flow
curves defined by Definition 3.2: for every $t\geq 0$ the gradient flow
$\mu_{t}$ at time $t$ is the pushforward of the initialization $\mu_{0}$ by a
flow-map. The proof relies on classical results from transport equation theory
[Amb08].

######  Proposition 3.2.

Assume $\phi$ is twice continuously differentiable and satisfies Assumptions
1, 2 and 3. Let $(\mu_{t})_{t\geq 0}$ be a gradient flow for the risk $L$ and
consider for every $t\geq 0$ the vector field:

$\displaystyle V_{t}:(s,\omega)\mapsto\left(0,\nabla
L[\mu_{t}](s,\omega)\right)=\left(0,\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{\mu_{t}}(s))^{\top}p_{\mu_{t},x,y}(s)\right)\in\mathbb{R}\times\Omega.$

Then for every $t\geq 0$ we have $\mu_{t}=(X_{t})_{\\#}\mu_{0}$ where $X_{t}$
is the flow-map solution of the ODE:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}X_{t}(s,\omega)=V_{t}(X_{t}(s,\omega)),\quad
X_{0}=\mathrm{Id}.$ (18)

###### Proof.

The existence and uniqueness of the flow-map $X_{t}$ for every $t\geq 0$
follows from the assumptions on $\phi$ (in particular linear growth and local-
Lipschitzness of $\mathrm{D}_{\omega}\phi$ w.r.t. $\omega$) and classical
theory of ODEs. The flow-map representation of $\mu_{t}$ then follows from
[Amb08, Thm.3.2] as for any initial value $(s,\omega)\in[0,1]\times\Omega$ the
set of curves solutions to the ODE is the singleton
$\\{(X_{t}(s,\omega))_{t\geq 0}\\}$. ∎

##### Consistency with the adjoint gradient flow

A case of particular interest for numerical applications is when the measure
$\mu$ is discretized and approximated at every time $s\in[0,1]$ by an
empirical distribution. Given $m\geq 0$ and $\theta=(\theta^{j})_{1\leq j\leq
m}\in L^{2}([0,1],\Omega)^{m}$ we define the associated empirical distribution
$\mu_{\theta}\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ by:

$\displaystyle\text{for $\mathrm{d}s$-a.e.
$s\in[0,1]$},\quad\mu_{\theta}(.|s)\coloneqq\frac{1}{m}\sum_{j=1}^{m}\delta_{\theta^{j}(s)}\mathrm{d}s.$
(19)

_i.e._ $\mu_{\theta}$ is the measure whose disintegration at any time
$s\in[0,1]$ is the empirical measure
$\frac{1}{m}\sum_{j=1}^{m}\delta_{\theta^{j}(s)}$. Then we note
$L(\theta)\coloneqq L(\mu_{\theta})$ the risk associated to $\theta$. In the
original work of [Che+18], the authors propose to train the Neural ODE
parameterized by $\theta$ and minimize $L(\theta)$ by performing gradient
descent for the _adjoint gradient_ defined as:

$\displaystyle\nabla_{\theta^{j}}L(\theta)\coloneqq\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\theta^{j},x(s))^{\top}p(s),$

where $x$ and $p$ are respectively the solutions of eqs. 3 and 10 for the
parameter $\mu_{\theta}$. One can observe that the adjoint gradient is the one
calculated by eq. 17 when $\mu=\mu_{\theta}$. Given sufficient regularity
assumptions on the feature map $\phi$, we have by Proposition 3.2 that
$(\mu_{t})_{t\geq 0}$ is a gradient flow in the sense of Definition 3.2 with
$\mu_{0}=\mu_{\theta_{0}}$ if and only if $\mu_{t}=\mu_{\theta_{t}}$ for every
$t\geq 0$ and $(\theta_{t})_{t\geq 0}$ is a gradient flow for the above
adjoint gradient.

###  3.3 Gradient flows as curves of maximal slope

There exists a large body of mathematical works devoted to the generalization
of the classical theory of gradient flows to functionals over metric spaces.
[AGS08] give an in-depth presentation of this theory. Complementary and more
synthetic presentations are given by [Amb+13, San17]. Based on those works, we
introduce here another definition of gradient flows for the risk $L$ which is
the one of _curves of maximal slope_ and show it coincides with the definition
from the previous section.

####  3.3.1 Curves of maximal slope in metric spaces

When $Z$ is a Euclidean space, and $f$ is a smooth function the gradient flow
of $f$ is defined as the solution of the ODE
$\frac{\mathrm{d}}{\mathrm{d}t}z_{t}=-\nabla f(z_{t})$. Then such a gradient
flow satisfies:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}f(z_{t})=-\|\nabla
f(z_{t})\|^{2}=-\frac{1}{2}\left(\|\frac{\mathrm{d}}{\mathrm{d}t}z_{t}\|^{2}+\|\nabla
f(z_{t})\|^{2}\right),$

whereas for any other smooth curve $(y_{t})$ we have by Young’s inequality:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}f(y_{t})=-\langle\nabla
f(y_{t}),\frac{\mathrm{d}}{\mathrm{d}t}y_{t}\rangle\geq-\frac{1}{2}\left(\|\frac{\mathrm{d}}{\mathrm{d}t}y_{t}\|^{2}+\|\nabla
f(y_{t})\|^{2}\right)$ (20)

with equality if and only if $\frac{\mathrm{d}}{\mathrm{d}t}y_{t}=-\nabla
f(y_{t})$. Hence, we see that imposing equality in the above inequality gives
a characterization of gradient flow curves in the Euclidean case. The
definition of _curves of maximal slope_ is based on the generalization of this
characterization to metric spaces. For example, a generalization of the
speed’s norm $\|\frac{\mathrm{d}}{\mathrm{d}t}z_{t}\|$ is given by the metric
derivative $\left|\frac{\mathrm{d}}{\mathrm{d}t}z_{t}\right|$. To give a sense
to the gradient’s norm $\|\nabla f(z)\|$ we need to introduce the notion of
_upper gradient_.

######  Definition 3.3 (Upper gradient [AGS08, Def.1.2.1]).

Let $(Z,d)$ be a complete metric space and $f:Z\to\mathbb{R}$ be a function.
The map $g:Z\to[0,+\infty]$ is a _upper gradient_ for $f$ if for every curve
absolutely continuous curve $(z_{t})_{t\in I}$ we have that $t\mapsto
g(z_{t})$ is measurable and:

$\displaystyle\left|f(z_{t_{1}})-f(z_{t_{2}})\right|\leq\int_{t_{1}}^{t_{2}}g(z_{t})\left|\frac{\mathrm{d}}{\mathrm{d}t}z_{t}\right|\mathrm{d}t,\quad\forall
t_{1}\leq t_{2}\in I.$

When no ambiguity an upper gradient of $f$ will simply be noted $|\nabla f|$.

Given an upper gradient for $f$, the definition of curves of maximal slope
consists in imposing equality in eq. 20.

######  Definition 3.4 (Curve of maximal slope [AGS08, Def.1.3.2]).

Let $(Z,d)$ be a complete metric space, $I\in\mathbb{R}$ be an interval and
$f:Z\to\mathbb{R}$ a function with $|\nabla f|$ an upper gradient for $f$. We
say that $(z_{t})_{t\in I}$ is a _curve a maximal slope_ for $f$ w.r.t.
$|\nabla f|$ if it satisfies:

  1. 1.

$(z_{t})_{t\in I}$ is locally absolutely continuous,

  2. 2.

the map $t\mapsto f(z_{t})$ is non-increasing,

  3. 3.

for $\mathrm{d}t$-a.e. $t\in I$ it holds
$\frac{\mathrm{d}}{\mathrm{d}t}f(z_{t})\leq-\frac{1}{2}\left(\left|\frac{\mathrm{d}}{\mathrm{d}t}z_{t}\right|^{2}+|\nabla
f|^{2}(z_{t})\right)$.

Moreover, if $\lim_{t\to\inf I}z_{t}=z$ exists then we say $(z_{t})$ is a
_curve of maximal slope starting at $z$_.

######  Remark 3.1 (About the various definitions of _curves of maximal
slope_).

There exist various definitions of the notion of curve of maximal slope in
metric spaces, see for example [Amb+13, Sec.4] for a discussion about the
various definitions and their relations. Our definition is the same as the one
used in [HM19, Def.2.12]. In particular, it implies the following _Energy
Dissipation Equality_ [HM19, Prop.2.14]:

$\displaystyle
f(z_{t_{1}})-f(z_{t_{2}})=\frac{1}{2}\int_{t_{1}}^{t_{2}}\left(\left|\frac{\mathrm{d}}{\mathrm{d}t}z_{t}\right|^{2}+|\nabla
f|^{2}(z_{t})\right)\mathrm{d}t,\quad\forall t_{1},t_{2}\in I.$ (EDE)

This definition differs from the one exposed in [SMP23, Def.2.2] (see also
[MS20, Def.4.4]) as the map $t\mapsto f(z_{t})$ need not be locally absolutely
continuous. The difference between these two definitions is discussed in
[SMP23, Rem.2.6] but observe that for our purpose _(i)_ the loss $L$ will be
shown to be locally Lipschitz in Corollary 3.1, hence implying that
$L(\mu_{t})$ is locally absolutely continuous, _(ii)_ the gradient norm
$\|\nabla L[\mu]\|_{L^{2}(\mu)}$ will be shown to be an upper gradient in
Proposition 3.4. For these reasons, we need not here make the distinction
between these two definitions and prefer weaker assumptions.

Note that there is _a priori_ no reasons for Definitions 3.2 and 3.4 to define
the same notion of gradient flow for the risk $L$. In particular, the first
definition uses the existence of the adjoint variable $p$ and thus some
regularity on $\phi$. In contrast, the second definition requires an upper
gradient which is yet unspecified for $L$. Taking appropriate assumptions on
the feature map $\phi$, we show in the rest of this section that the risk $L$
is sufficiently regular for the two definitions to coincide. This will be the
content of Theorem 3.1.

####  3.3.2 Curve of maximal slope for the risk $L$

Provided with Definition 3.4, we seek to characterize the curves of maximal
slope for the risk $L$ in the metric space
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. We first show the
$\mathrm{NODE}$’s output is a locally Lipschitz function of the parameter
$\mu\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$.

######  Assumption 2 (local Lipschitzness w.r.t. $\omega$).

Assume that $\phi:\Omega\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ is locally
Lipschitz w.r.t. $\omega$ with a Lipschitz constant that grows at most
linearly w.r.t. $\Omega$: for every $R\geq 0$ there exists a constant $C(R)$
s.t.:

$\displaystyle\forall x\in
B(0,R),\,\forall\omega,\omega^{\prime}\in\Omega,\quad\|\phi(\omega,x)-\phi(\omega^{\prime},x)\|\leq
C(R)(1+\max(\|\omega\|,\|\omega^{\prime}\|))\|\omega-\omega^{\prime}\|.$

######  Lemma 3.2 (local Lipschitzness of the flow).

Assume $\phi$ satisfies Assumptions 1 and 2 and consider some input
$x\in\mathbb{R}^{d}$. Then the map
$\mu\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)\mapsto(x_{\mu}(s))_{s\in[0,1]}\in\mathcal{C}([0,1],\mathbb{R}^{d})$
is locally Lipschitz. More precisely, for every $\mathcal{E}\geq 0$ there
exists a constant $C=C(\mathcal{E})$ s.t.:

$\displaystyle\sup_{s\in[0,1]}\|x_{\mu}(s)-x_{\mu^{\prime}}(s)\|\leq
Cd(\mu,\mu^{\prime}),$

for every
$\mu,\mu^{\prime}\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ with
$\mathcal{E}_{2}(\mu),\mathcal{E}_{2}(\mu^{\prime})\leq\mathcal{E}$. Moreover,
the constant $C$ can be chosen uniformly over $x$ in a compact set.

###### Proof.

Consider $x\in\mathbb{R}^{d}$, $\mathcal{E}\geq 0$ and $\mu,\mu^{\prime}$ such
as in the statement. We note $(x(s))_{s\in[0,1]}$ and
$(x^{\prime}(s))_{s\in[0,1]}$ the flow associated to $x$ and to the parameters
$\mu$ and $\mu^{\prime}$ respectively. Let $R\geq 0$ be such that $\|x\|\leq
R$. Then by Proposition 1.1 the trajectories $x,x^{\prime}$ are uniformly
bounded by some $R^{\prime}=R^{\prime}(R,\mathcal{E})$ . Then using eq. 2 we
have for every $s\in[0,1]$:

$\displaystyle\|x(s)-x^{\prime}(s)\|\leq$
$\displaystyle\|x(0)-x^{\prime}(0)\|+\|\int_{0}^{s}\int_{\Omega}\phi(\omega,x(r))\mathrm{d}\mu(\omega|r)\mathrm{d}r-\int_{0}^{s}\int_{\Omega}\phi(\omega,x^{\prime}(r))\mathrm{d}\mu^{\prime}(\omega|r)\mathrm{d}r\|$
$\displaystyle\leq$
$\displaystyle\|x(0)-x^{\prime}(0)\|+\int_{0}^{s}\int_{\Omega}\|\phi(\omega,x(r))-\phi(\omega,x^{\prime}(r))\|\mathrm{d}\mu(\omega|r)\mathrm{d}r$
$\displaystyle+\int_{0}^{s}\|\int_{\Omega}\phi(\omega,x^{\prime}(r))\mathrm{d}(\mu-\mu^{\prime})(\omega|r)\|\mathrm{d}r.$

For the first integral note that using the local Lipschitzness of $\phi$
w.r.t. $x$ in Assumption 1 we have for every $r\in[0,1]$:

$\displaystyle\int_{\Omega}\|\phi(\omega,x(r))-\phi(\omega,x^{\prime}(r))\|\mathrm{d}\mu(\omega|r)\leq
C_{1}\int_{\Omega}(1+\|\omega\|^{2})\mathrm{d}\mu(\omega|r)\|x(r)-x^{\prime}(r)\|,$

where $C_{1}=C_{1}(R,\mathcal{E})$. For the second integral, note that at
fixed $r\in[0,1]$, if $\gamma\in\Gamma_{o}(\mu(.|r),\mu^{\prime}(.|r))$ is a
(optimal) coupling between $\mu(.|r)$ and $\mu^{\prime}(.|r)$ then:

$\displaystyle\int_{\Omega}\phi(\omega,x^{\prime}(r))\mathrm{d}(\mu-\mu^{\prime})(\omega|r)=\int_{\Omega^{2}}(\phi(\omega,x^{\prime}(r))-\phi(\omega^{\prime},x^{\prime}(r)))\mathrm{d}\gamma(\omega,\omega^{\prime})\,.$

Using the local Lipschitzness of $\phi$ w.r.t. $\omega$ in Assumption 2 and
the optimality of $\gamma$:

$\displaystyle\|\int_{\Omega}\phi(\omega,x^{\prime}(r))\mathrm{d}(\mu-\mu^{\prime})(\omega|r)\|$
$\displaystyle\leq\int_{\Omega^{2}}\|\phi(\omega,x^{\prime}(r))-\phi(\omega^{\prime},x^{\prime}(r))\|\mathrm{d}\gamma(\omega,\omega^{\prime})$
$\displaystyle\leq\int_{\Omega^{2}}C_{2}(1+\max(\|\omega\|,\|\omega^{\prime}\|))\|\omega-\omega^{\prime}\|\mathrm{d}\gamma(\omega,\omega^{\prime})$
$\displaystyle\leq\sqrt{3}C_{2}\left(1+\mathcal{E}_{2}(\mu(.|r))+\mathcal{E}_{2}(\mu^{\prime}(.|r))\right)^{1/2}\mathcal{W}_{2}(\mu(.|r),\mu^{\prime}(.|r)),$

where $C_{2}=C_{2}(R,\mathcal{E})$. Integrating those inequalities gives by
Grönwall’s lemma that for $s\in[0,1]$:

$\displaystyle\|x(s)-x^{\prime}(s)\|\leq\sqrt{3}e^{C_{1}(1+\mathcal{E}_{2}(\mu))}C_{2}\int_{0}^{s}\left(1+\mathcal{E}_{2}(\mu(.|r))+\mathcal{E}_{2}(\mu^{\prime}(.|r))\right)^{1/2}\mathcal{W}_{2}(\mu(.|r),\mu^{\prime}(.|r))\mathrm{d}r\leq
C_{3}d(\mu,\mu^{\prime}),$

where $C_{3}=C_{3}(R,\mathcal{E})$. ∎

As an immediate corollary of the above proposition we get that, provided the
loss function $\ell$ is itself locally Lipschitz then the risk $L$ is also a
locally Lipschitz function of $\mu$.

######  Corollary 3.1 (local Lipschitzness of the risk).

Assume that $\phi$ satisfies Assumptions 1 and 2 and $\ell$ is locally
Lipschitz w.r.t. $x$. Then the risk map
$L:\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)\mapsto L(\mu)$ is
locally Lipschitz.

Assuming more regularity on the map $\phi$, one can express the first
variation $\delta x$ of the flow map with respect to a variation of the
parameter transported by a velocity field $v\in L^{2}(\mu)$.

######  Assumption 3 (Differentiability of $\phi$).

Assume that $\phi$ is continuously differentiable and s.t.

  1. 1.

$\mathrm{D}_{x}\phi$ grows at most quadratically with $\omega$: for every
$R\geq 0$ there exists a constant $C(R)$ s.t.

$\displaystyle\forall x\in
B(0,R),\,\forall\omega\in\Omega,\quad\|\mathrm{D}_{x}\phi(\omega,x)\|\leq
C(R)(1+\|\omega\|^{2}).$

  2. 2.

$\mathrm{D}_{\omega}\phi$ grows at most linearly with $\omega$: for every
$R\geq 0$ there exists a constant $C=C(R)$ s.t.

$\displaystyle\forall x\in
B(0,R),\,\forall\omega\in\Omega,\quad\|\mathrm{D}_{\omega}\phi(\omega,x)\|\leq
C(R)(1+\|\omega\|).$

######  Proposition 3.3.

Assume $\phi$ satisfies Assumptions 1, 2 and 3. Consider
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and a velocity field
$v:[0,1]\times\Omega\to\Omega$ in $L^{2}(\mu)$. For $t\in\mathbb{R}$ note
$\mu_{t}\coloneqq(\mathrm{Id}+t(0,v))_{\\#}\mu$. Then, for
$x\in\mathbb{R}^{d}$, $(x_{\mu_{t}})_{t\in\mathbb{R}}$ is differentiable in
$\mathcal{C}([0,1],\mathbb{R}^{d})$ at $t=0$ and $\delta
x\coloneqq\frac{\mathrm{d}}{\mathrm{d}t}x_{\mu_{t}}|_{t=0}$ is the solution
to:

$\displaystyle\forall s\in[0,1],\,\delta
x(s)=\int_{0}^{s}\mathrm{D}F_{\mu(.|r)}(x_{\mu}(r))\delta
x(r)\mathrm{d}r+\int_{0}^{s}\int_{\Omega}\mathrm{D}_{\omega}\phi(\omega,x_{\mu}(r))v(r,\omega)\mathrm{d}\mu(\omega|r)\mathrm{d}r.$
(21)

###### Proof.

First, thanks to Assumption 3, for $\nu\in\mathcal{P}_{2}(\Omega)$ the map
$F_{\nu}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is differentiable with
$\mathrm{D}F_{\nu}:x\mapsto\int_{\Omega}\mathrm{D}_{x}\phi(\omega,x)\mathrm{d}\nu(\omega)$.
Also, $\delta x$ is well-defined as the unique solution of eq. 21 and:

$\displaystyle\forall s\in[0,1],\quad\delta
x(s)=\int_{0}^{s}\int_{\Omega}\Phi_{\mu,x}(s)\Phi_{\mu,x}(r)^{-1}\mathrm{D}_{\omega}\phi(\omega,x(r))v(r,\omega)\mathrm{d}\mu(r,\omega).$

For simplicity, in the rest of the proof we will note $x_{t}\coloneqq
x_{\mu_{t}}$ for any $t\in\mathbb{R}$. Let us then show that $\delta x$ is the
derivative of $x_{t}$ at $t=0$. For $t\neq 0$, consider the normalized
increment:

$\displaystyle
z_{t}\coloneqq\frac{1}{t}(x_{t}-x_{0})\in\mathcal{C}([0,1],\mathbb{R}^{d}).$

Then we have by definition of $x_{t}$ and $x_{0}$ that for every $s\in[0,1]$:

$\displaystyle z_{t}(s)=$
$\displaystyle\frac{1}{t}\int_{0}^{s}\int_{\Omega}\phi(\omega,x_{t}(r))\mathrm{d}\mu_{t}(r,\omega)-\frac{1}{t}\int_{0}^{s}\int_{\Omega}\phi(\omega,x_{0}(r))\mathrm{d}\mu(r,\omega)$
$\displaystyle=$
$\displaystyle\frac{1}{t}\int_{0}^{s}\int_{\Omega}\left(\phi(\omega+tv(r,\omega),x_{t}(r))-\phi(\omega,x_{0}(r))\right)\mathrm{d}\mu(r,\omega)$
$\displaystyle=$
$\displaystyle\int_{0}^{s}\int_{\Omega}\left(\int_{0}^{1}\mathrm{D}_{x}\phi(\omega,x_{0}(r)+utz_{t}(r))\mathrm{d}u\right)\cdot
z_{t}(r)\mathrm{d}\mu(r,\omega)$
$\displaystyle+\int_{0}^{s}\int_{\Omega}\left(\int_{0}^{1}\mathrm{D}_{\omega}\phi(\omega+utv(r,\omega),x_{t}(r))\mathrm{d}u\right)\cdot
v(r,\omega)\mathrm{d}\mu(r,\omega).$

Hence $z_{t}$ is solution of the linear ODE
$z_{t}(s)=\int_{0}^{s}\left(A_{t}(r)\cdot z_{t}(r)+b_{t}(r)\right)\mathrm{d}r$
where we defined for $\mathrm{d}r$-a.e. $r\in[0,1]$:

$\displaystyle A_{t}(r)$
$\displaystyle\coloneqq\int_{\Omega}\int_{0}^{1}\mathrm{D}_{x}\phi(\omega,x_{0}(r)+utz_{t}(r))\mathrm{d}u\mathrm{d}\mu(\omega|r)$
$\displaystyle b_{t}(r)$
$\displaystyle\coloneqq\int_{\Omega}\int_{0}^{1}\mathrm{D}_{\omega}\phi(\omega+utv(r,\omega),x_{t}(r))\cdot
v(r,\omega)\mathrm{d}u\mathrm{d}\mu(\omega|r)$

and in order to prove that $z_{t}\to\delta x$ in
$\mathcal{C}([0,1],\mathbb{R}^{d})$ as $t\to 0$ it suffices to show that
$A_{t}$ and $b_{t}$ converge respectively in $L^{1}([0,1])$ to:

$\displaystyle
A(r)\coloneqq\int_{\Omega}\mathrm{D}_{x}\phi(\omega,x_{0}(r))\mathrm{d}\mu(\omega|r),\quad\text{and}\quad
b(r)\coloneqq\int_{\Omega}\mathrm{D}_{\omega}\phi(\omega,x_{0}(r))\cdot
v(r)\mathrm{d}\mu(\omega|r).$

Indeed, note that $d(\mu_{t},\mu)\leq t\|v\|_{L^{2}(\mu)}$ and the family
$(z_{t})_{t\in[-1,1]}$ is bounded $\mathcal{C}([0,1],\mathbb{R}^{d})$ by Lemma
3.2. Thus for $t\in\mathbb{R}$:

$\displaystyle\int_{0}^{1}\left|A_{t}(r)-A(r)\right|\mathrm{d}r\leq\int_{0}^{1}\int_{\Omega}\left|\int_{0}^{1}\mathrm{D}_{x}\phi(\omega,x_{0}(r)+utz_{t}(r))\mathrm{d}u-\mathrm{D}_{x}\phi(\omega,x_{0}(r))\right|\mathrm{d}\mu(r,\omega)\xrightarrow[t\to
0]{}0$

where Assumption 3 allows to bound the integrand by an integrable function and
to apply Lebesgue’s theorem, showing that $A_{t}\to A$ as $t\to 0$ in
$L^{1}([0,1])$. Similarly for $b_{t}$:

$\displaystyle\int_{0}^{1}\left|b_{t}(r)-b(r)\right|\mathrm{d}r\leq\int_{0}^{1}\int_{\Omega}\left|\int_{0}^{1}\mathrm{D}_{\omega}\phi(\omega+utv(r,\omega),x_{t}(r))\mathrm{d}u-\mathrm{D}_{\omega}\phi(\omega,x_{0}(r))\right|\|v(r,\omega)\|\mathrm{d}\mu(r,\omega)\xrightarrow[t\to
0]{}0.$

∎

A direct consequence of the previous result is the differentiability of the
flow map and consequently of the risk along absolutely continuous curves.

######  Corollary 3.2 (Differentiability of the flow).

Assume $\phi$ satisfies Assumptions 1, 2 and 3. Let $I\subset\mathbb{R}$ be an
interval and consider $(\mu_{t})_{t\in I}$ an absolutely continuous curve in
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ satisfying the continuity
equation:

$\displaystyle\partial_{t}\mu_{t}+\mathrm{div}((0,v_{t})\mu_{t})=0\quad\text{over
$I\times[0,1]\times\Omega$.}$

Consider some $x\in\mathbb{R}^{d}$. Then $(x_{\mu_{t}})_{t\in I}$ is an
absolutely continuous curve in $\mathcal{C}([0,1],\mathbb{R}^{d})$ and is
differentiable in $\mathcal{C}([0,1],\mathbb{R}^{d})$ for $\mathrm{d}t$-a.e.
$t\in I$ with $\delta x_{t}\coloneqq\frac{\mathrm{d}}{\mathrm{d}t}x_{\mu_{t}}$
the solution to:

$\displaystyle\forall s\in[0,1],\,\delta
x_{t}(s)=\int_{0}^{s}\mathrm{D}F_{\mu_{t}(.|r)}(x_{\mu_{t}}(r))\delta
x_{t}(r)\mathrm{d}r+\int_{0}^{s}\int_{\Omega}\mathrm{D}_{\omega}\phi(\omega,x_{\mu_{t}}(r))v_{t}(r,\omega)\mathrm{d}\mu_{t}(\omega|r)\mathrm{d}r.$
(22)

###### Proof.

For $t\in I$ we use the shortcut notation $x_{t}\coloneqq x_{\mu_{t}}$. The
fact that $(x_{t})_{t\in I}$ is absolutely continuous follows from Lemma 3.2
stating that the flow map is locally Lipschitz. To prove the result it hence
suffices to show that $\delta x_{t}$ is the derivative of $t\mapsto x_{t}$ in
$\mathcal{C}([0,1],\mathbb{R}^{d})$.

Note that, without loss of generality we can consider $v_{t}$ to be the
(uniquely defined) tangent velocity field of the curve $(\mu_{t})_{t\in I}$.
Indeed if $\tilde{v}_{t}$ is the tangent velocity field then we have by
Proposition 2.4 that in the sense of distributions:

$\displaystyle\mathrm{div}((0,v_{t}-\tilde{v}_{t})\mu_{t})=0.$

Hence for every $x\in\mathbb{R}^{d}$ and every $s\in[0,1]$:

$\displaystyle\int_{0}^{s}\int_{\Omega}\mathrm{D}_{\omega}\phi(\omega,x_{t}(r))v_{t}(r,\omega)\mathrm{d}\mu_{t}(r,\omega)=\int_{0}^{s}\int_{\Omega}\mathrm{D}_{\omega}\phi(\omega,x_{t}(r))\tilde{v}_{t}(r,\omega)\mathrm{d}\mu_{t}(r,\omega)$

and the definition of $\delta x_{t}$ stays unchanged. Then, assuming $v_{t}$
is the tangent velocity field to the curve $\mu_{t}$, we can consider a subset
$\Lambda\subset I$ of full Lebesgue measure such that the conclusions of Lemma
2.2 hold. For every $t\in\Lambda$ and every $h\neq 0$ consider
$\tilde{\mu}^{h}_{t}\coloneqq(\mathrm{Id}+h(0,v_{t}))_{\\#}\mu_{t}$ and
$\tilde{x}^{h}_{t}$ the associated flow. Then by Proposition 3.3:

$\displaystyle\|\frac{x_{t+h}-x_{t}}{h}-\delta
x_{t}\|_{\mathcal{C}([0,1])}\leq\|\frac{\tilde{x}^{h}_{t}-x_{t}}{h}-\delta
x_{t}\|_{\mathcal{C}([0,1])}+\|\frac{x_{t+h}-\tilde{x}^{h}_{t}}{h}\|_{\mathcal{C}([0,1])}\xrightarrow[h\to
0]{}0$

where the first term goes to $0$ by application of Proposition 3.3.The second
term also goes to $0$ by the fact that the flow map is locally Lipschitz, thus
$\|x_{t+h}-\tilde{x}^{h}_{t}\|_{\mathcal{C}([0,1])}\leq
Cd(\mu_{t+h},\tilde{\mu}^{h}_{t})$ for some constant $C$ and
$\frac{1}{h}\|x_{t+h}-\tilde{x}^{h}_{t}\|_{\mathcal{C}([0,1])}\to 0$ by Lemma
2.2. Note that, as $\Lambda\subset I$ is independent of $x$, it follows that
the curve $t\mapsto x_{t}$ is differentiable at every $t\in\Lambda$ for every
$x\in\mathbb{R}^{d}$. ∎

######  Corollary 3.3 (Differentiability of the loss).

Assume $\phi$ satisfies Assumptions 1, 2 and 3 and $\ell$ is continuously
differentiable. Let $I\subset\mathbb{R}$ be an interval and $(\mu_{t})_{t\in
I}$ be as in Corollary 3.2. Then $(L(\mu_{t}))_{t\in I}$ is absolutely
continuous and for almost every $t\in I$:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}L(\mu_{t})=\int_{[0,1]\times\Omega}\langle\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{\mu_{t}}(s))^{\top}p_{\mu_{t},x,y}(s),v_{t}(s,\omega)\rangle\mathrm{d}\mu_{t}(s,\omega).$

###### Proof.

First, the fact that $t\mapsto L(\mu_{t})$ is absolutely continuous follows
from the fact that $\mu\mapsto L(\mu)$ is locally Lipschitz, as shown in
Corollary 3.1. It remains to show the formula for its derivative.

For $t\in I$ and $(x,y)\in\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}}$ use the
shortcut notations $x_{t}\coloneqq x_{\mu_{t}}$, $p_{t}\coloneqq
p_{\mu_{t},x,y}$ and $\Phi_{t}\coloneqq\Phi_{\mu_{t},x}$. By the proof of
Corollary 3.2 we know that there exists a subset $\Lambda\subset I$ of full
Lebesgue measure such that for every $t\in\Lambda$, the map $t\mapsto x_{t}$
is differentiable at $t$ for every $x\in\mathbb{R}^{d}$. By Lebesgue theorem,
the map $t\mapsto L(\mu_{t})$ is differentiable at every $t\in\Lambda$ with:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}L(\mu_{t})=\mathbb{E}_{x,y}\langle\nabla_{x}\ell(x_{t}(1),y),\delta
x_{t}(1)\rangle,$

where, at fixed $x\in\mathbb{R}^{d}$, $\delta x_{t}$ verifies eq. 22 and is
given by

$\delta
x_{t}(1)=\int_{[0,1]\times\Omega}\Phi_{t}(1)\Phi_{t}(s)^{-1}\mathrm{D}_{\omega}\phi(\omega,x_{t}(s))v_{t}(s,\omega)\mathrm{d}\mu_{t}(s,\omega).$

Also, the adjoint variable $p_{t}$ is given by
$p_{t}(s)=\Phi_{t}(s)^{-\top}\Phi_{t}(1)^{\top}\nabla_{x}\ell(x_{t}(1),y)$.
Hence by inverting the integration order, we see that for $t\in\Lambda$:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}L(\mu_{t})=\mathbb{E}_{x,y}\langle\nabla_{x}\ell(x_{t}(1),y),\delta
x_{t}(1)\rangle=\int_{[0,1]\times\Omega}\langle\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{t}(s))^{\top}p_{t}(s),v_{t}(s,\omega)\rangle\mathrm{d}\mu_{t}(s,\omega).$

∎

Thanks to Corollary 3.2 (which can be seen as a chain-rule formula) one can
show that the gradient norm $\|\nabla L[\mu]\|_{L^{2}(\mu)}$ gives an upper
gradient for the risk $L$ in the sense of Definition 3.3. Moreover the
following Proposition 3.4 shows it corresponds to the notion of _local slope_
([AGS08, Def.1.2.4]). The result relies on the following lemma.

######  Lemma 3.3 (Continuity of the adjoint variable $p$).

Assume $\phi$ satisfies Assumptions 1, 2 and 3. Then, for fixed
$(x,y)\in\mathbb{R}^{d+d^{\prime}}$, the map $\mu\mapsto
p_{\mu,x,y}\in\mathcal{C}([0,1]\times\mathbb{R}^{d})$ is $d$-continuous on
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$.

###### Proof.

Let $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and consider a
sequence $(\mu_{n})_{n\geq 0}$ in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ s.t. $d(\mu_{n},\mu)\to
0$. Fix a pair $(x,y)\in\mathbb{R}^{d+d^{\prime}}$ and use the shortcuts
$x_{n}\coloneqq x_{\mu_{n}}$ (resp. $x\coloneqq x_{\mu}$) and $p_{n}\coloneqq
p_{\mu_{n},x,y}$ (resp. $p\coloneqq p_{\mu,x,y}$). By Lemma 3.2 we already
have $x_{n}\to x$ in $\mathcal{C}([0,1])$ and we show now that $p_{n}\to p$ in
$\mathcal{C}([0,1])$ using Ascoli’s theorem.

Remark that by the assumptions on $\phi$, all the trajectories $x$, $x_{n}$,
$p$ and $p_{n}$ stay in a bounded set $B(0,R)$ for some $R\geq 0$. Also, as
$\mu_{n}\to\mu$, we have that the sequence $(\mu_{n})$ has uniformly
integrable second moment and for every $\varepsilon>0$ we can find a $k\geq 0$
s.t.

$\int_{\|\omega\|\geq k}(1+\|\omega\|^{2})\mathrm{d}\mu,\,\sup_{n\geq
0}\int_{\|\omega\|\geq k}(1+\|\omega\|^{2})\mathrm{d}\mu_{n}\leq\varepsilon.$

Then for $n\geq 0$ and $s_{1}\leq s_{2}\in[0,1]$ we have by eq. 10 and the
assumptions on $\phi$:

$\displaystyle\|p_{n}(s_{2})-p_{n}(s_{1})\|$
$\displaystyle\leq\int_{s_{1}}^{s_{2}}\int_{\Omega}\|\mathrm{D}_{x}\phi(\omega,x_{n}(r))\|\|p_{n}(r)\|\mathrm{d}\mu_{n}(r,\omega)$
$\displaystyle\leq
C\int_{s_{1}}^{s_{2}}\int_{\Omega}(1+\|\omega\|^{2})\mathrm{d}\mu_{n}(r,\omega)$
$\displaystyle\leq C(\varepsilon+(1+k^{2})|s_{2}-s_{1}|),$

where $C=C(R)$. Hence the sequence $(p_{n})_{n\geq 0}$ is equicontinuous and,
up to a subsequence, we have $p_{n}\to\bar{p}\in\mathcal{C}([0,1])$. Let us
then show $\bar{p}=p$. Indeed using the initial condition we have for $n\geq
0$ and $s\in[0,1]$:

$\displaystyle
p_{n}(s)=\nabla_{x}\ell(x_{n}(1),y)+\int_{s}^{1}\int_{\Omega}\mathrm{D}_{x}\phi(\omega,x_{n}(r))^{\top}p_{n}(r)\mathrm{d}\mu_{n}(r,\omega).$

First we have
$\nabla_{x}\ell(x_{n}(1),y)\xrightarrow[n\to\infty]{}\nabla_{x}\ell(x(1),y)$.
Also, note that by the assumptions on $\phi$ we have
$\mathrm{D}_{x}\phi(\omega,x_{n}(r))^{\top}p_{n}(r)\leq C(1+\|\omega\|^{2})$
and
$\mathrm{D}_{x}\phi(\omega,x_{n}(r))^{\top}p_{n}(r)\to\mathrm{D}_{x}\phi(\omega,x(r))^{\top}\bar{p}(r)$
locally uniformly over $[0,1]\times\Omega$. Hence by the properties of
$\mathcal{W}_{2}$-convergence, we can take the limit in the above equation to
obtain:

$\displaystyle\bar{p}(s)=\nabla_{x}\ell(x(1),y)+\int_{s}^{1}\int_{\Omega}\mathrm{D}_{x}\phi(\omega,x(r))^{\top}\bar{p}(r)\mathrm{d}\mu(r,\omega),$

_i.e._ $\bar{p}=p$ by uniqueness of the solutions to eq. 10. ∎

######  Lemma 3.4 (Continuity of $\|\nabla L(\mu)\|_{L^{2}(\mu)}$ ).

Assume $\phi$ satisfies Assumptions 1, 2 and 3. Then the map
$\mu\mapsto\|\nabla L[\mu]\|_{L^{2}(\mu)}$ is $d$-continuous on
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$.

###### Proof.

Let $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and consider a
sequence $(\mu_{n})_{n\geq 0}$ in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ s.t. $d(\mu_{n},\mu)\to
0$. For an input $x\in\mathbb{R}^{d}$, note $x_{n}$ (resp. $x$) the flow
associated to $\mu_{n}$ (resp. $\mu$) and starting from $x$. Similarly
introduce the adjoint variables $(p_{n})_{n\geq 0}$ and $p$. Then by Lemma 3.2
and Lemma 3.3 we have that $x_{n}\to x$ and $p_{n}\to p$ in
$\mathcal{C}([0,1],\mathbb{R}^{d})$. As a consequence the sequence of
continuous maps

$f_{n}:(r,\omega)\mapsto\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{n}(r))^{\top}p_{n}(r)$

converges locally uniformly towards the map
$f:(r,\omega)\mapsto\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x(r))^{\top}p(r)$
and is uniformly bounded by a function of linear growth. As $d$-convergence
implies $\mathcal{W}_{2}$-convergence and by the properties of
$\mathcal{W}_{2}$-convergence ([Vil09, Thm.6.9]) this implies:

$\displaystyle\|\nabla
L[\mu_{n}]\|^{2}_{L^{2}(\mu_{n})}=\int_{[0,1]\times\Omega}\|f_{n}\|^{2}\mathrm{d}\mu_{n}\xrightarrow[n\to\infty]{}\int_{[0,1]\times\Omega}\|f\|^{2}\mathrm{d}\mu=\|\nabla
L[\mu]\|^{2}_{L^{2}(\mu)}.$

∎

######  Proposition 3.4 ($\|\nabla L(\mu)\|_{L^{2}(\mu)}$ is an upper-
gradient).

Assume $\phi$ satisfies Assumptions 1, 2 and 3. Let
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$, then $\|\nabla
L[\mu]\|_{L^{2}}(\mu)$ is the _local slope_ of the risk $L$ at $\mu$, that is:

$\displaystyle\|\nabla
L[\mu]\|_{L^{2}(\mu)}=\limsup_{\nu\to\mu}\frac{\left(L(\mu)-L(\nu)\right)^{+}}{d(\mu,\nu)}.$
(23)

Moreover, it is an upper-gradient in the sense of Definition 3.3.

###### Proof.

The last part of the result follows from Proposition 2.4, since if
$(\mu_{t})_{t\in I}$ is an absolutely continuous curve then it satisfies the
continuity equation with a vector field $v$ such that
$\|v_{t}\|_{L^{2}(\mu_{t})}\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|$
for a.e. $t\in I$. Hence by Corollary 3.3 and Cauchy-Schwarz we have:

$\displaystyle\forall t_{1}\leq t_{2}\in
I,\quad\left|L(\mu_{t_{1}})-L(\mu_{t_{2}})\right|\leq\int_{t_{1}}^{t_{2}}\|\nabla
L[\mu_{t}]\|_{L^{2}(\mu_{t})}\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|\mathrm{d}t.$

Let us then show eq. 23. Consider some parameter
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ and note $|\nabla
L|(\mu)$ the r.h.s. of eq. 23. Then for $\varepsilon>0$, by continuity of
$\|\nabla L[\mu]\|_{L^{2}(\mu)}$ (Lemma 3.4) and by definition of $|\nabla
L|(\mu)$ one can find a
$\nu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ s.t.:

$\displaystyle\frac{\left(L(\mu)-L(\nu)\right)^{+}}{d(\mu,\nu)}\geq\left|\nabla
L\right|(\mu)-\varepsilon,\quad\text{and}\quad\left|\|\nabla
L[\mu]\|_{L^{2}(\mu)}-\|\nabla
L[\nu^{\prime}]\|_{L^{2}(\nu^{\prime})}\right|\leq\varepsilon,\,\text{if
$d(\mu,\nu^{\prime})\leq d(\mu,\nu)$}.$

Consider $(\mu_{t})_{t\in[0,1]}$ a constant speed geodesics with endpoints
$\mu_{0}=\mu$ and $\mu_{1}=\nu$ (such a geodesic can easily be constructed by
similarity with classical Wasserstein geodesics, see [AGS08, Thm.7.2.2]). Then
by Proposition 2.4 the tangent velocity field $v$ of the curve $(\mu_{t})$
satisfies for $\mathrm{d}t$-a.e. $t\in[0,1]$,
$\|v_{t}\|_{L^{2}(\mu_{t})}\leq\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|=d(\mu,\nu)$
and using Corollary 3.3:

$\displaystyle L(\nu)$ $\displaystyle=L(\mu)+\int_{0}^{1}\langle\nabla
L[\mu_{t}],v_{t}\rangle_{L^{2}(\mu_{t})}\mathrm{d}t$ $\displaystyle\leq
L(\mu)+d(\mu,\nu)\int_{0}^{1}\|\nabla
L[\mu_{t}]\|_{L^{2}(\mu_{t})}\mathrm{d}t$ $\displaystyle\leq
L(\mu)+d(\mu,\nu)\|\nabla L[\mu]\|_{L^{2}(\mu)}+\varepsilon.$

Similarly we have $L(\mu)\leq L(\nu)+d(\mu,\nu)\|\nabla
L[\mu]\|_{L^{2}(\mu)}+\varepsilon$ and hence $\left|\nabla
L\right|(\mu)\leq\|\nabla L[\mu]\|_{L^{2}(\mu)}+\varepsilon$.

For the converse inequality consider for $t\in\mathbb{R}$ the parameter
$\mu_{t}=(\mathrm{Id}+t(0,\nabla L[\mu]))_{\\#}\mu$. Then, by Proposition 3.3
with $v=\nabla L[\mu]$, the map $t\mapsto L(\mu_{t})$ is differentiable at
$t=0$ and applying the same calculations as in Corollary 3.3:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}L[\mu_{t}]\Bigr{|}_{t=0}=\langle\nabla
L[\mu],v\rangle_{L^{2}(\mu)}=\|\nabla L[\mu]\|^{2}_{L^{2}(\mu)}.$

Hence observing that $d(\mu_{t},\mu)\leq t\|v\|_{L^{2}(\mu_{t})}$ we have
$\liminf_{t\to
0^{+}}\frac{(L(\mu_{t})-L(\mu))^{+}}{d(\mu_{t},\mu)}\geq\|\nabla
L[\mu]\|_{L^{2}(\mu)}$. ∎

As a consequence of the previous result, we will from now on only consider as
upper gradient of $L$ the one given for every
$\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ by:

$\displaystyle|\nabla L|(\mu)\coloneqq\|\nabla
L[\mu]\|_{L^{2}(\mu)}=\left(\int_{[0,1]\times\Omega}\|\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x(s))^{\top}p(s)\|^{2}\mathrm{d}\mu(s,\omega)\right)^{1/2}.$
(24)

Note that the vector field $\nabla L[\mu]$ was used in Definition 3.2 to
define the notion of _gradient flow_ whereas the upper-gradient $|\nabla
L|(\mu)$ is used in the Definition 3.4 of curves of maximal slope. The
following theorem is the main result of this section and shows these two
notions coincide.

######  Theorem 3.1.

Assume $\phi$ satisfies Assumptions 1, 2 and 3 and $\ell$ is smooth. Let
$I\subset\mathbb{R}$ be an open interval. Then a curve $(\mu_{t})_{t\in I}$ is
a _gradient flow_ in the sense of Definition 3.2 if and only if it is a _curve
of maximal slope_ for $L$ in the sense of Definition 3.4.

###### Proof.

_Part 3: Gradient flows are curves of maximal slope._

Let $(\mu_{t})_{t\in I}$ be a gradient flow for $L$ in the sense of Definition
3.2. Then $(\mu_{t})$ is a locally absolutely continuous curve satisfying the
continuity equation $\partial_{t}\mu_{t}+\mathrm{div}(v_{t}\mu_{t})=0$ with
$v_{t}=-\nabla L[\mu_{t}]$. Hence by Proposition 2.4 we have for a.e. $t\in
I$:

$\displaystyle\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|\leq\|v_{t}\|_{L^{2}(\mu_{t})}=\|\nabla
L[\mu_{t}]\|_{L^{2}(\mu_{t})}.$

Also, by Corollary 3.2, $(L(\mu_{t}))_{t\in I}$ is absolutely continuous with
for a.e. $t\in I$:

$\displaystyle-\frac{\mathrm{d}}{\mathrm{d}t}L(\mu_{t})=\langle v_{t},\nabla
L[\mu_{t}]\rangle_{L^{2}(\mu_{t})}=\|\nabla
L[\mu_{t}]\|^{2}_{L^{2}(\mu_{t})}.$

Thus recalling that $|\nabla L|(\mu)=\|\nabla L[\mu]\|_{L^{2}(\mu)}$ we get
Definition 3.4 by putting together the two previous equations.

_Part 4: Curves of maximal slope are gradient flows._

Let $(\mu_{t})_{t\in I}$ be a curve of maximal slope for $L$ in the sense of
Definition 3.4. Then in particular $(\mu_{t})_{t\in I}$ is locally abolutely
continuous in $(\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega),d)$ and by
Proposition 2.4 there exists a Borel velocity field
$v:I\times[0,1]\times\Omega\to\Omega$ such that $\mu$ satisfies the continuity
equation:

$\displaystyle\partial_{t}\mu_{t}+\mathrm{div}((0,v_{t})\mu_{t})=0,\quad\text{over
$I\times[0,1]\times\Omega$,}$

and such that the metric derivative satisfies
$|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}|\geq\|v_{t}\|_{L^{2}(\mu_{t})}$ for
a.e. $t\in I$. Hence it follows from Corollary 3.3 that $(L(\mu_{t}))_{t\in
I}$ is absolutely continuous and for a.e. $t\in I$:

$\displaystyle-\frac{\mathrm{d}}{\mathrm{d}t}L(\mu_{t})=-\langle v_{t},\nabla
L[\mu_{t}]\rangle.$

Using the EDE condition we thus have:

$\displaystyle-\langle v_{t},\nabla
L[\mu_{t}]\rangle\geq\frac{1}{2}(|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}|^{2}+|\nabla
L|^{2}(\mu_{t}))\geq\frac{1}{2}(\|v_{t}\|^{2}_{L^{2}(\mu_{t})}+\|\nabla
L[\mu_{t}]\|^{2}_{L^{2}(\mu_{t})})$

from which it follows by Young’s inequality that $v_{t}=-\nabla L[\mu_{t}]$ in
$L^{2}(\mu_{t})$ for a.e. $t\in I$. ∎

Note that, although it does not appear in Definition 3.2, the above
equivalence shows that if $(\mu_{t})_{t\in I}$ is a gradient flow for $L$ then
$\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|=\|\nabla
L[\mu_{t}]\|_{L^{2}(\mu_{t})}$ _i.e._ $\nabla L[\mu_{t}]$ is in fact the
(uniquely defined) tangent velocity field of the curve $(\mu_{t})_{t\in I}$.

###  3.4 Existence, uniqueness, and stability of gradient flow curves

We show here the well-posedness result for the gradient flow equation of the
risk $L$, namely we show the existence, uniqueness, and stability of gradient
flow curves starting from any initialization
$\mu_{0}\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$. For the
“existence” part we will rely on classical results from the theory of gradient
flows in metric spaces showing the convergence of proximal sequences towards a
curve known as _(Generalized) Minimising Movements_ [De ̵93]. For the
“uniqueness” part we will show that gradient flow trajectories are stable,
that is if two initializations $\mu_{0},\mu_{0}^{\prime}$ are close (in the
sense of the metric $d$), then the emanating gradient flow curves
$(\mu_{t})_{t\geq 0},(\mu_{t}^{\prime})_{t\geq 0}$ stay close in finite time.

####  3.4.1 Existence

We proceed to show the existence of gradient flow curves as defined in
Definition 3.2. For that purpose, we need a strengthening of Assumption 1.
Notably, Assumption A allows to show the flow map $\mu\mapsto x_{\mu}$ is
continuous for the topology of narrow convergence over
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$.

######  Assumption A.

For some $\alpha\in[1,2)$ we assume that:

  1. 1.

The feature map $\phi$ has $\alpha$-growth w.r.t. $\omega$, locally w.r.t.
$x$. That is for every compact $K\subset\mathbb{R}^{d}$ there exists a
constant $C=C(K)$ s.t:

$\displaystyle\forall x\in
K,\forall\omega\in\Omega,\quad\|\phi(\omega,x)\|\leq
C(1+\|\omega\|^{\alpha}).$

  2. 2.

The feature map $\phi$ is continuously differentiable and its differential
$\mathrm{D}_{x}\phi$ w.r.t. $x$ has $\alpha$-growth w.r.t. $\omega$, locally
w.r.t. $x$. That is for every compact $K\subset\mathbb{R}^{d}$ there exists a
constant $C=C(K)$ s.t.:

$\displaystyle\forall x\in
K,\forall\omega\in\Omega,\quad\|\mathrm{D}_{x}\phi(\omega,x)\|\leq
C(1+\|\omega\|^{\alpha}).$

######  Theorem 3.2 (Existence of curves of maximal slope).

Assume $\phi$ satisfies Assumptions 1, 2 and 3 and Assumption A. Let
$\mu_{0}\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. Then there
exists a curve of maximal slope $(\mu_{t})_{t\in[0,+\infty)}$ starting from
$\mu_{0}$ and
$\left(\left|\frac{\mathrm{d}}{\mathrm{d}t}\mu_{t}\right|\right)_{t\geq 0}\in
L^{2}_{loc}([0,+\infty))$.

###### Proof.

The result follows from the successive application of [AGS08, Thm.2.2.3 and
Thm.2.3.3], the first result ensuring the existence of _Generalized Minimizing
Movements_ and the second result stating that these curves are curves of
maximal slope for the local slope. The proof proceeds by verifying the
assumptions of these theorems. We consider here
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ provided with the topology
induced by the distance $d$ and with the (weaker) topology of narrow
convergence, noted $\tau$. Note that
$(\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega),d)$ is a complete metric
space (Proposition 2.3) and that the distance $d$ is $\tau$-lower-semi-
continuous (Lemma 2.1).

_Part 5: $d$-bounded sets are $\tau$-relatively compact._

This property is verified as $d$-bounded sets are tight and hence
$\tau$-relatively compact by Prokhorov’s theorem.

_Part 6: $L$ is $\tau$-continuous on $d$-bounded sets._

Let $(\mu_{n})$ be a $d$-bounded sequence in
$\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ such that
$\mu_{n}\xrightarrow{\tau}\mu$ for some
$\mu\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$ and let us show that
$L(\mu_{n})\to L(\mu)$. Take $x\in\mathbb{R}^{d}$ and note $x_{n}\coloneqq
x_{\mu_{n}}$ the flow trajectory starting from $x$ and associated to
$\mu_{n}$. By Lebesgue’s theorem, it suffices to show that $x_{n}(1)\to
x_{\mu}(1)$. Using Ascoli’s theorem, we will proceed by showing that $x_{n}\to
x_{\mu}$ in $\mathcal{C}([0,1],\mathbb{R}^{d})$.

From the $d$-boundedness and the proof of Proposition 1.1, it follows that the
trajectories $x_{n}$ stay in a compact set. Moreover given
$s_{1}<s_{2}\in[0,1]$ we have using the $\alpha$-growth assumption:

$\displaystyle\|x_{n}(s_{2})-x_{n}(s_{1})\|\leq\int_{[s_{1},s_{2}]\times\Omega}\|\phi(\omega,x_{n}(r))\|\mathrm{d}\mu_{n}(r,\omega)\leq\int_{[s_{1},s_{2}]\times\Omega}C(1+\|\omega\|^{\alpha})\mathrm{d}\mu_{n}(r,\omega).$

Also as the sequence $(\mu_{n})$ is $d$-bounded it has uniformly integrable
$\alpha$-moments. Given $\varepsilon>0$ we can thus find a $k\geq 0$ such
that, for every $n\geq 0$, $\int_{\|\omega\|\geq
k}C(1+\|\omega\|^{\alpha})\mathrm{d}\mu_{n}\leq\varepsilon$. Using this in the
previous inequality and the fact that the marginal of $\mu_{n}$ on $[0,1]$ is
the Lebesgue measure gives:

$\displaystyle\forall
s_{1}<s_{2}\in[0,1],\quad\|x_{n}(s_{2})-x_{n}(s_{1})\|\leq\varepsilon+C(1+k^{\alpha})(s_{2}-s_{1}).$

Thus the trajectories $(x_{n})$ are equicontinuous and, by Arzela-Ascoli’s
theorem, we have (up to a subsequence) that $x_{n}\to\bar{x}$ in
$\mathcal{C}([0,1],\mathbb{R}^{d})$.

Let us show that $\bar{x}=x_{\mu}$ is the flow generated by $\mu$. This will
conclude this part of the proof as it will imply $x_{n}(1)\to x_{\mu}(1)$ and
then $L(\mu_{n})\to L(\mu)$ by Lebesgue’s convergence theorem. Considering
$s\in[0,1]$ we have:

$\displaystyle x_{n}(s)=x$
$\displaystyle+\int_{[0,1]\times\Omega}\mathbbm{1}_{r\leq
s}\phi(\omega,x_{n}(r))\mathrm{d}\mu_{n}(r,\omega)$ $\displaystyle=x$
$\displaystyle+\int_{[0,1]\times\Omega}\mathbbm{1}_{r\leq
s}\phi(\omega,\bar{x}(r))\mathrm{d}\mu(r,\omega)$
$\displaystyle+\int_{[0,1]\times\Omega}\mathbbm{1}_{r\leq
s}\left(\phi(\omega,x_{n}(r))-\phi(\omega,\bar{x}(s))\right)\mathrm{d}\mu_{n}(r,\omega)$
(L1) $\displaystyle+\int_{[0,1]\times\Omega}\mathbbm{1}_{r\leq
s}\phi(\omega,\bar{x}(r))\mathrm{d}(\mu_{n}-\mu)(r,\omega).$ (L2)

In this last equality, we need to show that L1 and L2 vanish as $n\to\infty$.
In L1 we have that the integrand has $\alpha$-growth, hence given
$\varepsilon>0$ we have using the uniform integrability of
$\|\omega\|^{\alpha}$ on the sequence $(\mu_{n})$ that for every $n\geq 0$:

$\displaystyle\|\int\mathbbm{1}_{r\leq
s}\left(\phi(\omega,x_{n}(r))-\phi(\omega,\bar{x}(r))\right)\mathrm{d}\mu_{n}(r,\omega)\|\leq\varepsilon+\int\mathbbm{1}_{r\leq
s,\|\omega\|\leq
k}\|\phi(\omega,x_{n}(r))-\phi(\omega,\bar{x}(r))\|\mathrm{d}\mu_{n}(r,\omega).$

Then as $\phi$ is locally-Lipschitz w.r.t. $x$ and $x_{n}\to\bar{x}$ we have
that the integrand on the r.h.s. converges uniformly to 0 and hence:

$\displaystyle\limsup_{n\to\infty}\|\int_{[0,1]\times\Omega}\mathbbm{1}_{r\leq
s}\left(\phi(\omega,x_{n}(r))-\phi(\omega,\bar{x}(s))\right)\mathrm{d}\mu_{n}(r,\omega)\|\leq\varepsilon.$

In L2 the integrand is not continuous so we can’t simply apply the definition
of narrow convergence and need to leverage the fact that $(\mu_{n})$ is a
$d$-bounded sequence in $\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$.
Note that for every $(r,\omega)\in[0,1]\times\Omega$,
$\|\phi(\omega,\bar{x}(r))\|\leq C(1+\|\omega\|^{\alpha})$. Given
$\varepsilon>0$ and using the uniform integrability of $\|\omega\|^{\alpha}$
we can thus have a $k\geq 0$ such that:

$\displaystyle\sup_{n\geq 0}\int_{\|\omega\|\geq
k}C(1+\|\omega\|^{\alpha})\mathrm{d}\mu_{n},\,\int_{\|\omega\|\geq
k}C(1+\|\omega\|^{\alpha})\mathrm{d}\mu\,\leq\varepsilon.$

Then whenever $\delta>0$ we can find a continuous function
$\varphi:[0,1]\times\Omega\to\mathbb{R}^{d}$ such that
$\|\varphi(r,\omega)\|\leq C(1+\|\omega\|^{\alpha})$,
$\varphi(r,\omega)=\phi(\omega,\bar{x}(r))$ for every $r\leq s$ and
$\varphi(r,\omega)=0$ whenever $r\geq s+\delta$. Considering such a function
$\varphi$ we have for L2:

$\displaystyle\|\int\mathbbm{1}_{r\leq
s}\phi(\omega,\bar{x}(r))\mathrm{d}(\mu_{n}-\mu)(r,\omega)\|$
$\displaystyle\leq\|\int\varphi\mathrm{d}(\mu_{n}-\mu)\|+\int\|\mathbbm{1}_{r\leq
s}\phi(\omega,\bar{x}(r))-\varphi(r,\omega)\|\mathrm{d}(\mu_{n}+\mu)(r,\omega)$
$\displaystyle\leq\|\int\varphi\mathrm{d}(\mu_{n}-\mu)\|+4\varepsilon+C(1+k^{\alpha})\delta$

where we used the fact that, in the second term, the integrand has
$\alpha$-growth and is only non-zero for $r\in[s,s+\delta]$. Hence having
chosen $\delta$ sufficiently small gives:

$\displaystyle\limsup_{n\to\infty}\|\int\mathbbm{1}_{r\leq
s}\phi(\omega,\bar{x}(r))\mathrm{d}(\mu_{n}-\mu)(r,\omega)\|\leq\limsup_{n\to\infty}\|\int\varphi\mathrm{d}(\mu_{n}-\mu)\|+5\varepsilon\leq
5\varepsilon,$

where the first $\limsup$ is 0 by definition of narrow convergence. We have
thus shown that for every $s\in[0,1]$, taking the limit as $n\to\infty$:

$\displaystyle\bar{x}(s)=x+\int_{0}^{s}\phi(\omega,\bar{x}(r))\mathrm{d}\mu(\omega|r)\mathrm{d}r,$

_i.e._ $\bar{x}=x_{\mu}$ is the flow trajectory associated to $\mu$ and
starting from $x$.

_Part 7: $\|\nabla L(\mu)\|_{L^{2}(\mu)}$ is $\tau$-lower semi-continuous on
$d$-bounded subsets._

We previously considered the upper-gradient $|\nabla L|$ defined in eq. 24 as
$|\nabla L|(\mu)\coloneqq\|\nabla L[\mu]\|_{L^{2}(\mu)}$. However, [AGS08,
Thm.2.3.3] state that the obtained curve of maximal slope is a curve of
maximal slope for another definition of gradient, referred to as _relaxed
slope_. To show these two notions coincide here we show that the map
$\mu\mapsto\|\nabla L[\mu]\|_{L^{2}(\mu)}$ is $\tau$-lower semi-continuous on
$d$-bounded subsets (_c.f._ [AGS08, Rm.2.3.4]).

As before consider a $d$-bounded sequence $(\mu_{n})$ in
$\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$ which narrowly converges
towards some $\mu\in\mathcal{P}_{2}^{\mathrm{Leb}}([0,1]\times\Omega)$. Then
we previously have shown that for every $x\in\mathbb{R}^{d}$ we have
$x_{\mu_{n}}\to x_{\mu}$ in $\mathcal{C}([0,1],\mathbb{R}^{d})$. Proceeding
with similar arguments one could show the same for the adjoint variable $p$,
that is $p_{n}\coloneqq p_{\mu_{n},x,y}\to p_{\mu,x,y}$ in
$\mathcal{C}([0,1],\mathbb{R}^{d})$. Then using a generalization of Fatou’s
lemma with varying measures (_e.g._ [FKL20, Thm.2.4]) we have:

$\displaystyle\liminf_{n\to\infty}\|\nabla L[\mu_{n}]\|^{2}_{L^{2}(\mu_{n})}$
$\displaystyle=\liminf_{n\to\infty}\int_{[0,1]\times\Omega}\|\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega,x_{n}(r))^{\top}p_{n}(r)\|^{2}\mathrm{d}\mu_{n}(r,\omega)$
$\displaystyle\geq\int_{[0,1]\times\Omega}\liminf_{\begin{subarray}{c}n\to\infty\\\
(r^{\prime},\omega^{\prime})\to(r,\omega)\end{subarray}}\|\mathbb{E}_{x,y}\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{n}(r^{\prime}))^{\top}p_{n}(r^{\prime})\|^{2}\mathrm{d}\mu(r,\omega)$
$\displaystyle=\|\nabla L[\mu]\|^{2}_{L^{2}(\mu)},$

which is the desired property. ∎

####  3.4.2 Uniqueness

We present here a uniqueness result for solutions of the gradient flow
equation, which is the content of the following Theorem 3.3. The proof is
standard and relies on the lipschitzness of the gradient vector field $\nabla
L[\mu]$ w.r.t. the measure $\mu$. It uses the following Assumption B on the
feature map $\phi$ to ensure local lipschitzness of the adjoint variable $p$
(Lemma 3.5).

######  Assumption B.

Assume that $\phi$ is twice continuously differentiable with
$\mathrm{D}^{2}_{\omega,\omega}\phi$ uniformly bounded,
$\mathrm{D}^{2}_{\omega,x}\phi$ having linear growth and
$\mathrm{D}^{2}_{x,x}\phi$ having quadratic growth w.r.t. $\omega$. Namely,
for every $R\geq 0$ there exists a constant $C=C(R)$ s.t. for every
$x,x^{\prime}\in B(0,R)$ and every $\omega,\omega^{\prime}\in\Omega$:

$\displaystyle\|\mathrm{D}^{2}_{\omega,\omega}\phi(\omega,x)\|\leq
C(R),\quad\|\mathrm{D}^{2}_{\omega,x}\phi(\omega,x)\|\leq
C(R)(1+\|\omega\|),\quad\|\mathrm{D}^{2}_{x,x}\phi(\omega,x)\|\leq
C(R)(1+\|\omega\|^{2}).$

######  Theorem 3.3 (Uniqueness of curves of maximal slope).

Assume $\phi$ satisfies Assumptions 1, 2, 3 and B and that $\nabla_{x}\ell$ is
locally Lipschitz w.r.t. $x$. Let
$\mu_{0}\in\mathcal{P}^{\mathrm{Leb}}_{2}([0,1]\times\Omega)$. Then the
gradient flow for the risk $L$ starting from $\mu_{0}$, if it exists, is
unique.

###### Proof.

Let $(\mu_{t})_{t\geq 0}$ and $(\mu_{t}^{\prime})_{t\geq 0}$ be two gradient
flow curves for the risk $L$ starting from $\mu_{0}$. We will proceed to show
that $d(\mu_{t},\mu_{t}^{\prime})=0$ for every $t\geq 0$.

We use the shorter notations $v_{t}\coloneqq\nabla L[\mu_{t}]$,
$v_{t}^{\prime}\coloneqq\nabla L[\mu_{t}^{\prime}]$ to refer to the _tangent
vector fields_ of $\mu$ and $\mu^{\prime}$ respectively. Observe that the map
$t\mapsto d(\mu_{t},\mu_{t}^{\prime})^{2}$ is locally absolutely continuous
and by Lemma 2.3 its differential is given at almost every $t\geq 0$ by:

$\displaystyle\frac{\mathrm{d}}{\mathrm{d}t}d(\mu_{t},\mu_{t}^{\prime})^{2}=2\int_{0}^{1}\int_{\Omega}\langle\omega^{\prime}-\omega,v_{t}^{\prime}(s,\omega^{\prime})-v_{t}(s,\omega)\rangle\mathrm{d}\gamma_{t}(s,\omega,\omega^{\prime}),$

where $\gamma_{t}\in\Pi^{\mathrm{Leb}}_{o}(\mu_{t},\mu_{t}^{\prime})$ can be
any optimal coupling.

Let $T\geq 0$ and note
$\mathcal{E}\coloneqq\sup_{t\in[0,T]}\max(\mathcal{E}_{2}(\mu_{t}),\mathcal{E}_{2}(\mu_{t}^{\prime}))<\infty$.
Fix some $t\in[0,T]$ and consider $(x,y)$ in the support of the data
ditsribution $\mathcal{D}$ with the shortcuts $x_{t}\coloneqq x_{\mu_{t}}$,
$p_{t}\coloneqq p_{\mu_{t},x,y}$ and similarly $x_{t}^{\prime},p_{t}^{\prime}$
for $\mu_{t}^{\prime}$. As the data distribution has compact support, we have
that there exists some $R=R(\mathcal{E})$ such that
$\|x_{t}(s)\|,\|x_{t}^{\prime}(s)\|,\|p_{t}(s)\|,\|p_{t}^{\prime}(s)\|\leq R$.
Then using Lemmas 3.2 and 3.5 as well as the assumptions on $\phi$ we have for
every $(s,\omega,\omega^{\prime})\in[0,1]\times\Omega^{2}$:

$\|\mathrm{D}_{\omega}\phi(\omega,x_{t}(s))^{\top}p_{t}(s)-\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}^{\prime}(s))^{\top}p_{t}^{\prime}(s)\|\\\
\leq\|\mathrm{D}_{\omega}\phi(\omega,x_{t}(s))^{\top}p_{t}(s)-\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}(s))^{\top}p_{t}(s)\|\\\
+\|\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}(s))^{\top}p_{t}(s)-\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}^{\prime}(s))^{\top}p_{t}^{\prime}(s)\|\\\
\leq\|\mathrm{D}_{\omega}\phi(\omega,x_{t}(s))-\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}(s))\|\|p_{t}(s)\|\\\
+\|\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}(s))-\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}^{\prime}(s))\|\|p_{t}(s)\|\\\
+\|\mathrm{D}_{\omega}\phi(\omega^{\prime},x_{t}^{\prime}(s))\|\|p_{t}(s)-p_{t}^{\prime}(s)\|\\\
\leq
C_{1}\|\omega-\omega^{\prime}\|+C_{1}(1+\|\omega^{\prime}\|)d(\mu_{t},\mu_{t}^{\prime})+C_{1}(1+\|\omega^{\prime}\|)d(\mu_{t},\mu_{t}^{\prime}),$
(25)

with $C_{1}=C_{1}(\mathcal{E})$ some constant. Fixing some