# On the Global Convergence of Gradient Descent for multi-layer ResNets in the
mean-field regime

Zhiyan Ding, Shi Chen & Qin Li
Mathematics Department
University of Wisconsin-Madison
Madison, WI 53706 USA.
<EMAIL_ADDRESS>
&Stephen Wright
Department of Computer Sciences
University of Wisconsin-Madison
Madison, WI 53706 USA.
<EMAIL_ADDRESS>

###### Abstract

Finding the optimal configuration of parameters in ResNet is a nonconvex
minimization problem, but first order methods nevertheless find the global
optimum in the overparameterized regime. We study this phenomenon with mean-
field analysis, by translating the training process of ResNet to a gradient-
flow partial differential equation (PDE) and examining the convergence
properties of this limiting process. The activation function is assumed to be
$2$-homogeneous or partially $1$-homogeneous; the regularized ReLU satisfies
the latter condition. We show that if the ResNet is sufficiently large, with
depth and width depending algebraically on the accuracy and confidence levels,
first order optimization methods can find global minimizers that fit the
training data.

##  1 Introduction

Training of multi-layer neural networks (NN) requires us to find weights in
the network such that its outputs perfectly match the prescribed outputs for a
given set of training data. The usual approach is to formulate this problem as
a nonconvex minimization problem and solve it with a first-order optimization
method based on gradient descent (GD). Extensive computational experience
shows that in the overparametrized regime (where the total number of
parameters in the NN far exceeds the minimum number required to fit the
training data), GD methods run for sufficiently many iterations consistently
find a global minimum achieving the zero-loss property, that is, a perfect fit
to the training data.

What is the mechanism that allows GD to perform so well on this large-scale
nonconvex problem?

Part of the explanation is that in the overparametrized case, the parameter
space contains many global minima, and some evidence suggests that they are
distributed throughout the space, making it easier for the optimization
process to find one such solution. Many approaches have been taken to
characterize this phenomenon more rigorously, including landscape analysis,
the neural tangent kernel approach, and mean-field analysis. All such
viewpoints aim to give an idea of the structure and size of the NN required to
ensure global convergence.

Our approach in this paper is based on mean-field analysis and gradient-flow
analysis, the latter being the continuous and mean-field limit of GD. We will
examine residual neural networks (ResNets), and study how deep and wide a
ResNet needs to be to match the data with high accuracy and high confidence.
To relax the assumptions on the activation function as far as possible, we
follow the setup in (Chizat & Bach, 2018), which requires this function to be
either $2$-homogeneous or partially $1$-homogeneous. We show that both depth
and width of the NN depend algebraically on ${\epsilon}$ and $\eta$, which are
the accuracy and confidence levels, respectively.

Mean-field analysis translates the training process of the ResNet to a
gradient-flow partial differential equation (PDE). The training process
evolves weights on connections between neurons. When dealing with wide neural
networks, instead of tracing the evolution of each weight individually, one
can record the evolution of the full distribution of the weight configuration.
This perspective translates the coupled ordinary differential equation system
(ODE) that characterizes evolution of individual weights into a PDE (the
gradient-flow equation) characterizing the evolution of the distribution. The
parameters in the PDE naturally depend on the properties of the activation
functions. Gradient-flow analysis is used to show that the PDE drives the
solution to a point where the loss function becomes zero. We obtain our
results on zero-loss training of ResNet with GD by translating the zero-loss
property of the gradient-flow PDE back to the discrete-step setting.

This strategy of the proof was taken in an earlier paper (Ding et al., 2021)
where multi-layer ResNets were also analyzed. The main difference in this
current paper is that the assumptions on the activation function and the
initial training state for obtaining the global convergence are both much
relaxed. This paper adopts the setup from (Chizat & Bach, 2018) of minimal
Lipschitz continuity requirements on the activation function. Furthermore, the
paper (Ding et al., 2021) required a dense support condition to be satisfied
on the final parameter configuration has a support condition. This condition
is hard to justify in any realistic setting, and is discared from current
paper. Further details on these issues appear in Section 3.

We discuss the setup of the problem and formally derive the continuous and
mean-field limits in Section 2. In Section 3, we discuss related work,
identify our contribution and present the main theorem in its general terms.
After precise definitions and assumptions are specified in Section 4, we
present the two main ingredients in the proof strategy. The mean-field limit
is obtained by connecting the training process of the ResNet to a gradient-
flow PDE in Section 5, and the zero-loss property of the limiting PDE is
verified in Section 6. The main theorem is a direct corollary of Theorem 5.1
and Theorem 6.1 (or Theorem 6.2).

##  2 ResNet and gradient descent

The ResNet can be specified as follows:

$z_{l+1}(x)=z_{l}(x)+\frac{1}{ML}\sum^{M}_{m=1}f(z_{l}(x),\theta_{l,m})\,,\quad
l=0,1,\dots,L-1\,,$ (1)

where $M$ and $L$ are the width and depth, respectively;
$z_{0}(x)=x\in\mathbb{R}^{d}$ is the input data; and $z_{L}(x)$ is the output
from the last layer. The configuration of the NN is encoded in parameters
${\Theta}_{L,M}=\left\\{\theta_{l,m}\right\\}^{L-1,M}_{l=0,m=1}$, where each
parameter $\theta_{l,m}$ is a vector in $\mathbb{R}^{k}$ and
$f:\mathbb{R}^{d}\times\mathbb{R}^{k}\rightarrow\mathbb{R}^{d}$ is the
activation function. The formulation (1) covers“conventional” ResNets, which
have the specific form

$z_{l+1}(x)=z_{l}(x)+\frac{1}{ML}\sum^{M}_{m=1}U_{l,m}\sigma(W^{\top}_{l,m}z_{l}(x)+b_{l,m})\,,\quad
l=0,1,\dots,L-1\,,$

where $W_{l,m},U_{l,m}\in\mathbb{R}^{d}$, $b_{l.m}\in\mathbb{R}$, and $\sigma$
is the ReLU activation function. In this example, we have
$\theta_{l,m}=(W_{l,m},U_{l,m},b_{l,m})\in\mathbb{R}^{k}$, with $k=2d+1$.

Denote by $Z_{{\Theta}_{L,M}}(l;x)$ the output of the ResNet defined by (1).
(This quantity is the same as $z_{L}(x)$ defined above, but we use this
alternative notation to emphasize the dependece on parameters $\Theta_{L,M}$.)
The goal of training ResNet is to seek parameters $\Theta_{L,M}$ that minimize
the following mismatch or loss function:

$E({\Theta}_{L,M})=\mathbb{E}_{x\sim\mu}\left[\frac{1}{2}\left(g(Z_{{\Theta}_{L,M}}(L;x))-y(x)\right)^{2}\right]\,,$
(2)

where $g(x):\mathbb{R}^{d}\rightarrow\mathbb{R}$ is a given measuring
function, $y(x)\in\mathbb{R}$ is the label corresponding to $x$, and $\mu$ is
the probability from which the data $x$ is drawn.

Classical gradient descent updates the parameters according to the formula

$\Theta_{L,M}^{n+1}=\Theta_{L,M}^{n}-h\nabla_{{\Theta}}E(\Theta^{n}_{L,M})\,,$

where $h$ is the step length. In the limit as $h\to 0$, the updating process
can be characterized by the following ODE (Chizat & Bach, 2018, Def 2.2)
(rescaled by $L,M$):

$\frac{\,\mathrm{d}{\Theta}_{L,M}(s)}{\,\mathrm{d}s}=-ML\nabla_{{\Theta}}E({\Theta}_{L,M})\,,\quad\text{for}\;s\geq
0\,,$ (3)

where $s$ represents pseudo-time, the continuous analog of the discrete
stepping process.

###  2.1 The continuous limit and the mean-field limit

The continuous limit of (1) is obtained when the ResNet is infinitely deep,
with $L\to\infty$. By reparametrizing the indices $l=[0,\cdots,L-1]$ with the
continuous variable $t\in[0,1]$, we can view $z$ in (1) as a function in $t$
that satisfies a coupled ODE, with $1/L$ being the stepsize in $t$.
Accordingly, $\theta_{l,m}$ can be recast as $\theta_{m}(t=l/L)$, and denoting
$\Theta(t)=\\{\theta_{m}(t)\\}_{m=1}^{M}$, we can write the continuous limit
of (1) as

$\frac{\,\mathrm{d}z(t;x)}{\,\mathrm{d}t}=\frac{1}{M}\sum^{M}_{m=1}f(z(t;x),\theta_{m}(t))\,,\quad
t\in[0,1],\quad\text{with}\;z(0;x)=x\,.$ (4)

Extending (2), we define the cost functional $E$ as

$E(\Theta)=\mathbb{E}_{x\sim\mu}\left[\frac{1}{2}\left(g(Z_{\Theta}(1;x))-y(x)\right)^{2}\right]\,,$
(5)

where $Z_{\Theta}(t;x)$ solves (4) for a given collection $\Theta(t)$ of the
$M$ functions $\\{\theta_{m}(t)\\}$. Similar to (3), we can use GD to find the
configuration of $\Theta(t)$ that minimizes (5) by making $\Theta(t)$ flow in
the descending direction of $E(\Theta)$. Denote $s$ the pseudo-time of the
training process, and $\Theta(s;t)$ the collection of functions at the
training time $s$:

$\frac{\partial\Theta}{\partial s}=-M\left.\frac{\delta
E}{\delta\Theta}\right|_{\Theta(s;\cdot)}\,,\quad s>0,\quad t\in[0,1]$ (6)

where $\frac{\delta E}{\delta\Theta}$ is the functional derivative of $E$ with
respect to $\Theta$, and thus a list of $M$ functions of $t$ for every fixed
$s$.

The mean-field limit is obtained by making the ResNet infinitely wide, that
is, $M\to\infty$. Considering that the right hand side of (4) has the form of
an expectation, it approaches an integral in the limit, with respect to a
certain probability density. Denoting this PDF by
$\rho(\theta,t)\in\mathcal{C}([0,1];\mathcal{P}^{2})$111A collection of
probability distribution that is continuous in $t$ and has bounded second
moment in $\theta$ for all $t$. The definition is to be made rigorous in Def
4.1., and assuming that the $\theta_{m}$ are drawn from it, the ODE for $z$
translates to the following:

$\frac{\,\mathrm{d}z(t;x)}{\,\mathrm{d}t}=\int_{\mathbb{R}^{k}}f(z(t;x),\theta)\,\mathrm{d}\rho(\theta,t)\,,\quad
t\in[0,1]\quad\text{with}\;z(0;x)=x\,.$ (7)

Mimicking (5), we define the following cost function in the mean-field
setting:

$E(\rho)=\mathbb{E}_{x\sim\mu}\left[\frac{1}{2}\left(g(Z_{\rho}(1;x))-y(x)\right)^{2}\right]\,,$
(8)

where $Z_{\rho}(t;x)$ is the solution to (7) for a given $\rho$. Then, similar
to the gradient flow for $\Theta_{L,M}$ and $\Theta(t)$, the probability
distribution $\rho$ that encodes the configuration of $\theta$ flows in the
descending direction of $E(\rho)$ in pseudo-time $s$. Since $\rho(\theta,t,s)$
needs to be a probability density for all $s$ and $t$, its evolution in $s$ is
characterized by a gradient flow in the Wasserstein metric (Chizat & Bach,
2018; Lu et al., 2020; Ding et al., 2021):

$\frac{\partial\rho}{\partial
s}=\nabla_{\theta}\cdot\left(\rho\nabla_{\theta}\left.\frac{\delta
E}{\delta\rho}\right|_{\rho(\cdot,\cdot,s)}\right)\,,\quad
s>0,\;t\in[0,1]\,\quad\text{with}\quad\rho(\theta,t,0)=\rho_{\mathrm{ini}}(\theta,t)\,,$
(9)

where $\frac{\delta E}{\delta\rho}$ is the functional derivative with respect
to $\rho$, and thus a function of $(\theta,t)$ for every fixed $s$. Using the
classical calculus-of-variations method, this functional derivative can be
computed as:

$\left.\frac{\delta
E}{\delta\rho}\right|_{\rho}(\theta,t)=\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho}(t;x)f(Z_{\rho}(t;x),\theta)\right)\,,$
(10)

where $p_{\rho}(\cdot;x)$, parameterized by $x$, maps
$[0,1]\to\mathbb{R}^{d}$, and is a vector solution to the following ODE:

$\frac{\,\mathrm{d}p^{\top}_{\rho}}{\,\mathrm{d}t}=-p^{\top}_{\rho}\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho},\theta)\rho(\theta,t)\,\mathrm{d}\theta\,.$
(11)

with $p_{\rho}(t=1;x)=\left(g(Z_{\rho}(1;x))-y(x)\right)\nabla
g(Z_{\rho}(1;x))$. In the later sections, to emphasize the $s$ dependence, we
use $\frac{\delta E(\Theta(s))}{\delta\Theta}$ and $\frac{\delta
E(\rho(s))}{\delta\rho}$ to denote $\left.\frac{\delta
E}{\delta\Theta}\right|_{\Theta(s;\cdot)}$ and $\left.\frac{\delta
E}{\delta\rho}\right|_{\rho(\cdot,\cdot,s)}$ respectively. As a summary, to
update $\rho(\theta,t,s)$ to $\rho(\theta,t,s+\delta s)$ with an infinitesimal
$\delta s$, we solve (7) for $Z_{\rho}(t;x)$, using the given
$\rho(\theta,t,s)$, and compute $p_{\rho}$ using (11). This then allows us to
compute $\frac{\delta E(\rho(s))}{\delta\rho}(\theta,t)$ which, in turn,
yields $\rho(\theta,t,s+\delta s)$ from (9). In (11), $\partial_{z}f$ is a
$d\times d$ matrix that stands for the Jacobian of $f$ with respect to its $z$
argument.

##  3 Related work and Contribution

There is a vast literature addressing the overparameterization of DNN. Many
perspectives have been taken to justify the success of the application of the
first order (gradient descent) optimization methods, in this overparameterized
regime. We briefly review related works, and identify our contribution.

The earliest approach to understanding overparametrization was landscape
analysis, in which the countours of the nonconvex objective function were
studied to find which properties make it possible for a first order method to
converge to the optimizer. Different NN structures are then analyzed to see
which have these properties (Jin et al., 2017; Ge et al., 2015; Du et al.,
2017; Ge et al., 2018; Nguyen & Hein, 2018; Du & Lee, 2018; Soltanolkotabi et
al., 2019; Nguyen & Hein, 2017; Kawaguchi, 2016; Yun et al., 2018). This
approach naturally limits the types of DNN that can be “explained,” since most
DNN structures do not satisfy the required properties.

Another approach taken in the literature is related to the Neural Tangent
Kernel (NTK) regime, which is the regime in which the nonlinear problem is
reduced to a nearly linear model due to the confinement of the iterates to a
small region around the initial values. Insensitivity of the so-called Gram
matrix is evaluated in the limit of the number of weights (Allen-Zhu et al.,
2019; Du et al., 2019a; Zhang et al., 2019; Chatterji et al., 2021; Du et al.,
2019b; Jacot et al., 2018; Liu et al., 2020; Frei et al., 2019). The argument
is that zero-loss solutions are close to every point in the space, and one can
find an optimal point within a small region of the initial guess. The NTK
arguments are shown to work well in several real application problems, such as
the classification problem (Li & Liang, 2018; Zou et al., 2019). However, as
pointed out by (Ba et al., 2020; Wei et al., 2019; Fang et al., 2019), NTK
approximately views nonlinear DNN as a linear kernel model, a rather limited
description, so the estimates obtained through NTK might not be sharp. Indeed,
the empirical observation in (Allen-Zhu & Li, 2019; Arora et al., 2019) have
suggested that the kernel models are not as general as NN, and certain
(nonlinear) features of NN are not captured.

Finally, there is the mean-field limit perspective that we adopt in this
paper. The term “mean-field" indicates that in a system with a large ensemble
of particles, the field formed by averaging across all samples exerts a force
on each sample. Instead of tracing the trajectory of each sample, one can
characterize the evolution of the full distribution function that represents
the field. This idea originated in statistical physics, and is made rigorous
under the framework of kinetic theory. In training an overparametrized ResNet
context, a large number of weights evolve to decrease the cost function. In
the mean-field limit, the training process evolves the distribution function
of these weights. A significant advantage of the mean-field approach is that
once we derive a formula for the gradient flow, standard PDE techniques can be
adopted to describe the convergence behavior. This approach was taken in
(Araújo et al., 2019; Fang et al., 2019; Nguyen, 2019; Du et al., 2019a;
Chatterji et al., 2021; Chizat & Bach, 2018; Mei et al., 2018; Wojtowytsch,
2020; Lu et al., 2020; Sirignano & Spiliopoulos, 2021; 2020; E et al., 2020b;
Jabir et al., 2021). The case of a single hidden layer NN in the regime as
$M\to\infty$ is studied by Chizat & Bach (2018); Mei et al. (2018);
Wojtowytsch (2020), who justified the mean-field approach and demonstrated
convergence of the gradient flow process to a zero objective. In the multi-
layer case, Lu et al. (2020) showed the convergence of a PDE that can be
viewed as a modified version of the true gradient flow, hinting at convergence
of the real mean-field limit. Nguyen & Pham (2021) also gave the global
convergence of the mean-field limit of DNN for a certain class of NN
structures, but their work excludes such important practical NN structures as
ResNet. The work most closely related to ours is (Ding et al., 2021), but this
paper makes technical assumptions on $\rho_{\infty}$ and $f$ that restrict the
usefulness of the results, as we discuss below following the statement of
Theorem 3.1.

We note that in certain parameter regimes, the mean-field and NTK perspectives
can sometimes be unified; see (Chen et al., 2020).

We follow the roadmap of Chizat & Bach (2018); Ding et al. (2021), which shows
that the PDE (9) achieves the global minimum for which
$E(\rho(\theta,t,s=\infty))=0$, and that the gradient flow in the discrete
setting (3) can be closely approximated by the PDE, so that
$E(\Theta_{L,M}(s))\approx E(\rho(\cdot,\cdot,s))$. These two results together
show that $E(\Theta_{L,M}(s))\approx 0$ for pseudo-time $s$ sufficiently
large. Specifically, the two main tasks of the paper are as follows.

  * –

Task 1: We need to give a rigorous proof of the continuous and mean-field
limit. This will be stated in Theorem 5.1, to justify that for every fixed
$s<\infty$, when $M,L\to\infty$, $E(\Theta_{L,M}(s))\approx
E(\Theta(s;\cdot))\approx E(\rho(\cdot,\cdot,s))$. The dependence of these
approximations on $L$ and $M$ are made precise.

  * –

Task 2: We need to demonstrate the convergence to global minimum. This is
stated in Theorem 6.1 and 6.2, for two different cases. In both theorems, we
obtain the global convergence for the gradient flow, assuming certain
homogeneity and the Sard-type regularity for $f$. A weak assumption of the
initialization of $\rho_{\mathrm{ini}}$ is also imposed.

By combining these two, we obtain the main result of the paper.

###### Theorem 3.1

Let the conditions in Theorem 5.1 and 6.1 (or 6.2) hold. Then for any positive
$\epsilon$ and $\eta$, there exist positive constants $C_{0}$ depending on
$\rho_{\mathrm{ini}}(\theta,t),\epsilon$ and $C$ depending on
$\rho_{\mathrm{ini}}(\theta,t),s$ such that when

$s>C_{0}(\rho_{\mathrm{ini}}(\theta,t),\epsilon)\,,\quad
M>\frac{C(\rho_{\mathrm{ini}}(\theta,t),s)}{\epsilon^{2}\eta}\,,\quad
L>\frac{C(\rho_{\mathrm{ini}}(\theta,t),s)}{\epsilon}\,,$

we have

$\mathbb{P}\left(\left|E(\Theta_{L,M}(s))\right|\leq\epsilon\right)\geq
1-\eta\,,$

where $E$ is defined in (2) and $\Theta_{L,M}$ solves (3).

This theorem gives quantitative bounds for $M$ and $L$. The number of weights
required to reduce the cost function below $\epsilon$ is
$O(ML)=O(1/\epsilon^{3})$. The theorem also suggests that $L$ and $M$ are
independent parameters.

The results resonate with those obtained in (Chizat & Bach, 2018) for the
2-layer NN, and extend those in (Ding et al., 2021) greatly. Specifically,
compared with the results in (Chizat & Bach, 2018), where $\rho(\theta,s)$
follows a typical gradient flow in the probability space on $\theta$ in time
$s$ (Ambrosio et al., 2008), we have, at each training time $s$, a “list” of
probability measures $\rho(\theta,t)$ on $\theta$, for all $t$. The members of
this list are coupled, flowing together in $s$ in the descending direction of
the cost function $E$. New analytical estimates are developed to deal with
this non-traditional gradient flow.

Ding et al. (2021) takes a similar approach to ours, but their assumptions on
the support of $\rho(\theta,t,s=\infty)$ are quite strong: The limiting
probability measure $\rho(\theta,t,s=\infty)$ is assumed to have the full
support over $\theta$. The assumption greatly reduces the technical difficulty
of the proof, but it is impractical and hard to justify, thus preventing the
results from being of practical use. In the current paper, this support
condition is replaced by the well-accepted homogeneity condition adopted by
(Chizat & Bach, 2018). As a consequence, the structure of the gradient flow
must be examined closely to demonstrate convergence, requiring considerable
technical complications. Lu et al. (2020) also investigates gradient flow for
training multi-layer neural network, but the gradient flow structure is
modified for mathematical convenience. All blocks are integrated together,
making $\rho$ a probability measure over the full $(\theta,t)$-space. This
design is inconsistent with the structure of the ResNet design that we
investigate in this paper.

##  4 Notations, Assumptions, and Definitions

Throughout the paper we denote the collection of probability distributions
with bounded second moments as $\mathcal{P}^{2}(\mathbb{R}^{k})$, that is,
$\mathcal{P}^{2}(\mathbb{R}^{k})=\\{\rho:\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho(\theta)<\infty\\}$.
We assume certain regularity properties for the activation function $f$, the
measuring function $g$, the data $y$, and the input measure $\mu$, as follows.

######  Assumption 4.1 (Assumptions on $f$)

Let $f:\mathbb{R}^{d}\times\mathbb{R}^{k}\rightarrow\mathbb{R}^{d}$ be a
$\mathcal{C}^{2}$ function.

  * 1.

(linear growth) For all $x\in\mathbb{R}^{d},\theta\in\mathbb{R}^{k}$, there is
a constant $C_{1}$ such that

$\left|f\right|\leq C_{1}(|\theta|^{2}+1)(|x|+1)\,.$ (12)

  * 2.

(locally Lipschitz) For all $r>0$, and $|x|<r,\theta\in\mathbb{R}^{k}$, we
have for $C_{2}(r)$ monotonically increasing with respect to $r$ that the
following bounds hold:

$\left|\partial_{x}f\right|\leq
C_{2}(r)(|\theta|^{2}+1),\quad\left|\partial_{\theta}f\right|\leq
C_{2}(r)(|\theta|+1)\,.$ (13)

  * 3.

(local smoothness) There exists $k-1$ with $0<k_{1}\leq k$ with the following
property: Denoting $\theta=(\theta_{[1]},\theta_{[2]})$,
$r=\max\\{|x|,|\theta_{[1]}|\\}$, where $\theta_{[1]}\in\mathbb{R}^{k_{1}}$,
$\theta_{[2]}\in\mathbb{R}^{k-k_{1}}$, we have for $C_{3}(r)$ monotonically
increasing with respect to $r$ that the following bounds hold:

$\left|\partial^{i}_{x}\partial^{j}_{\theta}f(x,\theta)\right|\leq
C_{3}(r),\quad i+j=2,\ i,j\geq 0\,.$ (14)

When $k_{1}<k$, we have in addition that

$\max\\{|\partial_{x}f|,|f|\\}\leq C_{3}(r)\left(|\theta_{[2]}|+1\right),\
\left|\partial_{\theta_{[1]}}f\right|\leq
C_{3}(|x|)\left(|\theta_{[1]}|+1\right)\,.$ (15)

  * 4.

(universal kernel) The function set
$\left\\{h\middle|h=\int_{\mathbb{R}^{k}}f(x,\theta)\,\mathrm{d}\rho(\theta),\
\rho\in\mathcal{P}^{2}(\mathbb{R}^{k})\right\\}$ is dense in
$\mathcal{C}\left(|x|<R;\mathbb{R}^{d}\right)$ for all $R>0$.

######  Assumption 4.2 (Assumptions on data)

Let $g,y:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be $\mathcal{C}^{2}$ functions,
and let $\mu$ be the probability distribution of $x$. We assume the following.

  * 5.

$\mu(x)$ is compactly supported, meaning that there is $R_{\mu}>0$ such that
the support of $\mu$ is within a ball of size $R_{\mu}$ around the origin,
that is, $\mathrm{supp}(\mu)\subset\mathcal{B}_{R_{\mu}}(\vec{0})$.

  * 6.

$y(x)\in L^{\infty}_{\mathrm{loc}}(\mathbb{R}^{d})$, that is, $\sup_{|x|\leq
R}|y(x)|<\infty$.

  * 7.

$g(x)$ and $\nabla g(x)$ are Lipschitz continuous. Moreover, $|\nabla g(x)|$
has a positive lower bound, that is, $\inf_{x\in\mathbb{R}^{d}}|\nabla
g(x)|>0$.

We note that Assumption 4.1 admits many commonly used activation functions (E
et al., 2020a). One example of a function satisfying this assumption is
$f(x,\theta)=f(x,\theta_{[1]},\theta_{[2]},\theta_{[3]})=\theta_{[3]}\sigma(\theta_{[1]}x+\theta_{[2]})$,
where $\theta_{[1]}\in\mathbb{R}^{d\times d}$,
$\theta_{[2]}\in\mathbb{R}^{d}$, $\theta_{[3]}\in\mathbb{R}$ and $\sigma$ is a
component-wise regularized ReLU activation function, see Remark H.2.

We now build the metric on the function space for our solutions. Note that the
solution $\rho(\theta,t,s)$ to (9) is expected to be a continuous function in
$(t,s)$, and a distribution of $\theta$ for each $(t,s)$. For this non-
standard probability space, we first introduce the following metric.

###### Definition 4.1

222$\mathcal{C}([0,T];\mathcal{A})$ is defined to be the set of functions
$f(a,t)$ such that for any $t\in[0,T]$, $f(\cdot,t)\in\mathcal{A}$ and
$f(\cdot,t)$ is continuous with respect to $t$ under the metric defined on
$\mathcal{A}$. In Definition 4.1, we set $T=1$ and
$\mathcal{A}=\mathcal{P}^{2}$, where the natural metric on $\mathcal{A}$ is
$W_{2}$.

$\mathcal{C}([0,1];\mathcal{P}^{2})$ is a collection of continuous paths of
probability distribution $\nu(\theta,t)$ $(\theta\in\mathbb{R}^{k},t\in[0,1])$
where 1. $\nu(\cdot,t)\in\mathcal{P}^{2}(\mathbb{R}^{k})$ for every fixed
$t\in[0,1]$; 2. For any $t_{0}\in[0,1]$, $\lim_{t\rightarrow
t_{0}}W_{2}\left(\nu(\cdot,t),\nu(\cdot,t_{0})\right)=0$, where $W_{2}$ is the
classical Wasserstein-2 distance. The space
$\mathcal{C}([0,1];\mathcal{P}^{2})$ is equipped with the following metric
$d_{1}\left(\nu_{1},\nu_{2}\right)=\sup_{t}W_{2}(\nu_{1}(\cdot,t),\nu_{2}(\cdot,t))$.

Accordingly, $\mathcal{C}([0,\infty);\mathcal{C}([0,1];\mathcal{P}^{2}))$ is a
collection of continuous paths of probability distribution $\nu(\theta,t,s)$
(with $\theta\in\mathbb{R}^{k}$, $t\in[0,1]$, $s\in[0,\infty)$), where 1.
$\nu(\cdot,\cdot,s)\in\mathcal{C}([0,1];\mathcal{P}^{2})$ for every fixed
$s\in[0,\infty)$. 2. For any $s_{0}\in[0,\infty)$, $\lim_{s\rightarrow
s_{0}}d_{1}\left(\nu(\cdot,\cdot,s),\nu(\cdot,\cdot,s_{0})\right)=0$ (where
$d_{1}$ is defined above). The metric in
$\mathcal{C}([0,\infty);\mathcal{C}([0,1];\mathcal{P}^{2}))$ is defined by
$d_{2}\left(\nu_{1},\nu_{2}\right)=\sup_{t,s}W_{2}(\nu_{1}(\cdot,t,s),\nu_{2}(\cdot,t,s))$.

Since $\mathcal{P}^{2}$ is complete in $W_{2}$ distance,
$\mathcal{C}([0,1];\mathcal{P}^{2})$ and
$\mathcal{C}([0,\infty);\mathcal{C}([0,1];\mathcal{P}^{2}))$ are complete
metric spaces under $d_{1}$ and $d_{2}$ respectively also. To give a rigorous
justification of the mean-field limit, we use the particle representation of
$\rho(\theta,t,s)$. Thus, at least we need to assume that we can find a
stochastic process that is drawn from the initial condition
$\rho_{\mathrm{ini}}(\theta,t)$. We call such initial conditions admissible.

###### Definition 4.2

We call a continuous path of probability distribution
$\nu(\theta,t)\in\mathcal{C}([0,1];\mathcal{P}^{2})$ admissible if it has a
particle representation, namely there exists a continuous stochastic process
$\theta(t):[0,1]\rightarrow\mathbb{R}^{k}$ and $r>0$ such that for any
$t_{0}\in[0,1]$, we have

$\theta(t_{0})\sim\nu(\theta,t_{0})\,,\quad\lim_{t\rightarrow
t_{0}}\mathbb{E}\left(|\theta(t)-\theta(t_{0})|^{2}\right)=0,\quad|\theta(t_{0})|<r\,.$
(16)

Furthermore, $\nu(\theta,t)$ is called limit-admissible if its $M$-averaged
trajectory is bounded and Lipschitz with high probability. (See the rigorous
definition in Definition E.1.)

We note that without the dependence on $t$, probability distributions are
“admissible", in the sense that one can draw a sample from a given
distribution. In Appendix A, we show that if the initial condition
$\rho_{\mathrm{ini}}(\theta,t)$ is admissible, (9) has a unique solution
$\rho(\theta,t,s)$ that is admissible for each $s$.

The global convergence result depends on Sard-type regularity, defined as
follows.

###### Definition 4.3

Given a metric space $\Theta$, a differentiable function
$h:\Theta\rightarrow\mathbb{R}$, and a subset
$\widetilde{\Theta}\subset\Theta$, we say $h$ satisfies Sard-type regularity
in $\widetilde{\Theta}$ if the set of regular values333For a function
$h:\widetilde{\Theta}\rightarrow\mathbb{R}$, a regular value is a real number
$\alpha$ in the range of $h$ such that $h^{-1}(\alpha)$ is included in an open
set where $h$ is differentiable and where $\,\mathrm{d}h$ does not vanish. of
$h|_{\widetilde{\Theta}}$ is dense in its range, where
$h|_{\widetilde{\Theta}}:\widetilde{\Theta}\rightarrow\mathbb{R}$ is a
confinement of $h$ in $\widetilde{\Theta}$.

###### Remark 4.1

This regularity assumption is not a common one; we have adopted it from Chizat
& Bach (2018). This property is essentially that most of the points in the
range of $h$ lie in an open set and $h$ is locally monotonic. The assumption
is rather mild and can be satisfied by most commonly seen regular functions,
unless the function oscillates wildly.

##  5 Mean-field and continuous limit

In this section, we focus on the justification of mean-field and continuous-
limit result. This is to prove that $E(\rho(\cdot,\cdot,s))$,
$E(\Theta(s;\cdot))$, and $E(\Theta_{L,M}(s))$ are asymptotically close to
each other for every $s$. In the next section, we prove convergence of
$E(\rho(\cdot,\cdot,s))$ as $s\to\infty$.

To show the asymptotic equivalence of the three quantities, we need to compare
(3), (6), and (9), and take the measurement in $E$ according to (2), (5) and
(8).

###### Theorem 5.1

Suppose that Assumptions 4.1 and 4.2 are satisfied. Assume that
$\rho_{\mathrm{ini}}(\theta,t)$ is limit-admissible and
$\mathrm{supp}_{\theta}(\rho_{\mathrm{ini}}(\theta,t))\subset\\{\theta||\theta_{[1]}|\leq
R\\}$ with some $R>0$ for all $t\in[0,1]$. Let
$\\{\theta_{m}(0;t)\\}^{M}_{m=1}$ in (6) be $i.i.d.$ drawn from
$\rho_{\mathrm{ini}}(\theta,t)$. Let

  * •

$\Theta_{L,M}(s)=\\{\theta_{l,m}(s)\\}$ be the solution to (3) with initial
condition $\theta_{l,m}(s=0)=\theta_{m}\left(0;\frac{l}{L}\right)$;

  * •

$\theta_{m}(s;t)$ be the solution to (6) with the initial condition
$\theta_{m}(0;t)$;

  * •

$\rho(\theta,t,s)$ be the solution to (9) with initial condition
$\rho_{\mathrm{ini}}(\theta,t)$.

Then for any positive $\epsilon$, $\eta$, and $S$, there exists a constant
$C>0$ that depends on $\rho_{\mathrm{ini}}(\theta,t)$ and $S$ such that when

$M>\frac{C(\rho_{\mathrm{ini}}(\theta,t),S)}{\epsilon^{2}\eta}\,,\quad
L>\frac{C(\rho_{\mathrm{ini}}(\theta,t),S)}{\epsilon}\,,\quad s<S\,,$

we have:

$\min\\{\mathbb{P}\left(\left|E(\Theta_{L,M}(s))-E(\Theta(s;\cdot))\right|\leq\epsilon/2\right)\,,\mathbb{P}\left(\left|E(\Theta(s;\cdot))-E(\rho(\cdot,\cdot,s))\right|\leq\epsilon/2\right)\\}\geq
1-\tfrac{1}{2}\eta\,.$

It follows that

$\mathbb{P}\left(\left|E(\Theta_{L,M}(s))-E(\rho(\cdot,\cdot,s))\right|\leq\epsilon\right)\geq
1-\eta\,,\quad\forall s<S\,.$

Here $E(\Theta_{L,M}(s))$, $E(\Theta(s;\cdot))$, and $E(\rho(\cdot,\cdot,s))$
are defined in (2), (5), and (8), respectively.

The proof of this result appears in Appendix E. This theorem suggests that for
every fixed $S>0$, the gradient descent of $\Theta_{L,M}$ is approximately the
same as the gradient flow of $\rho(\theta,t)$, in the sense that the two costs
are close to each other with high probability, when $L$ and $M$ are
sufficiently large. The size of the ResNet depends negative-algebraically on
$\epsilon$ (the desired accuracy) and $\eta$ (the confidence of success). The
result translates the evolution (gradient descent) of $\Theta_{L,M}$ to the
evolution (gradient flow) of $\rho(\theta,t)$, and thus matches the zero-loss
property of the parameter configuration of a finite sized ResNet to its
limiting PDE, whose analysis can be performed with standard PDE tools.

The proof of Theorem 5.1 divides naturally into two components. We show that
for all $s<S$, $E(\rho(\cdot,\cdot,s))\approx E(\Theta(s;\cdot))$ and
$E(\Theta(s;\cdot))\approx E(\Theta_{L,M}(s))$ with high probability. The
former is obtained from mean-field limit theory, justifying that the particle
trajectory $\theta_{m}(t,s)$ follows $\rho(\theta,t,s)$ for all $t$ in pseudo-
time $s\in[0,S]$. The latter makes use of continuity in $t$ and traces the
differences between $\theta_{m}(\frac{l}{L})$ and $\theta_{l,m}$. These two
components of the proof are summarized in Theorems E.1 and E.2, respectively.
According to the formula of the Fréchet derivatives (10) and (11), the
estimates in these theorems naturally route through the boundedness of
$p_{\rho}$, $p_{\Theta}$, $p_{\Theta_{L,M}}$, and similarly
$Z_{\rho,\Theta,\Theta_{L,M}}$. It is technically demanding to derive these
bounds, but they are not surprising. We dedicate a large portion of the
appendix to addressing the well-posedness of these systems. See Appendices
A-D, where we show these equations have unique solutions with proper initial
conditions, along with the required bounds. Naturally, these estimates depend
on regularity of $f$, $g$, $y$, and $\mu$.

To gain some intuition for the equivalence between (6) and (9), we test them
on the same smooth function $h(\theta)$. To test (9), we multiply both sides
by $h$ and perform integration by parts to obtain
$\frac{\,\mathrm{d}}{\,\mathrm{d}s}\int_{\mathbb{R}^{k}}h\,\mathrm{d}\rho(\theta)=-\int_{\mathbb{R}^{k}}\nabla_{\theta}h\nabla_{\theta}\frac{\delta
E(\rho(s))}{\delta\rho}\,\mathrm{d}{\rho}$. This is to say
$\frac{\,\mathrm{d}}{\,\mathrm{d}s}\mathbb{E}(h)=\mathbb{E}\left(\nabla_{\theta}h\nabla_{\theta}\frac{\delta
E(\rho(s))}{\delta\rho}\right)$. Testing $h$ on (6) also gives the same
formula. Supposing that $\rho=\frac{1}{M}\sum_{m=1}^{M}\delta_{\theta_{m}}$,
we have
$\frac{\,\mathrm{d}}{\,\mathrm{d}s}\mathbb{E}(h)=\frac{1}{M}\sum_{m=1}^{M}\nabla_{\theta}h(\theta_{m})\frac{\,\mathrm{d}}{\,\mathrm{d}s}\theta_{m}=-\sum_{m=1}^{M}\nabla_{\theta}h(\theta_{m})\frac{\delta
E}{\delta\theta_{m}}$. The right hand side is also
$\mathbb{E}\left(\nabla_{\theta}h\nabla_{\theta}\frac{\delta
E(\rho(s))}{\delta\rho}\right)$ if and only if $M\frac{\delta
E(\Theta(s;\cdot))}{\delta\theta_{m}}=\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta_{m},t)$. This will be shown to hold true in
Appendix F.

##  6 Convergence to global minimizer

After translating the study of $\Theta_{L,M}$ to the study of
$\rho(\theta,t)$, this section presents results on when and how $E(\rho)$
converges to zero loss by examining the conditions for the global convergence
of (9). We first identify the property of global minimum.

###### Proposition 6.1

Suppose that $\rho\in\mathcal{C}([0,1];\mathcal{P}^{2})$ has $E(\rho)>0$. Then
for any $t_{0}\in[0,1]$, there exists a measure $\nu(\theta)$ of
$\mathbb{R}^{k}$ such that $\int_{\mathbb{R}^{k}}\,\mathrm{d}\nu(\theta)=0$
and

$\int_{\mathbb{R}^{k}}\left.\frac{\delta
E}{\delta\rho}\right|_{\rho}(\theta,t_{0})\,\mathrm{d}\nu(\theta)<0\,.$ (17)

See Appendix H.1 for the proof of this result. At the stationary point of the
cost function, $\frac{\delta E}{\delta\rho}(\theta,t_{0})=0$, then there is no
$\nu$ satisfying (17), so $E(\rho)$ is necessarily trivial. Our task now
becomes to check under what conditions on $f$ we can have $\frac{\delta
E}{\delta\rho}$ becoming $0$ as $s\to\infty$. We give two possibilities, both
requiring a separation initialization and certain homogeneities. In the first
case, we require $f$ to be 2-homogeneous, and the result is collected in
Section 6.1. In the second case, we require $f$ to be partially 1-homogeneous;
see Section 6.2.

###  6.1 The 2-homogeneous case

The results in this section are obtained under the following assumption on
$f$. (Functions can be designed to satisfy this assumption easily; see Remark
H.1.)

###### Assumption 6.1

The function $f(x,\theta):\mathbb{R}^{d}\times\mathbb{R}^{k}$ is
$2$-homogeneous, meaning that $f(x,\lambda\theta)=\lambda^{2}f(x,\theta)$ for
all $(x,\theta,\lambda)\in\mathbb{R}^{d}\times\mathbb{R}^{k}\times\mathbb{R}$.

###### Theorem 6.1

Let Assumption 6.1 and conditions of Theorem 5.1 hold true with $k_{1}=k$. Let
$\rho_{\infty}(\theta,t)$ be the long-time limit of (9), that is,
$\rho(\theta,t,s)$ converges to $\rho_{\infty}(\theta,t)$ in
$\mathcal{C}([0,1];\mathcal{P}^{2})$ as $s\rightarrow\infty$. Then
$E(\rho_{\infty})=0$ if the following hold:

  * •

Separation initialization: There exists $t_{0}\in[0,1]$ such that
$\rho_{\mathrm{ini}}(\theta,t_{0})$ separates444We say that a set $C$
separates the sets $A$ and $B$ if any continuous path in with endpoints in $A$
and $B$ intersects $C$. the spheres $r_{a}\mathbb{S}^{k-1}$ and
$r_{b}\mathbb{S}^{k-1}$, for some $0<r_{a}<r_{b}$.

  * •

Sard-type regularity: $\left.\frac{\delta
E}{\delta\rho}\right|_{\rho_{\infty}}(\theta,t_{0})$ satisfies the Sard-type
regularity condition in $\mathbb{S}^{k-1}$.

A proof appears in Appendix H.2. A corollary of this theorem, when combined
with Theorem 5.1 with $k_{1}=k$, gives Theorem 3.1, the main result of our
paper.

###  6.2 The partially 1-homogeneous case

The following assumption is used in this section. (Functions that satisfy this
assumption include regularized ReLU; see Remark H.2.)

###### Assumption 6.2

Let $\theta=(\theta_{[1]},\theta_{[2]})$ with $\theta_{[1]}\in\mathbb{R}$ and
$\theta_{[2]}\in\mathbb{R}^{k-1}$:

  * 1.

(partially 1-homogeneous in $\theta$) $f$ can be written as
$f(x,\theta)=f(x,\theta_{[1]},\theta_{[2]})=\theta_{[1]}\widehat{f}(x,\theta_{[2]})$,

  * 2.

(locally bounded and smooth) For any $r>0$, $\widehat{f}(x,\theta_{[2]})$ is
bounded and Lipschitz with Lipschitz continuous differential for
$(x,\theta_{[2]})\in\mathcal{B}_{r}(\vec{0})\times\mathbb{R}^{k-1}$.

When Assumption 6.2 holds true, Assumption 4.1 part 3 can be satisfied with
$k_{1}=1$. Then, we introduce the main result in this section.:

###### Theorem 6.2

Let Assumption 6.2 and conditions of Theorem 5.1 hold true with $k_{1}=1$. Let
$\rho_{\infty}(\theta,t)$ be the limit of (9) as $s\to\infty$. Then
$E(\rho_{\infty})=0$ if the following hold:

  * •

Separation initialization: There exists $t_{0}\in[0,1]$ such that
$\rho_{\mathrm{ini}}(\theta_{[1]},\theta_{[2]},t_{0})$ separates the spheres
$\\{-r_{0}\\}\times\mathbb{R}^{k-1}$ and $\\{r_{0}\\}\times\mathbb{R}^{k-1}$
for some $r_{0}>0$, where $\theta_{[1]},\theta_{[2]}$ are defined by
Assumption 4.1 item 3 with $k_{1}=1$.

  * •

Sard-type regularity: $\left.\frac{\delta
E}{\delta\rho}\right|_{\rho_{\infty}}((1,\theta_{[2]}),t_{0}):\mathbb{R}^{k-1}\rightarrow\mathbb{R}$
satisfies the Sard-type regularity condition.

  * •

For any $\tilde{\rho}\in\mathcal{C}([0,1],\mathcal{P}^{2})$, define
$H_{r,\tilde{\rho}}\left(\widetilde{\theta}_{[2]}\right)=\left.\frac{\delta
E(\rho)}{\delta\rho}\right|_{\tilde{\rho}}\left(1,r\widetilde{\theta}_{[2]},t_{0}\right)$
where $\theta_{[2]}=r\tilde{\theta}_{[2]}$ with $r=|\theta_{2}|$ and
$\tilde{\theta}_{[2]}\in\mathbb{S}^{k-2}$. Suppose that $H_{r,\tilde{\rho}}$
converges in $\mathcal{C}^{1}(\mathbb{S}^{k-2})$ as $r\rightarrow\infty$ to a
function $H_{\infty,\tilde{\rho}}$. Furthermore, assume that
$H_{\infty,\rho_{\infty}}$ satisfies the Sard-type regularity condition in
$\mathbb{S}^{k-2}$ and that the intersection of regular values of
$H_{\infty,\rho_{\infty}}$ and $\left.\frac{\delta
E(\rho)}{\delta\rho}\right|_{\rho_{\infty}}(1,\theta_{[2]},t_{0})$ is also
dense in the intersection of their range.

The proof is found in Appendix H.3. This result, combined with Theorem 5.1 in
the case of $k_{1}=1$, gives the main result in Theorem 3.1. The assumptions
in Theorem 6.2 are rather technical and seemingly tedious. However, we note
that only the first two assumptions — $1$-homogeneous and the separation
assumption — are crucial. The third and fourth assumptions concern the
regularity of the Fréchet derivative of $\rho$; they are rather mild and serve
to exclude wildly oscillating functions; see Remark 4.1. Most commonly seen
functions are regular enough that these two assumptions are satisfied.

##  7 Ethics Statement

This work does not present any foreseeable societal consequence.

##  8 Reproducibility Statement

The notations, assumptions and definitions are clarified in Section 4. And all
proofs appear in the Appendix.

## References

  * Allen-Zhu & Li (2019) Z. Allen-Zhu and Y. Li.  What can ResNet learn efficiently, going beyond kernels?  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * Allen-Zhu et al. (2019) Z. Allen-Zhu, Y. Li, and Z. Song.  A convergence theory for deep learning via over-parameterization.  In _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 242–252. PMLR, 09–15 Jun 2019.
  * Ambrosio et al. (2008) L. Ambrosio, N. Gigli, and G. Savaré.  _Gradient flows: in metric spaces and in the space of probability measures_.  Springer Science and Business Media. Birkhäuser Basel, 2008.
  * Araújo et al. (2019) D. Araújo, R. Oliveira, and D. Yukimura.  A mean-field limit for certain deep neural networks.  _arXiv/1906.00193_ , 2019.
  * Arora et al. (2019) S. Arora, S. Du, W. Hu, Z. Li, R. Salakhutdinov, and R. Wang.  On exact computation with an infinitely wide neural net.  In _NeurIPS_ , 2019.
  * Ba et al. (2020) J. Ba, M. Erdogdu, T. Suzuki, D. Wu, and T. Zhang.  Generalization of two-layer neural networks: An asymptotic viewpoint.  In _ICLR_ , 2020.
  * Chatterji et al. (2021) N. Chatterji, P. Long, and P. Bartlett.  When does gradient descent with logistic loss interpolate using deep networks with smoothed relu activations?  _arXiv/2102.04998_ , 2021.
  * Chen et al. (2020) Z. Chen, Y. Cao, Q. Gu, and T. Zhang.  A generalized neural tangent kernel analysis for two-layer neural networks.  In _Advances in Neural Information Processing Systems_ , volume 33, pp. 13363–13373, 2020.
  * Chizat & Bach (2018) L. Chizat and F. Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  In _Advances in Neural Information Processing Systems_ , volume 31, 2018.
  * Ding et al. (2021) Z. Ding, S. Chen, Q. Li, and S. Wright.  Overparameterization of deep resnet: zero loss and mean-field analysis.  _arXiv/2105.14417_ , 2021.
  * Du & Lee (2018) S. Du and J. Lee.  On the power of over-parametrization in neural networks with quadratic activation.  In _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 1329–1338. PMLR, 10–15 Jul 2018.
  * Du et al. (2017) S. Du, C. Jin, J. Lee, M. Jordan, A. Singh, and B. Póczos.  Gradient descent can take exponential time to escape saddle points.  In _Advances in Neural Information Processing Systems_ , volume 30, 2017.
  * Du et al. (2019a) S. Du, J. Lee, H. Li, L. Wang, and X. Zhai.  Gradient descent finds global minima of deep neural networks.  In _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 1675–1685, 09–15 Jun 2019a.
  * Du et al. (2019b) S. Du, X. Zhai, B. Póczos, and A. Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In _International Conference on Learning Representations_ , 2019b.
  * E et al. (2020a) W. E, C. Ma, S. Wojtowytsch, and L. Wu.  Towards a mathematical understanding of neural network-based machine learning: what we know and what we don’t.  _arXiv/2009.10713_ , 2020a.
  * E et al. (2020b) W. E, C. Ma, and L. Wu.  Machine learning from a continuous viewpoint, i.  _Science China Mathematics_ , 63(11):2233–2266, Sep 2020b.
  * Fang et al. (2019) C. Fang, Y. Gu, W. Zhang, and T. Zhang.  Convex formulation of overparameterized deep neural networks.  _arXiv/1911/07626_ , 2019.
  * Frei et al. (2019) S. Frei, Y. Cao, and Q. Gu.  Algorithm-dependent generalization bounds for overparameterized deep residual networks.  In _NeurIPS_ , 2019.
  * Ge et al. (2015) R. Ge, F. Huang, C. Jin, and Y. Yuan.  Escaping from saddle points — online stochastic gradient for tensor decomposition.  In _Proceedings of The 28th Conference on Learning Theory_ , volume 40 of _Proceedings of Machine Learning Research_ , pp. 797–842. PMLR, 03–06 Jul 2015.
  * Ge et al. (2018) R. Ge, J. Lee, and T. Ma.  Learning one-hidden-layer neural networks with landscape design.  In _International Conference on Learning Representations_ , 2018.
  * Jabir et al. (2021) J. Jabir, D. Šiška, and Ł. Szpruch.  Mean-field neural odes via relaxed optimal control.  _arxiv/1912.05475_ , 2021.
  * Jacot et al. (2018) A. Jacot, F. Gabriel, and C. Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Advances in Neural Information Processing Systems_ , volume 31, 2018.
  * Jin et al. (2017) C. Jin, R. Ge, P. Netrapalli, S. Kakade, and M. Jordan.  How to escape saddle points efficiently.  In _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 1724–1732. PMLR, 06–11 Aug 2017.
  * Kawaguchi (2016) K. Kawaguchi.  Deep learning without poor local minima.  In _Advances in Neural Information Processing Systems_ , volume 29, 2016.
  * Li & Liang (2018) Y. Li and Y. Liang.  Learning overparameterized neural networks via stochastic gradient descent on structured data.  In _Proceedings of the 32nd International Conference on Neural Information Processing Systems_ , NIPS’18, pp. 8168–8177, 2018.
  * Liu et al. (2020) C. Liu, L. Zhu, and M. Belkin.  On the linearity of large non-linear models: when and why the tangent kernel is constant.  In _Advances in Neural Information Processing Systems_ , volume 33, pp. 15954–15964, 2020.
  * Lu et al. (2020) Y. Lu, C. Ma, Y. Lu, J. Lu, and L. Ying.  A mean field analysis of deep ResNet and beyond: Towards provably optimization via overparameterization from depth.  In _Proceedings of the 37th International Conference on Machine Learning_ , volume 119, pp. 6426–6436, 13–18 Jul 2020.
  * Mei et al. (2018) S. Mei, A. Montanari, and P. M. Nguyen.  A mean field view of the landscape of two-layer neural networks.  _Proceedings of the National Academy of Sciences_ , 115(33):E7665–E7671, 2018.
  * Nguyen (2019) P. M. Nguyen.  Mean field limit of the learning dynamics of multilayer neural networks.  _arXiv/1902.02880_ , 2019.
  * Nguyen & Pham (2021) P. M. Nguyen and H. Pham.  A rigorous framework for the mean field limit of multilayer neural networks.  _arxiv/2001.11443_ , 2021.
  * Nguyen & Hein (2017) Q. Nguyen and M. Hein.  The loss surface of deep and wide neural networks.  In _Proceedings of the 34th International Conference on Machine Learning_ , volume 70 of _Proceedings of Machine Learning Research_ , pp. 2603–2612. PMLR, 06–11 Aug 2017.
  * Nguyen & Hein (2018) Q. Nguyen and M. Hein.  Optimization landscape and expressivity of deep cnns.  In _ICML_ , 2018.
  * Sirignano & Spiliopoulos (2020) J. Sirignano and K. Spiliopoulos.  Mean field analysis of neural networks: A law of large numbers.  _SIAM Journal on Applied Mathematics_ , 80(2):725–752, 2020.
  * Sirignano & Spiliopoulos (2021) J. Sirignano and K. Spiliopoulos.  Mean field analysis of deep neural networks.  _Mathematics of Operations Research_ , 2021.  doi: 10.1287/moor.2020.1118.
  * Soltanolkotabi et al. (2019) M. Soltanolkotabi, A. Javanmard, and J. Lee.  Theoretical insights into the optimization landscape of over-parameterized shallow neural networks.  _IEEE Trans. Inf. Theor._ , 65(2):742–769, February 2019.  ISSN 0018-9448.
  * Wei et al. (2019) C. Wei, J. Lee, Q. Liu, and T. Ma.  Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel.  In _NeurIPS_ , 2019.
  * Wojtowytsch (2020) S. Wojtowytsch.  On the convergence of gradient descent training for two-layer relu-networks in the mean field regime.  _arXiv/2005/13530_ , 2020.
  * Yun et al. (2018) C. Yun, S. Sra, and A. Jadbabaie.  Global optimality conditions for deep neural networks.  In _International Conference on Learning Representations_ , 2018.
  * Zhang et al. (2019) H. Zhang, D. Yu, M. Yi, W. Chen, and T. Liu.  Convergence theory of learning over-parameterized resnet: A full characterization.  _arXiv/1903.07120_ , 2019.
  * Zou et al. (2019) D. Zou, Y. Cao, D. Zhou, and Q. Gu.  Gradient descent optimizes over-parameterized deep relu networks.  _Machine Learning_ , 109:467–492, 2019.

## Appendix: Introduction

The appendix contains proofs and supporting analysis for the theorems in the
main text. We start by proving well-posedness of the ResNet ODE and the
gradient flow. We then justify the continuous and mean-field limit. Finally,
we prove the global convergence of the gradient-flow PDE. The sections are
organized as follows.

  1. Appendix A:

Well-posedness. We summarize the well-posedness of the continuous limit ODE
(4), the mean-field limit ODE (7), and the associated gradient flows (3), (6),
and (9). The results are collected in Theorem A.1.

  2. Appendix B:

This section collects the detailed proof for the well-posedness of the ResNet
ODE in its continuous limit (4), and the mean-field limit (7), and finalizes
the proof of Theorem A.2.

  3. Appendix C:

This section prepares some a-priori estimates for showing the well-posedness
of the gradient flow equations.

  4. Appendix D:

This section contains a detailed proof for the well-posedness of gradient flow
equations (3) and (9) and finalizes the proofs of Theorem A.3 and A.4.

  5. Appendix E-G:

Proof of Theorem 5.1, the continuous and mean-field limit. Section E lays out
the structure of the proof, Section F shows the continuous limit, and Section
G shows the mean-field limit.

  6. Appendix H:

Proof of Theorem 6.1 and 6.2: Global convergence of the gradient flow.

The analytical core of the paper lies in Appendices H and E, which describe
properties of the gradient flow PDEs and explain why the gradient descent
method for ResNet can be explained by these equations. The technical results
of Appendix B-D can be skipped by readers who are interested to proofs of the
main results.

Throughout, we denote by $C(\cdot)$ a generic constant whose value depends
only on its arguments. The precise value of this constant may change each time
it is invoked. Throughout Appendices A-G, we assume that Assumption 4.1 holds
for some $0<k_{1}\leq k$.

##  Appendix A Well-posedness result

In this section, we show the well-posedness of ODEs (4), (7) and gradient
flows (3), (6), and (9).

###### Theorem A.1

The following claims hold.

  1. –

Well-posedness of ODE:

     * •

If $\\{\theta_{m}(t)\\}_{m=1}^{M}$ is continuous, then (4) has a unique
$\mathcal{C}^{1}$solution.

     * •

If $\rho\in\mathcal{C}([0,1];\mathcal{P}^{2})$, then (7) has a unique
$\mathcal{C}^{1}$ solution.

  2. –

Well-posedness of gradient flow:

     * •

(3) has a unique solution.

     * •

If $\\{\theta_{m}(0;t)\\}_{m=1}^{M}$ is continuous, then (6) has a unique
solution $\\{\theta_{m}(s;t)\\}_{m=1}^{M}$ that is continuous in $(s,t)$.

     * •

If $\rho_{\mathrm{ini}}(\theta,t)$ is admissible and
$\mathrm{supp}_{\theta}(\rho_{\mathrm{ini}}(\theta,t))\subset\\{\theta||\theta_{[1]}|\leq
R\\}$ with some $R>0$ for all $t\in[0,1]$, then (9) has a unique solution
$\rho(\theta,t,s)$ in
$\mathcal{C}([0,\infty);\mathcal{C}([0,1];\mathcal{P}^{2}))$ with initial
condition $\rho_{\mathrm{ini}}(\theta,t)$. Furthermore, for each $s$,
$\rho(\theta,t,s)$ is admissible.

Note that the well-posedness of (4) and (6) are direct corollaries of that of
(7), (3), and (9) (the corresponding continuous versions), according to Remark
A.1, A.2. Thus, we merely prove well-posedness on the continuous level, in
Theorems A.2, A.3, and A.4. Specifically,

  * •

Theorem A.2 (Appendix A.1) shows the well-posedness of the dynamical system
for $z$;

  * •

Theorems A.3 and A.4 (Appendix A.2) justify the well-posedness of the gradient
flow of the parameter configuration.

###  A.1 Well-posedness of the OID (7)

As $L$ and $M$ approach $\infty$, $z$ satisfies the ordinary-integral equation
(7). We justify that this differential equation is well-posed, in the sense
that the solution is unique and stable.

###### Theorem A.2

Suppose that Assumption 4.1 holds with and that $x$ is in the support of
$\mu$, then (7) has a unique $\mathcal{C}^{1}$ solution. Specifically, for
$\rho_{1},\rho_{2}\in\mathcal{C}([0,1];\mathcal{P}^{2})$, we have

$\left|Z_{\rho_{1}}(t;x)\right|\leq C(\mathcal{L}_{1})\,,$ (18)

and for all $t\in[0,1]$:

$\left|Z_{\rho_{1}}(t;x)-Z_{\rho_{2}}(t;x)\right|\leq
C(\mathcal{L}_{1},\mathcal{L}_{2})d_{1}(\rho_{1},\rho_{2})\,,$ (19)

where $\mathcal{L}_{i}$ are the second moments of $\rho_{i}$, that is,

$\mathcal{L}_{i}=\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{i}(\theta,t)\,\mathrm{d}t,\quad
i=1,2.$ (20)

We leave the proof to Appendix B. Besides the well-posedness result, the
theorem also suggests that a small perturbation to $\rho$ is reflected
linearly in $Z_{\rho}$, the solution to (7). This means that a small
perturbation in the parameterization of the ResNet leads to only a small
perturbation to the ResNet output.

###### Remark A.1

Although we do not directly show the well-posedness of (4), it follow
immediately from Theorem A.2. One way to make this connection is to
reformulate the discrete probability distribution as

$\rho^{\mathrm{dis}}(\theta,t)=\frac{1}{M}\sum^{M}_{m=1}\delta_{\theta_{m}(t)}(\theta)\,,$

where $\Theta(t)=\\{\theta_{m}(t)\\}_{m=1}^{M}$ is the list of trajectories.
Since $\theta_{m}(t)$ is continuous in $t$, we have
$\rho^{\mathrm{dis}}(\theta,t)\in\mathcal{C}([0,1];\mathcal{P}^{2})$. Since

$\frac{1}{M}\sum^{M}_{m=1}f(z(t;x),\theta_{m}(t))=\int_{\mathbb{R}^{k}}f(z(t;x),\theta)\,\mathrm{d}\rho^{\mathrm{dis}}(\theta,t)\,,$

using Theorem A.2, (4) has a unique $\mathcal{C}^{1}$ solution when
$\Theta(t)$ is continuous.

###  A.2 Well-posedness of the gradient flow

The gradient flow of the parameterization is also well-posed, both in the
discrete setting and the continuous mean-field limit. In the discrete setting,
we have the following result.

###### Theorem A.3

(3) has a unique solution.

Further, (9) characterizes the dynamics of the continuous mean-field limit of
the parameter configuration, and is also well-posed.

###### Theorem A.4

If $\rho_{\mathrm{ini}}(\theta,t)$ is admissible and
$\mathrm{supp}_{\theta}(\rho_{\mathrm{ini}}(\theta,t))\subset\\{\theta||\theta_{[1]}|\leq
R\\}$ with some $R>0$ for all $t\in[0,1]$. Then (9) has a unique solution
$\rho(\theta,t,s)$ in
$\mathcal{C}([0,\infty);\mathcal{C}([0,1];\mathcal{P}^{2}))$ that is
admissible for each $s$. Further, we have

$\frac{\,\mathrm{d}E(\rho(\cdot,\cdot,s))}{\,\mathrm{d}s}\leq 0\,.$ (21)

The proofs of these two theorems can be found in Appendix D.

###### Remark A.2

Using the same argument as in Remark A.1, calling

$\rho^{\mathrm{dis}}(\theta,t,s)=\frac{1}{M}\sum^{M}_{m=1}\delta_{\theta_{m}(s;t)}(\theta)\,,$
(22)

the well-posedness of (6) follows immediately from Theorem A.4. Furthermore,
according to the definition (5) and (8), we have

$E\left(\rho^{\mathrm{dis}}(\cdot,\cdot,s)\right)=E\left(\Theta(s;\cdot)\right)\,.$

As a consequence, if $\Theta(s;t)$ satisfies (6), then $\rho^{\mathrm{dis}}$
satisfies (9), and vice versa. The well-posedness result in Theorem A.4 for
(9) then can be extended to justify well-posedness of (6).

##  Appendix B Proof of Theorem A.2

This section contains the proof of Theorem A.2. We rewrite (7) as follows:

$\frac{\,\mathrm{d}Z_{\rho}(t;x)}{\,\mathrm{d}t}=F(Z_{\rho},t)\,,\quad\forall
t\in[0,1]\quad\text{with}\quad z(0;x)=x\,,$ (23)

where for a given $\rho\in\mathcal{C}([0,1];\mathcal{P}^{2})$ we use the
notation:

$F(z,t)=\int_{\mathbb{R}^{k}}f(z,\theta)\,\mathrm{d}\rho(\theta,t)\,.$

The proof of Theorem A.2 relies on the classical Lipschitz condition for the
well-posedness of an ODE.

Proof [Proof of Theorem A.2] Since
$\rho\in\mathcal{C}([0,1];\mathcal{P}^{2})$, we have a constant $C$ such that

$\sup_{0\leq t\leq
1}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho(\theta,t)<C<\infty\,.$

For any $t\in[0,1]$, using (12) from Assumption 4.1 equation, we have

$\displaystyle\left|F(z,t)\right|\leq\int_{\mathbb{R}^{k}}\left|f(z_{1},\theta)\right|\,\mathrm{d}\rho(\theta,t)\leq
C_{1}(|z|+1)\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho(\theta,t)\,.$
(24)

To show the boundedness result (18), we multiply (23) by $Z_{\rho_{1}}(t;x)$
and use (24) to obtain

$\displaystyle\frac{\,\mathrm{d}|Z_{\rho_{1}}(t;x)|^{2}}{\,\mathrm{d}t}\leq$
$\displaystyle
2C_{1}\left(|Z_{\rho_{1}}|^{2}+|Z_{\rho_{1}}|\right)\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho_{1}(\theta,t)$
$\displaystyle\leq$ $\displaystyle
4C_{1}\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho_{1}(\theta,t)\left(|Z_{\rho_{1}}(t;x)|^{2}+1\right)\,.$

Using Grönwall’s inequality, we have

$\displaystyle|Z_{\rho_{1}}(t;x)|$
$\displaystyle\leq\exp\left(2C_{1}\left(\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho_{1}\,\mathrm{d}t+1\right)\right)(|x|+1)$
$\displaystyle\leq\exp\left(2C_{1}\left(\mathcal{L}_{1}+1\right)\right)(|x|+1)\,,$

where
$\mathcal{L}_{1}=\int_{0}^{1}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho_{1}\,\mathrm{d}t$.
Since $x\in\text{supp}\,\mu$ (so that $|x|<R$), we have (18). This gives us an
a-priori estimate of (23).

Next, using (13) from Assumption 4.1, we have

$|f(z_{1},\theta)-f(z_{2},\theta)|\leq
C_{2}(|z_{1}|+|z_{2}|)\left(|\theta|^{2}+1\right)|z_{1}-z_{2}|\,.$ (25)

Then, using boundedness of the second moment of $\theta$, we have

$\displaystyle\left|F(z_{1},t)-F(z_{2},t)\right|$
$\displaystyle\leq\left|\int_{\mathbb{R}^{k}}\left(f(z_{1},\theta)-f(z_{2},\theta)\right)\,\mathrm{d}\rho(\theta,t)\right|$
$\displaystyle\leq
C_{2}(|z_{1}|+|z_{2}|)\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho(\theta,t)|z_{1}-z_{2}|$
$\displaystyle<C_{2}(|z_{1}|+|z_{2}|)(C+1)|z_{1}-z_{2}|\,,$

which implies that $F(z,t)$ is locally Lipschitz in $z$ for all $t\in[0,1]$.
Combining this with the a-priori estimate in (18), classical ODE theory
implies that (23) has a unique $\mathcal{C}^{1}$ solution.

To prove the stability result (19), we define $Z_{\rho_{i}}$ as in (23)
parameterized by $\rho_{i}\in\mathcal{C}([0,1];\mathcal{P}^{2})$, and denote

$\Delta(t;x)=Z_{\rho_{1}}(t;x)-Z_{\rho_{2}}(t;x)\,.$

Then by subtracting the two equations, we obtain

$\displaystyle\frac{\,\mathrm{d}|\Delta(t;x)|^{2}}{\,\mathrm{d}t}$ (26)
$\displaystyle=2\left\langle\Delta(t;x),\int_{\mathbb{R}^{k}}f(Z_{\rho_{1}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)-\int_{\mathbb{R}^{k}}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{2}(\theta,t)\right\rangle$
$\displaystyle=2\left\langle\Delta(t;x),\underbrace{\int_{\mathbb{R}^{k}}f(Z_{\rho_{1}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)-\int_{\mathbb{R}^{k}}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)}_{\mathrm{(I)}}\right\rangle$
$\displaystyle\quad+2\left\langle\Delta(t;x),\underbrace{\int_{\mathbb{R}^{k}}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)-\int_{\mathbb{R}^{k}}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{2}(\theta,t)}_{\mathrm{(II)}}\right\rangle\,.$

We now bound $\mathrm{(I)}$ and $\mathrm{(II)}$. For $\mathrm{(I)}$, we have
using (13) and (25) that

$\displaystyle\left|\mathrm{(I)}\right|\leq$ $\displaystyle
2C_{2}\left(|Z_{\rho_{1}}(t;x)|+|Z_{\rho_{2}}(t;x)|\right)|\Delta(t;x)|\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho_{1}(\theta,t)$
(27) $\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{1},\mathcal{L}_{2})|\Delta(t;x)|\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho_{1}(\theta,t)\,,$

where the second inequality comes from (18). For $\mathrm{(II)}$, we denote
the particle representation $\theta_{1}(t)\sim\rho_{1}(\theta,t)$ and
$\theta_{2}(t)\sim\rho_{2}(\theta,t)$ such that
$\left(\mathbb{E}\left(|\theta_{1}-\theta_{2}|^{2}\right)\right)^{1/2}=W_{2}(\rho_{1}(\cdot,t),\rho_{2}(\cdot,t))$.
Then

$\displaystyle\left|\textrm{(II)}\right|$
$\displaystyle\leq\mathbb{E}\left(\left|f(Z_{\rho_{2}}(t;x),\theta_{1})-f(Z_{\rho_{2}}(t;x),\theta_{2})\right|\right)$
(28) $\displaystyle\leq
C_{2}(Z_{\rho_{2}}(t;x))\mathbb{E}\left((|\theta_{1}|+|\theta_{2}|+1)|\theta_{1}-\theta_{2}|\right)$
$\displaystyle\leq
C(\mathcal{L}_{2})\mathbb{E}\left((|\theta_{1}|+|\theta_{2}|+1)|\theta_{1}-\theta_{2}|\right)$
$\displaystyle\leq
C(\mathcal{L}_{2})\left(\mathbb{E}\left(|\theta_{1}|^{2}+|\theta_{2}|^{2}+1\right)\right)^{1/2}\left(\mathbb{E}\left(|\theta_{1}-\theta_{2}|^{2}\right)\right)^{1/2}$
$\displaystyle\leq
C(\mathcal{L}_{2})\left(\mathbb{E}\left(|\theta_{1}|^{2}+|\theta_{2}|^{2}+1\right)\right)^{1/2}W_{2}(\rho_{1}(\cdot,t),\rho_{2}(\cdot,t))$
$\displaystyle\leq
C(\mathcal{L}_{2})\left(\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)+\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{2}(\theta,t)+1\right)^{1/2}d_{1}(\rho_{1},\rho_{2})\,,$

where we use mean-value theorem and (13) in the first inequality to obtain

$\displaystyle\left|f(Z_{\rho_{2}}(t;x),\theta_{1})-f(Z_{\rho_{2}}(t;x),\theta_{2})\right|\leq$
$\displaystyle|\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta_{1}+\lambda\theta_{2})||\theta_{1}-\theta_{2}|$
$\displaystyle\leq$ $\displaystyle
C_{2}(Z_{\rho_{2}}(t;x))(|\theta_{1}|+|\theta_{2}|+1)|\theta_{1}-\theta_{2}|,$

for some $\lambda\in[0,1]$. We use (18) in the second inequality of (28).

Plugging (27)-(28) into (26) and using Hölder’s inequality, we obtain that

$\displaystyle\frac{\,\mathrm{d}|\Delta(t;x)|^{2}}{\,\mathrm{d}t}\leq$
$\displaystyle
C(\mathcal{L}_{1},\mathcal{L}_{2})|\Delta(t;x)|^{2}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho_{1}(\theta,t)$
$\displaystyle+C(\mathcal{L}_{2})|\Delta(t;x)|\left(\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)+\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{2}(\theta,t)+1\right)^{1/2}d_{1}(\rho_{1},\rho_{2})$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{1},\mathcal{L}_{2})\left(\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)+\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{2}(\theta,t)+1\right)\left(|\Delta(t;x)|^{2}+d^{2}_{1}(\rho_{1},\rho_{2})\right)\,,$

where we used Young’s inequality in the last line. Since $|\Delta(0;x)|=0$, we
complete the proof of (19) using Grönwall’s inequality.

##  Appendix C A-priori estimation of the cost function

Some a-priori estimates are necessary in the proof for the main theorems. We
first consider the case when $f$ satisfies only Assumption 4.1 with
$0<k_{1}\leq k$. (Better a-priori estimates can be obtained when $f$ also
satisfies the homogeneity properties of Sections 6.1 or 6.2.)

###  C.1 A-priori estimate for general $f$

According to (10), the Fréchet derivative can be computed, similarly to (Lu et
al., 2020), as follows:

$\frac{\delta
E(\rho)}{\delta\rho}(\theta,t)=\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho}(t;x)f(Z_{\rho}(t;x),\theta)\right)\,,$
(29)

where $p_{\rho}(t;x)$ is the solution to the following ODE:

$\left\\{\begin{aligned} \frac{\partial p^{\top}_{\rho}}{\partial
t}&=-p^{\top}_{\rho}\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho}(t;x),\theta)\,\mathrm{d}\rho(\theta,t)\,,\\\
p_{\rho}(1;x)&=\left(g(Z_{\rho}(1;x))-y(x)\right)\nabla
g(Z_{\rho}(1;x))\,.\end{aligned}\right.$ (30)

We now show that $p_{\rho}$ is Lipschitz continuous with respect to $\rho$.

###### Lemma C.1

Suppose that $x$ is in the support of $\mu$. Suppose that
$\rho_{1},\rho_{2}\in\mathcal{C}([0,1];\mathcal{P}^{2})$ and $p_{\rho_{1}}$,
$p_{\rho_{2}}$ are the corresponding solutions of (30). Denote
$\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ as in (20) and

$\mathcal{R}=\min_{r}\left\\{r\middle|\mathrm{supp}(\rho_{1})\cup\mathrm{supp}(\rho_{2})\subset\left\\{\theta\middle|\left|\theta_{[1]}\right|<r\right\\}\right\\}\,.$

Then the following two bounds are satisfied:

$\left|p_{\rho_{1}}(t;x)\right|\leq C(\mathcal{L}_{1})\,,$ (31)

and

$\left|p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\right|\leq
C(\mathcal{R},\mathcal{L})\,d_{1}(\rho_{1},\rho_{2})\,,$ (32)

where $\mathcal{L}=\max\\{\mathcal{L}_{1},\mathcal{L}_{2}\\}$.

Proof  From (13) in Assumption 4.1,

$\displaystyle\left|\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{1}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)\right|\leq$
$\displaystyle
C(Z_{\rho_{1}}(t;x))\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho_{1}(\theta,t)$
(33) $\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{1})\int_{\mathbb{R}^{k}}(|\theta|^{2}+1)\,\mathrm{d}\rho_{1}(\theta,t)\,,$

where we use (18) in the second inequality. It follows from the initial
condition of (30) that

$p_{\rho_{1}}(1;x)\leq C(|Z_{\rho_{1}}(1,x)|+1)\leq C(\mathcal{L}_{1})\,,$

where we use Assumption 4.1 in the first inequality and (18) in the second
inequality. Noting that (30) is a linear equation, (31) follows naturally when
we combine (33) with the inequality above.

To prove (32), we define

$\Delta(t;x)=p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\,.$

For $t=1$, with $x\in\text{supp}\,\mu$, we have

$\displaystyle|\Delta(1;x)|$
$\displaystyle=|p_{\rho_{1}}(1;x)-p_{\rho_{2}}(1;x)|$ (34)
$\displaystyle=\left|\left(g(Z_{\rho_{1}}(1;x))-y(x)\right)\nabla
g(Z_{\rho_{1}}(1;x))-\left(g(Z_{\rho_{2}}(1;x))-y(x)\right)\nabla
g(Z_{\rho_{2}}(1;x))\right|$ $\displaystyle\leq
C(\mathcal{L})|Z_{\rho_{1}}(1;x)-Z_{\rho_{2}}(1;x)|$ $\displaystyle\leq
C(\mathcal{L})d_{1}(\rho_{1},\rho_{2})\,,$

where we use Assumption 4.1, (18), and $|x|<{R_{\mu}}$ in the first inequality
and (19) in the second inequality. The following ODE is satisfied by $\Delta$:

$\frac{\partial\Delta^{\top}(t;x)}{\partial
t}=-\Delta^{\top}(t;x)\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{1}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)+p^{\top}_{\rho_{2}}(t;x)D_{\rho_{1},\rho_{2}}(t;x)\,,$
(35)

where

$D_{\rho_{1},\rho_{2}}(t;x)=\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{2}(\theta,t)-\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{1}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)\,.$
(36)

To show the boundedness of $D_{\rho_{1},\rho_{2}}(t;x)$, we follow the same
strategy as that for (26). By splitting into two terms, we obtain

$\displaystyle|D_{\rho_{1},\rho_{2}}(t;x)|\leq$
$\displaystyle\underbrace{\left|\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{2}(\theta,t)-\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)\right|}_{\mathrm{(I)}}$
(37)
$\displaystyle+\underbrace{\left|\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{2}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)-\int_{\mathbb{R}^{k}}\partial_{z}f(Z_{\rho_{1}}(t;x),\theta)\,\mathrm{d}\rho_{1}(\theta,t)\right|}_{\mathrm{(II)}}\,.$

The bound of $\mathrm{(II)}$ relies on part 3 of Assumption 4.1: Because the
supports of $\rho_{1}$, $\rho_{2}$ are contained in
$\left\\{\theta\middle|\left|\theta_{[1]}\right|<\mathcal{R}\right\\}$ and $Z$
is bounded by (18), we have

$|\mathrm{(II)}|\leq
C(\mathcal{R},\mathcal{L})\left|Z_{\rho_{1}}(t;x)-Z_{\rho_{2}}(t;x)\right|\leq
C(\mathcal{R},\mathcal{L})d_{1}(\rho_{1},\rho_{2})\,,$ (38)

where we use (19) in the second inequality.

To bound (I), we use the particle representation
$\theta_{1}\sim\rho_{1}(\theta,t)$ and $\theta_{2}\sim\rho_{2}(\theta,t)$ such
that
$\left(\mathbb{E}\left(|\theta_{1}-\theta_{2}|^{2}\right)\right)^{1/2}=W_{2}(\rho_{1}(\cdot,t),\rho_{2}(\cdot,t))$.
We then have

(I)
$\displaystyle\leq\mathbb{E}\left(\left|\partial_{z}f(Z_{\rho_{2}}(t;x),\theta_{1})-\partial_{z}f(Z_{\rho_{2}}(t;x),\theta_{2})\right|\right)$
(39) $\displaystyle\leq
C(\mathcal{R},\mathcal{L})\mathbb{E}\left(|\theta_{1}-\theta_{2}|\right)$
$\displaystyle\leq
C(\mathcal{R},\mathcal{L})\left(\mathbb{E}\left(|\theta_{1}-\theta_{2}|^{2}\right)\right)^{1/2}$
$\displaystyle\leq C(\mathcal{R},\mathcal{L})d_{1}(\rho_{1},\rho_{2}),$

where we use the mean-value theorem, part 3 of Assumption 4.1 with
$|\theta_{1,[1]}|<\mathcal{R}$ and $|\theta_{2,[1]}|<\mathcal{R}$, (18) in the
second inequality. Substituting (38) and (39) into (37), we obtain

$|D_{\rho_{1},\rho_{2}}(t;x)|\leq
C(\mathcal{R},\mathcal{L})d_{1}(\rho_{1},\rho_{2})\,.$

By substituting this bound into (35) and using (33), we have

$\frac{\,\mathrm{d}|\Delta(t;x)|^{2}}{\,\mathrm{d}t}\leq
C(\mathcal{R},\mathcal{L})\left(|\Delta(t;x)|^{2}+d^{2}_{1}(\rho_{1},\rho_{2})\right).$
(40)

The result (32) follows from the initial condition (34) and Grönwall’s
inequality.

The second lemma concerns the continuity of $\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}$.

###### Lemma C.2

Suppose that $\rho,\rho_{1},\rho_{2}\in\mathcal{C}([0,1];\mathcal{P}^{2})$.
Define

$\mathcal{L}=\max_{1\leq i\leq
3}\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho_{i}(\theta,t)\,\mathrm{d}t,\quad\mathcal{L}^{\sup}=\max_{1\leq
i\leq
3}\sup_{t\in[0,1]}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho_{i}(\theta,t)\,,$

and

$\mathcal{R}=\min_{r}\left\\{r\middle|\mathrm{supp}(\rho)\cup\mathrm{supp}(\rho_{1})\cup\mathrm{supp}(\rho_{2})\subset\\{\theta||\theta_{[1]}|<r\\}\right\\}\,.$

Then for any
$(\theta,t),(\theta_{1},t_{1}),(\theta_{2},t_{2})\in\mathbb{R}^{k}\times[0,1]$
and $s>0$, the following properties hold.

  * •

Boundedness:

$\left|\nabla_{\theta}\frac{\delta E(\rho)}{\delta\rho}(\theta,t)\right|\leq
C(\mathcal{L})(|\theta|+1),\quad\left|\partial_{\theta_{[1]}}\frac{\delta
E(\rho)}{\delta\rho}(\theta,t)\right|\leq C(\mathcal{L})(|\theta_{[1]}|+1)$
(41)

  * •

Lipschitz continuity in $\theta$ and $t$: There exists
$Q_{1}:\mathbb{R}^{2}\rightarrow\mathbb{R}^{+}$ that depends increasingly on
both arguments such that

$\displaystyle\left|\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta_{1},t_{1})-\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta_{2},t_{2})\right|$ (42) $\displaystyle\leq$
$\displaystyle
Q_{1}\left(\mathcal{L},\max_{i=1,2}(|\theta_{i,[1]}|)\right)|\theta_{1}-\theta_{2}|+Q_{1}\left(\mathcal{L}^{\sup},\max_{i=1,2}(|\theta_{i,[1]}|)\right)(|\theta_{2}|+1)|t_{1}-t_{2}|\,,$

  * •

Lipschitz continuity in $\rho$. There exists
$Q:\mathbb{R}^{2}\rightarrow\mathbb{R}^{+}$ that increasingly depends on both
arguments such that

$\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta,t)-\nabla_{\theta}\frac{\delta
E(\rho_{2})}{\delta\rho}(\theta,t)\right|\leq
Q(\mathcal{L},|\theta_{[1]}|,\mathcal{R})(1+|\theta|)d_{1}(\rho_{1},\rho_{2})\,,$
(43)

where $d_{1}$ is defined in Definition 4.1.

Proof  To prove the first bound of (41), we restate (29) as follows

$\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta,t)=\mathbb{E}_{x\sim\mu}\left(p_{\rho}^{\top}(t;x)\partial_{\theta}f(Z_{\rho}(t;x),\theta)\right)\,,$

from which it follows that

$\left|\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta,t)\right|\leq\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho}(t;x),\theta)\right|\left|p_{\rho}(t;x)\right|\right)\leq
C(\mathcal{L})(|\theta|+1)\,,$

where we use (13), (18), and (31) in the second inequality. To prove the
second bound in (41), we use the bound of
$\left|\partial_{\theta_{[1]}}f(x,\theta)\right|$ according to Assumption 4.1
part 3 to obtain

$\left|\partial_{\theta_{[1]}}\frac{\delta
E(\rho)}{\delta\rho}(\theta,t)\right|\leq\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta_{[1]}}f(Z_{\rho}(t;x),\theta)\right|\left|p_{\rho}(t;x)\right|\right)\leq
C(\mathcal{L})(|\theta_{[1]}|+1)\,,$

To prove (42), we assume $t_{1}>t_{2}$ without loss of generality, and use the
triangle inequality to obtain

$\displaystyle\left|\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta_{1},t_{1})-\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta_{2},t_{2})\right|$ (44)
$\displaystyle\leq\underbrace{\left|\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho}(t_{1};x)\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{1})-p^{\top}_{\rho}(t_{1};x)\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{2})\right)\right|}_{\mathrm{(I)}}$
$\displaystyle\quad+\underbrace{\left|\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho}(t_{1};x)\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{2})-p^{\top}_{\rho}(t_{1};x)\partial_{\theta}f(Z_{\rho}(t_{2};x),\theta_{2})\right)\right|}_{\mathrm{(II)}}$
$\displaystyle\quad+\underbrace{\left|\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho}(t_{1};x)\partial_{\theta}f(Z_{\rho}(t_{2};x),\theta_{2})-p^{\top}(t_{2};x)\partial_{\theta}f(Z_{\rho}(t_{2};x),\theta_{2})\right)\right|}_{\mathrm{(III)}}\,.$

To bound (I), we use the mean-value theorem, Assumption 4.1 (14), and (18) to
obtain

$\displaystyle|\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{1})-\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{2})|$
$\displaystyle\leq|\partial^{2}_{\theta}f(Z_{\rho}(t_{1};x),(1-\lambda)\theta_{1}+\lambda\theta_{2})||\theta_{1}-\theta_{2}|$
$\displaystyle\leq
C\left(\mathcal{L},\max_{i=1,2}(|\theta_{i,[1]}|)\right)|\theta_{1}-\theta_{2}|\,.$

For (II), we note first that

$\displaystyle\left|Z_{\rho}(t_{1};x)-Z_{\rho}(t_{2};x)\right|$
$\displaystyle\leq\left|\int^{t_{1}}_{t_{2}}\int_{\mathbb{R}^{k}}f(Z_{\rho}(t;x),\theta)\,\mathrm{d}\rho(\theta,t)\,\mathrm{d}t\right|$
$\displaystyle\leq
C(\mathcal{L})\left|\int^{t_{1}}_{t_{2}}\int_{\mathbb{R}^{k}}\left(\left|\theta\right|^{2}+1\right)\,\mathrm{d}\rho(\theta,t)\,\mathrm{d}t\right|$
$\displaystyle\leq C(\mathcal{L})(\mathcal{L}^{\sup}+1)|t_{1}-t_{2}|$
$\displaystyle\leq C(\mathcal{L}^{\sup})|t_{1}-t_{2}|\,,$

where we use Assumption 4.1 (12) together with (18) and
$\mathcal{L}\leq\mathcal{L}^{\sup}$. This bound implies that

$\displaystyle\mathrm{(II)}$
$\displaystyle\leq\mathbb{E}_{x\sim\mu}\left(|p^{\top}_{\rho}(t_{1};x)|\left|\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{2})-\partial_{\theta}f(Z_{\rho}(t_{2};x),\theta_{2})\right|\right)$
$\displaystyle\leq
C(\mathcal{L})\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho}(t_{1};x),\theta_{2})-\partial_{\theta}f(Z_{\rho}(t_{2};x),\theta_{2})\right|\right)$
$\displaystyle\leq
C(\mathcal{L}^{\sup},|\theta_{2,[1]}|)\mathbb{E}_{x\sim\mu}\left(\left|Z_{\rho}(t_{1};x)-Z_{\rho}(t_{2};x)\right|\right)$
$\displaystyle\leq C(\mathcal{L}^{\sup},|\theta_{2,[1]}|)|t_{1}-t_{2}|\,,$

where we use (31) in the first inequality, Assumption 4.1 (14) for $f$, and
(18) in the second inequality.

To bound term (III), we again use boundedness of $Z_{\rho}$, $p_{\rho}$ and
the Lipschitz condition of $f$ to obtain

$\displaystyle|p_{\rho}(t_{1};x)-p_{\rho}(t_{2};x)|$
$\displaystyle\leq\left|\int^{t_{1}}_{t_{2}}\int_{\mathbb{R}^{k}}p^{\top}_{\rho}(t;x)\partial_{z}f(Z_{\rho}(t;x),\theta)\,\mathrm{d}\rho(\theta,t)\,\mathrm{d}t\right|$
$\displaystyle\leq
C(\mathcal{L})\left|\int^{t_{1}}_{t_{2}}\int_{\mathbb{R}^{k}}(\left|\theta\right|^{2}+1)\,\mathrm{d}\rho(\theta,t)\,\mathrm{d}t\right|$
$\displaystyle\leq C(\mathcal{L})(\mathcal{L}^{\sup}+1)|t_{1}-t_{2}|$
$\displaystyle\leq C(\mathcal{L}^{\sup})|t_{1}-t_{2}|\,,$

so that

$\mathrm{(III)}\leq\mathbb{E}_{x\sim\mu}\left(|\partial_{\theta}f(Z_{\rho}(t_{2};x),\theta_{2})||p_{\rho}(t_{1};x)-p_{\rho}(t_{2};x)|\right)\leq
C(\mathcal{L}^{\sup})(|\theta_{2}|+1)|t_{1}-t_{2}|\,,$

where we use Assumption 4.1 (13) in the second inequality. By substituting
these bounds into (44), we complete the proof of (42).

Finally, to prove (43), we recall the definition of the Fréchet derivative, to
obtain

$\displaystyle\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta,t)-\nabla_{\theta}\frac{\delta
E(\rho_{2})}{\delta\rho}(\theta,t)\right|$ $\displaystyle\leq$
$\displaystyle\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho_{1}}(t;x),\theta)p_{\rho_{1}}(t;x)-\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta)p_{\rho_{2}}(t;x)\right|\right)$
$\displaystyle\leq$
$\displaystyle\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho_{1}}(t;x),\theta)-\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta)\right||p_{\rho_{1}}(t;x)|\right)$
$\displaystyle+\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta)\right|\left|p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\right|\right)$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L},|\theta_{[1]}|)\mathbb{E}_{x\sim\mu}\left(\left|Z_{\rho_{1}}(t;x)-Z_{\rho_{2}}(t;x)\right|\right)+C(\mathcal{L})(|\theta|+1)\mathbb{E}_{x\sim\mu}\left(\left|p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\right|\right)$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L},|\theta_{[1]}|,\mathcal{R})(1+|\theta|)d_{1}(\rho_{1},\rho_{2})\,,$

where we use Assumption 4.1 for $f$ with (18) and (31) in the third
inequality. In the last inequality, we use (19), (32).

###  C.2 A-priori estimate under the homogeneous assumption

In Theorem 6.1 and 6.2, 2-homogeneity or partial 1-homogeneity are assumed.
When these properties hold, we can sharpen the estimates obtained in the
previous section. We summarize our results here.

###### Lemma C.3

Suppose that Assumption 4.1 holds, then there exists a constant $C_{3}(r)$
depending increasingly on $r$ such that for any $r>0$ with $|x|<r$ and
$\theta\in\mathbb{R}^{k}$, we have the following.

  * •

If $f$ is 2-homogeneous (Assumption 6.1), then

$\left|\partial^{2}_{x}f\right|\leq
C_{3}(r)\left(|\theta|^{2}+1\right),\quad\left|\partial_{x}\partial_{\theta}f\right|\leq
C_{3}(r)\left(|\theta|+1\right),\quad\left|\partial^{2}_{\theta}f\right|\leq
C_{3}(r)\,.$ (45)

  * •

If $f$ is partially 1-homogeneous (Assumption 6.2), then

$\left|\partial^{2}_{x}f\right|\leq
C_{3}(r)(|\theta|+1),\quad\left|\partial_{x}\partial_{\theta}f\right|\leq
C_{3}(r)(|\theta|+1),\quad\left|\partial^{2}_{\theta}f\right|\leq
C_{3}(r)(|\theta|+1)\,,$ (46)

where $|\cdot|$ denotes the Frobenius norm.

Proof  When $f$ satisfies Assumption 4.1 and Assumption 6.2, (46) can be
obtained from direct calculation. When $f$ is 2-homogeneous in $\theta$, we
have

$f(x,\theta)=|\theta|^{2}f\left(x,\theta/|\theta|\right)=f_{1}(\theta)f_{2}(x,\theta)\,,$

where
$f_{1}(\theta)=|\theta|^{2},f_{2}(x,\theta)=f\left(x,\theta/|\theta|\right)$.
Naturally, with product rule, the derivatives of $f$ becomes the products of
derivatives of $f_{1}$ and $f_{2}$. We then obtain (45), noting that when
$|\theta|>1$, we have

$\left|\partial^{2}_{x}f_{1}\right|=0,\quad\left|\partial_{x}\partial_{\theta}f_{1}\right|\leq
0,\quad\left|\partial_{x}f_{1}\right|=0,\quad\left|\partial_{\theta}f_{1}\right|\leq
2|\theta|,\quad\left|\partial^{2}_{\theta}f_{1}\right|\leq 2k,$

and

$\left|\partial^{2}_{x}f_{2}\right|=C(r),\quad\left|\partial_{x}\partial_{\theta}f_{2}\right|\leq\frac{C(r)}{|\theta|},\quad\left|\partial_{x}f_{2}\right|=C(r),\quad\left|\partial_{\theta}f_{2}\right|\leq\frac{C(r)}{|\theta|},\quad\left|\partial^{2}_{\theta}f_{2}\right|\leq\frac{C(r)}{|\theta|^{2}}\,.$

The estimates above allow us to improve the stability results in Lemmas C.1
and C.2.

###### Lemma C.4

Suppose that Assumption 4.1 holds and let
$\rho_{1},\rho_{2}\in\mathcal{C}([0,1];\mathcal{P}^{2})$. Define
$\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ as in (20) and

$\mathcal{L}^{\sup}_{1}=\sup_{t\in[0,1]}\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)\,.$

Then we have the following results.

  * •

If $f$ either satisfies Assumption 6.1 or Assumption 6.2, we have the
following properties.

    * 1.

Stability of $p_{\rho}$:

$\left|p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\right|\leq
C(\mathcal{L}_{1},\mathcal{L}_{2})d_{1}(\rho_{1},\rho_{2})\,.$ (47)

    * 2.

Lipschitz continuity in $\rho$:

$\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta,t)-\nabla_{\theta}\frac{\delta
E(\rho_{2})}{\delta\rho}(\theta,t)\right|\leq
C(\mathcal{L}_{1},\mathcal{L}_{2})d_{1}(\rho_{1},\rho_{2})\left(|\theta|+1\right)\,,$
(48)

where $d_{1}$ is defined in (4.1).

  * •

Lipschitz continuity in $\theta$ and $t$: If $f$ satisfies Assumption 6.1,
then for any
$(\theta_{1},t_{1}),(\theta_{2},t_{2})\in\mathbb{R}^{k}\times[0,1]$, we have

$\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta_{1},t_{1})-\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta_{2},t_{2})\right|\leq
C(\mathcal{L}_{1})|\theta_{1}-\theta_{2}|+C(\mathcal{L}^{\sup}_{1})(|\theta_{2}|+1)|t_{1}-t_{2}|\,.$
(49)

  * •

Lipschitz continuity in $\theta$ and $t$: If $f$ satisfies Assumption 6.2,
then for any
$(\theta_{1},t_{1}),(\theta_{2},t_{2})\in\mathbb{R}^{k}\times[0,1]$, we have

$\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta_{1},t_{1})-\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta_{2},t_{2})\right|\leq
C(\mathcal{L}^{\sup}_{1})(|\theta_{1}|+|\theta_{2}|+1)(|\theta_{1}-\theta_{2}|+|t_{1}-t_{2}|)\,,$
(50)

###### Remark C.1

We note that in comparing (32) with (47) and (42)-(43) with (48)-(50), the
main differences are the dependence of the bounds on $\theta$. The new
estimates have explicit (and mild) dependence on $\theta$.

Proof  First, to prove (47), we let
$\Delta(t;x)=p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)$, and recall (35) and (36).
Using the boundedness of $Z_{\rho}$ in (18) and (19), and calling (45) (or
(46)), we have from (36) that

$|D_{\rho_{1},\rho_{2}}(t;x)|\leq
C(\mathcal{L}_{1},\mathcal{L}_{2})\left(\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)+\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{2}(\theta,t)+1\right)d_{1}(\rho_{1},\rho_{2})\,.$

By substituting into (35), and using (31), (33), and Hölder’s inequality, we
have

$\displaystyle\frac{\,\mathrm{d}|\Delta(t;x)|^{2}}{\,\mathrm{d}t}\leq$
$\displaystyle
C(\mathcal{L}_{1},\mathcal{L}_{2})\left(\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)+\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{2}(\theta,t)+1\right)|\Delta(t;x)|^{2}$
$\displaystyle+C(\mathcal{L}_{1},\mathcal{L}_{2})\left(\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{1}(\theta,t)+\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho_{2}(\theta,t)+1\right)d^{2}_{1}(\rho_{1},\rho_{2})\,.$

The result (47) follows from the initial condition (34) together with
Grönwall’s inequality.

Next, to prove (48), we have

$\displaystyle\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta,t)-\nabla_{\theta}\frac{\delta
E(\rho_{2})}{\delta\rho}(\theta,t)\right|$ $\displaystyle\leq$
$\displaystyle\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho_{1}}(t;x),\theta)p_{\rho_{1}}(t;x)-\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta)p_{\rho_{2}}(t;x)\right|\right)$
$\displaystyle\leq$
$\displaystyle\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho_{1}}(t;x),\theta)-\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta)\right||p_{\rho_{1}}(t;x)|\right)$
$\displaystyle+\mathbb{E}_{x\sim\mu}\left(\left|\partial_{\theta}f(Z_{\rho_{2}}(t;x),\theta)\right|\left|p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\right|\right)$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{1},\mathcal{L}_{2})(|\theta|+1)\left(\mathbb{E}_{x\sim\mu}\left(\left|Z_{\rho_{1}}(t;x)-Z_{\rho_{2}}(t;x)\right|\right)+\mathbb{E}_{x\sim\mu}\left(\left|p_{\rho_{1}}(t;x)-p_{\rho_{2}}(t;x)\right|\right)\right)$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{1},\mathcal{L}_{2})d_{1}(\rho_{1},\rho_{2})\left(|\theta|+1\right)\,,$

where we also use (13), (18), (31), (45) (or (46)) in the second inequality
and (19) and (47) in the final inequality.

Finally, to prove (49) and (50), we have as in (44) that

$\displaystyle\left|\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta_{1},t_{1})-\nabla_{\theta}\frac{\delta
E(\rho_{1})}{\delta\rho}(\theta_{2},t_{2})\right|$ (51) $\displaystyle\leq$
$\displaystyle\underbrace{\left|\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho_{1}}(t_{1};x)\partial_{\theta}f(Z_{\rho_{1}}(t_{1};x),\theta_{1})-p^{\top}_{\rho_{1}}(t_{1};x)\partial_{\theta}f(Z_{\rho_{1}}(t_{1};x),\theta_{2})\right)\right|}_{\mathrm{(I)}}$
$\displaystyle+\underbrace{\left|\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho_{1}}(t_{1};x)\partial_{\theta}f(Z_{\rho_{1}}(t_{1};x),\theta_{2})-p^{\top}_{\rho_{1}}(t_{1};x)\partial_{\theta}f(Z_{\rho_{1}}(t_{2};x),\theta_{2})\right)\right|}_{\mathrm{(II)}}$
$\displaystyle+\underbrace{\left|\mathbb{E}_{x\sim\mu}\left(p^{\top}_{\rho_{1}}(t_{1};x)\partial_{\theta}f(Z_{\rho_{1}}(t_{2};x),\theta_{2})-p^{\top}_{\rho_{1}}(t_{2};x)\partial_{\theta}f(Z_{\rho_{1}}(t_{2};x),\theta_{2})\right)\right|}_{\mathrm{(III)}}\,.$

The boundedness of the three terms above relies on Assumption 4.1 and (45) (or
(46)).

To bound $\mathrm{(I)}$, if $f$ is $2$-homogeneous, we have

$|\mathrm{(I)}|\leq C(\mathcal{L}_{1})|\theta_{1}-\theta_{2}|\,,$

where we use (18), (31), and $\partial^{2}_{\theta}f\leq C(|z|)$. If $f$
satisfies Assumption 6.2, we have

$|\mathrm{(I)}|\leq
C(\mathcal{L}_{1})(|\theta_{1}|+|\theta_{2}|+1)|\theta_{1}-\theta_{2}|\,,$

where we use (18), (31), and $\partial^{2}_{\theta}f(z,\theta)\leq
C(|z|)(|\theta|+1)$.

The bounds of (II) and (III) are same for both homogeneity assumptions and
similar to the proof of Lemma C.2. For (II), we have

$|\partial_{\theta}f(Z_{\rho_{1}}(t_{1};x),\theta_{2})-\partial_{\theta}f(Z_{\rho_{1}}(t_{2};x),\theta_{2})|\leq
C(\mathcal{L}_{1})(|\theta_{2}|+1)|Z_{\rho_{1}}(t_{1};x)-Z_{\rho_{1}}(t_{2};x)|$

and

$\left|Z_{\rho_{1}}(t_{1};x)-Z_{\rho_{1}}(t_{2};x)\right|\leq
C(\mathcal{L}_{1})(\mathcal{L}^{\sup}_{1}+1)|t_{1}-t_{2}|\,.$

For (III), we also use

$|\partial_{\theta}f(Z_{\rho_{1}}(t_{2};x),\theta_{2})|\leq
C(\mathcal{L}_{1})(|\theta_{2}|+1)$

and

$|p_{\rho_{1}}(t_{1};x)-p_{\rho_{1}}(t_{2};x)|\leq
C(\mathcal{L}_{1})(\mathcal{L}^{\sup}_{1}+1)|t_{1}-t_{2}|\,.$

These estimates, together with $\mathcal{L}_{1}\leq\mathcal{L}^{\sup}_{1}$ and
(51), prove the claims (49) and (50).

##  Appendix D Well posedness of gradient flow

Theorem A.4 is about well-posedness of the gradient flow equation (9) in the
mean-field limit, while Theorem A.3 shows the corresponding result in the
discrete setting. The proof for the two are similar. We first show the mean-
field limit well-posedness, Theorem A.4.

###  D.1 Proof of Theorem A.4

We use the fixed-point argument. To do so, we first define a subset of
$\mathcal{C}([0,S];\mathcal{C}([0,1];\mathcal{P}^{2}))$ with compact support
measures, as follows:

$\displaystyle\Omega_{S}=$
$\displaystyle\left\\{\phi(\theta,t,s)\in\mathcal{C}([0,S];\mathcal{C}([0,1];\mathcal{P}^{2}))\middle|\exists
r>0,\ \mathrm{supp}(\phi)\subset\\{\theta||\theta_{[1]}|<r\\},\
\forall(t,s)\in[0,1]\times[0,S]\right\\}$
$\displaystyle\cap\left\\{\phi(\theta,t,s)\in\mathcal{C}([0,S];\mathcal{C}([0,1];\mathcal{P}^{2}))\middle|\phi(\theta,t,0)=\rho_{\mathrm{ini}}(\theta,t)\right\\}$

For any $\phi(\theta,t,s)\in\Omega_{S}$ with
$\phi(\theta,t,0)=\rho_{\mathrm{ini}}(\theta,t)$, we define a map

$\varphi=\mathcal{T}_{S}(\phi):\Omega_{S}\rightarrow\Omega_{S}$ (52)

where $\varphi$ solves

$\left\\{\begin{aligned} &\frac{\partial\varphi(\theta,t,s)}{\partial
s}=\nabla_{\theta}\cdot\left(\varphi(\theta,t,s)\nabla_{\theta}\frac{\delta
E(\phi(s))}{\delta\rho}(\theta,t)\right)\,,\\\
&\varphi(\theta,t,0)=\rho_{\mathrm{ini}}(\theta,t)\,.\end{aligned}\right.$
(53)

The strategy is to show this map is a contraction map, so that there is a
fixed point in the set $\Omega_{S}$, which is then the solution to equation
(9).

The proof of Theorem A.4 is divided into three steps:

  * Step 1:

We show $\mathcal{T}_{S}$ is well-defined map from $\Omega_{S}$ to
$\Omega_{S}$.

  * Step 2:

We give a bound of
$d_{2}(\mathcal{T}_{S}(\phi_{1}),\mathcal{T}_{S}(\phi_{2}))$ in terms of
$d_{2}(\phi_{1}\,,\phi_{2})$. One can then tune $S$ to ensure
$\mathcal{T}_{S}$ is a contraction map, meaning there is a unique fixed point
$\phi^{\ast}$ so that $\phi^{\ast}=\mathcal{T}_{S}(\phi^{\ast})$, and thus
$\phi^{\ast}$ solves (53) for $s<S$.

  * Step 3:

We extend the local solution to a global solution.

##### Step 1.

According to the definition of (53), for a fixed $\phi(\theta,t,s)$, let

$\begin{cases}\frac{\,\mathrm{d}\theta_{\phi}(s;t)}{\,\mathrm{d}s}=-\nabla_{\theta}\frac{\delta
E(\phi(s))}{\delta\rho}\left(\theta_{\phi}(s;t),t\right)\,,\\\
\theta_{\phi}(0;t)\sim\rho_{\mathrm{ini}}(\theta,t)\,.\end{cases}$ (54)

Then $\theta_{\phi}\sim\varphi=\mathcal{T}_{S}(\phi)$. To show the existence
of $\varphi$ amounts to showing the wellposedness of (54).

According to Lemma C.2 (41) and (42), the force $\nabla_{\theta}\frac{\delta
E}{\delta\phi(s)}(\cdot,t)$ has a linear growth and is locally Lipschitz.
Classical ODE theory then suggests there is a unique solution for $s\in[0,S]$,
which depends continuously on the initial value $\theta(0;t)$.

Denoting

$\displaystyle\mathcal{L}_{S,\phi}$ $\displaystyle=\sup_{0\leq s\leq
S}\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi(\theta,t,s)\,\mathrm{d}t,$
(55) $\displaystyle\mathcal{L}^{\sup}_{\mathrm{ini}}$
$\displaystyle=\sup_{0\leq t\leq
1}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho_{\mathrm{ini}}(\theta,t),$
$\displaystyle\mathcal{L}^{\sup}_{S,\phi}$ $\displaystyle=\sup_{0\leq t\leq
1,0\leq s\leq S}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi(\theta,t,s)$
$\displaystyle\mathcal{R}_{\mathrm{ini}}$
$\displaystyle=\inf_{r>0}\left\\{\mathrm{supp}(\rho_{\mathrm{ini}}(\theta,t))\subset\\{\theta||\theta_{[1]}|<r\\},\
\forall t\in[0,1]\right\\},$ $\displaystyle\mathcal{R}_{S,\phi}$
$\displaystyle=\inf_{r>0}\left\\{\mathrm{supp}(\phi(\theta,t,s))\subset\\{\theta||\theta_{[1]}|<r\\},\
\forall(t,s)\in[0,1]\times[0,S]\right\\}\,,$

we have the following proposition.

###### Proposition D.1

Suppose that $\theta_{\phi}(s;t)$ solves (54) and $\phi\in\Omega_{S}$. Then
for any $(t_{1},s_{1}),(t_{2},s_{2})\in[0,1]\times[0,S]$, we have

$\displaystyle\mathbb{E}\left(\left|\theta_{\phi}(s_{1};t_{1})\right|^{2}\right)$
$\displaystyle\leq\exp(SC(\mathcal{L}_{S,\phi}))\left(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\right)\,,\
|\theta_{\phi,[1]}(s_{1};t_{1})|$
$\displaystyle\leq\exp(SC(\mathcal{L}_{S,\phi}))\left(\mathcal{R}_{\mathrm{ini}}+1\right)$
(56)

and

$\displaystyle\mathbb{E}\left(\left|\theta_{\phi}(s_{1};t_{1})-\theta_{\phi}(s_{2};t_{2})\right|^{2}\right)$
(57) $\displaystyle\leq
C\left(\mathcal{L}^{\sup}_{S,\phi},\mathcal{R}_{\mathrm{ini}},S\right)\left(\mathbb{E}\left(\left|\theta_{\phi}(0;t_{1})-\theta_{\phi}(0;t_{2})\right|^{2}\right)+|t_{1}-t_{2}|^{2}+|s_{1}-s_{2}|^{2}\right)\,.$

Proof  To prove the first bound in (56), we multiply (54) on both sides by
$\theta_{\phi}$ and use boundedness of the forcing term (41) to obtain

$\frac{\,\mathrm{d}\left|\theta_{\phi}(s;t_{1})\right|^{2}}{\,\mathrm{d}s}\leq
C(\mathcal{L}_{S,\phi})(|\theta_{\phi}(s;t_{1})|^{2}+1)\,.$

Using Grönwall’s inquality together with
$\mathbb{E}\left(|\theta_{\phi}(0;t_{1})|^{2}\right)\leq\mathcal{L}^{\sup}_{\mathrm{ini}}$,
we obtain the first bound in (56). To prove the second bound in (56), we use
the bound of $\left|\partial_{\theta_{[1]}}\frac{\delta
E(\rho)}{\delta\rho}(\theta,t)\right|$ according to Lemma C.2 (41):

$\frac{\,\mathrm{d}\left|\theta_{\phi,[1]}(s;t_{1})\right|^{2}}{\,\mathrm{d}s}\leq
C(\mathcal{L}_{S,\phi})(\left|\theta_{\phi,[1]}(s;t_{1})\right|^{2}+1)\,.$

Using Grönwall’s inquality together with
$|\theta_{\phi,[1]}(0;t_{1})|\leq\mathcal{R}_{\mathrm{ini}}$, we obtain the
second bound in (56).

To show (57), we first write

$\displaystyle\mathbb{E}\left(\left|\theta_{\phi}(s_{1};t_{1})-\theta_{\phi}(s_{2};t_{2})\right|^{2}\right)$
(58) $\displaystyle\leq
2\underbrace{\mathbb{E}\left(\left|\theta_{\phi}(s_{1};t_{1})-\theta_{\phi}(s_{1};t_{2})\right|^{2}\right)}_{\textrm{(I)}}+2\underbrace{\mathbb{E}\left(\left|\theta_{\phi}(s_{1};t_{2})-\theta_{\phi}(s_{2};t_{2})\right|^{2}\right)}_{\textrm{(II)}}\,,$

then bound the two terms (I) and (II) separately. For (I), we use (54), the
second bound of (56), and Lemma C.2 (42)

$\displaystyle\frac{\,\mathrm{d}\left|\theta_{\phi}(s;t_{1})-\theta_{\phi}(s;t_{2})\right|^{2}}{\,\mathrm{d}s}$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{S,\phi},\mathcal{R}_{\mathrm{ini}})\left|\theta_{\phi}(s;t_{1})-\theta_{\phi}(s;t_{2})\right|^{2}+C(\mathcal{L}^{\sup}_{S,\phi},\mathcal{R}_{\mathrm{ini}})(|\theta_{\phi}(s_{1};t_{1})|^{2}+|\theta_{\phi}(s_{1};t_{2})|^{2}+1)|t_{1}-t_{2}|^{2}\,,$

Using the first bound in (56) and Grönwall’s inequality, this implies that

$\displaystyle\mathbb{E}\left(\left|\theta_{\phi}(s_{1};t_{1})-\theta_{\phi}(s_{1};t_{2})\right|^{2}\right)\leq
C\left(\mathcal{L}^{\sup}_{S,\phi},\mathcal{R}_{\mathrm{ini}},S\right)\left(\mathbb{E}\left(\left|\theta_{\phi}(0;t_{1})-\theta_{\phi}(0;t_{2})\right|^{2}\right)+|t_{1}-t_{2}|^{2}\right)\,.$
(59)

For (II), we obtain an estimate by integrating (54) from $s_{1}$ to $s_{2}$
and using the boundedness of $\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}$ in (41). From the Grönwall inequality, we have

$\left|\theta_{\phi}(s_{1};t_{2})-\theta_{\phi}(s_{2};t_{2})\right|\leq
C\left(\mathcal{L}_{S,\phi}\right)\left(\int^{s_{2}}_{s_{1}}|\theta_{\phi}(s;t_{2})|\,\mathrm{d}s+|s_{1}-s_{2}|\right)\,.$

Using first bound in (56) and Hölder’s inequality, this implies

$\mathbb{E}\left|\theta_{\phi}(s_{1};t_{2})-\theta_{\phi}(s_{2};t_{2})\right|^{2}\leq
C\left(\mathcal{L}_{S,\phi},\mathcal{L}^{\sup}_{\mathrm{ini}},S\right)|s_{1}-s_{2}|^{2}\,$
(60)

and thus combining with (59) and substituting in (58), we complete the proof.

An immediate corollary of Proposition D.1 is that the map
$\mathcal{T}_{S}(\cdot)$ is well defined.

###### Corollary D.1

For every fixed $S>0$, the map $\mathcal{T}_{S}$ is well defined. That is, for
any $\phi\in\Omega_{S}$, one can find
$\varphi=\mathcal{T}_{S}(\phi)\in\Omega_{S}$ as the unique solution of (53).
In particular, for any $(t,s)\in[0,1]\times[0,S]$, we have

$\displaystyle\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\varphi(\theta,t,s)\leq\exp(SQ_{1}(\mathcal{L}_{S,\phi}))\left(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\right)\,,$
(61)
$\displaystyle\mathrm{supp}_{\theta}(\varphi(\theta,t,s))\subset\left\\{\theta\middle||\theta_{[1]}|\leq\exp(SQ_{1}(\mathcal{L}_{S,\phi}))\left(\mathcal{R}_{\mathrm{ini}}+1\right)\right\\}\,,$

where $Q_{1}:\mathbb{R}_{+}\rightarrow\mathbb{R}_{+}$ is depends only on
$\mathcal{L}_{S,\phi}$ and is an increasing function of its argument.

Proof  For fixed $(t,s)\in[0,1]\times[0,S]$, define $\varphi(\theta,t,s)$ as
the distribution of $\theta_{\phi}(s;t)$. Using classical stochastic theory
(Ambrosio et al., 2008, Prop 8.1.8), $\varphi(\theta,t,s)$ is a solution to
(53). The estimate of the support is a consequence of (56). Finally, using
(16) and (57), we obtain that

$\lim_{(t,s)\rightarrow(t_{0},s_{0})}W_{2}(\varphi(\cdot,t,s),\varphi(\cdot,t_{0},s_{0}))\leq\lim_{(t,s)}\left(\mathbb{E}\left(\left|\theta_{\phi}(s;t)-\theta_{\phi}(s_{0};t_{0})\right|^{2}\right)\right)^{1/2}=0\,,$

which proves the continuity in $s$ and $t$, so that
$\varphi\in\mathcal{C}([0,S];\mathcal{C}([0,1];\mathcal{P}^{2}))$. By
combining all the factors above, we conclude that $\varphi\in\Omega_{S}$.

##### Step 2.

We show now that $\mathcal{T}_{S}$ is a contraction map for $S$ sufficiently
small.

###### Proposition D.2

For any $\phi_{1},\phi_{2}\in\Omega_{S}$, we have

$d_{2}(\mathcal{T}_{S}(\phi_{1}),\mathcal{T}_{S}(\phi_{2}))\leq
SQ_{2}(\mathcal{L}_{S},\mathcal{R}_{S},S)d_{2}(\phi_{1},\phi_{2}),$ (62)

where $Q_{2}:\mathbb{R}^{3}\rightarrow\mathbb{R}_{+}$ is an increasing
function and
$\mathcal{L}_{S}=\max\\{\mathcal{L}_{S,\phi_{1}}\,,\mathcal{L}_{S,\phi_{2}}\\}$,
$\mathcal{R}_{S}=\max\\{\mathcal{R}_{S,\phi_{1}}\,,\mathcal{R}_{S,\phi_{2}}\\}$,
with $\mathcal{L}_{S,\phi},\mathcal{R}_{S,\phi}$ defined in (55).

Proof  Denote by $\theta_{\phi_{i}}(s;t)$ the solutions to (54) with
$\phi=\phi_{i}$ using the same initial data, that is,

$\theta_{\phi_{1}}(0;t)=\theta_{\phi_{2}}(0;t)\,.$

As in the previous subsection, we translate the study of $\varphi_{i}$ to the
study of $\theta_{\phi_{i}}$. Defining

$\Delta_{t}(s)=|\theta_{\phi_{1}}(s;t)-\theta_{\phi_{2}}(s;t)|\,,$

we have according to the definition of Wasserstein distance that

$d_{2}(\mathcal{T}_{S}(\phi_{1}),\mathcal{T}_{S}(\phi_{2}))\leq\sup_{(t,s)\in[0,1]\times[0,S]}\mathbb{E}\left(\Delta^{2}_{t}(s)\right)\,.$

Using (54), we obtain

$\displaystyle\frac{\,\mathrm{d}(\Delta_{t}(s))^{2}}{\,\mathrm{d}s}$
$\displaystyle\leq 4(\Delta_{t}(s))^{2}+4\left|\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{1}},t)-\nabla_{\theta}\frac{\delta
E(\phi_{2}(s))}{\delta\rho}(\theta_{\phi_{2}},t)\right|^{2}$ (63)
$\displaystyle\leq 4(\Delta_{t}(s))^{2}+8\left|\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{1}},t)-\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{2}},t)\right|^{2}$
$\displaystyle+8\left|\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{2}},t)-\nabla_{\theta}\frac{\delta
E(\phi_{2}(s))}{\delta\rho}(\theta_{\phi_{2}},t)\right|^{2}\,.$

The second term on the right hand side involves the continuity addressed in
Lemma C.2 (42)

$\left|\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{1}},t)-\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{2}},t)\right|^{2}\leq
C(\mathcal{L}_{S},\mathcal{R}_{S},S)(\Delta_{t}(s))^{2}\,,$

where we use the second bound in (56) to substitute the constant in (42). Then
the last term of (63) involves the continuity discussed in (43). In
particular, we have

$\displaystyle\left|\nabla_{\theta}\frac{\delta
E(\phi_{1}(s))}{\delta\rho}(\theta_{\phi_{2}},t)-\nabla_{\theta}\frac{\delta
E(\phi_{2}(s))}{\delta\rho}(\theta_{\phi_{2}},t)\right|$ $\displaystyle\leq$
$\displaystyle
C(\mathcal{R}_{S},|\theta_{\phi_{2},[1]}|)(1+|\theta_{\phi_{2}}|)d_{1}(\phi_{1},\phi_{2})\leq
C(\mathcal{L}_{S},\mathcal{R}_{S},S)(1+|\theta_{\phi_{2}}|)d_{2}(\phi_{1},\phi_{2})\,,$

where $d_{2}$ is defined in Definition 4.1 and in the second inequality we use
the second bound in (56) to substitute the constant in (43).

By substituting in (63), we obtain

$\frac{\,\mathrm{d}(\Delta_{t}(s))^{2}}{\,\mathrm{d}s}\leq
C(\mathcal{L}_{S},\mathcal{R}_{S},S)\left[(\Delta_{t}(s))^{2}+(1+|\theta_{\phi_{2}}|^{2})d^{2}_{2}(\phi_{1},\phi_{2})\right]\,,$

which implies

$\frac{\,\mathrm{d}\left(\mathbb{E}(\Delta_{t}(s))^{2}\right)}{\,\mathrm{d}s}\leq
C(\mathcal{L}_{S},\mathcal{R}_{S},S)\left[\left(\mathbb{E}(\Delta_{t}(s))^{2}\right)+d^{2}_{2}(\phi_{1},\phi_{2})\right]\,,$

where we use the first inequality in (56). From the Grönwall inequality, there
exists $Q_{2}:\mathbb{R}^{3}\rightarrow\mathbb{R}_{+}$ is an increasing
function such that

$\displaystyle\mathbb{E}\left((\Delta_{t}(s))^{2}\right)\leq
SQ_{2}(\mathcal{L}_{S},\mathcal{R}_{S},S)d^{2}_{2}(\phi_{1},\phi_{2})\,,$

completing the proof.

To apply the contraction mapping theorem, we need to verify two conditions in
order to show that there exists a fixed point
$\phi^{\ast}=\mathcal{T}_{S}(\phi^{\ast})$:

  * •

There is a closed subset in $\Omega_{S}$ such that $\mathcal{T}_{S}$ maps this
set to itself.

  * •

$\mathcal{T}_{S}$ is a contraction map in this subset.

For the closed subset, we define

$\displaystyle B_{\rho_{0}}=$
$\displaystyle\left\\{\phi\in\Omega_{S}\middle|\mathrm{supp}(\phi(t,s))\subset\\{\theta||\theta_{[1]}|\leq
4(\mathcal{R}_{\mathrm{ini}}+1)\\},\quad\forall(t,s)\in[0,1]\times[0,S]\right\\}$
(64)
$\displaystyle\cap\left\\{\phi\in\Omega_{S}\middle|\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi(\theta,t,s)\leq
4\left(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\right),\quad\forall(t,s)\in[0,1]\times[0,S]\right\\}\,.$

We now claim that for small enough $S$, $\mathcal{T}_{S}$ is a contraction map
in $B_{\rho_{0}}$.

###### Proposition D.3

Suppose that $S$ is small enough that

$\displaystyle\exp(SQ_{1}(4(\mathcal{L}^{\sup}_{\mathrm{ini}}+1)))\left(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\right)$
$\displaystyle\leq 4(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\,,$
$\displaystyle\exp(SQ_{1}(4(\mathcal{L}^{\sup}_{\mathrm{ini}}+1)))\left(\mathcal{R}_{\mathrm{ini}}+1\right)$
$\displaystyle\leq 4(\mathcal{R}_{\mathrm{ini}}+1)\,,$ $\displaystyle
SQ_{2}(4(\mathcal{L}^{\sup}_{\mathrm{ini}}+1),4(\mathcal{R}_{\mathrm{ini}}+1),S)$
$\displaystyle<\frac{1}{2}\,,$

where $Q_{1}$ and $Q_{2}$ are defined in Corollary D.1 and Proposition D.2,
respectively. Then we have the following.

  * •

If $\phi\in B_{\rho_{0}}$, then $\mathcal{T}_{S}(\phi)\in B_{\rho_{0}}$, that
is, for any $(t,s)\in[0,1]\times[0,S]$, we have

$\mathrm{supp}(\mathcal{T}_{S}(\phi)(t,s))\subset\\{\theta||\theta_{[1]}|\leq
4(\mathcal{R}_{\mathrm{ini}}+1)\\}\,,$ (65)

and

$\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\mathcal{T}_{S}(\phi)(\theta,t,s)\leq
4\left(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\right)\,.$ (66)

  * •

$\mathcal{T}_{S}$ is a contraction map in this subset, meaning that for any
$\phi_{1},\phi_{2}\in B_{\rho_{0}}$, we have

$d_{2}(\mathcal{T}_{S}(\phi_{1}),\mathcal{T}_{S}(\phi_{2}))<\frac{1}{2}d_{2}(\phi_{1},\phi_{2})\,.$
(67)

Proof  First, using Corollary D.1 (61) and noticing $\mathcal{L}_{S,\phi}\leq
4\left(\mathcal{L}^{\sup}_{\mathrm{ini}}+1\right)$, we prove (65), (66). Then,
using (62) with (65) and (66), we have

$d_{2}(\mathcal{T}_{S}(\phi_{1}),\mathcal{T}_{S}(\phi_{2}))\leq
SQ_{2}(4(\mathcal{L}^{\sup}_{\mathrm{ini}}+1),4(\mathcal{R}_{\mathrm{ini}}+1),S)<\frac{1}{2}d_{2}(\phi_{1},\phi_{2}),$

which proves (67).

Using the contraction mapping theorem, we can obtain directly that
$\mathcal{T}_{S}(\phi)$ has a fixed point in $B_{\rho_{0}}$ when $S$ is small
enough.

###### Corollary D.2

If $S$ satisfies conditions in Proposition D.3, then there exists a unique
$\phi^{\ast}(\theta,t,s)\in B_{\rho_{0}}\subset\Omega_{S}$ such that
$\phi^{\ast}(\theta,t,s)$ is a solution to (9) with initial condition
$\rho_{\mathrm{ini}}(\theta,t)$.

This is a direct consequence of the application of contraction mapping
theorem.

Finally, we prove that the cost function decreases along the flow.

###### Lemma D.1

Suppose that
$\phi^{\ast}(\theta,t,s)\in\mathcal{C}([0,S];\mathcal{C}([0,1];\mathcal{P}^{2}))$
solves (9) with initial condition $\rho_{\mathrm{ini}}(\theta,t)$. Then for
$0<s<S$, we have

$\frac{\,\mathrm{d}E(\phi^{\ast}(\theta,t,s))}{\,\mathrm{d}s}=-\int^{1}_{0}\int_{\mathbb{R}^{k}}\left|\nabla_{\theta}\frac{\delta
E(\phi^{\ast}(s))}{\delta\rho}(\theta,t)\right|^{2}\,\mathrm{d}\phi^{\ast}(\theta,t,s)\,\mathrm{d}t\leq
0\,.$ (68)

Proof  Denote by $\theta^{\ast}(s;t)$ the associated path, meaning that
$\theta^{\ast}(s;t)$ solves (54) with $\phi=\phi^{\ast}$, then
$\theta^{\ast}\sim\phi^{\ast}$, meaning the distribution of $\theta^{\ast}$ is
$\phi^{\ast}$. According to (9), we obtain using a change of variable that

$\frac{\,\mathrm{d}E(\phi^{\ast}(\theta,t,s))}{\,\mathrm{d}s}=-\int^{1}_{0}\int_{\mathbb{R}^{k}}\left|\nabla_{\theta}\frac{\delta
E(\phi^{\ast}(s))}{\delta\rho}(\theta,t)\right|^{2}\,\mathrm{d}\phi^{\ast}(\theta,t,s)\,\mathrm{d}t\leq
0\,,$ (69)

which proves the result. We note that the derivation in (69) is formal. A
rigorous proof can be found in (Ding et al., 2021, Appendix I).

##### Step 3.

In this final step of the proof, we extend the local solution from Corollary
D.2 to a global solution. Lemma D.1 shows that the formula of
$\frac{\,\mathrm{d}E}{\,\mathrm{d}s}$, so we can then use this formula to
improve the bound for the support of the solution (61). This improvement will
be shown in the following corollary. This improved estimate helps in extending
the local solution to the global solution.

###### Corollary D.3

For fixed $S$ satisfying the condition in Proposition D.3, denote by
$\phi^{\ast}(\theta,t,s)\in\mathcal{C}([0,S];\mathcal{C}([0,1];\mathcal{P}^{2}))$
the solution to (9) with initial condition $\rho_{\mathrm{ini}}(\theta,t)$.
Then for any $(t,s)\in[0,1]\times[0,S]$, we have

$\displaystyle\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi^{*}(\theta,t,s)$
$\displaystyle\leq
C(S,\mathcal{R}_{\mathrm{ini}},\mathcal{L}^{\sup}_{\mathrm{ini}})\,,$ (70)
$\displaystyle\mathrm{supp}(\phi^{*}(t,s))\subset\left\\{\theta\middle||\theta_{[1]}|\right.$
$\displaystyle\leq\left.C(S,\mathcal{R}_{\mathrm{ini}},\mathcal{L}^{\sup}_{\mathrm{ini}})\right\\}\,,$

where the quantity $C$ depends only on $S$, $\mathcal{R}_{\mathrm{ini}}$, and
$\mathcal{L}^{\sup}_{\mathrm{ini}})$.

Proof  According to (61), it suffices to prove

$\mathcal{L}_{S,\phi^{\ast}}=\sup_{0\leq s\leq
S}\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi^{\ast}(\theta,t,s)\,\mathrm{d}t\leq
C(S,\mathcal{R}_{\mathrm{ini}},\mathcal{L}^{\sup}_{\mathrm{ini}})\,.$

Denote by $\theta^{\ast}(s;t)$ the particle representation of $\phi^{\ast}$,
meaning that $\theta^{\ast}(s;t)$ solves (54) with $\phi=\phi^{\ast}$. Since
$\theta^{\ast}(s;t)\sim\phi^{\ast}(s;t)$,

$\mathcal{L}_{S,\phi^{\ast}}=\sup_{0\leq s\leq
S}\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi^{\ast}(\theta,t,s)\,\mathrm{d}t=\sup_{0\leq
s\leq
S}\int^{1}_{0}\mathbb{E}\left(|\theta^{\ast}(s;t)|^{2}\right)\,\mathrm{d}t\,.$

Using (54), we obtain that

$\frac{\,\mathrm{d}|\theta^{\ast}(s;t)|^{2}}{\,\mathrm{d}s}\leq|\theta^{\ast}(s;t)|\left|\nabla_{\theta}\frac{\delta
E(\phi^{\ast}(s))}{\delta\rho}\left(\theta^{*}(s;t),t\right)\right|\,,$

which gives

$\displaystyle\frac{\,\mathrm{d}\int^{1}_{0}\mathbb{E}\left(|\theta^{\ast}(s;t)|^{2}\right)dt}{\,\mathrm{d}s}$
$\displaystyle\leq$
$\displaystyle\left(\int^{1}_{0}\mathbb{E}\left(|\theta^{\ast}(s;t)|^{2}\right)dt\right)^{1/2}\left(\int^{1}_{0}\mathbb{E}\left(\left|\nabla_{\theta}\frac{\delta
E(\phi^{\ast}(s))}{\delta\rho}\left(\theta^{*}(s;t),t\right)\right|^{2}\right)dt\right)$
$\displaystyle=$
$\displaystyle\left(\int^{1}_{0}\mathbb{E}\left(|\theta^{\ast}(s;t)|^{2}\right)dt\right)^{1/2}\left|\frac{\,\mathrm{d}E(\phi^{\ast}(\theta,t,s))}{\,\mathrm{d}s}\right|\,,$

where we use the Hölder inequality and (68) from Lemma D.1 in the last
equality. Since

$\displaystyle\int^{1}_{0}\mathbb{E}\left(|\theta^{\ast}(0;t)|^{2}\right)dt$
$\displaystyle\leq\mathcal{L}^{\sup}_{\mathrm{ini}},$
$\displaystyle\int^{S}_{0}\left|\frac{\,\mathrm{d}E(\phi^{\ast}(\theta,t,s))}{\,\mathrm{d}s}\right|\,\mathrm{d}s$
$\displaystyle\leq
E(\rho_{\mathrm{ini}}(\theta,t))-E(\phi^{\ast}(\theta,t,S))\leq
C(\mathcal{L}^{\sup}_{\mathrm{ini}})\,,$

we obtain

$\mathcal{L}_{S,\phi^{\ast}}=\sup_{0\leq u\leq
s}\int^{1}_{0}\mathbb{E}\left(|\theta^{\ast}(u;t)|^{2}\right)dt\leq
C(S,\mathcal{L}^{\sup}_{\mathrm{ini}})$

by Grönwall’s inequality. This proves (70).

By contrast with (61), this estimate removes the dependence of the bound on
$\mathcal{L}_{S,\phi}$. This improvement is important because it relaxes the
fixed-point argument from the dependence on the initial guess $\phi$.

We are now ready to prove Theorem A.4.

Proof  [Proof of Theorem A.4] From Corollary D.2, let $S_{1}$ be a constant
satisfying the conditions in Proposition D.3. Then there is a local solution
$\phi^{\ast}\in C([0,S_{1}];\mathcal{C}([0,1];\mathcal{P}^{2}))$ to (9).

We now denote by $S^{\ast}$ the largest time within which the solution exists,
where we denote this solution by $\phi^{\ast}\in
C\left([0,S^{\ast});\mathcal{C}([0,1];\mathcal{P}^{2})\right)$. We aim to show
that $S^{\ast}=\infty$. According to Corollary D.3 (70), for any $s<S^{\ast}$
and $t\in[0,1]$, we have

$\displaystyle\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\phi^{*}(\theta,t,s)\leq
C(S^{*},\mathcal{R}_{\mathrm{ini}},\mathcal{L}^{\sup}_{\mathrm{ini}})\,,$
$\displaystyle\mathrm{supp}(\phi^{*}(t,s))\subset\left\\{\theta\middle||\theta_{[1]}|\leq
C(S^{*},\mathcal{R}_{\mathrm{ini}},\mathcal{L}^{\sup}_{\mathrm{ini}})\right\\}\,,$

Define $\mathcal{R}^{*}=\mathcal{R}_{S^{*},\phi^{*}}$,
$\mathcal{L}^{*}=\mathcal{L}^{\sup}_{S^{*},\phi^{*}}$ according to (55). Since
$\mathcal{R}^{\ast},\mathcal{L}^{\ast}<\infty$, let us choose
$\Delta_{S^{\ast}}$ small enough to satisfy

$\exp(\Delta_{S^{\ast}}Q_{1}(4(\mathcal{L}^{*}+1)))\left(\mathcal{L}^{*}+1\right)\leq
4(\mathcal{L}^{*}+1),\quad\exp(\Delta_{S^{\ast}}Q_{1}(4(\mathcal{L}^{*}+1)))\left(\mathcal{R}^{*}+1\right)\leq
4(\mathcal{R}^{*}+1)\,,$

and

$\Delta_{S^{\ast}}Q_{2}(4(\mathcal{L}^{*}+1),4(\mathcal{R}^{*}+1),\Delta_{S^{\ast}})\leq\frac{1}{2}\,,$

If $S^{\ast}$ is finite, then, using Proposition D.3 and Corollary D.2, we can
further extend $\phi^{\ast}$ to be supported on
$C\left([0,S^{\ast}+\Delta_{S^{\ast}});\mathcal{C}([0,1];\mathcal{P}^{2})\right)$,
giving a contradiction. If follows that $S^{\ast}=\infty$, as desired.

Finally, (21) is a direct result of Lemma D.1.

###  D.2 Proof of Theorem A.3

This section is dedicated to Theorem A.3 — we show the well posedness of the
gradient flow in the finite-layer case. We rewrite the gradient of (2) as
follows

$\frac{\partial
E({\Theta}_{L,M})}{\partial\theta_{l,m}}=\frac{1}{ML}\mathbb{E}_{x\sim\mu}\left(\partial_{\theta}f(Z_{\Theta_{L,M}}(l;x),\theta_{l,m})p_{\Theta_{L,M}}(l;x)\right)\,,$
(71)

where $p_{\Theta_{L,M}}(l;x)$ solves:

$\left\\{\begin{aligned}
p^{\top}_{\Theta_{L,M}}(l;x)&=p^{\top}_{\Theta_{L,M}}(l+1;x)\left(I+\frac{1}{ML}\sum^{M}_{m=1}\partial_{z}f\left(Z_{\Theta_{L,M}}(l+1;x),\theta_{l+1,i}\right)\right)\\\
p_{\Theta_{L,M}}(L-1;x)&=\left(g(Z_{\Theta_{L,M}}(L;x))-y(x)\right)\nabla
g(Z_{\Theta_{L,M}}(L;x)),\end{aligned}\right.\,,$ (72)

for $0\leq l\leq L-2$. We unify the space in a similar fashion to Definition
4.1.

###### Definition D.1

${\Theta}_{L,M}=\left\\{\theta_{l,m}\right\\}^{L-1,M}_{l=0,m=1}\in
L^{\infty}_{L,M}$ if and only if

$\sup_{l,m}|\theta_{l,m}|<\infty\,.$

The metric in $L^{\infty}_{L,M}$ is defined as

$d_{1,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)=\max_{l}\left(\frac{1}{M}\sum^{M}_{m=1}|\theta_{l,m}-\widetilde{\theta}_{l,m}|^{2}\right)^{1/2}\,.$

###### Definition D.2

For $s\geq 0$, we have
$\Theta_{L,M}(s)=\left\\{\theta_{l,m}(s)\right\\}^{L-1,M}_{l=0,m=1}\in\mathcal{C}([0,\infty);L^{\infty}_{L,M})$
if and only if

  * 1.

For fixed $s\in[0,\infty)$, $\Theta_{L,M}(s)\in L^{\infty}_{L,M}$.

  * 2.

For any $s_{0}\in[0,\infty)$,

$\lim_{s\rightarrow
s_{0}}d_{1,L,M}\left(\Theta_{L,M}(s),\Theta_{L,M}(s_{0})\right)=0\,,$

where $d_{1,L,M}$ is defined in Definition D.1.

The metric in $\mathcal{C}([0,\infty);L^{\infty}_{L,M})$ is defined by

$d_{2,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)=\sup_{s}d_{1,L,M}({\Theta}_{L,M}(s),\widetilde{\Theta}_{L,M}(s))\,.$

Theorem A.3 is to say that the solution to (3) is unique in
$\mathcal{C}([0,\infty);L^{\infty}_{L,M})$ if $\Theta_{L,M}(0)\in
L^{\infty}_{L,M}$.

Before proving the theorem, prepare some a-priori estimates of
$Z_{\Theta_{L,M}}$ and $p_{\Theta_{L,M}}$.

###### Lemma D.2

Suppose that Assumption 4.1 holds and that $x$ is in the support of $\mu$. Let

${\Theta}_{L,M}=\left\\{{\theta}_{l,m}\right\\}^{L-1,M}_{l=0,m=1}\,,\quad\text{and}\quad\widetilde{\Theta}_{L,M}=\left\\{\widetilde{\theta}_{l,m}\right\\}^{L-1,M}_{l=0,m=1}\,,$

and denote

$\displaystyle\mathcal{L}_{\Theta_{L,M}}$
$\displaystyle=\frac{1}{LM}\sum^{L-1}_{l=0}\sum^{M}_{m=1}|\theta_{l,m}|^{2},$
$\displaystyle\mathcal{L}_{\widetilde{\Theta}_{L,M}}$
$\displaystyle=\frac{1}{LM}\sum^{L-1}_{l=0}\sum^{M}_{m=1}|\widetilde{\theta}_{l,m}|^{2},$
$\displaystyle\mathcal{R}_{L,M}$
$\displaystyle=\sup_{l,m}\left\\{|\theta_{l,m}|,\left|\widetilde{\theta}_{l,m}\right|\right\\}\,.$

Then for $0\leq l\leq L-1$, we have the following properties:

  * •

Boundedness in $Z_{\Theta_{L,M}}$:

$\left|Z_{\Theta_{L,M}}(l+1;x)\right|\leq C(\mathcal{L}_{\Theta_{L,M}})\,,$
(73)

  * •

Lipschitz in $Z_{\Theta_{L,M}}$:

$\left|Z_{\Theta_{L,M}}(l+1;x)-Z_{\widetilde{\Theta}_{L,M}}(l+1;x)\right|\leq
C\left(\mathcal{L}_{\Theta_{L,M}},\mathcal{L}_{\widetilde{\Theta}_{L,M}}\right)d_{1,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)\,,$
(74)

  * •

Boundedness in $p_{\Theta_{L,M}}$:

$\left|p_{\Theta_{L,M}}(l;x)\right|\leq C(\mathcal{L}_{\Theta_{L,M}})\,,$ (75)

  * •

Lipschitz in $p_{\Theta_{L,M}}$:

$\left|p_{\Theta_{L,M}}(l;x)-p_{\widetilde{\Theta}_{L,M}}(l;x)\right|\leq
C\left(\mathcal{R}_{L,M}\right)d_{1,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)\,.$
(76)

Proof  From (1) and (12) we obtain

$\displaystyle\left(\left|Z_{\Theta_{L,M}}(l+1;x)\right|+1\right)$
$\displaystyle\leq
C_{1}\left(1+\frac{1}{LM}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+1)\right)(\left|Z_{\Theta_{L,M}}(l;x)\right|+1)$
$\displaystyle\leq
C_{1}\exp\left(\frac{1}{LM}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+1)\right)(\left|Z_{\Theta_{L,M}}(l;x)\right|+1)\,,$

which proves (73) by iteration on $l$.

From (13) and (73) we obtain

$\frac{1}{ML}\sum^{M}_{m=1}|\partial_{z}f\left(Z_{\Theta_{L,M}}(l+1;x),\theta_{l,m}\right)|\leq\frac{C(\mathcal{L}_{\Theta_{L,M}})}{ML}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+1)\,,$

which by (72) implies

$|p_{\Theta_{L,M}}(l;x)|\leq\left(1+\frac{C(\mathcal{L}_{\Theta_{L,M}})}{ML}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+1)\right)|p_{\Theta_{L,M}}(l+1;x)|\,.$

From this bound, together with $|p_{\Theta_{L,M}}(x,L-1)|\leq
C|Z_{\Theta_{L,M}}(L;x)|\leq C(\mathcal{L}_{\Theta_{L,M}})$, we prove (75) by
iteration on $l$.

To prove (74), we subtract the two updating formulas and split the estimate to
obtain

$\displaystyle\left|Z_{\Theta_{L,M}}(l+1;x)-Z_{\widetilde{\Theta}_{L,M}}(l+1;x)\right|$
$\displaystyle\leq\left(1+\frac{C(\mathcal{L}_{\Theta_{L,M}})}{ML}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+1)\right)\left|Z_{\Theta_{L,M}}(l;x)-Z_{\widetilde{\Theta}_{L,M}}(l;x)\right|$
$\displaystyle\quad+\frac{C(\mathcal{L}_{\Theta_{L,M}},\mathcal{L}_{\widetilde{\Theta}_{L,M}})}{L}\left(\frac{1}{M}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+|\widetilde{\theta}_{l,m}|^{2}+1)\right)d_{1,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)\,,$

where we use (1) together with the bounds (13), and (73). Noting that
$|Z_{\Theta_{L,M}}(0;x)-Z_{\widetilde{\Theta}_{L,M}}(0;x)|=0$, we prove (74)
by iteration on $l$.

Finally, for (76), we subtract two equations in the form of (72), and use
(72)-(75) together with Lipschitz continuity to obtain

$\displaystyle\left|p_{\Theta_{L,M}}(l;x)-p_{\widetilde{\Theta}_{L,M}}(l;x)\right|$
(77)
$\displaystyle\leq\left|(p_{\Theta_{L,M}}(l+1;x)-p_{\widetilde{\Theta}_{L,M}}(l+1;x))^{\top}\left(I+\frac{1}{ML}\sum^{M}_{m=1}\partial_{z}f\left(Z_{\Theta_{L,M}}(l+1;x),\theta_{l+1,m}\right)\right)\right|$
$\displaystyle+\left|p^{\top}_{\widetilde{\Theta}_{L,M}}(l+1;x)\left(\frac{1}{ML}\sum^{M}_{m=1}\left(\partial_{z}f\left(Z_{\Theta_{L,M}}(l+1;x),\theta_{l+1,i}\right)-\partial_{z}f\left(Z_{\widetilde{\Theta}_{L,M}}(l+1;x),\widetilde{\theta}_{l+1,m}\right)\right)\right)\right|$
$\displaystyle\leq\left(1+\frac{C(\mathcal{R}_{L,M})}{ML}\sum^{M}_{m=1}(|\theta_{l,m}|^{2}+1)\right)\left|p_{\Theta_{L,M}}(l+1;x)-p_{\widetilde{\Theta}_{L,M}}(l+1;x)\right|$
$\displaystyle+\frac{C(\mathcal{R}_{L,M})}{L}d_{1,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)\,.$

The initial data is also controlled, as follows:

$\displaystyle|p_{\Theta_{L,M}}(L-1;x)-p_{\widetilde{\Theta}_{L,M}}(L-1;x)|$
$\displaystyle\leq C|Z_{\Theta_{L,M}}(L;x)-Z_{\widetilde{\Theta}_{L,M}}(L;x)|$
$\displaystyle\leq
C(\mathcal{R}_{L,M})d_{1,L,M}\left({\Theta}_{L,M},\widetilde{\Theta}_{L,M}\right)\,.$

By combining this with (77), we prove (76) by iteration on $l$.

Lemma D.2 resembles Theorem A.2 and Lemma C.1. These estimates allow us to
prove Theorem A.3. Since the proof strategy is exactly the same, we omit
details. Essentially we define a map

$\widetilde{\Theta}(s)=\mathcal{T}^{L,M}_{S}(\Theta^{\prime}(s)):\mathcal{C}([0,\infty);L^{\infty}_{L,M})\rightarrow\,,\mathcal{C}([0,\infty);L^{\infty}_{L,M})\,,$

where $\widetilde{\Theta}(s)$ solves:

$\frac{\,\mathrm{d}\widetilde{\Theta}(s)}{\,\mathrm{d}s}=-ML\nabla_{\Theta}E(\Theta^{\prime}(s))\,,\quad\text{for}\;s\geq
0\,,$

where $\Theta$ defines the forcing term. The estimates above provide all the
ingredients to show the map is well-defined, and for a small enough $S$, the
map is also contracting, leading to the uniqueness of the solution to (3).
Similar to Lemma D.1, one can also show
$\frac{\,\mathrm{d}E}{\,\mathrm{d}s}=-ML|\nabla_{{\Theta}}E({\Theta}_{L,M})|^{2}$,
improving the estimates and removing the constants’ dependence on the initial
guess. This extends the local solution to the global one, as done in Step 3
for the continuous case.

##  Appendix E Proof of Theorem 5.1

Theorem 5.1 links the cost defined by $\Theta_{L,M}(s)$ with that defined by
$\rho(\theta,t,s)$ for all $s$. The continuous and mean-field limits are
obtained, with both $L$ and $M$ sent to infinity. We decompose this result
into two parts, discussing mean-field and continuous limits separately.

We start with the full definition of "limit-admissible" for a distribution
$\rho$.

###### Definition E.1

For an admissible $\rho(\theta,t)$, we say $\rho(\theta,t)$ is limit-
admissible if the average of a large number of particle presentations is
bounded and Lipschitz with high probability. That is, for an admissible
$\rho(\theta,t)$, there are two constants $C_{3}$ and $C_{4}$, both greater
than $\sup_{t\in[0,1]}\int_{\mathbb{R}^{k}}|\theta|^{2}d\rho(\theta,t)$ such
that, for any $M$ stochastic process presentation
$\\{\theta_{m}(t)\\}^{M}_{m=1}$ that are $i.i.d.$ drawn from $\rho(\theta,t)$,
the following properties are satisfied for any $\eta>0$ and
$M>\frac{C_{3}}{\eta}$:

  * 1.

Second moment boundedness in time:

$\mathbb{P}\left(\sup_{t\in[0,1]}\frac{1}{M}\sum^{M}_{m=1}|\theta_{m}(t)|^{2}\leq
C_{4}\right)\geq 1-\eta\,.$ (78)

  * 2.

For all $L>0$, we have

$\mathbb{P}\left(\frac{1}{M}\sum^{L-1}_{l=0}\sum^{M}_{m=1}\int^{\frac{l+1}{L}}_{\frac{l}{L}}\left|\theta_{m}(t)-\theta_{m}\left(\frac{l}{L}\right)\right|^{2}\,\mathrm{d}t\leq\frac{C_{4}}{L^{2}}\right)\geq
1-\eta\,.$ (79)

We now state the two theorems that play complementary parts in Theorem 5.1.
The first theorem addresses the limit in $M$ under the assumption that
$L=\infty$. This is the mean-field part of the analysis.

###### Theorem E.1

Let Assumptions 4.1 and 4.2 hold with some $0<k_{1}\leq k$. Assume that
$\rho_{\mathrm{ini}}(\theta,t)$ is limit-admissible and
$\mathrm{supp}_{\theta}(\rho_{\mathrm{ini}}(\theta,t))\subset\\{\theta||\theta_{[1]}|\leq
R\\}$ with some $R>0$ for all $t\in[0,1]$. Suppose that
$\\{\theta_{m}(0;t)\\}^{M}_{m=1}$ are $i.i.d$ drawn from
$\rho_{\mathrm{ini}}(\theta,t)$. Suppose in addition that

  * •

$\rho(\theta,t,s)$ solves (9) with the initial condition
$\rho_{\mathrm{ini}}(\theta,t)$, and

  * •

$\theta_{m}(s;t)$ solves (6) with the initial condition $\theta_{m}(0;t)$.

Then for any $\epsilon,\eta,S>0$, there exists a constant
$C(\rho_{\mathrm{ini}}(\theta,t),S)>0$ depending on
$\rho_{\mathrm{ini}}(\theta,t)$, $S$ such that when
$M>\frac{C(\rho_{\mathrm{ini}}(\theta,t),S)}{\epsilon^{2}\eta}$, we have

$\mathbb{P}\left(\left|E(\Theta(s;\cdot))-E(\rho(\cdot,\cdot,s))\right|\leq\epsilon\right)\geq
1-\eta\,,\quad\forall s<S\,.$

Proof  See Appendix F.
The conclusion of this result suggests that for a $1-\eta$ confidence of an
$\epsilon$ accuracy, $M$ grows polynomially with respect to $1/\epsilon$ and
$1/\eta$.

The second result considers the convergence of the parameter configuration for
the discrete ResNet (1) to that for the continuous ResNet (4) as $L\to\infty$.
This is the continuous-limit part of the analysis.

###### Theorem E.2

Let Assumptions 4.1 and 4.2 hold with some $0<k_{1}\leq k$. Assume that
$\rho_{\mathrm{ini}}(\theta,t)$ is limit-admissible and
$\mathrm{supp}_{\theta}(\rho_{\mathrm{ini}}(\theta,t))\subset\\{\theta||\theta_{[1]}|\leq
R\\}$ with some $R>0$ for all $t\in[0,1]$. Suppose that
$\\{\theta_{m}(0;t)\\}^{M}_{m=1}$ are $i.i.d$ drawn from
$\rho_{\mathrm{ini}}(\theta,t)$. Suppose in addition that

  * •

$\theta_{m}(s;t)$ solves (6) with initial condition $\theta_{m}(0;t)$,

  * •

$\theta_{l,m}(s)$ solves (3) with initial condition
$\theta_{m}\left(0;\frac{l}{L}\right)$ .

Then for any $\epsilon,\eta,S>0$, there exists a constant
$C(\rho_{\mathrm{ini}}(\theta,t),S)>0$ depending on
$\rho_{\mathrm{ini}}(\theta,t),S$ such that when
$M\geq\frac{C(\rho_{\mathrm{ini}}(\theta,t),S)}{\eta}$ and
$L\geq\frac{C(\rho_{\mathrm{ini}}(\theta,t),S)}{\epsilon}$, we have for all
$s<S$ that

$\mathbb{P}\left(\left|E(\Theta(s;\cdot))-E(\Theta_{L,M}(s))\right|\leq\epsilon\right)\geq
1-\eta\,.$

Proof  See Appendix G.
This theorem shows that when the width is large enough, then with high
probability, in the whole training process with $s<S$, the difference between
the loss functions defined by the discrete ResNet and its continuous
counterpart decreases to $0$ as $L\rightarrow\infty$.

##  Appendix F Convergence to the mean-field PDE

This section is dedicated to mean-field analysis and the proof of Theorem E.1.
The intuition of this theorem is largely aligned with many other mean-field
results, as demonstrated in (Ding et al., 2021).

As argued in Section 5, to show “equivalence" between (6) and (9), we can test
them on the same smooth function $h(\theta)$. Testing (9) on amounts to
multiplying $h$ on both sides of the equation by $h$. From integration by
parts we have

$\frac{\,\mathrm{d}}{\,\mathrm{d}s}\int_{\mathbb{R}^{k}}h\,\mathrm{d}\rho(\theta)=-\int_{\mathbb{R}^{k}}\nabla_{\theta}h\nabla_{\theta}\frac{\delta
E(\rho(s))}{\delta\rho}\,\mathrm{d}{\rho}\,,$

that is,

$\frac{\,\mathrm{d}}{\,\mathrm{d}s}\mathbb{E}(h)=\mathbb{E}\left(\nabla_{\theta}h\nabla_{\theta}\frac{\delta
E(\rho(s))}{\delta\rho}\right)\,.$

To test (6) on $h$, we let $\rho=\frac{1}{M}\sum_{m=1}^{M}\delta_{\theta_{m}}$
and obtain

$\frac{\,\mathrm{d}}{\,\mathrm{d}s}\mathbb{E}(h)=\frac{1}{M}\sum_{m=1}^{M}\nabla_{\theta}h(\theta_{m})\frac{\,\mathrm{d}}{\,\mathrm{d}s}\theta_{m}=-\sum_{m=1}^{M}\nabla_{\theta}h(\theta_{m})\frac{\delta
E}{\delta\theta_{m}}\,.$

We see that (9) and (6) are equivalent when tested by $h$, if and only if the
right hand sides of the two equations above are the same, that is,

$M\frac{\delta E}{\delta\theta_{m}}=\nabla_{\theta}\frac{\delta
E(\rho)}{\delta\rho}(\theta_{m},t)\,.$ (80)

This claim can be established from the definitions of the Fréchet derivatives
for $\frac{\delta E(\rho)}{\delta\rho}$ and $\frac{\delta
E}{\delta\theta_{m}}$; see (Ding et al., 2021, Lemma 33).

To give a quantitative estimate on how quickly (6) converges to of (9), we
utilize the particle method, a classical strategy for the mean-field limit. We
sketch the proof here and will it more rigorous in the following subsections.
We make use of two particle systems. In one system, the particles evolve
themselves, while in the second, the particles are moved forward according to
the underlying field constructed by the limit. In our situation, the former
particle system consists of the $M$ stochastic processes
$\\{\theta_{m}(s;t)\\}$ that descend according to $E(\Theta(s;\cdot))$. The
latter particle system will be termed
$\widetilde{\Theta}(s;t)=\\{\widetilde{\theta}_{m}(s;t)\\}_{m=1}^{M}$; it
descends according to $E(\rho(\cdot,\cdot,s))$, the limiting cost function.
Essentially, we prove

$E(\Theta(s;\cdot))\approx E\left(\widetilde{\Theta}(s;t)\right)\approx
E(\rho(\cdot,\cdot,s))\,.$

The latter approximation arises roughly from the law of large numbers, but the
former needs to be proved rigorously by tracing the two different evolving
ODEs.

To be more specific, let $\rho(\theta,t,s)$ be the solution to (9) with
admissible initial conditions $\rho_{\mathrm{ini}}(\theta,t)$, and let
$\theta_{m}(s;t)$ be the solution to (6) with initial conditions
$\\{\theta_{m}(t,0)\\}^{M}_{m=1}$ that are $i.i.d$ drawn from
$\rho_{\mathrm{ini}}(\theta,t)$. Using the definition of $E$ in (8), we have

$\displaystyle\left|E(\rho(\cdot,\cdot,s))-E(\Theta(s;\cdot))\right|$ (81)
$\displaystyle\leq$
$\displaystyle\mathbb{E}_{x\sim\mu}\left[\left|\frac{1}{2}\left(g(Z_{\rho(s)}(1;x))-y(x)\right)^{2}-\frac{1}{2}\left(g(Z_{{\Theta(s)}}(1;x))-y(x)\right)^{2}\right|\right]$
$\displaystyle\leq$
$\displaystyle\mathbb{E}_{x\sim\mu}\left(\left|g(Z_{\rho(s)}(1;x))-g(Z_{{\Theta(s)}}(1;x))\right|\left(\left|g(Z_{\rho(s)}(1;x))+g(Z_{{\Theta(s)}}(1;x))\right|+|y(x)|\right)\right)$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{s})\left|g(Z_{\rho(s)}(1;x))-g(Z_{{\Theta(s)}}(1;x))\right|$
$\displaystyle\leq$ $\displaystyle
C(\mathcal{L}_{s})\left|Z_{\rho(s)}(1;x)-Z_{{\Theta(s)}}(1;x)\right|\,,$

where
$\mathcal{L}_{s}=\max\left\\{\int^{1}_{0}\int_{\mathbb{R}^{k}}|\theta|^{2}\,\mathrm{d}\rho(\theta,t,s)\,\mathrm{d}t,\frac{1}{M}\int^{1}_{0}\sum^{M}_{m=1}\left|\theta_{m}(s;t)\right|^{2}\,\mathrm{d}t\right\\}$.
In this derivation, we used the Lipschitz property of $g$, the boundedness of
$y$ (required in Assumptions 4.1-4.2), and the boundedness of $Z_{\rho}$ and
$Z_{{\Theta(s)}}$. Boundedness of $Z_{\rho}$ was shown in Theorem A.2, while
the bound for $Z_{\Theta(s)}$ will be addressed in Lemma F.3. The constant $C$
depends on the support of $\rho_{\mathrm{ini}}$, as well as on $s$ and the
Lipschitz constant of $g$. It follows from (81) that to control
$E(\rho(\cdot,\cdot,s))-E(\Theta(s;\cdot))$, we need to control

$\left|Z_{\rho(s)}(1;x)-Z_{{\Theta(s)}}(1;x)\right|\,.$ (82)

To do so, we employ the particle method and invent a new particle system.

According to (80), we can reformulate the original particle system as

$\frac{\,\mathrm{d}\theta_{m}(s;t)}{\,\mathrm{d}s}=-\nabla_{\theta}\frac{\delta
E(\rho^{\mathrm{dis}}(s))}{\delta\rho}(\theta_{m}(s;t))\,,\quad\forall(t,s)\in[0,1]\times[0,\infty)\,,$
(83)

where we denote

$\rho^{\mathrm{dis}}(\theta,t,s)=\frac{1}{M}\sum^{M}_{m=1}\delta_{\theta_{m}(s;t)}(\theta)\,.$
(84)

We invent a new system that follows the underlying flow governed by the limit.
Define $\widetilde{\Theta}(s)=\\{\widetilde{\theta}_{m}(s;t)\\}_{m=1}^{M}$,
where $\widetilde{\theta}_{m}$ solves

$\frac{\,\mathrm{d}\widetilde{\theta}_{m}(s;t)}{\,\mathrm{d}s}=-\nabla_{\theta}\frac{\delta
E(\rho(s))}{\delta\rho}\left(\widetilde{\theta}_{m}(s;t)\right)\,,\quad\forall(t,s)\in[0,1]\times[0,\infty)\,,$
(85)

with initial condition

$\widetilde{\theta}_{m}(0;t)=\theta_{m}(0;t)\,.$ (86)

As a consequence, we have $\widetilde{\theta}_{m}(s;t)\sim\rho(\theta,t,s)$
for all $(t,s)$. The corresponding ensemble distribution is

$\widetilde{\rho}^{\mathrm{dis}}(\theta,t,s)=\frac{1}{M}\sum^{M}_{m=1}\delta_{\widetilde{\theta}_{m}(s;t)}(\theta)\,.$
(87)

Now we have available particle system $\Theta(s)=\\{\theta_{m}(s)\\}$, a newly
invented particle system
$\widetilde{\Theta}(s)=\\{\widetilde{\theta}_{m}(s)\\}$ and the mean-field
flow $\rho$. Accordingly, there are three versions of $Z$: $Z_{\rho(s)}$ that
solves (7) using $\rho(s)$ and $Z_{\Theta(s)}$ and $Z_{\widetilde{\Theta}(s)}$
that solve (4) using $\Theta(s)$ and $\widetilde{\Theta}(s)$, respectively. We
use the following relabelling for convenience:

$Z_{s}=Z_{\rho(s)}\,,\quad
Z^{\mathrm{dis}}_{s}=Z_{\Theta(s)}\,,\quad\widetilde{Z}^{\mathrm{dis}}_{s}=Z_{\widetilde{\Theta}(s)}\,.$
(88)

Similarly, there are three sets of $p$: $p_{\rho(s)}$,
$p_{\rho^{\mathrm{dis}}(s)}$, and $p_{\widetilde{\rho}^{\mathrm{dis}}(s)}$
that solve (11) using $\rho(s)$, $\rho^{\mathrm{dis}}(s)$, and
$\widetilde{\rho}^{\mathrm{dis}}(s)$, respectively. We relabel similarly to
(88) and write

$p_{s}=p_{\rho(s)}\,,\quad
p^{\mathrm{dis}}_{s}=p_{\rho^{\mathrm{dis}}(s)}\,,\quad\widetilde{p}^{\mathrm{dis}}_{s}=p_{\widetilde{\rho}^{\mathrm{dis}}(s)}\,.$
(89)

Since $\widetilde{\Theta}(s)$ serves as a bridge, we translate the control of
(82) to:

$\left|Z_{s}(t;x)-Z^{\mathrm{dis}}_{s}(t;x)\right|\leq\left|Z_{s}(t;x)-\widetilde{Z}^{\mathrm{dis}}_{s}(t;x)\right|+\left|\widetilde{Z}^{\mathrm{dis}}_{s}(t;x)-Z_{s}^{\mathrm{dis}}(t;x)\right|\,.$
(90)

Bounding $\left|Z_{s}(t;x)-\widetilde{Z}^{\mathrm{dis}}_{s}(t;x)\right|$ can
be done using the law of large numbers. Bounding
$\left|\widetilde{Z}^{\mathrm{dis}}_{s}(t;x)-Z_{s}^{\mathrm{dis}}(t;x)\right|$
translates to controlling
$\sum_{m}\int_{0}^{1}|\theta_{m}(s;t)-\widetilde{\theta}_{m}(s;t)|\,\mathrm{d}{t}$,
for which we will evaluate the difference between equations (83) and (85).
Since these two equations have the same initial data, the difference between
$\theta_{m}$ and $\widetilde{\theta}_{m}$ can then be controlled when the
right-hand side forcing terms are close.

This approach divides the proof naturally into two components. In Section F.1,
we give the rigorous bound of (90), while in Section F.2, we trace the
evolution of the difference
$\sum_{m}\int_{0}^{1}|\theta_{m}(s;t)-\widetilde{\theta}_{m}(s;t)|\,\mathrm{d}{t}$
in $s$, thus finalizing the proof for Theorem E.1.

###  F.1 Stability in the mean-field regime

Here we discuss control of the two terms on the right-hand side of (90).
Recall that $\Theta(s)$ and $\widetilde{\Theta}(s)$ satisfy (83) and (85), and
the two corresponding ensemble distribution are defined in (84) and (87),
respectively. For any $S>0$, define

$\displaystyle\mathcal{L}^{\sup}_{S}$ $\displaystyle=\sup_{0\leq t\leq 1,0\leq
s\leq