# Generalization of Scaled Deep ResNets in the Mean-Field Regime

Yihang Chen
EPFL
<EMAIL_ADDRESS>&Fanghui Liu
University of Warwick
<EMAIL_ADDRESS>Yiping Lu
New York University
<EMAIL_ADDRESS>&Grigorios G. Chrysos
University of Wisconsin-Madison
<EMAIL_ADDRESS>&Volkan Cevher
EPFL
<EMAIL_ADDRESS>
###### Abstract

Despite the widespread empirical success of ResNet, the generalization
properties of deep ResNet are rarely explored beyond the lazy training regime.
In this work, we investigate _scaled_ ResNet in the limit of infinitely deep
and wide neural networks, of which the gradient flow is described by a partial
differential equation in the large-neural network limit, i.e., the _mean-
field_ regime. To derive the generalization bounds under this setting, our
analysis necessitates a shift from the conventional time-invariant Gram matrix
employed in the lazy training regime to a time-variant, distribution-dependent
version. To this end, we provide a global lower bound on the minimum
eigenvalue of the Gram matrix under the mean-field regime. Besides, for the
traceability of the dynamic of Kullback-Leibler (KL) divergence, we establish
the linear convergence of the empirical error and estimate the upper bound of
the KL divergence over parameters distribution. Finally, we build the uniform
convergence for generalization bound via Rademacher complexity. Our results
offer new insights into the generalization ability of deep ResNet beyond the
lazy training regime and contribute to advancing the understanding of the
fundamental properties of deep neural networks.

##  1 Introduction

Deep neural networks (DNNs) have achieved great success empirically, a notable
illustration of which is ResNet (He et al., 2016), a groundbreaking network
architecture with skip connections. One typical way to theoretically
understand ResNet (e.g., optimization, generalization), is based on the neural
tangent kernel (NTK) tool (Jacot et al., 2018). Concretely, under proper
assumptions, the training dynamics of ResNet can be described by a fixed
kernel function (NTK). Hence, the global convergence and generalization
guarantees can be given via NTK and the benefits of residual connection can be
further demonstrated by spectral properties of NTK (Hayou et al., 2019; Huang
et al., 2020; Hayou et al., 2021; Tirer et al., 2022). However, the NTK
analysis requires the parameters of ResNet to not move much during training
(which is called _lazy training_ or kernel regime (Chizat et al., 2019;
Woodworth et al., 2020; Barzilai et al., 2022)). Accordingly, the NTK analysis
fails to describe the true non-linearity of ResNet. Beyond the NTK analysis
thus received great attention in the deep learning theory community.

One typical approach is, _mean-field_ based analysis, which allows for
unrestricted movement of the parameters of DNNs during training, (Wei et al.,
2019; Woodworth et al., 2020; Ghorbani et al., 2020; Yang & Hu, 2021; Akiyama
& Suzuki, 2022; Ba et al., 2022; Mahankali et al., 2023; Chen et al., 2022;
Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Nguyen, 2019; Sirignano &
Spiliopoulos, 2020b). The training dynamics can be formulated as an
optimization problem over the weight space of probability measures by studying
suitable scaling limits. For deep ResNet, taking the limit of infinite depth
is naturally connected to a continuous neural ordinary differential equation
(ODE) (Sonoda & Murata, 2019; Weinan, 2017; Lu et al., 2018; Haber & Ruthotto,
2017; Chen et al., 2018), which makes the optimization guarantees of deep
ResNets feasible under the mean-field regime (Lu et al., 2020; Ma et al.,
2020; Ding et al., 2021; 2022; Barboni et al., 2022). The obtained
optimization results indicate that infinitely deep and wide ResNets can easily
fit the training data with random labels, i.e., _global convergence_. While
previous works have obtained promising optimization results for deep ResNets,
there is a notable absence of generalization analysis, which is essential for
theoretically understanding why deep ResNet can generalize well _beyond the
lazy training regime_. Accordingly, this naturally raises the following
question:

_Can we build a generalization analysis of trained Deep ResNets in the mean-
field setting?_

We answer this question by providing a generalization analysis framework of
_scaled_ deep ResNet in the mean-field regime, where the word _scaled_ denotes
a scaling factor on deep ResNet. To this end, we consider an infinite width
and infinite depth ResNet parameterized by two measures: $\nu$ over the
feature encoder and $\tau$ over the output layer, respectively. By proving the
condition number of the optimization dynamics of deep ResNets parameterized by
$\tau$ and $\nu$ is lower bounded, we obtain the global linear convergence
guarantees, aiming to derive the Kullback–Leibler (KL) divergence of such two
measures between initialization and after training. Based on the KL divergence
results, we, therefore, build the uniform convergence result for
generalization via Rademacher complexity under the mean-field regime,
obtaining the convergence rate at $\mathcal{O}(1/\sqrt{n})$ given $n$ training
data.

Our contributions are summarized as below:

  * •

The paper provides the minimum eigenvalue estimation (lower bound) of the Gram
matrix of the gradients for deep ResNet parameterized by the ResNet encoder’s
parameters and MLP predictor’s parameters.

  * •

The paper proves that the KL divergence of feature encoder $\tau$ and output
layer $\nu$ can be bounded by a constant (depending only on network
architecture parameters) during the training.

  * •

This paper builds the connection between the Rademacher complexity result and
KL divergence, and then derives the convergence rate $\mathcal{O}(1/\sqrt{n})$
for generalization.

Our theoretical analysis provides an in-depth understanding of the global
convergence under minimal assumptions, sheds light on the KL divergence of
network measures before and after training, and builds the generalization
guarantees under the mean-field regime, matching classical results in the lazy
training regime (Allen-Zhu et al., 2019; Du et al., 2019b). We expect that our
analysis opens the door to generalization analysis for feature learning and
look forward to which adaptive features can be learned from the data under the
mean-field regime.

##  2 Related Work

In this section, we briefly introduce the large width/depth ResNets in an ODE
formulation, NTK analysis, and mean-field analysis for ResNets.

###  2.1 Infinite-width, infinite-depth ResNet, ODE

The limiting model of deep and wide ResNets can be categorized into three
class, by either taking the width or depth to infinity: a) the infinite depth
limit to the ODE/SDE model (Sonoda & Murata, 2019; Weinan, 2017; Lu et al.,
2018; Chen et al., 2018; Haber & Ruthotto, 2017; Marion et al., 2023); b)
infinite width limit, (Hayou et al., 2021; Hayou & Yang, 2023; Frei et al.,
2019); c) infinite depth width ResNets, mean-field ODE framework under the
infinite depth width limit (Li et al., 2021; Lu et al., 2020; Ding et al.,
2021; 2022; Barboni et al., 2022)

In this work, we are particularly interested in mean-field ODE formulation.
The deep ResNets modeling by mean-field ODE formulation stems from (Lu et al.,
2020), in which every residual block is regarded as a particle and the target
is changed to optimize over the empirical distribution of particles. Sander et
al. (2022) discusses the rationale behind such equivalence between discrete
dynamics and continuous ODE for ResNet under certain cases. The global
convergence is built under a modified cost function (Ding et al., 2021), and
further improved by removing the regularization term on the cost function
(Ding et al., 2022). However, the analysis in (Ding et al., 2022) requires
more technical assumptions about the limiting distribution. Barboni et al.
(2022) show a local linear convergence by parameterizing the network with
RKHS. However, the radius of the ball containing parameters relies on the
$N$-universality and is difficult to estimate. Our work requires minimal
assumptions under a proper scaling of the network parameters and the design of
architecture, and hence foresters optimization and generalization analyses.
There is a concurrent works (Marion et al., 2023) that studies the implicit
regularization of ResNets converging to ODEs, but the employed technique is
different from ours and the generalization analysis in their work is missing.

###  2.2 NTK analysis for deep ResNet

Jacot et al. (2018) demonstrate that the training process of wide neural
networks under gradient flow can be effectively described by the Neural
Tangent Kernel (NTK) as the network’s width (denoted as ’$M$’) tends towards
infinity under the NTK scaling (Du et al., 2019b). During the training, the
NTK remains unchanged and thus the theoretical analyses of neural networks can
be transformed to those of kernel methods. In this case, the optimization and
generalization properties of neural networks can be well controlled by the
minimum eigenvalue of NTK (Cao & Gu, 2019; Nguyen et al., 2021). Regarding
ResNets, the architecture we are interested in, the NTK analysis is also valid
(Tirer et al., 2022; Huang et al., 2020; Belfer et al., 2021), as well as an
algorithm-dependent bound (Frei et al., 2019) in the lazy training regime.
Compared with kernelized analysis of the wide neural network, we do not rely
on the convergence to a fixed kernel as the width approaches infinity to
perform the convergence and generalization analysis. Instead, under the mean
field regime, the so-called kernel falls into a time-varying, measure-
dependent version. We also remark that, for a ResNet with an infinite depth
but constant width, the global convergence is given by Cont et al. (2022)
beyond the lazy training regime by studying the evolution of gradient norm. To
our knowledge, this is the first work that analyzes the (varying) kernel
eigenvalue of infinite-width/depth ResNet beyond the NTK regime in terms of
optimization and generalization.

###  2.3 Mean-field Analysis

Under suitable scaling limits (Mei et al., 2018; Rotskoff & Vanden-Eijnden,
2018; Sirignano & Spiliopoulos, 2020a), as the number of neurons goes to
infinity, _i.e._ $M\to\infty$, neural networks work in the mean-field limit.
In this setting, the training dynamics of neural networks can be formulated as
an optimization problem over the distribution of neurons. A notable benefit of
the mean-field approach is that, after deriving a formula for the gradient
flow, conventional PDE methods can be utilized to characterize convergence
behavior, which enables both nonlinear feature learning and global convergence
(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; 2020b; E et al., 2020;
Jabir et al., 2021).

In the case of the two-layer neural network, Chizat & Bach (2018); Mei et al.
(2018); Wojtowytsch (2020); Chen et al. (2020); Barboni et al. (2022) justify
the mean-field approach and demonstrate the convergence of the gradient flow
process, achieving the zero loss. For the wide shallow neural network, Chen et
al. (2022) proves the linear convergence of the training loss by virtue of the
Gram matrix. In the multi-layer case, Lu et al. (2020); Ding et al. (2021;
2022) translate the training process of ResNet to a gradient-flow partial
differential equation (PDE) and showed that with depth and width depending
algebraically on the accuracy and confidence levels, first-order optimization
methods can be guaranteed to find global minimizers that fit the training
data.

In terms of the generalization of a trained neural network under the mean-
field regime, current results are limited to two-layer neural networks. For
example, Chen et al. (2020) provides a generalized NTK framework for two-layer
neural networks, which exhibits a "kernel-like" behavior. Chizat & Bach (2020)
demonstrate that the limits of the gradient flow of two-layer neural networks
can be characterized as a max-margin classifier in a certain non-Hilbertian
space. Our work, instead, focuses on deep ResNets in the mean-field regime and
derives the generalization analysis framework.

##  3 From Discrete to Continuous ResNet

In this section, we present the problem setting of our deep ResNets for binary
classification under the infinite depth and width limit, which allows for
parameter evolution of ResNets. Besides, several mild assumptions are
introduced for our proof.

###  3.1 Problem setting

For an integer $L$, we use the shorthand $[L]=\left\\{1,2,\dots,L\right\\}$.
Let $\mathcal{X}\subseteq\mathbb{R}^{d}$ be a compact metric space and
$\mathcal{Y}\subseteq\mathbb{R}$. We assume that the training set
$\mathcal{D}_{n}=\\{(\bm{x}_{i},y_{i})\\}_{i=1}^{n}$ is drawn from an unknown
distribution $\mu$ on $\mathcal{X}\times\mathcal{Y}$, and $\mu_{X}$ is the
marginal distribution of $\mu$ over $\mathcal{X}$. The goal of our supervised
learning task is to find a hypothesis (i.e., a ResNet used in this work)
$f:\mathcal{X}\rightarrow\mathcal{Y}$ such that $f(\bm{x};\bm{\Theta})$
parameterized by $\bm{\Theta}$ is a good approximation of the label
$y\in\mathcal{Y}$ corresponding to a new sample $\bm{x}\in\mathcal{X}$. In
this paper, we consider a binary classification task, denoted by minimizing
the expected risk, let $\ell_{0-1}(f,y)=\mathbbm{1}\\{yf<0\\}$,

$\min_{\bm{\Theta}}~{}{L}_{0-1}(\bm{\Theta}):=\mathbb{E}_{(\bm{x},y)\sim\mu}~{}\ell_{0-1}(f(\bm{x};\bm{\Theta}),y)\;.$

Note that the $0-1$ loss is non-convex and non-smooth, and thus difficult for
optimization. One standard way in practice for training is using a _surrogate_
loss for empirical risk minimization (ERM), normally convex and smooth (or at
least continuous), e.g., the hinge loss, the cross-entropy loss.
Interestingly, the squared loss, originally used for regression, can be also
applied for classification with good statistical properties in terms of
robustness and calibration error, as systematically discussed in Hui & Belkin
(2020); Hu et al. (2022). Therefore, we employ the squared loss in ERM for
training, let $\ell(f,y)=\frac{1}{2}(y-f)^{2}$,

$\min_{\bm{\Theta}}~{}\widehat{L}(\bm{\Theta}):=\frac{1}{n}\sum_{i=1}^{n}\ell(f(\bm{x}_{i};\bm{\Theta}),y_{i})=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}\ell(f(\bm{x};\bm{\Theta}),y(\bm{x}))\,,$
(1)

where $\mathcal{D}_{n}$ is the empirical measure of $\mu_{X}$ over
$\\{\bm{x}_{i}\\}_{i=1}^{n}$ and note that $y$ is a function of $\bm{x}$.

We call the probability measure $\rho\in\mathcal{P}^{2}$ if $\rho$ has the
finite second moment, and $\rho\in\mathcal{C}(\mathcal{P}^{2};[0,1])$ if
$\rho^{s}\in\mathcal{P}^{2},\forall s\in[0,1]$. For
$\rho_{1},\rho_{2}\in\mathcal{P}^{2}$, the 2-Wasserstein distance is denoted
by $\mathcal{W}_{2}(\rho_{1},\rho_{2})$; and for
$\rho_{1},\rho_{2}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$, we define
$\mathcal{W}_{2}(\rho_{1},\rho_{2}):=\sup_{s\in[0,1]}\mathcal{W}_{2}(\rho_{1}^{s},\rho_{2}^{s})$.

###  3.2 ResNets in the infinite depth and width limit

The continuous formulation of ResNets is a recent approach that uses
differential equations to model the behavior of the ResNet. This formulation
has the advantage of enabling continuous analysis of the ODE, which can make
the analysis of ResNets easier (Lu et al., 2020; Ding et al., 2021; 2022;
Barboni et al., 2022). We firstly consider the following ResNet (He et al.,
2016) of depth $L$ can be formulated as
$\bm{z}_{0}(\bm{x})=\bm{x}\in\mathbb{R}^{d}$, and

$\begin{split}\bm{z}_{l+1}(\bm{x})&=\bm{z}_{l}(\bm{x})+\frac{\alpha}{ML}\sum_{m=1}^{M}\bm{\sigma}(\bm{z}_{l}(\bm{x}),\bm{\bm{\theta}}_{l,m})\in\mathbb{R}^{d},\quad
l\in[L-1]\,,\\\
f_{\bm{\Omega}_{K},\bm{\Theta}_{L,M}}(\bm{x})&=\frac{\beta}{K}\sum_{k=1}^{K}h(\bm{z}_{L},\bm{\omega}_{k})\in\mathbb{R}\,,\end{split}$
(2)

where $\bm{x}\in\mathbb{R}^{d}$ is the input data,
$\alpha,\beta\in\mathbb{R}_{+}$ are the scaling factors.
$\bm{\Theta}_{L,M}=\\{\bm{\theta}_{l,m}\in\mathbb{R}^{k_{\nu}}\\}_{l=0,m=0}^{L-1,M}$
is the parameters of the ResNet encoder
$\bm{\sigma}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ ( activation functions
are implicitly included into $\bm{\sigma}$), and
$\bm{\Omega}_{K}=\\{\bm{\omega}_{k}\in\mathbb{R}^{k_{\tau}}\\}_{k=1}^{K}$ is
the parameters of the predictor $h:\mathbb{R}^{d}\rightarrow\mathbb{R}$. We
introduce a trainable MLP parametrized by $\bm{\omega}$ in the end, which is
different from the fixed linear predictor in Lu et al. (2020); Ding et al.
(2021; 2022). We make the assumptions on the choices of activation function
$\bm{\sigma}$ and predictor $h$ later in 3.3. Different scaling of
$\alpha,\beta$ leads to different training schemes. Note that setting
$\alpha=\sqrt{M},\beta=\sqrt{K}$ corresponds to the standard scaling in the
NTK regime (Du et al., 2019b), while setting $\alpha=\beta=1$ corresponds to
the classical mean field analysis (Mei et al., 2018; Rotskoff & Vanden-
Eijnden, 2018; Lu et al., 2020; Ding et al., 2022). We will keep
$\alpha,\beta$ as a hyperparameter in our theoretical analysis and determine
the choice of $\alpha,\beta$ in future discussions. Besides, the scaling $1/L$
is necessary to derive the neural ODE limit (Marion et al., 2022), which has
been supported by the empirical observations from Bachlechner et al. (2021);
Marion et al. (2023).

We then introduce the infinitely deep and wide ResNet which is known as the
mean-field limit of deep ResNet (Lu et al., 2020; Ma et al., 2020; Ding et
al., 2022).

##### Infinite Depth

To be specific, we re-parametrize the indices $l\in[L]$ in Eq. 2 with
$s=\frac{l}{L}\in[0,1]$. We view $\bm{z}$ in Eq. 2 as a function in $s$ that
satisfies a coupled ODE, with $1/L$ being the stepsize. Accordingly, we write
$\bm{\theta}_{m}(s):=\bm{\theta}_{m}(l/L)=\bm{\theta}_{l,m}$, and
$\bm{\Theta}_{M}(s)=\\{\bm{\theta}_{m}(s)\\}_{m=1}^{M}$. The continuous limit
of Eq. 2 by taking $L\rightarrow\infty$ is

$\displaystyle\frac{\mathrm{d}\bm{z}(\bm{x},s)}{\mathrm{d}s}=\frac{\alpha}{M}\sum_{m=1}^{M}\bm{\sigma}(\bm{z}(\bm{x},s),\bm{\theta}_{m}(s))=\alpha\int_{\mathbb{R}^{k_{\nu}}}\bm{\sigma}(\bm{z}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{M}(\bm{\theta},s),\quad\bm{z}(\bm{x},0)=\bm{x}\,,$
(3)

where the discrete probability $\nu_{M}(\bm{\theta},s)$ is defined as
$\nu_{M}(\bm{\theta},s):=\frac{1}{M}\sum_{i=1}^{M}\delta_{\bm{\theta}_{m}(s)}(\bm{\theta})$.
Accordingly, the empirical risk in Eq. (1) can be written as

$\displaystyle\widehat{L}(\bm{\Omega}_{K},\bm{\Theta}_{M}):=\mathbb{E}_{\bm{x}\sim{{\mathcal{D}}_{n}}}\
\ell(f_{\bm{\Omega}_{K},\bm{\Theta}_{M}}(\bm{x}),y(\bm{x}))\,.$ (4)

##### Infinite Width

The mean-field limit is obtained by considering a ResNet of infinite width,
i.e. $M\to\infty$. Denoting the limiting density of $\nu_{M}(\bm{\theta},s)$
by $\nu(\bm{\theta},s)\in\mathcal{C}(\mathcal{P}^{2};[0,1])$, Eq. 3 can be
written as

$\frac{\mathrm{d}\bm{z}(\bm{x},s)}{\mathrm{d}s}=\alpha\cdot\int_{\mathbb{R}^{k_{\nu}}}\bm{\sigma}(\bm{z}(\bm{x},s),\bm{\theta})\mathrm{d}\nu(\bm{\theta},s),\quad
s\in[0,1],\quad\bm{z}(\bm{x},0)=\bm{x}\,.$ (5)

We denote the solution of Eq. 5 as $\bm{Z}_{\nu}(\bm{x},s)$. Besides, we also
take the infinite width limit in the final layer, i.e. $K\to\infty$, and then
the limiting density of $\bm{\omega}$ is $\tau(\bm{\omega})$. The whole
network can be written as

$\displaystyle
f_{{\tau,\nu}}(\bm{x}):=\beta\cdot\int_{\mathbb{R}^{k_{\tau}}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})\mathrm{d}\tau(\bm{\omega})\,,$
(6)

and the empirical loss in Eq. (1) can be defined as:

$\displaystyle\widehat{L}(\tau,\nu):=\mathbb{E}_{\bm{x}\sim{{\mathcal{D}}}_{n}}\
\ell(f_{\tau,\nu}(\bm{x}),y(\bm{x}))\,.$ (7)

####  3.2.1 Parameter Evolution

In the discrete ResNet (2), consider minimizing the empirical loss
$\widehat{L}(\bm{\Omega}_{K},\bm{\Theta}_{L,M})$ with an infinitesimally small
learning rate, the updating process can be characterized by the particle
gradient flow, see Definition 2.2 in Chizat & Bach (2018):

$\displaystyle\frac{\mathrm{d}\bm{\Omega}_{K}(t)}{\mathrm{d}t}$
$\displaystyle=-K\nabla_{\bm{\Omega}_{k}}\widehat{L}(\bm{\Omega}_{K}(t),\bm{\Theta}_{L,M}(t)),$
(8) $\displaystyle\frac{\mathrm{d}\bm{\Theta}_{L,M}(t)}{\mathrm{d}t}$
$\displaystyle=-LM\nabla_{\bm{\Theta}_{L,M}}\widehat{L}(\bm{\Omega}_{K}(t),\bm{\Theta}_{L,M}(t))\,,$
(9)

where $t$ is the rescaled pseudo-time, which amounts to assigning a
$\frac{1}{K}$ or $\frac{1}{LM}$ mass to each particle, and is convenient to
take the many-particle limit.

In the continuous ResNet (5), we use the gradient flow in the Wasserstein
metric to characterize the evolution of ${\tau,\nu}$ (Chizat & Bach, 2018).
The evolution of the final layer distribution $\tau(\bm{\omega})$ can be
characterized as

$\displaystyle\frac{\partial\tau}{\partial
t}(\bm{\omega},t)=\nabla_{\bm{\omega}}\cdot\left(\tau(\bm{\omega},t)\nabla_{\bm{\omega}}\frac{\delta\widehat{L}({\tau,\nu})}{\delta\tau}(\bm{\omega},t)\right),\quad
t\geq 0\,,$ (10)

where

$\displaystyle\frac{\delta\widehat{L}({\tau,\nu})}{\delta\tau}(\bm{\omega})$
$\displaystyle=\mathbb{E}_{\bm{x}\sim{{\mathcal{D}_{n}}}}[\beta\cdot(f_{{\tau,\nu}}(\bm{x})-y(\bm{x}))\cdot
h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})]\,.$ (11)

In addition, the evolution of the ResNet layer $\nu(\bm{\theta},s)$ can be
characterized as

$\displaystyle\frac{\partial\nu}{\partial
t}(\bm{\theta},s,t)=\nabla_{\bm{\theta}}\cdot\left(\nu(\bm{\theta},s,t)\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau,\nu})}{\delta\nu}(\bm{\theta},s,t)\right),\quad
t\geq 0\,.$ (12)

From the results in Lu et al. (2020); Ding et al. (2022; 2021), we can compute
the functional derivative as follows:

$\displaystyle\frac{\delta\widehat{L}({\tau,\nu})}{\delta\nu}(\bm{\theta},s)$
$\displaystyle=\mathbb{E}_{\bm{x}\sim{{\mathcal{D}_{n}}}}[\beta\cdot(f_{{\tau,\nu}}(\bm{x})-y(\bm{x}))\cdot\bm{\omega}^{\top}\frac{\partial\bm{Z}_{\nu}(\bm{x},1)}{\partial\bm{Z}_{\nu}(\bm{x},s)}\frac{\delta\bm{Z}_{\nu}(\bm{x},s)}{\delta\nu}(\bm{\theta},s)]$
(13)
$\displaystyle=\mathbb{E}_{\bm{x}\sim{{\mathcal{D}_{n}}}}[\beta\cdot(f_{{\tau,\nu}}(\bm{x})-y(\bm{x}))\cdot{\bm{p}}^{\top}_{\nu}(\bm{x},s)\cdot\alpha\cdot\bm{\sigma}(\bm{Z}_{\nu}(\bm{x},s),\bm{\theta})]\,,$
(14)

where ${\bm{p}}_{\nu}\in\mathbb{R}^{k_{\nu}}$, parameterized by
$\bm{x},s,\nu$, is the solution to the following adjoint ODE, with initial
condition dependent on $\tau$:

$\displaystyle\frac{\mathrm{d}{\bm{p}}^{\top}_{\nu}}{\mathrm{d}s}(\bm{x},s)$
$\displaystyle=-\alpha\cdot{\bm{p}}^{\top}_{\nu}(\bm{x},s)\int_{\mathbb{R}^{k_{\nu}}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu}(\bm{x},s),\bm{\theta})\mathrm{d}\nu(\bm{\theta},s)\,,$
(15) $\displaystyle{\bm{p}}_{\nu}^{\top}(\bm{x},1)$
$\displaystyle=\int_{\mathbb{R}^{k_{\tau}}}\nabla_{\bm{z}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})\mathrm{d}\tau(\bm{\omega})\,.$
(16)

For the linear ODE (15), we can directly obtain the explicit formula,
${\bm{p}}_{\nu}^{\top}(\bm{x},s)={\bm{p}}^{\top}_{\nu}(\bm{x},1)\bm{q}_{\nu}(\bm{x},s)$,
where $\bm{q}_{\nu}(\bm{x},s)$ is the exponentially scaling matrix defined in
Eq. 17. The correctness of solution (17) can be verified by taking the
gradient w.r.t. $s$ at both sides.

$\displaystyle\bm{q}_{\nu}(\bm{x},s)=\exp\left(\alpha\int_{s}^{1}\int_{\mathbb{R}^{k_{\nu}}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu}(\bm{x},s^{\prime}),\bm{\theta})\mathrm{d}\nu(\bm{\theta},s^{\prime})\right).$
(17)

###  3.3 Assumptions

In the following, we use the upper subscript for ResNet ODE layer $s\in[0,1]$,
and the lower subscript for training time $t\in[0,+\infty)$. For example,
$\tau_{t}(\bm{\omega}):=\tau(\bm{\omega},t)$, and
$\nu_{t}^{s}(\bm{\theta}):=\nu(\bm{\theta},s,t)$. First, we assume the
boundedness and second moment of the dataset by 3.1.

######  Assumption 3.1 (Assumptions on data).

We assume that for $\bm{x}_{i}\neq\bm{x}_{j}\sim\mu_{X}$, the following holds
with probability 1,

$\displaystyle\|\bm{x}_{i}\|_{2}=1,|y(\bm{x}_{i})|\leq
1,\langle\bm{x}_{i},\bm{x}_{j}\rangle\leq C_{\max}<1\,,\forall i,j\in[n]\,.$

Remark: The assumption, i.e., $\bm{x}_{i},\bm{x}_{j}$ being not parallel, is
attainable and standard in the analysis of neural networks (Du et al., 2019b;
Zhu et al., 2022).

Second, we adopt the standard Gaussian initialization for distribution $\tau$
and $\nu$.

######  Assumption 3.2 (Assumption on initialization).

The initial distribution ${\tau_{0},\nu_{0}}$ is standard Gaussian:
$(\tau_{0},\nu_{0})(\bm{\omega},\bm{\theta},s)\propto\exp\left(-\frac{\|\bm{\omega}\|_{2}^{2}+\|\bm{\theta}\|_{2}^{2}}{2}\right),\forall
s\in[0,1]$.

Next, we adopt the following assumption on activation $\bm{\sigma},h$ in terms
of formulation and smoothness. The widely used activation functions, such as
Sigmoid, Tanh, satisfy this assumption.

######  Assumption 3.3 (Assumptions on activation $\bm{\sigma},h$).

Let $\bm{\theta}:=(\bm{u},\bm{w},b)\in\mathbb{R}^{k_{\nu}}$, where
$\bm{u},\bm{w}\in\mathbb{R}^{k_{\nu}},b\in\mathbb{R}$, i.e. $k_{\nu}=2d+1$;
$\bm{\omega}:=(a,\bm{w},b)\in\mathbb{R}^{k_{\tau}}$, where
$\bm{w}\in\mathbb{R}^{k_{\nu}},a,b\in\mathbb{R}$, i.e. $k_{\tau}=d+2$.

For any $\bm{z}\in\mathbb{R}^{k_{\nu}}$, we assume

$\displaystyle\bm{\sigma}(\bm{z},\bm{\theta})=\bm{u}\sigma_{0}(\bm{w}^{\top}\bm{z}+b),\quad
h(\bm{z},\bm{\omega})=a\sigma_{0}(\bm{w}^{\top}\bm{z}+b),\quad\sigma_{0}:\mathbb{R}\to\mathbb{R}.$
(18)

In addition, we have the following assumption on $\sigma_{0}$.
$|\sigma_{0}(x)|\leq C_{1}\max(|x|,1),|\sigma_{0}^{\prime}(x)|\leq
C_{1},|\sigma_{0}^{\prime\prime}(x)|\leq C_{1}$, and let $\mu_{i}(\sigma_{0})$
be the $i$-th Hermite coefficient of $\sigma_{0}$.

Based on our description of the evolution of deep ResNets and standard
assumptions, we are ready to present our main results on optimization and
generalization in the following section.

##  4 Main results

In this section, we derive a quantitative estimation of the convergence rate
of optimizing the ResNet. Our main results are three-fold: a) minimum
eigenvalue estimation of the Gram matrix during the training dynamics which
controls the training speed; b) a quantitative estimation of KL divergence
between the weight destruction of trained ResNet with initialization; c)
Rademacher complexity generalization guarantees for the trained ResNet.

###  4.1 Gram Matrix and Minimum Eigenvalue

The training dynamics is governed by the Gram matrix of the coordinate tangent
vectors to the functional derivatives. In this section, we bound the minimum
eigenvalue of the Gram matrix of the gradients through the whole training
dynamics, which controls the convergence of gradient flow.

In the lazy training regime (Jacot et al., 2018), the Gram matrix converges
pointwisely to the NTK as the width approaches infinity. Hence one only needs
to bound the Gram matrix’s minimum eigenvalue at initialization, and the
global convergence rate under gradient descent is bounded by the minimum
eigenvalue of NTK. In our setting under the mean-field regime, we also need
similar Gram matrix/matrices to aid our proof on the training dynamics, but we
do not rely on their convergence to the NTK when the width approaches
infinity. Instead, we consider the Gram matrix of the limiting mean-field
model Eq. 6.

For the ResNet parameter distribution $\nu$, we define one Gram matrix
$\bm{G}_{1}({\tau,\nu})\in\mathbb{R}^{n\times n}$ by

$\displaystyle\bm{G}_{1}({\tau,\nu})=\int_{0}^{1}\bm{G}_{1}({\tau,\nu},s)\mathrm{d}s,\quad\bm{G}_{1}({\tau,\nu},s)=\mathbb{E}_{\bm{\theta}\sim\nu(\cdot,s)}\bm{J}_{1}(\tau,\nu,\bm{\theta},s)\bm{J}_{1}(\tau,\nu,\bm{\theta},s)^{\top}\,,$
(19)

where the row vector of $\bm{J}_{1}$ is defined as

$\displaystyle\left(\bm{J}_{1}(\tau,\nu,\bm{\theta},s)\right)_{i,\cdot}={\bm{p}}^{\top}_{\nu}(\bm{x}_{i},s)\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu}(\bm{x}_{i},s),\bm{\theta}),\quad
1\leq i\leq n\,,$

where the dependence on $\tau$ on the right side of the equality is from the
initial condition ${\bm{p}}_{\nu}^{\top}(\bm{x},1)$. We also define the Gram
matrix for the MLP parameter distribution $\tau$,
$\bm{G}_{2}({\tau,\nu})\in\mathbb{R}^{n\times n}$ by

$\displaystyle\bm{G}_{2}({\tau,\nu})=\mathbb{E}_{\bm{\omega}\sim\tau(\cdot)}\bm{J}_{2}(\nu,\bm{\omega})\bm{J}_{2}(\nu,\bm{\omega})^{\top}\,,$
(20)

where the row vector of $\bm{J}_{2}$ is defined as

$\displaystyle\left(\bm{J}_{2}(\nu,\bm{\omega})\right)_{i,\cdot}=\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega}),\quad
1\leq i\leq n\,.$

We characterize the training dynamics of the neural networks by the following
theorem (the proof deferred to Section C.1), which demonstrates the
relationship between the gradient flow of the loss and those of functional
derivatives.

######  Theorem 4.1.

The training dynamics of $\widehat{L}({\tau_{t},\nu_{t}})$ can be written as:

$\displaystyle\frac{\mathrm{d}\widehat{L}({\tau_{t},\nu_{t}})}{\mathrm{d}t}=-\int_{0}^{1}\int_{\mathbb{R}^{k_{\nu}}}\left\|\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t},\nu_{t}})}{\delta\nu_{t}}(\bm{\theta},s)\right\|_{2}^{2}\mathrm{d}\nu_{t}(\bm{\theta},s)-\int_{\mathbb{R}^{k_{\nu}}}\left\|\nabla_{\bm{\omega}}\frac{\delta\widehat{L}({\tau_{t},\nu_{t}})}{\delta\tau_{t}}(\bm{\omega})\right\|_{2}^{2}\mathrm{d}\tau_{t}(\bm{\omega})\,.$

From the definition of functional derivatives
$\frac{\delta\widehat{L}({\tau_{t},\nu_{t}})}{\delta\nu_{t}}(\bm{\theta},s)$
and
$\frac{\delta\widehat{L}({\tau_{t},\nu_{t}})}{\delta\tau_{t}}(\bm{\omega})$,
we immediately obtain Proposition 4.2, an extension of Theorem 4.1, which
demonstrates that the training dynamics can be controlled by the corresponding
Gram matrices.

######  Proposition 4.2.

Let
$\bm{b}_{t}=(f_{{\tau_{t},\nu_{t}}}(\bm{x}_{1})-y(\bm{x}_{1}),\cdots,f_{{\tau_{t},\nu_{t}}}(\bm{x}_{n})-y(\bm{x}_{n}))$,
using the Gram matrix defined in Eq. 19 and 20, the training dynamics of
$\widehat{L}({\tau_{t},\nu_{t}})$ can be written as:

$\displaystyle\frac{\mathrm{d}\widehat{L}({\tau_{t},\nu_{t}})}{\mathrm{d}t}=-\frac{\beta^{2}}{n^{2}}\bm{b}_{t}^{\top}(\alpha^{2}\bm{G}_{1}({\tau_{t},\nu_{t}})+\bm{G}_{2}({\tau_{t},\nu_{t}}))\bm{b}_{t}\,.$

Our analysis mainly relies on the minimum eigenvalue of the Gram matrix, which
is commonly used in the analysis of overparameterized neural network (Arora et
al., 2019; Chen et al., 2020). The minimum eigenvalue of the Gram matrix
controls the convergence rate of the gradient descent.

We remark that the Gram matrix $\bm{G}_{1}(\tau_{t},\nu_{t})$ is always
positive semi-definite for any $t\geq 0$, and
$\bm{G}_{1}(\tau_{0},\nu_{0})=\mathbf{0}_{n\times n}$. Therefore, we only need
to bound the minimum eigenvalue of $\bm{G}_{2}(\tau_{t},\nu_{t})$. First, we
present such result under initialization, i.e., the lower bound of
$\lambda_{\min}(\bm{G}_{2}(\tau_{0},\nu_{0}))$ by the following lemma. The
proof is deferred to Section C.2.

######  Lemma 4.3.

Under 3.1, 3.2, 3.3, there exist a constant $\Lambda:=\Lambda(d)$, only
depending on the dimension $d$, such that
$\lambda_{\min}[\bm{G}(\tau_{0},\nu_{0})]$ is lower bounded by

$\displaystyle\lambda_{0}:=\lambda_{\min}(\bm{G}(\tau_{0},\nu_{0}))\geq\lambda_{\min}(\bm{G}_{2}(\tau_{0},\nu_{0}))\geq\Lambda(d)\,.$

Remark: Using the stability of the ODE model, we derive the KL divergence by
virtue of the structure of the ResNet, build the lower bound of
$\lambda_{\min}(\bm{G}_{2})$, and prove the global convergence. In fact, our
results, e.g., global convergence, and KL divergence can also depend on
$\bm{G}_{1}$ by taking
$\Lambda(t):=\alpha^{2}\lambda_{\min}[\bm{G}_{1}(t)]+\lambda_{\min}(\bm{G}_{2})$
in Lemma C.5. Due to $\lambda_{\min}[\bm{G}_{1}(t)]\geq 0$ for any $t$, we
only use $\lambda_{\min}(\bm{G}_{2})\geq\Lambda(d)$ in Lemma 4.3 for proof
simplicity. Our model degenerates to a two-layer neural network if the
residual part is removed (can be regarded as an identity mapping).

Second, for $\tau,\nu$ different from initialization $\tau_{0},\nu_{0}$, we
first prove in the finite time $t<t_{\max}$, we have the minimum eigenvalue is
lower bounded $\lambda_{\min}(\bm{G}_{2}(\tau_{t},\nu_{t}))\geq\lambda_{0}/2$.
In the next, we choose a proper scaling of $\alpha,\beta$, such that
$t_{\max}=\infty$, so that we obtain a global guarantee.

The proof is deferred to Section C.2.

######  Lemma 4.4.

There exists $r_{\max}$, such that, for
$\nu\in\mathcal{C}(\mathcal{P}^{2};[0,1])$ and $\tau\in\mathcal{P}^{2}$
satisfying
$\max\\{\mathcal{W}_{2}(\nu,\nu_{0}),\mathcal{W}_{2}(\tau,\tau_{0})\\}\leq
r_{\max}$, we have
$\lambda_{\min}(\bm{G}_{2}(\tau,\nu))\geq\frac{\lambda_{0}}{2}$.

Remark: The radius is defined as
$r_{\max}:=\min\Big{\\{}\sqrt{d},\frac{\Lambda(d)}{4nC_{\bm{G}}(d,\alpha)}\Big{\\}}$,
where $\Lambda(d)$ is defined in Lemma 4.3 and $C_{\bm{G}}(d,\alpha)$ is some
constant depending on $d$ and $\alpha$, used for the uniform estimation of
$\bm{G}_{2}(\tau,\nu)$ around its initialization, refer to Lemma C.2. We
detail this in Section C.3.

######  Definition 4.5.

Define

$\displaystyle t_{\max}:=\sup\\{t_{0},\ {\rm s.t.}\forall
t\in[0,t_{0}],\max\\{\mathcal{W}_{2}(\nu_{t},\nu_{0}),\mathcal{W}_{2}(\tau_{t},\tau_{0})\\}\leq
r_{\max}\\}\,,$

where $r_{\max}$ is defined in Lemma 4.4.

###  4.2 KL divergence between Trained network and Initialization

Based on our previous results on the minimum eigenvalue of the Gram matrix, we
are ready to prove the global convergence of the empirical loss over the
weight distributions $\tau$ and $\nu$ of ResNets, and well control the KL
divergence of them before and after training. The proofs in this subsection
are deferred to Section C.4.

We first present the gradient flow of the KL divergence of the parameter
distribution.

######  Lemma 4.6.

The dynamics of the KL divergence ${\rm KL}(\tau_{t}\|\tau_{0}),{\rm
KL}(\nu_{t}\|\nu_{0})$ through training can be characterize by

$\displaystyle\frac{\mathrm{d}{\rm KL}(\tau_{t}\|\tau_{0})}{\mathrm{d}t}:=$
$\displaystyle-\int_{\mathbb{R}^{k_{\tau}}}\left(\nabla_{\bm{\omega}}\frac{\delta{\rm
KL}(\tau_{t}\|\tau_{0})}{\delta\tau_{t}}\right)\cdot\left(\nabla_{\bm{\omega}}\frac{\delta\widehat{L}(\tau_{t},\nu_{t})}{\delta\tau_{t}}\right)\mathrm{d}\tau_{t}(\bm{\omega})\,,$
$\displaystyle\frac{\mathrm{d}{\rm KL}(\nu_{t}\|\nu_{0})}{\mathrm{d}t}:=$
$\displaystyle-\int_{\mathbb{R}^{k_{\nu}}\times[0,1]}\left(\nabla_{\bm{\theta}}\frac{\delta{\rm
KL}(\nu_{t}^{s}\|\nu_{0}^{s})}{\delta\nu_{t}^{s}}\right)\cdot\left(\nabla_{\bm{\theta}}\frac{\delta\widehat{L}(\tau_{t},\nu_{t})}{\delta\nu_{t}}\right)\mathrm{d}\nu_{t}(\bm{\theta},s)\,.$

Since the evolution of $\tau_{t},\nu_{t}$ is continuous, we define the
notation $t_{\max}>0$ in the following, such that the minimum eigenvalue of
$\bm{G}_{2}(\tau_{t},\nu_{t})$ can be lower bounded for $t<t_{\max}$. In our
later proof, we will demonstrate that $t_{\max}=\infty$ can be achieved under
proper $\alpha$ and $\beta$.

Combining Lemma 4.4 and Definition 4.5, we immediately obtain that
$\lambda_{\min}(\bm{G}(\tau_{t},\nu_{t}))\geq\lambda_{\min}(\bm{G}_{2}(\tau_{t},\nu_{t}))\geq\lambda_{0}/2$
for $t<t_{\max}$.

By choosing certain $\alpha,\beta$, we could prove $t_{\max}=\infty$, which
leads to the bound ${\rm KL}(\tau_{t}\|\tau_{0})$, and ${\rm
KL}(\nu_{t}\|\nu_{0})$ uniformly for all $t>0$. (_c.f._ Theorem 4.7).

######  Theorem 4.7.

Assume the PDE (10) has solution $\tau_{t}\in\mathcal{P}^{2}$, and the PDE
(12) has solution $\nu_{t}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$. Under
Assumption 3.1, 3.2, 3.3, for some constant $C_{\rm KL}$ dependent on
$d,\alpha$, taking $\bar{\beta}:=\frac{\beta}{n}>\frac{4\sqrt{C_{\rm
KL}(d,\alpha)}}{\Lambda r_{\max}}$, the following results hold for all
$t\in[0,\infty)$:

$\displaystyle\widehat{L}(\tau_{t},\nu_{t})\leq
e^{-\frac{\beta^{2}\Lambda}{2n}t}\widehat{L}(\tau_{0},\nu_{0}),\quad{\rm
KL}(\tau_{t}\|\tau_{0})\leq\frac{C_{\rm
KL}(d,\alpha)}{\Lambda^{2}\bar{\beta}^{2}},\quad{\rm
KL}(\nu_{t}\|\nu_{0})\leq\frac{C_{\rm
KL}(d,\alpha)}{\Lambda^{2}\bar{\beta}^{2}}\,.$

We also derive a lower bound for the KL divergence. In Lemma C.10, we have
that the average movement of the KL divergence is in the same order as the
change in output layers.

###  4.3 Rademacher Complexity Bound

After our previous estimates on the minimum eigenvalue of the Gram matrix as
well as the KL divergence, we are ready to build the generalization bound for
such trained mean-field ResNets. The proofs in this subsection are deferred to
Section C.5.

Before we start the proof, we introduce some basic notations of Rademacher
complexity. Let $\mathcal{D}_{X}=\\{\bm{x}_{i}\\}_{i=1}^{n}$ be the training
dataset, and $\eta_{1},\cdots,\eta_{n}$ be an i.i.d. family of Rademacher
variables taking values $\pm 1$ with equal probability. For any function set
$\mathcal{H}$, the global Rademacher complexity is defined as
$\mathcal{R}_{n}(\mathcal{H}):=\mathbb{E}\left[\sup_{h\in\mathcal{H}}\frac{1}{n}\sum_{i=1}^{n}\eta_{i}h(\bm{x}_{i})\right]$.

Let $\mathcal{F}=\left\\{f_{\tau,\nu}(\bm{x})=\beta\cdot\int
h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})\mathrm{d}\tau(\bm{\omega})\right\\}$ be
the function class of infinite wide infinite depth ResNet defined in Section
3. We consider the following function class of infinite wide infinite depth
ResNets whose KL divergence to the initial distribution is upper bounded by
some $r>0$: $\mathcal{F}_{\rm KL}(r)=\left\\{f_{\tau,\nu}\in\mathcal{F}:{\rm
KL}(\tau\|\tau_{0})\leq r,{\rm KL}(\nu\|\nu_{0})\leq r\right\\}$. The
Rademacher complexity of $\mathcal{F}_{\rm KL}(r)$ is given by the following
lemma.

######  Lemma 4.8.

Under 3.3, if $r\leq r_{0}=O(1/\sqrt{n})$, the Rademacher complexity of
$\mathcal{F}_{\rm KL}(r)$ can be bounded by $\mathcal{R}_{n}(\mathcal{F}_{\rm
KL}(r))\lesssim\beta\sqrt{r/n}$, where $\lesssim$ hides the constant
dependence on $d,\alpha$.

Now we consider the generalization error of the 0-1 classification problem,

######  Theorem 4.9 (Generalization).

Assume $\tau_{y}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$ and
$\nu_{y}\in\mathcal{P}^{2}$ be the ground truth distributions, such that,
$y(\bm{x})=\mathbb{E}_{\bm{\omega}\sim\tau_{y}}h(\bm{Z}_{\nu_{y}}(\bm{x},1),\bm{\omega})$.
Under the 3.1, 3.2 and 3.3, we set $\beta>\Omega(\sqrt{n})$. For any
$\delta>0$, with probability at least $1-\delta$, the following bound holds:

$\displaystyle\mathbb{E}_{\bm{x}\sim\mu_{X}}\ell_{0-1}(f_{\tau_{\star},\nu_{\star}}(\bm{x}),y(\bm{x}))\lesssim
1/\sqrt{n}+6\sqrt{\log(2/\delta)/{2n}},$

where $\lesssim$ hides the constant dependence on $d,\alpha$.

Remark: Our results of $O(1/\sqrt{n})$ matches the standard generalization
error in the NTK regime (Du et al., 2019b). However, in contrast to setting
$\alpha=\sqrt{M},\beta=\sqrt{K}$ in Eq. 2 as the NTK regime, we directly
analyze the ResNet in the limiting infinite width depth model in Eq. 6, and
select proper choice of $\alpha,\beta$ independent of the width. We also
validate our theoretical results by some numerical experiments in Section C.6.

##  5 Conclusion

In this paper, we build the generalization bound for trained deep results
beyond the NTK regime under mild assumptions. Our results demonstrate that the
KL divergence between the distribution of parameters after training and
initialization of an infinitely width and deep ResNet can be controlled via
lower bounding the eigenvalue of the Gram matrix during training. Under some
stronger data assumptions, e.g., $k$-sparse parity problem (Suzuki et al.,
2023), we may ensure that the limiting distribution of deep ResNet moves far
away from its initialization in terms of KL divergence, which cannot be
derived under the current setting. We leave it as the future work.

##  6 Acknowledgement

This work was carried out in the EPFL LIONS group. This work was supported by
Hasler Foundation Program: Hasler Responsible AI (project number 21043), the
Army Research Office and was accomplished under Grant Number W911NF-24-1-0048,
and Swiss National Science Foundation (SNSF) under grant number 200021_205011.
Corresponding authors: Fanghui Liu and Yihang Chen.

## References

  * Akiyama & Suzuki (2022) Shunta Akiyama and Taiji Suzuki.  Excess risk of two-layer relu neural networks in teacher-student settings and its superiority to kernel methods.  _arXiv preprint arXiv:2205.14818_ , 2022.
  * 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.
  * 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) Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In _International Conference on Machine Learning_ , pp. 322–332. PMLR, 2019.
  * Ba et al. (2022) Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang.  High-dimensional asymptotics of feature learning: How one gradient step improves the representation.  _Advances in Neural Information Processing Systems_ , 35:37932–37946, 2022.
  * Bachlechner et al. (2021) Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley.  Rezero is all you need: Fast convergence at large depth.  In _Uncertainty in Artificial Intelligence_ , pp. 1352–1361. PMLR, 2021.
  * Barboni et al. (2022) Raphaël Barboni, Gabriel Peyré, and François-Xavier Vialard.  On global convergence of resnets: From finite to infinite width using linear parameterization.  _Advances in Neural Information Processing Systems_ , 35:16385–16397, 2022.
  * Barzilai et al. (2022) Daniel Barzilai, Amnon Geifman, Meirav Galun, and Ronen Basri.  A kernel perspective of skip connections in convolutional networks.  _arXiv preprint arXiv:2211.14810_ , 2022.
  * Belfer et al. (2021) Yuval Belfer, Amnon Geifman, Meirav Galun, and Ronen Basri.  Spectral analysis of the neural tangent kernel for deep residual networks.  _arXiv preprint arXiv:2104.03093_ , 2021.
  * Cao & Gu (2019) Yuan Cao and Quanquan Gu.  Generalization bounds of stochastic gradient descent for wide and deep neural networks.  In _Advances in Neural Information Processing Systems_ , volume 32, 2019.
  * 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. (2018) Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  _Advances in neural information processing systems_ , 31, 2018.
  * Chen et al. (2022) Zhengdao Chen, Eric Vanden-Eijnden, and Joan Bruna.  On feature learning in neural networks with global convergence guarantees.  _arXiv preprint arXiv:2204.10782_ , 2022.
  * Chen et al. (2020) Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang.  A generalized neural tangent kernel analysis for two-layer neural networks.  _Advances in Neural Information Processing Systems_ , 33: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.
  * Chizat & Bach (2020) Lenaic Chizat and Francis Bach.  Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss.  In _Conference on Learning Theory_ , pp. 1305–1338. PMLR, 2020.
  * Chizat et al. (2019) Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  _Advances in neural information processing systems_ , 32, 2019.
  * Cont et al. (2022) Rama Cont, Alain Rossier, and RenYuan Xu.  Convergence and implicit regularization properties of gradient descent for deep residual networks.  _arXiv preprint arXiv:2204.07261_ , 2022.
  * Ding et al. (2021) Zhiyan Ding, Shi Chen, Qin Li, and Stephen Wright.  On the global convergence of gradient descent for multi-layer resnets in the mean-field regime.  _arXiv preprint arXiv:2110.02926_ , 2021.
  * Ding et al. (2022) Zhiyan Ding, Shi Chen, Qin Li, and Stephen J Wright.  Overparameterization of deep resnet: Zero loss and mean-field analysis.  _J. Mach. Learn. Res._ , 23:48–1, 2022.
  * Donsker & Varadhan (1975) Monroe D Donsker and SR Srinivasa Varadhan.  Asymptotic evaluation of certain markov process expectations for large time, i.  _Communications on Pure and Applied Mathematics_ , 28(1):1–47, 1975.
  * 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. (2020) W. E, C. Ma, and L. Wu.  Machine learning from a continuous viewpoint, i.  _Science China Mathematics_ , 63(11):2233–2266, Sep 2020.
  * 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.
  * Ghorbani et al. (2020) Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  When do neural networks outperform kernel methods?  _Advances in Neural Information Processing Systems_ , 33:14820–14830, 2020.
  * Haber & Ruthotto (2017) Eldad Haber and Lars Ruthotto.  Stable architectures for deep neural networks.  _Inverse problems_ , 34(1):014004, 2017.
  * Hayou & Yang (2023) Soufiane Hayou and Greg Yang.  Width and depth limits commute in residual networks.  In _International Conference on Machine Learning_ , 2023.
  * Hayou et al. (2019) Soufiane Hayou, Arnaud Doucet, and Judith Rousseau.  Exact convergence rates of the neural tangent kernel in the large depth limit.  _arXiv preprint arXiv:1905.13654_ , 2019.
  * Hayou et al. (2021) Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau.  Stable resnet.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 1324–1332. PMLR, 2021.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Identity mappings in deep residual networks.  In _Computer Vision–ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11–14, 2016, Proceedings, Part IV 14_ , pp. 630–645. Springer, 2016.
  * Hu et al. (2022) Tianyang Hu, Jun Wang, Wenjia Wang, and Zhenguo Li.  Understanding square loss in training overparametrized neural network classifiers.  _Advances in Neural Information Processing Systems_ , 35:16495–16508, 2022.
  * Huang et al. (2020) Kaixuan Huang, Yuqing Wang, Molei Tao, and Tuo Zhao.  Why do deep residual networks generalize better than deep feedforward networks?—a neural tangent kernel perspective.  In _Advances in Neural Information Processing Systems_ , volume 33, pp. 2698–2709, 2020.
  * Hui & Belkin (2020) Like Hui and Mikhail Belkin.  Evaluation of neural architectures trained with square loss vs cross-entropy in classification tasks.  _arXiv preprint arXiv:2006.07322_ , 2020.
  * 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.
  * Li et al. (2021) Mufan Li, Mihai Nica, and Dan Roy.  The future is log-Gaussian: ResNets and their infinite-depth-and-width limit at initialization.  In _Advances in Neural Information Processing Systems_ , pp. 7852–7864, 2021.
  * 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.
  * Lu et al. (2018) Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong.  Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations.  In _International Conference on Machine Learning_ , pp. 3276–3285. PMLR, 2018.
  * Ma et al. (2020) Chao Ma, Lei Wu, et al.  Machine learning from a continuous viewpoint, i.  _Science China Mathematics_ , 63(11):2233–2266, 2020.
  * Mahankali et al. (2023) Arvind Mahankali, Jeff Z Haochen, Kefan Dong, Margalit Glasgow, and Tengyu Ma.  Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time.  _arXiv preprint arXiv:2306.16361_ , 2023.
  * Marion et al. (2022) Pierre Marion, Adeline Fermanian, Gérard Biau, and Jean-Philippe Vert.  Scaling resnets in the large-depth regime.  _arXiv preprint arXiv:2206.06929_ , 2022.
  * Marion et al. (2023) Pierre Marion, Yu-Han Wu, Michael E. Sander, and Gérard Biau.  Implicit regularization of deep residual networks towards neural odes, 2023.
  * 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 et al. (2021) Quynh Nguyen, Marco Mondelli, and Guido F Montufar.  Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks.  In _International Conference on Machine Learning_ , pp. 8119–8129. PMLR, 2021.
  * Nguyen & Mondelli (2020) Quynh N Nguyen and Marco Mondelli.  Global convergence of deep networks with one wide layer followed by pyramidal topology.  _Advances in Neural Information Processing Systems_ , 33:11961–11972, 2020.
  * Otto & Villani (2000) Felix Otto and Cédric Villani.  Generalization of an inequality by talagrand and links with the logarithmic sobolev inequality.  _Journal of Functional Analysis_ , 173(2):361–400, 2000.
  * Poli et al. (2021) Michael Poli, Stefano Massaroli, Atsushi Yamashita, Hajime Asama, Jinkyoo Park, and Stefano Ermon.  Torchdyn: implicit models and neural numerical methods in pytorch.  In _Neural Information Processing Systems, Workshop on Physical Reasoning and Inductive Biases for the Real World_ , volume 2, 2021.
  * Polyanskiy & Wu (2016) Yury Polyanskiy and Yihong Wu.  Wasserstein continuity of entropy and outer bounds for interference channels.  _IEEE Transactions on Information Theory_ , 62(7):3992–4002, 2016.
  * Rotskoff & Vanden-Eijnden (2018) Grant Rotskoff and Eric Vanden-Eijnden.  Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Sander et al. (2022) Michael Sander, Pierre Ablin, and Gabriel Peyré.  Do residual neural networks discretize neural ordinary differential equations?  In _Advances in Neural Information Processing Systems_ , pp. 36520–36532, 2022.
  * Sirignano & Spiliopoulos (2020a) J. Sirignano and K. Spiliopoulos.  Mean field analysis of neural networks: A central limit theorem.  _Stochastic Processes and their Applications_ , 130(3):1820–1852, 2020a.
  * Sirignano & Spiliopoulos (2020b) 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, 2020b.
  * 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.
  * Sonoda & Murata (2019) Sho Sonoda and Noboru Murata.  Transport analysis of infinitely deep neural network.  _The Journal of Machine Learning Research_ , 20(1):31–82, 2019.
  * Suzuki et al. (2023) Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda.  Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond.  In _Advances in Neural Information Processing Systems_ , 2023.
  * Tirer et al. (2022) Tom Tirer, Joan Bruna, and Raja Giryes.  Kernel-based smoothness analysis of residual networks.  In _Mathematical and Scientific Machine Learning_ , pp. 921–954. PMLR, 2022.
  * Wei et al. (2019) Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma.  Regularization matters: Generalization and optimization of neural nets vs their induced kernel.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Weinan (2017) Ee Weinan.  A proposal on machine learning via dynamical systems.  _Communications in Mathematics and Statistics_ , 1(5):1–11, 2017.
  * 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.
  * Woodworth et al. (2020) Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro.  Kernel and rich regimes in overparametrized models.  In _Conference on Learning Theory_ , pp. 3635–3673. PMLR, 2020.
  * Yang & Hu (2021) Greg Yang and Edward J Hu.  Tensor programs iv: Feature learning in infinite-width neural networks.  In _International Conference on Machine Learning_ , pp. 11727–11737. PMLR, 2021.
  * Zhu et al. (2022) Zhenyu Zhu, Fanghui Liu, Grigorios Chrysos, and Volkan Cevher.  Generalization properties of nas under activation and skip connection search.  _Advances in Neural Information Processing Systems_ , 35:23551–23565, 2022.

##  Appendix A Overview of Appendix

We give a brief overview of the appendix here.

  * •

Appendix B. In Section B.1, we prove some lemmas that will be useful. In
Section B.2, we provide the estimation of the activation function
$\bm{\sigma}$. In Section B.3, we provide the prior estimation of
$\bm{Z}_{\nu},\bm{p}_{\nu}$.

  * •

Appendix C In Section C.1, we prove the gradient flow of
$\frac{\mathrm{d}\widehat{L}}{\mathrm{d}t}$ and $\frac{\mathrm{d}{\rm
KL}}{\mathrm{d}t}$. In Section C.2, we bound the minimal eigenvalue at
initialization. In Section C.3, we bound the perturbation of minimal
eigenvalue. In Section C.4, we bound the KL divergence in finite time, and
choose scaling parameters to prove the main results in Theorem 4.7. In Section
C.5, we bound the KL divergence and provided the generalization bound.

  * •

Section C.6 We provide the experimental verification.

##  Appendix B Useful Estimations

###  B.1 Useful Lemmas

######  Lemma B.1 (2-Wasserstein continuity for functions of quadratic growth,
Proposition 1 in Polyanskiy & Wu (2016)).

Let ${\mu},\nu$ be two probability measures on $\mathbb{R}^{d}$ with finite
second moments, and let $g:\mathbb{R}^{d}\to\mathbb{R}$ be a $\mathcal{C}_{1}$
function obeying

$\displaystyle\|\nabla g(w)\|_{2}\leq c_{1}\|w\|_{2}+c_{2},\forall
w\in\mathbb{R}^{d}\;,$

for some constants $c_{1}>0$ and $c_{2}\geq 0$. Then

$\displaystyle|\mathbb{E}_{w\sim{\mu}}g(w)-\mathbb{E}_{w\sim\nu}g(w)|\leq(c_{1}\sigma+c_{2})\mathcal{W}_{2}({\mu},\nu)\;,$

where
$\sigma^{2}=\max\\{\mathbb{E}_{w\sim\mu}\|w\|^{2}_{2},\mathbb{E}_{w\sim\nu}\|w\|_{2}^{2}\\}$

######  Lemma B.2 (Corollary 2.1 in Otto & Villani (2000)).

The probability measure
$\nu_{0}(\bm{\theta})\propto\exp(-\frac{\|\bm{\theta}\|_{2}^{2}}{2})$
satisfies following Talagrand inequality (in short $T(\frac{1}{2})$) for any
$\nu\in\mathcal{P}^{2}(\mathbb{R}^{k_{\nu}})$

$\displaystyle\mathcal{W}_{2}^{2}(\nu,\nu_{0})\leq 4{\rm KL}(\nu\|\nu_{0}).$

######  Lemma B.3 (Donsker-Varadhan representation (Donsker & Varadhan,
1975)).

Let $\mu,\lambda$ be probability measures on a measurable space $(X,\Sigma)$.
For any bounded, $\Sigma$-measurable functions $\Phi:X\to\mathbb{R}$:

$\displaystyle\int_{X}\Phi\mathrm{d}\mu\leq{\rm
KL}(\mu\|\lambda)+\log\int_{X}\exp(\Phi)\mathrm{d}\lambda.$

###  B.2 Estimation of $\bm{\sigma}$

######  Lemma B.4 (Boundedness of $\bm{\sigma}(\bm{z},\bm{\theta})$).

Under 3.3, for $\bm{x}\in\mathbb{R}^{d},\bm{\theta}\in\mathbb{R}^{k}$, we have

$\displaystyle\|\bm{\sigma}(\bm{z},\bm{\theta})\|_{2}\leq
C_{\bm{\sigma}}(\|\bm{z}\|_{2}+1)(\|\bm{\theta}\|^{2}_{2}+1),$ (21)
$\displaystyle\|\nabla_{\bm{z}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}\leq
C_{\bm{\sigma}}(\|\bm{\theta}\|_{2}^{2}+1),$ (22)
$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}\leq
C_{\bm{\sigma}}(\|\bm{z}\|_{2}+1)(\|\bm{\theta}\|_{2}+1),$ (23)
$\displaystyle\|\Delta_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{2}\leq
C_{\bm{\sigma}}(\|\bm{z}\|_{2}^{2}+1)(\|\bm{\theta}\|_{2}+1),$ (24)
$\displaystyle\|\nabla_{\bm{\theta}}(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\cdot\bm{\theta})\|_{F}\leq
C_{\bm{\sigma}}(\|\bm{\theta}\|_{2}+1)(\|\bm{z}\|_{2}+1),$ (25)
$\displaystyle\|\nabla_{\bm{\theta}}\Delta_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}\leq
C_{\bm{\sigma}}(\|\bm{\theta}\|_{2}+1)(\|\bm{z}\|_{2}^{3}+1),$ (26)

where $\Delta$ is the Laplace operator. Let
$\bm{\sigma}=(\sigma_{i})_{i=1}^{d}$,
$(\nabla_{\bm{z}}\bm{\sigma})_{ij}=\nabla_{z_{j}}\sigma_{i}$,
$(\nabla_{\bm{\theta}}\bm{\sigma})_{ij}=\nabla_{\theta_{j}}\sigma_{i}$,
$(\Delta_{\bm{\theta}}\bm{\sigma})_{i}=\Delta_{\bm{\theta}}\sigma_{i}$,
$(\nabla_{\bm{\theta}}\bm{\sigma}\cdot\bm{\theta})_{ij}=(\nabla_{\bm{\theta}}\bm{\sigma})_{ij}\theta_{j}$.

###### Proof of Lemma B.4.

We prove the relations directly in the following:

$\displaystyle\|\bm{\sigma}(\bm{z},\bm{\theta})\|_{2}$
$\displaystyle=\|\bm{u}\sigma_{0}(\bm{w}^{\top}\bm{z}+b)\|_{2}\leq
C_{1}\|\bm{u}\|_{2}|\bm{w}^{\top}\bm{z}+b|\leq
C_{1}(\|\bm{z}\|_{2}+1)(\|\bm{\theta}\|_{2}^{2}+1),$ (27)
$\displaystyle\|\nabla_{\bm{z}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}$
$\displaystyle\leq\|\bm{u}\|_{2}|\bm{w}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}+b)|\leq\|\bm{u}\|_{2}\cdot
C_{1}\|\bm{w}\|_{2}\leq C_{1}(\|\bm{\theta}\|_{2}^{2}+1).$ (28)

We can write
$\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\in\mathbb{R}^{d\times k}$
by

$\displaystyle(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}))_{ij}=\begin{cases}\sigma_{0}(\bm{w}^{\top}\bm{z}+b)&\quad
j=i,\\\ 0,&\quad j\neq i,1\leq j\leq d,\\\
u_{i}z_{j-d}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}+b)&\quad d+1\leq j\leq
2d,\\\ u_{i}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}+b)&\quad
j=2d+1.\end{cases}$ (29)

Therefore,

$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}^{2}$
$\displaystyle=\sum_{i=1}^{d}\left[\sigma_{0}^{2}(\bm{w}^{\top}\bm{z}+b)+u_{i}^{2}(\sigma^{\prime}_{0}(\bm{w}^{\top}\bm{z}+b))^{2}\left(1+\|\bm{z}\|_{2}^{2}\right)\right]$
$\displaystyle\leq
dC_{1}^{2}(\|\bm{w}\|_{2}\|\bm{z}\|_{2}+b)^{2}+C_{1}^{2}\|u\|_{2}^{2}(1+\|\bm{z}\|_{2}^{2})$
$\displaystyle\leq
2dC_{1}^{2}(\|\bm{w}\|_{2}^{2}\|\bm{z}\|_{2}^{2}+b^{2})+C_{1}^{2}\|u\|_{2}^{2}(1+\|\bm{z}\|_{2}^{2})$
$\displaystyle\leq
2dC_{1}^{2}(1+\|\bm{z}\|_{2}^{2})(\|\bm{w}\|_{2}^{2}+\|\bm{u}\|_{2}^{2}+b^{2}+1)$
$\displaystyle=2dC_{1}^{2}(1+\|\bm{z}\|_{2}^{2})(1+\|\bm{\theta}\|_{2}^{2})\leq
2dC_{1}^{2}(1+\|\bm{z}\|)^{2}(1+\|\bm{\theta}\|_{2})^{2}.$

Therefore,

$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}\leq\sqrt{2d}C_{1}(1+\|\bm{z}\|)(1+\|\bm{\theta}\|_{2}).$
(30)

For $i\in[d]$

$\displaystyle|\Delta_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})_{i}|=\left|u_{i}(1+\|\bm{z}\|_{2}^{2})\sigma^{\prime\prime}_{0}(\bm{w}^{\top}\bm{z}+b)\right|\leq
C_{1}\cdot u_{i}(\|\bm{z}\|_{2}^{2}+1).$ (31)

Therefore,

$\displaystyle\|\Delta_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{2}\leq
C_{1}(\|\bm{z}\|_{2}^{2}+1)\cdot\sqrt{\sum_{i=1}^{d}u_{i}^{2}}\leq
C_{1}(\|\bm{z}\|_{2}^{2}+1)\cdot(\|\theta\|_{2}+1).$ (32)

By Eq. 29, we obtain

$\displaystyle\langle\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})_{i,\cdot},\bm{\theta}\rangle=u_{i}\sigma_{0}(\bm{w}^{\top}\bm{z}+b)+u_{i}(\bm{w}^{\top}\bm{z}+b)\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}+b),\quad
1\leq i\leq d.$

Hence,

$\displaystyle(\nabla_{\bm{\theta}}\langle\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})_{i,\cdot},\bm{\theta}\rangle)_{j}=\begin{cases}\sigma_{0}(\bm{w}^{\top}\bm{z}+b)+(\bm{w}^{\top}\bm{z}+b)\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}+b)&\quad
j=i,\\\ 0&\quad j\neq i,1\leq j\leq d,\\\
u_{i}z_{j-d}\left(\sigma_{0}^{\prime}(y)+(y\sigma_{0}^{\prime}(y))^{\prime}\right)|_{y=\bm{w}^{\top}\bm{z}+b}&\quad
d+1\leq j\leq 2d,\\\
u_{i}\left(\sigma_{0}^{\prime}(y)+(y\sigma_{0}^{\prime}(y))^{\prime}\right)|_{y=\bm{w}^{\top}\bm{z}+b}&\quad
j=2d+1.\end{cases}$

By 3.3, we have

$\displaystyle\|\nabla_{\bm{\theta}}\langle\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})_{i,\cdot},\bm{\theta}\rangle\|_{2}^{2}$
$\displaystyle=(\sigma_{0}(\bm{w}^{\top}\bm{z}+b)+(\bm{w}^{\top}\bm{z}+b)\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}+b))^{2}$
$\displaystyle+(u_{i}\left(\sigma_{0}^{\prime}(y)+(y\sigma_{0}^{\prime}(y))^{\prime}\right)|_{y=\bm{w}^{\top}\bm{z}+b})^{2}(1+\|\bm{z}\|_{2}^{2})$
$\displaystyle\leq[2C_{1}(\|\bm{\theta}\|_{2}+1)(\|\bm{z}\|_{2}+1)]^{2}+4u_{i}^{2}C_{1}^{2}(1+\|\bm{z}\|_{2}^{2})$

Hence,

$\displaystyle\|\nabla_{\bm{\theta}}\langle\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}),\bm{\theta}\rangle\|_{F}^{2}$
$\displaystyle=\sum_{i=1}^{d}\|\nabla_{\bm{\theta}}\langle\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})_{i,\cdot},\bm{\theta}\rangle\|_{2}^{2}$
$\displaystyle\leq
4dC_{1}^{2}(\|\bm{\theta}\|_{2}+1)^{2}(\|\bm{z}\|_{2}+1)^{2}+4\|u\|_{2}^{2}C_{1}^{2}(\|\bm{z}\|_{2}+1)^{2}$
$\displaystyle\leq(4d+4)C_{1}^{2}(\|\bm{\theta}\|_{2}+1)^{2}(\|\bm{z}\|_{2}+1)^{2}$
(33)

For the last part, by Eq. 31,

$\displaystyle\nabla_{\bm{\theta}}\Delta_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})_{ij}=\begin{cases}(1+\|\bm{z}\|_{2}^{2})\sigma^{\prime\prime}_{0}(\bm{w}^{\top}\bm{z}+b)&\quad
j=i\\\ 0&\quad j\neq i,1\leq j\leq d\\\
u_{i}z_{j-d}(1+\|\bm{z}\|_{2}^{2})\bm{\sigma}^{\prime\prime\prime}(\bm{w}^{\top}\bm{z}+b)&\quad
d+1\leq j\leq 2d\\\
u_{i}(1+\|\bm{z}\|_{2}^{2})\bm{\sigma}^{\prime\prime\prime}(\bm{w}^{\top}\bm{z}+b)&\quad
j=2d+1\end{cases}$

Therefore,

$\displaystyle\|\nabla_{\bm{\theta}}\Delta_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta})\|_{F}$
$\displaystyle\leq
C_{1}(1+\|\bm{z}\|_{2}^{2})\sqrt{\sum_{i=1}^{d}1+u_{i}^{2}(1+\|\bm{z}\|_{2}^{2})}$
$\displaystyle\leq\sqrt{d}C_{1}(\|\bm{z}\|_{2}^{2}+1)^{1.5}(\|\bm{\theta}\|_{2}+1)\leq
3\sqrt{d}C_{1}(\|\bm{\theta}\|_{2}+1)(\|\bm{z}\|_{2}^{3}+1)$ (34)

The last inequality is from that, for $x>0$

$\displaystyle
x^{3}+1=x^{3}+\frac{1}{2}+\frac{1}{2}\geq\frac{3}{2^{\frac{2}{3}}}x,\quad
x^{3}+1=1+\frac{x^{3}}{2}+\frac{x^{3}}{2}\geq\frac{3}{2^{\frac{2}{3}}}x^{2},$

then, we have

$\displaystyle(1+x^{2})^{\frac{3}{2}}\leq(1+x^{2})(1+x)=1+x+x^{2}+x^{3}\leq(1+x^{3})(1+\frac{2^{\frac{5}{3}}}{3})<3(1+x^{3}).$

From Eq. 27, Eq. 28, Eq. 30, Eq. 32, Eq. 33, and Eq. 34, taking
$C_{\bm{\sigma}}^{1}=4\sqrt{d}C_{1}$, the proof is finished. We defer the
definition of $C_{\bm{\sigma}}$ later. ∎

######  Lemma B.5 (Stability of $\bm{\sigma}(\bm{z},\bm{\theta})$).

Under 3.3, for $\bm{x}\in\mathbb{R}^{d},\bm{\theta}\in\mathbb{R}^{k}$, we have

$\displaystyle\|\bm{\sigma}(\bm{z}_{1},\bm{\theta})-\bm{\sigma}(\bm{z}_{2},\bm{\theta})\|_{2}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{\theta}\|_{2}^{2}+1)\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$
(35)
$\displaystyle\|\nabla_{\bm{z}}\bm{\sigma}(\bm{z}_{1},\bm{\theta})-\nabla_{\bm{z}}\bm{\sigma}(\bm{z}_{2},\bm{\theta})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{\theta}\|_{2}^{2}+1)\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$
(36)
$\displaystyle\|\nabla_{\bm{z}}\bm{\sigma}(\bm{z},\bm{\theta}_{1})-\nabla_{\bm{z}}\bm{\sigma}(\bm{z},\bm{\theta}_{2})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{\theta}_{1}\|_{2}+\|\bm{\theta}_{2}\|_{2}+1)\|\bm{\theta}_{1}-\bm{\theta}_{2}\|_{2}$
(37)
$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{1})-\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{2})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{z}\|_{2}+1)\|\bm{\theta}_{1}-\bm{\theta}_{2}\|_{2}$
(38)
$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z}_{1},\bm{\theta})-\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z}_{2},\bm{\theta})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{\theta}\|_{2}^{2}+1)\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$
(39)

###### Proof of Lemma B.5.

By the mean-value theorem, we have there exists $\epsilon\in[0,1]$

$\displaystyle\|\bm{\sigma}(\bm{z}_{1},\bm{\theta})-\bm{\sigma}(\bm{z}_{2},\bm{\theta})\|_{2}$
$\displaystyle\leq\|\nabla_{\bm{z}}\bm{\sigma}(\bm{z}_{1}+\epsilon(\bm{z}_{2}-\bm{z}_{1}),\bm{\theta})\|_{F}\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$
$\displaystyle\leq
C_{\bm{\sigma}}^{1}\cdot(\|\bm{\theta}\|_{2}^{2}+1)\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$

Denote by $\bm{\theta}=(\bm{u},\bm{w},b)$, we have

$\displaystyle\|\nabla_{\bm{z}}\sigma(\bm{z}_{1},\bm{\theta})-\nabla_{\bm{z}}\sigma(\bm{z}_{2},\bm{\theta})\|_{F}$
$\displaystyle\leq\|\bm{u}\bm{w}^{\top}(\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}_{1}+b)-\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}_{2}+b))\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}^{1}\cdot(\|\bm{\theta}\|_{2}^{2}+1)\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$

and

$\displaystyle\|\nabla_{\bm{z}}\bm{\sigma}(\bm{z},\bm{\theta}_{1})-\nabla_{\bm{z}}\bm{\sigma}(\bm{z},\bm{\theta}_{2})\|_{F}\leq\|\bm{u}_{1}\bm{w}_{1}^{\top}\sigma_{0}^{\prime}(\bm{w}_{1}^{\top}\bm{z}+b_{1})-\bm{u}_{2}\bm{w}_{2}^{\top}\sigma_{0}^{\prime}(\bm{w}_{2}^{\top}\bm{z}+b_{2})\|_{F}$
$\displaystyle\leq$ $\displaystyle
C_{\bm{\sigma}}^{1}\|\bm{u}_{1}\bm{w}_{1}^{\top}-\bm{u}_{2}\bm{w}_{2}^{\top}\|_{F}\leq
C_{\bm{\sigma}}^{1}\|(\bm{u}_{1}-\bm{u}_{2})(\bm{w}_{1}-\bm{w}_{2})^{\top}+(\bm{u}_{1}-\bm{u}_{2})\bm{w}_{2}^{\top}+\bm{u}_{1}(\bm{w}_{1}-\bm{w}_{2})^{\top}\|_{F}$
$\displaystyle\leq$ $\displaystyle
2C_{\bm{\sigma}}^{1}(\|\bm{\theta}_{1}\|_{2}+\|\bm{\theta}_{2}\|_{2}+1)\|\bm{\theta}_{1}-\bm{\theta}_{2}\|_{2}$

In the next, we have

$\displaystyle(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{1})-\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{2}))_{ij}$
$\displaystyle=$
$\displaystyle\begin{cases}\sigma_{0}(\bm{w}_{1}^{\top}\bm{z}+b_{1})-\sigma_{0}(\bm{w}_{2}^{\top}\bm{z}+b_{2})&\quad
j=i,\\\ 0,&\quad j\neq i,1\leq j\leq d,\\\
u_{i}^{1}z_{j-d}\sigma_{0}^{\prime}(\bm{w}_{1}^{\top}\bm{z}+b_{1})-u_{i}^{2}z_{j-d}\sigma_{0}^{\prime}(\bm{w}_{2}^{\top}\bm{z}+b_{2})&\quad
d+1\leq j\leq 2d,\\\
u_{i}^{1}\sigma_{0}^{\prime}(\bm{w}_{1}^{\top}\bm{z}+b_{1})-u_{i}^{2}\sigma_{0}^{\prime}(\bm{w}_{2}^{\top}\bm{z}+b_{2})&\quad
j=2d+1.\end{cases}$

and then,

$\displaystyle(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{1}))_{ij}-(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{2}))_{ij}|=\begin{cases}C_{\bm{\sigma}}^{1}\cdot\|\bm{w}_{1}-\bm{w}_{2}\|_{2}\|\bm{z}\|_{2}&\quad
j=i,\\\ 0,&\quad j\neq i,1\leq j\leq d,\\\
C_{\bm{\sigma}}^{1}\cdot|u_{i}^{1}-u_{i}^{2}|z_{j-d}&\quad d+1\leq j\leq
2d,\\\ C_{\bm{\sigma}}^{1}\cdot|u_{i}^{1}-u_{i}^{2}|&\quad j=2d+1.\end{cases}$

Therefore,

$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{1})-\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{2})\|_{F}\leq\sqrt{2d}C_{\bm{\sigma}}^{1}(\|\bm{z}\|_{2}+1)\cdot\|\bm{\theta}_{1}-\bm{\theta}_{2}\|_{2}.$

Similarly, we have

$\displaystyle(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z}_{1},\bm{\theta}))_{ij}-(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z}_{2},\bm{\theta}))_{ij}$
$\displaystyle=$
$\displaystyle\begin{cases}\sigma_{0}(\bm{w}^{\top}\bm{z}_{1}+b)-\sigma_{0}(\bm{w}^{\top}\bm{z}_{2}+b)&\quad
j=i,\\\ 0,&\quad j\neq i,1\leq j\leq d,\\\
u_{i}z_{j-d}^{1}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}_{1}+b_{1})-u_{i}z_{j-d}^{2}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}_{2}+b)&\quad
d+1\leq j\leq 2d,\\\
u_{i}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}_{1}+b)-u_{i}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{z}_{2}+b)&\quad
j=2d+1,\end{cases}$

and then

$\displaystyle(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{1}))_{ij}-(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z},\bm{\theta}_{2}))_{ij}|=\begin{cases}C_{\bm{\sigma}}^{1}\cdot\|\bm{w}\|_{2}\|\bm{z}_{1}-\bm{z}_{2}\|_{2}&\quad
j=i,\\\ 0,&\quad j\neq i,1\leq j\leq d,\\\ C_{\bm{\sigma}}^{1}\cdot
u_{i}|z_{j-d}^{1}-z_{j-d}^{2}|&\quad d+1\leq j\leq 2d,\\\
C_{\bm{\sigma}}^{1}\cdot
u_{i}\|\bm{z}_{1}-\bm{z}_{2}\|_{2}\|\bm{w}\|_{2}&\quad j=2d+1.\end{cases}$

Therefore,

$\displaystyle\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z}_{1},\bm{\theta})-\nabla_{\bm{\theta}}\bm{\sigma}(\bm{z}_{2},\bm{\theta})\|_{F}\leq\sqrt{d}C_{\bm{\sigma}}^{1}(\|\bm{\theta}\|_{2}^{2}+1)\|\bm{z}_{1}-\bm{z}_{2}\|_{2}$

taking $C_{\bm{\sigma}}^{2}=\sqrt{2d}C_{\bm{\sigma}}^{1}$, the proof is
finished. ∎

Combined with the estimation in the proofs of Lemma B.4 and Lemma B.5, we let

$\displaystyle C_{\bm{\sigma}}:=6d\cdot C_{1}$ (40)

###  B.3 Prior Estimation of ODE

The Lemma B.6 and Lemma B.7 establishes the boundedness and stability of
$\bm{Z}_{\nu}$ and $\bm{p}_{\nu}$ with respect to $\nu$.

######  Lemma B.6 (Boundedness and Stability of $\bm{Z}_{\nu}$).

Suppose that 3.3 holds and that $\bm{x}$ is in the support of $\mathcal{X}$.
Suppose that $\nu_{1},\nu_{2}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$ and
$\bm{Z}_{\nu_{1}}$, $\bm{Z}_{\nu_{2}}$ are the corresponding unique solutions
in Eq. 5.Then the following two bounds are satisfied for all $s\in[0,1]$:

$\left\|\bm{Z}_{\nu_{1}}(\bm{x},s)\right\|_{2}\leq
C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2};\alpha),$

and

$\left\|\bm{Z}_{\nu_{1}}(\bm{x},s)-\bm{Z}_{\nu_{2}}(\bm{x},s)\right\|_{2}\leq
C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2}),$

where $C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)$
is a constant depending only on
$\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2}$ and $\alpha$, and for
$\nu\in\mathcal{C}(\mathcal{P}^{2};[0,1])$, we denote by
$\|\nu\|^{2}_{\infty}:=\sup_{s\in[0,1]}\mathbb{E}_{\bm{\theta}\sim\nu(\cdot,s)}\|\bm{\theta}\|_{2}^{2}<\infty$.

###### Proof of Lemma B.6.

We firstly demonstrate that Eq. 5 has a unique $\mathcal{C}_{1}$ solution and
then prove the boundedness of $\bm{Z}_{\nu}$ under different probability
measures.

By Lemma B.5, we have

$\displaystyle\left\|\int_{\mathbb{R}^{k}}(\bm{\sigma}(\bm{z}_{1},\bm{\theta})-\bm{\sigma}(\bm{z}_{2},\bm{\theta}))\mathrm{d}\nu_{1}(\bm{\theta},s)\right\|_{2}$
$\displaystyle\leq
C_{\bm{\sigma}}\|\bm{z}_{1}-\bm{z}_{2}\|_{2}\int_{\mathbb{R}^{k}}(\|\bm{\theta}\|_{2}^{2}+1)\mathrm{d}\nu_{1}(\bm{\theta},s)$
$\displaystyle\leq
C_{\bm{\sigma}}\|\bm{z}_{1}-\bm{z}_{2}\|_{2}(\|\nu_{1}\|_{\infty}^{2}+1)\,,$
(41)

which implies that
$\int_{\mathbb{R}^{k}}\bm{\sigma}(\bm{z}_{1},\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)$
is locally Lipschitz. Combining this with the a-priori estimate, the ODE
theory implies that Eq. 5 has a unique $\mathcal{C}_{1}$ solution.

In the next, we aim to prove the boundedness of $\bm{Z}_{\nu}$. For any
$s\in[0,1]$, by Eq. 21 in Lemma B.4, we have

$\displaystyle\left\|\int_{\mathbb{R}^{k}}\bm{\sigma}(\bm{z},\bm{\theta})\mathrm{d}{\nu_{1}}(\bm{\theta},s)\right\|_{2}\leq\int_{\mathbb{R}^{k}}\|\bm{\sigma}(\bm{z},\bm{\theta})\|_{2}\mathrm{d}{\nu_{1}}(\bm{\theta},s)\leq
C_{\bm{\sigma}}(\|\bm{z}\|_{2}+1)\int_{\mathbb{R}^{k}}(\|\bm{\theta}\|_{2}^{2}+1)\mathrm{d}{\nu_{1}}(\bm{\theta},s)\,.$

To prove the boundedness of $\bm{Z}_{{\nu_{1}}}$, using Eq. 5 and Lemma B.4,
we have

$\displaystyle\frac{\mathrm{d}\|\bm{Z}_{{\nu_{1}}}(\bm{x},s)\|_{2}^{2}}{\mathrm{d}s}$
$\displaystyle=2\bm{Z}^{\top}_{{\nu_{1}}}(\bm{x},s)\frac{\mathrm{d}\bm{Z}_{{\nu_{1}}}(\bm{x},s)}{\mathrm{d}s}$
$\displaystyle\leq 2\alpha
C_{\bm{\sigma}}(\|\bm{Z}_{\nu_{1}}(\bm{x},s)\|_{2}^{2}+\|\bm{Z}_{\nu_{1}}(\bm{x},s)\|_{2})\int_{\mathbb{R}^{k}}(\|\bm{\theta}\|_{2}^{2}+1)\mathrm{d}{\nu_{1}}(\bm{\theta},s)$
$\displaystyle\leq 4\alpha
C_{\bm{\sigma}}(\|\bm{Z}_{\nu_{1}}(\bm{x},s)\|_{2}^{2}+1)\int_{\mathbb{R}^{k}}(\|\bm{\theta}\|_{2}^{2}+1)\mathrm{d}{\nu_{1}}(\bm{\theta},s).$

By Grönwall’s inequality, and $\bm{Z}_{\nu_{1}}(\bm{x},0)=\bm{x}$, we have

$\displaystyle\|\bm{Z}_{{\nu_{1}}}(\bm{x},s)\|_{2}$
$\displaystyle\leq\exp\left(2\alpha
C_{\bm{\sigma}}\left(\int_{0}^{1}\int_{\mathbb{R}^{k}}\|\bm{\theta}\|_{2}^{2}\mathrm{d}{\nu_{1}}(\bm{\theta},s)+1\right)\right)(\|\bm{x}\|_{2}+1)$
$\displaystyle\leq\exp(2\alpha
C_{\bm{\sigma}}(\|{\nu_{1}}\|_{\infty}^{2}+1))(\|\bm{x}\|_{2}+1).$

By 3.1, $\|\bm{x}\|_{2}\leq 1$, we thus have a priori estimate of
$Z_{{\nu_{1}}}$. Let
$C_{\bm{Z}}^{1}(\|\nu_{1}\|_{\infty}^{2};\alpha):=2\exp(2\alpha
C_{\bm{\sigma}}(\|{\nu_{1}}\|_{\infty}^{2}+1))$, we have
$\left\|\bm{Z}_{\nu_{1}}(\bm{x},s)\right\|_{2}\leq
C_{\bm{Z}}^{1}(\|\nu_{1}\|_{\infty}^{2};\alpha)$.

In the next, to estimate the difference under different measures $\nu_{1}$ and
$\nu_{2}$, define

$\displaystyle\bm{\delta}(\bm{x},s)=\bm{Z}_{\nu_{1}}(\bm{x},s)-\bm{Z}_{\nu_{2}}(\bm{x},s)\,,$

and we can easily obtain

$\displaystyle\frac{\mathrm{d}\|\bm{\delta}(\bm{x},s)\|_{2}^{2}}{\mathrm{d}s}$
$\displaystyle=2\alpha\left\langle\bm{\delta}(\bm{x},s),\int_{\mathbb{R}^{k}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)-\int_{\mathbb{R}^{k}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{2}(\bm{\theta},s)\right\rangle$
$\displaystyle:=2\alpha\left\langle\bm{\delta}(\bm{x},s),{\tt(A)}+{\tt(B)}\right\rangle,$
(42)

where, by Eq. 41, we have

$\displaystyle\|{\tt(A)}\|_{2}:=\left\|\int_{\mathbb{R}^{k}}(\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})-\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta}))\mathrm{d}\nu_{1}(\bm{\theta},s)\right\|_{2}\leq
C_{\bm{\sigma}}\|\bm{\delta}(\bm{x},s)\|_{2}(\|\nu_{1}\|_{\infty}^{2}+1)\,,$

and

$\displaystyle{\tt(B)}:=\int_{\mathbb{R}^{k}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)-\int_{\mathbb{R}^{k}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{2}(\bm{\theta},s)\,.$

By Lemma B.1 and
$\|\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\|_{F}\leq
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\|_{2}+1)(\|\bm{\theta}\|_{2}+1)$,
we can bound ${\tt(B)}$ by

$\displaystyle\|{\tt(B)}\|_{2}$ $\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\|_{2}+1)\cdot(\|\bm{\theta}\|_{2}+1)\cdot\mathcal{W}_{2}(\nu_{1}^{s},\nu_{2}^{s})$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\|_{2}+1)\cdot(\sqrt{\max\\{\|\nu_{1}^{s}\|_{2}^{2},\|\nu_{2}^{s}\|_{2}^{2}\\}}+1)\cdot\mathcal{W}_{2}(\nu_{1}^{s},\nu_{2}^{s})$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(C_{\bm{Z}}^{1}(\|\nu_{2}\|_{\infty}^{2};\alpha)+1)\cdot(\sqrt{\|\nu_{1}\|_{\infty}^{2}+\|\nu_{2}\|_{\infty}^{2}}+1)\mathcal{W}_{2}(\nu_{1},\nu_{2}).$

Plugging the estimate of ${\tt(A)}$ and ${\tt(B)}$ into Eq. 42, we have

$\displaystyle\frac{\mathrm{d}\|\bm{\delta}(\bm{x},s)\|_{2}^{2}}{\mathrm{d}s}$
$\displaystyle\leq 2\alpha
C_{\bm{\sigma}}\big{(}\|\bm{\delta}(\bm{x},s)\|_{2}^{2}(\|\nu_{1}\|_{\infty}^{2}+1)$
$\displaystyle+\|\bm{\delta}(\bm{x},s)\|_{2}(C_{\bm{Z}}^{1}(\|\nu_{1}\|_{\infty}^{2};\alpha)+1)(\sqrt{\|\nu_{1}\|_{\infty}^{2}+\|\nu_{2}\|_{\infty}^{2}}+1)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})\big{)}$
$\displaystyle\leq 2\alpha
C_{\bm{\sigma}}(\|\bm{\delta}(\bm{x},s)\|_{2}^{2}+\mathcal{W}_{2}^{2}(\nu_{1},\nu_{2}))(\sqrt{\|\nu_{1}\|_{\infty}^{2}+\|\nu_{2}\|_{\infty}^{2}}+1)^{2}(C_{\bm{Z}}^{1}(\|\nu_{1}\|_{\infty}^{2};\alpha)+1)^{2}\,.$

Since $\bm{\delta}(\bm{x},0)=0$, by Grönwall’s inequality, we have, $\forall
s\in[0,1]$,

$\displaystyle\|\bm{\delta}(\bm{x},s)\|_{2}\leq(\exp(\alpha
C_{\bm{\sigma}})-1)\cdot(\sqrt{\|\nu_{1}\|_{\infty}^{2}+\|\nu_{2}\|_{\infty}^{2}}+1)(C_{\bm{Z}}^{1}(\|\nu_{1}\|_{\infty}^{2};\alpha)+1)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2}),$

which concludes the proof. ∎

######  Lemma B.7 (Boundedness and Stability of $\bm{p}_{\nu}$).

Suppose that 3.3 holds and that $\bm{x}$ is in the support of $\mathcal{X}$.
Suppose that $\nu_{1},\nu_{2}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$ and
${\bm{p}}_{\nu_{1}}(\bm{x},s)$, ${\bm{p}}_{\nu_{2}}(\bm{x},s)$ are defined in
Eq. 15. Then the following three bounds are satisfied for all $s\in[0,1]$:

$\displaystyle\left\|{\bm{p}}_{\nu_{1}}(\bm{x},s)\right\|_{2}\leq
C_{\bm{p}}(\|\nu_{1}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha),$ (43)

and

$\displaystyle\left\|{\bm{p}}_{\nu_{1}}(\bm{x},s)-{\bm{p}}_{\nu_{2}}(\bm{x},s)\right\|_{2}\leq
C_{\bm{p}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})\,,$
(44)

where
$C_{\bm{p}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)$
is a constant depending only on
$\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2}$ and
$\alpha$, and for $\tau\in\mathcal{P}^{2}$, we denote by
$\|\tau\|^{2}_{2}:=\mathbb{E}_{\bm{\omega}\sim\tau(\cdot)}\|\tau\|_{2}^{2}<\infty$.

###### Proof of Lemma B.7.

For any $s\in[0,1]$, by Eq. 15 and the estimation of
$\nabla_{\bm{z}}\bm{\sigma}$ in Eq. 22, we have

$\displaystyle\frac{\mathrm{d}\|{\bm{p}}_{\nu_{1}}(\bm{x},s)\|_{2}^{2}}{\mathrm{d}s}$
$\displaystyle=2\frac{\mathrm{d}{\bm{p}}^{\top}_{\nu_{1}}(\bm{x},s)}{\mathrm{d}s}{\bm{p}}_{\nu_{1}}(\bm{x},s)$
$\displaystyle\leq
2\alpha\|{\bm{p}}^{\top}_{\nu_{1}}(\bm{x},s)\|_{2}^{2}\cdot\left\|\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)\right\|_{F}$
$\displaystyle\leq 2\alpha
C_{\bm{\sigma}}\|{\bm{p}}^{\top}_{\nu_{1}}(\bm{x},s)\|_{2}^{2}\int_{\mathbb{R}^{k}}(\|\bm{\theta}\|_{2}^{2}+1)\mathrm{d}\nu_{1}(\bm{\theta},s)\,.$

It follows from the estimation of $\nabla_{\bm{z}}\bm{\sigma}$ in Eq. 22,

$\displaystyle\|{\bm{p}}_{\nu_{1}}^{\top}(\bm{x},1)\|_{2}=\left\|\int_{\mathbb{R}^{k_{\tau}}}\nabla_{\bm{z}}h(\bm{Z}_{\nu_{1}}(\bm{x},1),\bm{\omega})\mathrm{d}\tau(\bm{\omega})\right\|_{2}\leq
C_{\bm{\sigma}}\cdot(\|\tau\|_{2}^{2}+1)\,.$

Therefore, by the Grönwall’s inequality

$\displaystyle\|{\bm{p}}_{\nu_{1}}(\bm{x},s)\|_{2}$ $\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\tau\|_{2}^{2}+1)\cdot\exp\left(\alpha
C_{\bm{\sigma}}\int_{0}^{1}\int_{\mathbb{R}^{k}}(\|\bm{\theta}\|_{2}^{2}+1)\mathrm{d}\nu_{1}(\bm{\theta},s)\right)$
$\displaystyle\leq C(\|\nu_{1}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)\,.$

In the next, we deal with Eq. 44, define

$\displaystyle\bm{\delta}_{2}(\bm{x},s):={\bm{p}}_{\nu_{1}}(\bm{x},s)-{\bm{p}}_{\nu_{2}}(\bm{x},s)\,,$

we have (taking $s=1$)

$\displaystyle\|\bm{\delta}_{2}(\bm{x},1)\|_{2}$
$\displaystyle=\left\|\int_{\mathbb{R}^{k_{\tau}}}\nabla_{\bm{z}}h(\bm{Z}_{\nu_{1}}(\bm{x},1),\bm{\omega})-\nabla_{\bm{z}}h(\bm{Z}_{\nu_{2}}(\bm{x},1),\bm{\omega})\mathrm{d}\tau(\bm{\omega})\right\|_{2}$
$\displaystyle\leq
C_{\bm{\sigma}}(\|\tau\|_{2}^{2}+1)\cdot\|\bm{Z}_{\nu_{1}}(\bm{x},1)-\bm{Z}_{\nu_{2}}(\bm{x},1)\|_{2}$
$\displaystyle\leq
C(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})\,.$

The following ODE is satisfied by $\bm{\delta}_{2}(\bm{x},s)$ by Eq. 15,

$\displaystyle\frac{\partial\bm{\delta}_{2}^{\top}(\bm{x},s)}{\partial
s}=-\alpha\cdot\bm{\delta}_{2}^{\top}(\bm{x},\bm{\omega}.s)\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)+\alpha\cdot{\bm{p}}_{\nu_{2}}(\bm{x},s)^{\top}\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s)\,,$

with

$\displaystyle\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s):=\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{2}(\bm{\theta},s)-\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)\,.$

Furthermore, we also split $\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s)$ as

$\displaystyle\|\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s)\|_{F}$
$\displaystyle\leq\underbrace{\left\|\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{2}(\bm{\theta},s)-\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)\right\|_{F}}_{\tt(A)}$
$\displaystyle+\underbrace{\left\|\int_{\mathbb{R}^{k}}\left(\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta})-\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\right)\mathrm{d}\nu_{1}(\bm{\theta},s)\right\|_{F}}_{\tt(B)}\,.$

Clearly, ${\tt(B)}$ can be estimated by

$\displaystyle{\tt(B)}\leq
C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)\cdot
C_{\bm{\sigma}}\cdot(\|\nu_{1}\|_{\infty}^{2}+1)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})\,.$

To estimate ${\tt(A)}$, denote
$\pi^{\star}_{\nu}\in\Pi(\nu_{1}^{s},\nu_{2}^{s})$ such that
$\mathbb{E}_{(\bm{\theta}_{1},\bm{\theta}_{2})\sim\pi^{\star}_{\nu}}\|\bm{\theta}_{1}-\bm{\theta}_{2}\|_{2}^{2}=\mathcal{W}_{2}^{2}(\nu_{1}^{s},\nu_{2}^{s})$,
by Lemma B.5, we then have

$\displaystyle{\tt(A)}$
$\displaystyle\leq\mathbb{E}_{(\bm{\theta}_{1},\bm{\theta}_{2})\sim\pi^{\star}_{\nu}}\|\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta}_{2})-\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{2}}(\bm{x},s),\bm{\theta}_{1})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot\sqrt{3\mathbb{E}_{(\bm{\theta}_{1},\bm{\theta}_{2})\sim\pi^{\star}_{\nu}}\|\bm{\theta}_{1}\|_{2}^{2}+\|\bm{\theta}_{2}\|_{2}^{2}+1}\cdot\sqrt{\mathbb{E}_{(\bm{\theta}_{1},\bm{\theta}_{2})\sim\pi^{\star}_{\nu}}\|\bm{\theta}_{1}-\bm{\theta}_{2}\|_{2}^{2}}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot\sqrt{3(\|\nu_{1}\|_{\infty}^{2}+\|\nu_{2}\|_{\infty}^{2}+1)}\cdot\mathcal{W}_{2}(\nu_{1},\nu_{0})\,.$

Combining the estimate of ${\tt(A)}$ and ${\tt(B)}$, we have

$\displaystyle\|\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s)\|_{F}\leq
C(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})\,.$

Accordingly, we are ready to estimate $\bm{\delta}_{2}(\bm{x},s)$.

$\displaystyle\frac{\mathrm{d}\|\bm{\delta}_{2}(\bm{x},s)\|_{2}^{2}}{\mathrm{d}s}$
$\displaystyle=2\left(-\alpha\cdot\delta^{\top}(\bm{x},s)\int_{\mathbb{R}^{k}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{1}(\bm{\theta},s)+\alpha\cdot{\bm{p}}_{\nu_{2}}(\bm{x},s)^{\top}\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s)\right)\bm{\delta}_{2}(\bm{x},s)$
$\displaystyle\leq
2\alpha\left(\|\bm{\delta}_{2}(\bm{x},s)\|_{2}^{2}\left(1+\int_{\mathbb{R}^{k}}\|\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{1}}(\bm{x},s),\bm{\theta})\|_{F}\mathrm{d}\nu_{1}(\bm{\theta},s)\right)+\|{\bm{p}}_{\nu_{2}}(\bm{x},s)^{\top}\bm{D}_{\nu_{1},\nu_{2}}(\bm{x},s)\|_{2}^{2}\right)$
$\displaystyle\leq\alpha\|\bm{\delta}_{2}(\bm{x},s)\|_{2}^{2}(1+C_{\bm{\sigma}}\cdot(1+\|\nu_{1}\|_{\infty}^{2}))+2\alpha[C(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)]\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})^{2}$
$\displaystyle\leq
C(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)\cdot(\|\bm{\delta}_{2}(\bm{x},s)\|_{2}^{2}+\mathcal{W}_{2}(\nu_{1},\nu_{2})^{2})\,.$

By the Grönwall’s inequality, $\|\bm{\delta}_{2}(\bm{x},s)\|_{2}\leq
C_{\bm{p}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2},\|\tau\|_{2}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})$,
and the proof is finished. ∎

##  Appendix C Main Results

###  C.1 Gradient Flow

###### Proof of Theorem 4.1.

To prove Theorem 4.1, we need to estimate

$\displaystyle\widehat{L}({\tau_{t},\nu_{t}})-\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})=\frac{1}{2}\mathbb{E}_{x\sim\mathcal{D}_{n}}[(\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x})-y(\bm{x}))^{2}-(\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})-y(\bm{x}))^{2}]$
$\displaystyle=$
$\displaystyle\mathbb{E}_{x\sim\mathcal{D}_{n}}(\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})-y(\bm{x}))(\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x})-\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x}))+o(|\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x})-\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})|),$

by $(a+\epsilon)^{2}-a^{2}=2a\epsilon+o(|\epsilon|)$, where $o(\cdot)$ denotes
the higher order of the error term. Then, we estimate
$\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x})-\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})$,

$\displaystyle\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x})-\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})=\beta\cdot\int_{\mathbb{R}^{k_{\tau}}}h(\bm{Z}_{{\nu_{t}}}(\bm{x},1),\bm{\omega})\mathrm{d}{\tau_{t}}(\bm{\omega})-h(\bm{Z}_{{\nu_{t_{0}}}}(\bm{x},1),\bm{\omega})\mathrm{d}{\tau_{t_{0}}}(\bm{\omega})$
$\displaystyle=$
$\displaystyle\beta\cdot\left(\underbrace{\int_{\mathbb{R}^{k_{\tau}}}h(\bm{Z}_{{\nu_{t}}}(\bm{x},1),\bm{\omega})(\mathrm{d}{\tau_{t}}(\bm{\omega})-\mathrm{d}{\tau_{t_{0}}}(\bm{\omega}))}_{\tt(A)}+\underbrace{\int_{\mathbb{R}^{k_{\tau}}}(h(\bm{Z}_{{\nu_{t}}}(\bm{x},1),\bm{\omega})-h(\bm{Z}_{{\nu_{t_{0}}}}(\bm{x},1),\bm{\omega}))\mathrm{d}{\tau_{t_{0}}}(\bm{\omega})}_{\tt(B)}\right)\,.$

We estimate $\bm{Z}_{\nu_{t}}(\bm{x},s)$ in the following, in which we assume
$\bm{\theta}_{t}^{s}\sim\nu_{t}(\cdot,s),\bm{\theta}_{t_{0}}^{s}\sim\nu_{t_{0}}(\cdot,s)$
in the expectation. Similar to the derivation in Lu et al. (2020); Ding et al.
(2022), we have

$\displaystyle\frac{1}{\alpha}\cdot\frac{\mathrm{d}(\bm{Z}_{\nu_{t}}-\bm{Z}_{\nu_{t_{0}}})(\bm{x},s)}{\mathrm{d}s}=\mathbb{E}\
(\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta}_{t}^{s})-\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s}))$
$\displaystyle=$ $\displaystyle\mathbb{E}\
(\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta}_{t}^{s})-\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t}^{s}))+\mathbb{E}\
(\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t}^{s})-\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s}))$
$\displaystyle=$ $\displaystyle\mathbb{E}\
\partial_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t}^{s})(\bm{Z}_{\nu_{t}}(\bm{x},s)-\bm{Z}_{\nu_{t_{0}}}(\bm{x},s))+\mathbb{E}\
\partial_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})(\bm{\theta}_{t}^{s}-\bm{\theta}_{t_{0}}^{s})+o(|t-t_{0}|)$
$\displaystyle=$ $\displaystyle\mathbb{E}\
\partial_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})(\bm{Z}_{\nu_{t}}(\bm{x},s)-\bm{Z}_{\nu_{t_{0}}}(\bm{x},s))+\mathbb{E}\
\partial_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})(\bm{\theta}_{t}^{s}-\bm{\theta}_{t_{0}}^{s})+o(|t-t_{0}|)$
$\displaystyle=$ $\displaystyle\mathbb{E}\
\partial_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})(\bm{Z}_{\nu_{t}}(\bm{x},s)-\bm{Z}_{\nu_{t_{0}}}(\bm{x},s))$
$\displaystyle-$ $\displaystyle\mathbb{E}\
\partial_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\mathrm{d}\nu_{t_{0}}}(\bm{\theta}_{t_{0}}^{s},s)(t-t_{0})+o(|t-t_{0}|)$

We therefore have, by the definition of $\bm{q}_{\nu}$ in Eq. 17,

$\displaystyle(\bm{Z}_{\nu_{t}}-\bm{Z}_{\nu_{t_{0}}})(\bm{x},1)$
$\displaystyle=$
$\displaystyle-\int_{0}^{1}\bm{q}_{v_{t_{0}}}(\bm{x},s)\cdot\mathbb{E}\
\left(\alpha\partial_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\mathrm{d}\nu_{t_{0}}}(\bm{\theta},s)\right)\cdot(t-t_{0})\mathrm{d}s+o(|t-t_{0}|)$

and then we have
$\|(\bm{Z}_{\nu_{t}}-\bm{Z}_{\nu_{t_{0}}})(\bm{x},1)\|_{2}=O(|t-t_{0}|)$.
Using this fact and by the evolution of $\tau_{t}$ in Eq. 10, we estimate (A),

$\displaystyle{\tt(A)}$
$\displaystyle=\int_{\mathbb{R}^{k_{\tau}}}h(\bm{Z}_{\nu_{t_{0}}}(\bm{x},1),\bm{\omega})(\mathrm{d}{\tau_{t}}(\bm{\omega})-\mathrm{d}{\tau_{t_{0}}}(\bm{\omega}))$
$\displaystyle+\int_{\mathbb{R}^{k_{\tau}}}(h(\bm{Z}_{\nu_{t}}(\bm{x},1),\bm{\omega})-h(\bm{Z}_{\nu_{t_{0}}}(\bm{x},1),\bm{\omega}))(\mathrm{d}{\tau_{t}}(\bm{\omega})-\mathrm{d}{\tau_{t_{0}}}(\bm{\omega}))$
$\displaystyle=-\int_{\mathbb{R}^{k_{\tau}}}h(\bm{Z}_{\nu_{t_{0}}}(\bm{x},1),\bm{\omega})\nabla_{\bm{\omega}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\delta\tau_{t_{0}}}(\bm{\omega})(t-t_{0})\mathrm{d}\tau_{t_{0}}(\bm{\omega})+o(|t-t_{0}|)$

We can also estimate (B), in which $h$ is hidden in the definition of
$\bm{p}_{\nu}$ in Eq. 17,

$\displaystyle{\tt(B)}=\int_{\mathbb{R}^{k_{\tau}}}\bm{p}_{\nu_{t_{0}}}(\bm{x},1)^{\top}(\bm{Z}_{{\nu_{t}}}-\bm{Z}_{\nu_{t_{0}}})(\bm{x},1)\mathrm{d}\tau_{t_{0}}(\bm{\omega})+o(|t-t_{0}|)$
$\displaystyle=-\int_{\mathbb{R}^{k_{\tau}}}\bm{p}_{\nu_{t_{0}}}(\bm{x},s)^{\top}$
$\displaystyle\left(\mathbb{E}\
\alpha\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta}_{t_{0}}^{s})\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\mathrm{d}\nu_{t_{0}}}(\bm{\theta}_{t_{0}}^{s},s)\cdot(t-t_{0})\right)\mathrm{d}\tau_{t_{0}}(\bm{\omega})+o(|t-t_{0}|)$
$\displaystyle=-\int_{\mathbb{R}^{k_{\tau}}\times\mathbb{R}^{k_{\nu}}\times[0,1]}\Big{(}\bm{p}_{\nu_{t_{0}}}(\bm{x},s)^{\top}\cdot\alpha\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta})\cdot$
$\displaystyle\qquad\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\delta\nu_{t_{0}}}(\bm{\theta},s)\Big{)}\mathrm{d}{\nu_{t_{0}}}(\bm{\theta},s)\cdot(t-t_{0})+o(|t-t_{0}|)\,.$

Combine the estimation of ${\tt(A)}$ and ${\tt(B)}$, we have

$\displaystyle\widehat{L}({\tau_{t},\nu_{t}})-\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}\beta(\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})-y(\bm{x}))({\tt(A)}+{\tt(B)})$
$\displaystyle=$
$\displaystyle-\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}\beta(\widehat{f}_{{\tau_{t_{0}},\nu_{t_{0}}}}(\bm{x})-y(\bm{x}))\int_{\mathbb{R}^{k_{\tau}}\times\mathbb{R}^{k}\times[0,1]}\mathrm{d}\tau_{t_{0}}(\bm{\omega})\mathrm{d}\nu_{t_{0}}(\bm{\theta},s)(t-t_{0})$
$\displaystyle\cdot$
$\displaystyle\left(\bm{Z}_{{\nu_{t_{0}}}}^{\top}(\bm{x},1)\nabla_{\bm{\omega}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\delta\tau_{t_{0}}}(\bm{\omega})\right.$
$\displaystyle+$
$\displaystyle\left.\bm{p}_{\nu_{t_{0}}}^{\top}(\bm{x},s)\cdot\alpha\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t_{0}}}(\bm{x},s),\bm{\theta})\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t_{0}},\nu_{t_{0}}})}{\delta\nu_{t_{0}}}(\bm{\theta},s)+o(|t-t_{0}|)\right)$
$\displaystyle=$
$\displaystyle-\mathbb{E}_{\bm{\omega}\sim\tau_{t_{0}},(\bm{\theta},s)\sim\nu_{t_{0}}}\left(\left\|\nabla_{\bm{\theta}}\frac{\delta\widehat{L}(\tau,\nu)}{\delta\nu}(\bm{\theta},s)\right\|_{2}^{2}+\left\|\nabla_{\bm{\omega}}\frac{\delta\widehat{L}(\tau,\nu)}{\delta\tau}(\bm{\omega})\right\|_{2}^{2}\right)(t-t_{0})+o(|t-t_{0}|)\,,$

from the definition of functional gradient in Eq. 11 and Eq. 14. In all, the
theorem is proved. ∎

###### Proof of Proposition 4.2.

We expand the function derivative:

$\displaystyle\int_{\mathbb{R}^{k_{\tau}}\times\mathbb{R}^{k_{\nu}}\times[0,1]}\left\|\nabla_{\bm{\theta}}\frac{\delta\widehat{L}({\tau_{t},\nu_{t}})}{\delta\nu_{t}}(\bm{\theta},s)\right\|_{2}^{2}\mathrm{d}\tau_{t}(\bm{\omega})\mathrm{d}\nu_{t}(\bm{\theta},s)$
$\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{k_{\tau}}\times\mathbb{R}^{k_{\nu}}\times[0,1]}\frac{\beta^{2}\alpha^{2}}{n^{2}}\sum_{i,j=1}^{n}(\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x}_{i})-y(\bm{x}_{i}))(\widehat{f}_{{\tau_{t},\nu_{t}}}(\bm{x}_{j})-y(\bm{x}_{j}))$
$\displaystyle\cdot{\bm{p}}^{\top}_{\nu_{t}}(\bm{x}_{i},s)\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x}_{i},s),\bm{\theta})\nabla_{\bm{\theta}}^{\top}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x}_{j},s),\bm{\theta}){\bm{p}}_{\nu_{t}}(\bm{x}_{j},s)\mathrm{d}\tau_{t}(\bm{\omega})\mathrm{d}\nu_{t}(\bm{\theta},s)$
$\displaystyle=$
$\displaystyle\frac{\alpha^{2}\beta^{2}}{n^{2}}\bm{b}_{t}^{\top}\bm{G}_{1}(\tau_{t},\nu_{t})\bm{b}_{t}\,,$

and similarly,

$\displaystyle\int_{\mathbb{R}^{k_{\tau}}}\left\|\nabla_{\bm{\omega}}\frac{\delta\widehat{L}({\tau_{t},\nu_{t}})}{\delta\tau_{t}}(\bm{\omega})\right\|_{2}^{2}\mathrm{d}\tau_{t}(\bm{\omega})=\frac{\beta^{2}}{n^{2}}\bm{b}_{t}^{\top}\bm{G}_{2}(\tau_{t},\nu_{t})\bm{b}_{t}.$

In all, the lemma is proved. ∎

###### Proof of Lemma 4.6.

We use the expansion of the gradient flow:

$\displaystyle\frac{\mathrm{d}{\rm KL}(\tau_{t}\|\tau_{0})}{\mathrm{d}t}$
$\displaystyle=\int_{\mathbb{R}^{k_{\tau}}}\frac{\delta{\rm
KL}(\tau_{t}\|\tau_{0})}{\delta\tau_{t}}\frac{\mathrm{d}\tau_{t}}{\mathrm{d}t}\mathrm{d}\bm{\omega}=\int_{\mathbb{R}^{k_{\tau}}}\frac{\delta{\rm
KL}(\tau_{t}\|\tau_{0})}{\delta\tau_{t}}\nabla\cdot\left(\tau_{t}(\bm{\omega})\nabla\frac{\delta\widehat{L}(\tau_{t},\nu_{t})}{\delta\tau_{t}}\right)\mathrm{d}\bm{\omega}$
$\displaystyle=-\int_{\mathbb{R}^{k_{\tau}}}\tau_{t}(\bm{\omega})\left(\nabla\frac{\delta{\rm
KL}(\tau_{t}\|\tau_{0})}{\delta\tau_{t}}\right)\left(\nabla\frac{\delta\widehat{L}(\tau_{t},\nu_{t})}{\delta\tau_{t}}\right)\mathrm{d}\bm{\omega}.$

Similarly, we have

$\displaystyle\frac{\mathrm{d}{\rm
KL}(\nu_{t}^{s}\|\nu_{0}^{s})}{\mathrm{d}t}$
$\displaystyle=\int_{\mathbb{R}^{k_{\nu}}}\frac{\delta{\rm
KL}(\nu_{t}^{s}\|\nu_{t}^{0})}{\delta\nu_{t}^{s}}\frac{\mathrm{d}\nu_{t}^{s}}{\mathrm{d}t}\mathrm{d}\bm{\theta}$
$\displaystyle=\int_{\mathbb{R}^{k_{\nu}}}\frac{\delta{\rm
KL}(\nu_{t}^{s}\|\nu_{t}^{0})}{\delta\nu_{t}^{s}}\nabla_{\bm{\theta}}\cdot\left(\nu_{t}^{s}(\bm{\theta})\nabla_{\bm{\theta}}\frac{\delta\widehat{L}(\tau_{t},\nu_{t})}{\delta\nu_{t}}(\bm{\theta},s)\right)\mathrm{d}\bm{\theta}$
$\displaystyle=-\int_{\mathbb{R}^{k_{\nu}}}\nu_{t}^{s}(\bm{\theta})\left(\nabla_{\bm{\theta}}\frac{\delta{\rm
KL}(\nu_{t}^{s}\|\nu_{t}^{0})}{\delta\nu_{t}^{s}}\right)\left(\nabla_{\bm{\theta}}\frac{\delta\widehat{L}(\tau_{t},\nu_{t})}{\delta\nu_{t}}(\bm{\theta},s)\right)\mathrm{d}\bm{\theta}.$

Therefore, the proof is completed. ∎

###  C.2 Minimum Eigenvalue at Initialization

###### Proof of Lemma 4.3.

In the proof, similar to 3.3, we assume
$\bm{\theta}=(\bm{u},\bm{w},b)\in\mathbb{R}^{2d+1},\bm{\omega}=(a,\bm{w},b)\in\mathbb{R}^{d+2},\bm{u},\bm{w}\in\mathbb{R}^{d},a,b\in\mathbb{R}$.
At initialization, we notice that
$\nu_{0}(\bm{\theta},s)\propto\exp\left(-\frac{\|\bm{\theta}\|_{2}^{2}}{2}\right)$,
and
$\tau_{0}(\bm{\omega})\propto\exp\left(-\frac{\|\bm{\omega}\|_{2}^{2}}{2}\right)$
are standard Gaussian. Since the distribution of $\bm{u},a$ is symmetric, and
independent from other parts of $\bm{\theta}$, $\bm{\omega}$ respectively, we
have

$\displaystyle\frac{\mathrm{d}\bm{Z}_{\nu_{0}}(\bm{x},s)}{\mathrm{d}s}$
$\displaystyle=\int_{\mathbb{R}^{k_{\nu}}}\bm{u}^{\top}\sigma_{0}(\bm{w}^{\top}\bm{Z}_{\nu_{0}}(\bm{x},s)+b)\mathrm{d}\nu_{0}(\bm{u},\bm{w},b,s)$
$\displaystyle=\int_{\mathbb{R}^{d}}\bm{u}^{\top}\mathrm{d}\nu_{0}(\bm{u})\int_{\mathbb{R}^{k_{\nu}-d}}\sigma_{0}(\bm{w}^{\top}\bm{Z}_{\nu_{0}}(\bm{x},s)+b)\mathrm{d}\nu_{0}(\bm{w},b,s)=\bm{0},\forall
s\in[0,1],$
$\displaystyle\frac{\mathrm{d}{\bm{p}}^{\top}_{\nu_{0}}}{\mathrm{d}s}(\bm{x},s)$
$\displaystyle=-\alpha\cdot{\bm{p}}^{\top}_{\nu_{0}}(\bm{x},s)\int_{\mathbb{R}^{k_{\nu}}}\nabla_{\bm{z}}\bm{\sigma}(\bm{Z}_{\nu_{0}}(\bm{x},s),\bm{\theta})\mathrm{d}\nu_{0}(\bm{\theta},s)$
$\displaystyle=-\alpha\cdot{\bm{p}}^{\top}_{\nu_{0}}(\bm{x},s)\int_{\mathbb{R}^{k_{\nu}}}\bm{u}\bm{w}^{\top}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{Z}_{\nu_{0}}(\bm{x},s)+b)\mathrm{d}\nu_{0}(\bm{u},\bm{w},b,s)=\bm{0},$
$\displaystyle\bm{p}_{\nu_{0}}(\bm{x},1)$
$\displaystyle=\int_{\mathbb{R}^{k_{\tau}}}\nabla_{\bm{z}}^{\top}h(\bm{Z}_{\nu_{0}}(\bm{x},1),\bm{\omega})\mathrm{d}\tau_{0}(\bm{\omega})$
$\displaystyle=\int_{\mathbb{R}^{k_{\tau}}}a\bm{w}\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{Z}_{\nu_{0}}(\bm{x},1)+b)\tau_{0}(a,\bm{w},b)=\bm{0}.$

From the first two equations, we have

$\displaystyle\bm{Z}_{\nu_{0}}(\bm{x},s)=\bm{x},\forall
s\quad{\bm{p}}_{\nu_{0}}(\bm{x},s)={\bm{p}}_{\nu_{0}}(\bm{x},1)=0$

By the definition of $\bm{G}_{2}(\tau_{0},\nu_{0})$, we have

$\displaystyle\bm{G}_{2}(\tau_{0},\sigma_{0})$
$\displaystyle=\mathbb{E}_{(a,\bm{w},b)\sim\mathcal{N}(0,I)}\nabla_{\bm{\omega}}h(\bm{X},\bm{\omega})\nabla_{\bm{\omega}}^{\top}h(\bm{X},\bm{\omega})$
$\displaystyle=\mathbb{E}_{(a,\bm{w},b)\sim\mathcal{N}(0,I)}(\bm{\sigma}_{0}((\bm{X},\mathbbm{1})(\bm{w},b)),a\bm{\sigma}_{0}^{\prime}((\bm{X},\mathbbm{1})(\bm{w},b)),\bm{\sigma}_{0}^{\prime}((\bm{X},\mathbbm{1})(\bm{w},b)))$
$\displaystyle\quad\quad(\bm{\sigma}_{0}((\bm{X},\mathbbm{1})(\bm{w},b)),a\bm{\sigma}_{0}^{\prime}((\bm{X},\mathbbm{1})(\bm{w},b)),\bm{\sigma}_{0}^{\prime}((\bm{X},\mathbbm{1})(\bm{w},b)))^{\top},$
$\displaystyle\geq\mathbb{E}_{(a,\bm{w},b)\sim\mathcal{N}(0,I)}\bm{\sigma}_{0}((\bm{X},\mathbbm{1})(\bm{w},b))\bm{\sigma}_{0}((\bm{X},\mathbbm{1})(\bm{w},b))^{\top}\,.$

Let $\bar{\bm{x}}=(\bm{x},1)$, by 3.1, the cosine similarity of
$\bar{\bm{x}}_{i}$ and $\bar{\bm{x}}_{j}$ is no larger than $(1+C_{\max})/2$.
Then we bound $\lambda_{\min}(\bm{G}^{(2)})$:

$\begin{split}\lambda_{\min}(\bm{G}^{(2)})&\geq\lambda_{\min}\bigg{(}\mathbb{E}_{\bm{w}\sim\mathcal{N}(0,\mathbb{I}_{d+1})}[\sigma_{1}(\bm{\bar{X}w})\sigma_{1}(\bm{\bar{X}w})^{\top}]\bigg{)}\\\
&=\lambda_{\min}\bigg{(}\sum_{s=0}^{\infty}\mu_{s}(\sigma_{1})^{2}\bigcirc_{i=1}^{s}(\bm{\bar{X}\bar{X}}^{\top})\bigg{)}\quad\text{\cite[citep]{(\@@bibref{AuthorsPhrase1Year}{nguyen2020global}{\@@citephrase{,
}}{}, Lemma D.3)}}\\\
&\geq\mu_{r}(\sigma_{1})^{2}\lambda_{\min}(\bigcirc_{i=1}^{r}\bm{\bar{X}\bar{X}}^{\top})\quad\bigg{(}\mbox{taking~{}}r\geq\frac{2\log(2n)}{1-C_{\text{max}}}\bigg{)}\\\
&\geq\mu_{r}(\sigma_{1})^{2}\bigg{(}\min_{i\in[n]}\left\|\bar{\bm{x}}_{i}\right\|_{2}^{2r}-(n-1)\max_{i\neq
j}\left|\left\langle\bar{\bm{x}}_{i},\bar{\bm{x}}_{j}\right\rangle\right|^{r}\bigg{)}\quad\text{[Gershgorin
circle theorem]}\\\
&\geq\mu_{r}(\sigma_{1})^{2}\bigg{(}1-(n-1)\left(\frac{1+C_{\max}}{2}\right)^{r}\bigg{)},\\\
&\geq\mu_{r}(\sigma_{1})^{2}\bigg{(}1-(n-1)\bigg{(}1-\frac{\log(2n)}{r}\bigg{)}^{r}\bigg{)}\\\
&\geq\mu_{r}(\sigma_{1})^{2}\bigg{(}1-(n-1)\exp(-\log(2n))\bigg{)}\\\
&\geq\mu_{r}(\sigma_{1})^{2}/2\,,\end{split}$

where the last inequality holds by the fact that
$\big{(}1-\frac{\log(2n)}{r}\big{)}^{r}$ is an increasing function of $r$.

∎

###  C.3 Perturbation of Minimum Eigenvalue

In this section, we analyze the minimum eigenvalue of the Gram matrix.

######  Lemma C.1.

The perturbation of $\bm{G}_{2}(\tau,\nu)$ can be upper bounded in the
following, for any $i,j\in[n]$,

$\displaystyle|\bm{G}_{2}(\tau,\nu)-\bm{G}_{2}(\tau_{0},\nu_{0})|_{i,j}$
$\displaystyle\leq
C_{\bm{G}}(\|\tau\|_{2}^{2},\|\nu\|_{\infty}^{2};d,\alpha)(\mathcal{W}_{2}(\tau,\tau_{0})+\mathcal{W}_{2}(\nu,\nu_{0})),$

where $\bm{G}_{2}$ is defined in Section 4.1, and $\tau_{0},\nu_{0}$ satisfies
3.2.

###### Proof of Lemma C.1.

We deal with $\bm{G}_{2}(\tau,\nu)$ in an element-wise way. Let
$(\bm{\omega},\bm{\omega}_{0})\sim{\pi^{\star}_{\tau}}$ be the optimal
coupling of $\mathcal{W}_{2}(\tau,\tau_{0})$, the difference can be estimated
by

$\displaystyle|\bm{G}_{2}(\tau,\nu)-\bm{G}_{2}(\tau_{0},\nu_{0})|_{i,j}$
$\displaystyle\leq$
$\displaystyle\mathbb{E}|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})\nabla_{\bm{\omega}}^{\top}h(\bm{Z}_{\nu}(\bm{x}_{j},1),\bm{\omega})-\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{0}}(\bm{x}_{i},1),\bm{\omega}_{0})\nabla_{\bm{\omega}}^{\top}h(\bm{Z}_{\nu_{0}}(\bm{x}_{j},1),\bm{\omega}_{0})|$
$\displaystyle\leq$
$\displaystyle\underbrace{\mathbb{E}|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})(\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x}_{j},1),\bm{\omega})-\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{0}}(\bm{x}_{j},1),\bm{\omega}_{0}))^{\top}|}_{\tt(A)}$
$\displaystyle+$
$\displaystyle\underbrace{\mathbb{E}|(\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})-\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{0}}(\bm{x}_{i},1),\bm{\omega}_{0}))\nabla_{\bm{\omega}}^{\top}h(\bm{Z}_{\nu_{0}}(\bm{x}_{j},1),\bm{\omega}_{0})|}_{\tt(B)}\,.$

We then estimate ${\tt(A)}$ and ${\tt(B)}$ separately. The term ${\tt(A)}$
involves

$\displaystyle\|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{0}}(\bm{x},1),\bm{\omega}_{0})\|_{2}$
$\displaystyle\leq$
$\displaystyle\|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega}_{0})\|_{2}+\|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega}_{0})-\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{0}}(\bm{x},1),\bm{\omega}_{0})\|_{2}$
$\displaystyle\leq$ $\displaystyle
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu}(\bm{x},1)\|_{2}+1)\cdot\|\bm{\omega}-\bm{\omega}_{0}\|_{2}+C_{\bm{\sigma}}\cdot(\|\bm{\omega}_{0}\|_{2}^{2}+1)\cdot\|\bm{Z}_{\nu}(\bm{x},1)-\bm{Z}_{\nu_{0}}(\bm{x},1)\|_{2}$
$\displaystyle\leq$ $\displaystyle
C_{\bm{\sigma}}\cdot((C_{\bm{Z}}(\|\nu\|_{\infty}^{2};\alpha)+1)\cdot\|\bm{\omega}-\bm{\omega}_{0}\|_{2}+(\|\bm{\omega}_{0}\|_{2}^{2}+1)\cdot
C_{\bm{Z}}(\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};\alpha)\mathcal{W}_{2}(\nu,\nu_{0}))\,,$

where we use Lemmas B.5 and B.6 in our proof. Besides, the term ${\tt(B)}$
involves

$\displaystyle\|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})\|_{2}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu}(\bm{x},1)\|_{2}+1)\cdot(\|\bm{\omega}\|_{2}+1)$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(C_{\bm{Z}}(\|\nu\|_{\infty}^{2};\alpha)+1)\cdot(\|\bm{\omega}\|_{2}+1).$

Therefore, by $\mathbb{E}\|\bm{\omega}_{0}\|_{2}^{4}=3k_{\tau}=3(d+2)$.

$\displaystyle{\rm(A)+(B)}\leq
C(\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};\alpha)\mathbb{E}(\|\bm{\omega}-\bm{\omega}_{0}\|_{2}+(\|\bm{\omega}_{0}\|_{2}^{2}+1)\mathcal{W}_{2}(\nu,\nu_{0}))(\|\bm{\omega}\|_{2}+\|\bm{\omega}_{0}\|_{2}+2)$
$\displaystyle\leq$ $\displaystyle
C(\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};\alpha)\mathbb{E}(\|\bm{\omega}-\bm{\omega}_{0}\|_{2}+(\|\bm{\omega}_{0}\|_{2}^{2}+1)\mathcal{W}_{2}(\nu,\nu_{0}))(\|\bm{\omega}\|_{2}^{2}+\|\bm{\omega}_{0}\|_{2}^{2}+4)$
$\displaystyle\leq$ $\displaystyle
C(\|\tau\|_{2}^{2},\|\tau_{0}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};d,\alpha)(\mathbb{E}\|\bm{\omega}-\bm{\omega}_{0}\|_{2}(\|\bm{\omega}\|_{2}^{2}+\|\bm{\omega}_{0}\|_{2}^{2}+4)+\mathcal{W}_{2}(\nu,\nu_{0}))$
$\displaystyle\leq$ $\displaystyle
C(\|\tau\|_{2}^{2},\|\tau_{0}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};d,\alpha)[(\mathbb{E}\|\bm{\omega}-\bm{\omega}_{0}\|_{2}^{2})^{\frac{1}{2}}(\mathbb{E}(\|\bm{\omega}\|_{2}^{2}+\|\bm{\omega}_{0}\|_{2}^{2}+4)^{2})^{\frac{1}{2}}+\mathcal{W}_{2}(\nu,\nu_{0})]$
$\displaystyle\leq$ $\displaystyle
C(\|\tau\|_{2}^{2},\|\tau_{0}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};d,\alpha)(\mathcal{W}_{2}(\tau,\tau_{0})+\mathcal{W}_{2}(\nu,\nu_{0}))\,,$

since
$\mathbb{E}\|\bm{\omega}-\bm{\omega}_{0}\|_{2}^{2}=(\mathcal{W}_{2}(\tau,\tau_{0}))^{2}$,
by the definition of optimal coupling. Since
$\|\tau_{0}\|_{2}^{2}=d+2,\|\nu_{0}\|_{\infty}^{2}=2d+1$, we can drop
dependence of $C$ on $\|\tau_{0}\|_{2}^{2},\|\nu_{0}\|_{\infty}^{2}$ and
replace them by $d$. In all, the lemma is proved. Specifically, we could set

$\displaystyle C_{\bm{G}}(\|\tau\|_{2}^{2},\|\nu\|_{\infty}^{2};d,\alpha)$
$\displaystyle:=$ $\displaystyle
16(d+1)C_{\bm{\sigma}}^{2}(C_{\bm{Z}}(\|\nu\|_{\infty}^{2};\alpha)+1)+C_{\bm{Z}}(\|\nu\|_{\infty}^{2},2d+1;\alpha))^{2}(\|\tau\|_{2}^{2}+d+1)$

∎

######  Lemma C.2.

If $\nu\in\mathcal{C}(\mathcal{P}^{2};[0,1]),\tau\in\mathcal{P}^{2}$,
$\mathcal{W}_{2}(\nu,\nu_{0})\leq\sqrt{d}$, and
$\mathcal{W}_{2}(\tau,\tau_{0})\leq\sqrt{d}$, we have $\forall i,j\in[n]$,

$\displaystyle|\bm{G}_{2}(\tau,\nu)-\bm{G}_{2}(\tau_{0},\nu_{0})|_{i,j}$
$\displaystyle\leq
C_{\bm{G}}(d,\alpha)\cdot(\mathcal{W}_{2}(\nu_{t},\nu_{0})+\mathcal{W}_{2}(\tau_{t},\tau_{0}))$

###### Proof of Lemma C.2.

For any $s\in[0,1]$, let
$(\bm{\theta}^{s},\bm{\theta}_{0}^{s})\sim{\pi^{\star}_{\nu^{s}}}$ be the
optimal coupling of $\mathcal{W}_{2}(\nu^{s},\nu_{0}^{s})$.

$\displaystyle\|\nu^{s}\|_{2}^{2}=\mathbb{E}\|\bm{\theta}^{s}\|_{2}^{2}\leq
2\mathbb{E}(\|\bm{\theta}^{s}-\bm{\theta}_{0}^{s}\|_{2}^{2}+\|\bm{\theta}_{0}^{s}\|_{2}^{2})=2\mathcal{W}_{2}^{2}(\nu^{s},\nu_{0}^{s})+2(2d+1)\leq
6d+2.$

where the last inequality holds, since
$\mathcal{W}_{2}(\nu^{s},\nu_{0}^{s})\leq\mathcal{W}_{2}(\nu,\nu_{0}),\forall
s\in[0,1]$.

We also let $(\bm{\omega},\bm{\omega}_{0})\sim{\pi^{\star}_{\tau}}$ be the
optimal coupling of $\mathcal{W}_{2}(\tau,\tau_{0})$.

$\displaystyle\|\tau\|_{2}^{2}=\mathbb{E}\|\bm{\omega}\|_{2}^{2}\leq
2\mathbb{E}(\|\bm{\omega}-\bm{\omega}_{0}\|_{2}^{2}+\|\bm{\omega}_{0}\|_{2}^{2})=2\mathcal{W}_{2}^{2}(\tau,\tau_{0})+2(d+2)\leq
4(d+1).$

Therefore, $\|\nu\|_{\infty}\leq\sqrt{6d+2},\|\tau\|_{2}\leq 2\sqrt{d+1}$.

By Lemma C.1, replacing $\|\tau\|_{2}^{2},\|\nu\|_{\infty}^{2}$ with their
upper bound w.r.t. $d$ in the definition of
$C_{\bm{G}}(\|\tau\|_{2}^{2},\|\nu\|_{\infty}^{2};d,\alpha)$, there exist
$C_{\bm{G}}(d,\alpha)$ satisfying Lemma C.2. ∎

For ease of description, we restate Lemma 4.4 with more details here.

######  Lemma C.3.

If $\nu\in\mathcal{C}(\mathcal{P}^{2};[0,1]),\tau\in\mathcal{P}^{2}$,
$\mathcal{W}_{2}(\nu,\nu_{0})\leq r$ and $\mathcal{W}_{2}(\tau,\tau_{0})\leq
r$, we have

$\displaystyle\lambda_{\min}(\bm{G}_{2}(\tau,\nu))\geq\frac{\Lambda}{2},\textit{with
}r:=r_{\max}(d,\alpha)=\min\Big{\\{}\sqrt{d},\frac{\Lambda}{4nC_{\bm{G}}(d,\alpha)}\Big{\\}}\,.$

where $\Lambda$ is defined in Lemma 4.3, and $C_{\bm{G}}(d,\alpha)$ is defined
in Lemma C.2.

###### Proof of Lemma C.3.

By Lemma C.2, let

$\displaystyle
r:=\min\Big{\\{}\sqrt{d},\frac{\Lambda}{4nC_{\bm{G}}(d,\alpha)}\Big{\\}}\,,$

By Lemma 4.3, we have $\forall i,j\in[n]$,

$\displaystyle|\bm{G}_{2}(\tau,\nu)-\bm{G}_{2}(\tau_{0},\nu_{0})|_{i,j}$
$\displaystyle\leq
C_{\bm{G}}(d;\alpha)\cdot(\mathcal{W}_{2}(\nu,\nu_{0})+\mathcal{W}_{2}(\tau,\tau_{0}))\leq\frac{\Lambda}{2n}.$

By the standard matrix perturbation bounds, we have

$\displaystyle\lambda_{\min}(\bm{G}_{2}(\tau,\nu))$
$\displaystyle\geq\lambda_{\min}(\bm{G}_{2}(\tau_{0},\nu_{0}))-\|\bm{G}_{2}(\tau,\nu)-\bm{G}_{2}(\tau_{0},\nu_{0})\|_{2}$
$\displaystyle\geq\lambda_{\min}(\bm{G}_{2}(\tau_{0},\nu_{0}))-n\|\bm{G}_{2}(\tau,\nu)-\bm{G}_{2}(\tau_{0},\nu_{0})\|_{\infty,\infty}$
$\displaystyle\geq\frac{\Lambda}{2}.$

∎

###  C.4 Estimation of KL divergence.

Inspired by Lemma C.3, we propose the following definition:

######  Definition C.4.

Define

$\displaystyle t_{\max}:=\sup\\{t_{0},{\rm s.t.}\forall
t\in[0,t_{0}],\max\\{\mathcal{W}_{2}(\nu_{t},\nu_{0}),\mathcal{W}_{2}(\tau_{t},\tau_{0})\\}\leq
r_{\max}\\},$

where $r_{\max}$ is defined in Lemma 4.4.

By Lemma C.2 and Lemma C.3, for $t\leq t_{\max}$, we have
$\max\\{\|\nu_{t}\|_{\infty}^{2},\|\tau_{t}\|_{2}^{2}\\}=O(d)$, and
$\lambda_{\min}(\bm{G}_{2}(\tau_{t},\nu_{t}))\geq\frac{\Lambda}{2}$.

We first prove the linear convergence of empirical loss under finite time.

######  Lemma C.5.

Assume the PDE (10) has solution $\tau_{t}\in\mathcal{P}^{2}$, and the PDE
(12) has solution $\nu_{t}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$. Under 3.1,
3.2, 3.3, for all $t\in[0,t_{\max})$, we have

$\widehat{L}(\tau_{t},\nu_{t})\leq
e^{-\frac{\beta^{2}\Lambda}{2n}t}\widehat{L}(\tau_{0},\nu_{0}),\quad{\rm
KL}(\tau_{t}\|\tau_{0})\leq\frac{C_{\rm
KL}(d,\alpha)}{\Lambda^{2}\bar{\beta}^{2}},\quad{\rm
KL}(\nu_{t}\|\nu_{0})\leq\frac{C_{\rm
KL}(d,\alpha)}{\Lambda^{2}\bar{\beta}^{2}}\,.$ (45)

###### Proof of Lemma C.5.

Please see Lemma C.6, Lemma C.7, and Lemma C.8. ∎

######  Lemma C.6.

Assume $\tau_{t},\nu_{t}$ is the solution to PDE (10) and (12), we have for
$t<t_{\max}$,

$\displaystyle\widehat{L}(\tau_{t},\nu_{t})\leq
e^{-\frac{\beta^{2}\Lambda}{2n}t}\widehat{L}(\tau_{0},\nu_{0})\,,$

where $\Lambda$ is defined in Lemma 4.3.

###### Proof of Lemma C.6.

By Lemma C.3, for $t<t_{\max}$,
$\lambda_{\min}(\bm{G}(\tau_{t},\nu_{t}))\geq\frac{\Lambda}{2}$,

$\displaystyle\frac{\partial\widehat{L}(\nu_{t},\tau_{t})}{\partial
t}=-\frac{\beta^{2}}{n^{2}}\bm{b}_{t}^{\top}(\alpha^{2}\bm{G}_{1}(\tau_{t},\nu_{t})+\bm{G}_{2}(\tau_{t},\nu_{t}))\bm{b}_{t}\leq-\frac{\beta^{2}\Lambda}{2n}\widehat{L}(\tau_{t},\nu_{t})\leq-\frac{\beta^{2}\Lambda}{2n}\widehat{L}(\tau_{0},\nu_{0})\,.$

Therefore, we have

$\displaystyle\widehat{L}(\tau_{t},\nu_{t})\leq
e^{-\frac{\beta^{2}\Lambda}{2n}t}\widehat{L}(\tau_{0},\nu_{0})\,.$

∎

######  Lemma C.7.

Assume the PDE (12) has solution
$\nu_{t}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$, and the PDE (10) has solution
$\tau_{t}\in\mathcal{P}^{2}$. Under 3.3, 3.1, 3.2, then for all
$t\in[0,t_{\max})$, the following results hold,

$\displaystyle{\rm
KL}(\nu_{t}^{s}\|\nu_{0}^{s})\leq\frac{1}{\Lambda^{2}\bar{\beta}^{2}}C_{\rm
KL}(d,\alpha),\quad\forall s\in[0,1].$

###### Proof of Lemma C.7.

By Gaussian initialization of $\nu_{0}^{s}$,
$\log\nu_{0}^{s}(\bm{\theta})=-\frac{\|\bm{\theta}\|_{2}^{2}}{2}+C$, and we
thus have $\nabla_{\bm{\theta}}\frac{\partial{\rm
KL}(\nu_{t}^{s}||\nu_{0}^{s})}{\partial\nu_{t}^{s}}=\nabla_{\bm{\theta}}\log\nu_{t}^{s}+\bm{\theta}$.
Combining this with Lemma 4.6 and Eq. 14, we have

$\displaystyle\frac{\partial{\rm KL}(\nu_{t}^{s}||\nu_{0}^{s})}{\partial t}$
$\displaystyle=-\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}\beta\cdot(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))$
$\displaystyle\cdot\alpha\int_{\mathbb{R}^{k_{\tau}}}\nabla_{\bm{\theta}}({\bm{p}}^{\top}_{\nu_{t}}(\bm{x},s)\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta}))\cdot(\nabla_{\bm{\theta}}\log\nu_{t}^{s}+\bm{\theta})\mathrm{d}\nu_{t}^{s}(\bm{\theta}).$

Define

$J_{t}^{s}(\bm{x},\bm{\theta}):=-\alpha\cdot{\bm{p}}_{\nu_{t}}(\bm{x},s)^{\top}\Big{(}\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta})\cdot\bm{\theta}-\Delta_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta})\Big{)}.$

By the definition of $J_{t}^{s}(\bm{x},\bm{\theta})$, and integration by
parts, we have

$\displaystyle\frac{\partial{\rm KL}(\nu_{t}^{s}||\nu_{0}^{s})}{\partial
t}=\beta\cdot\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}[(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))\mathbb{E}_{\bm{\theta}\sim\nu_{t}^{s}}J_{t}^{s}(\bm{x},\bm{\theta})].$

We can obtain the gradient of $J_{t}^{s}$ w.r.t. $\bm{\theta}$,

$\displaystyle\nabla_{\bm{\theta}}J_{t}^{s}(\bm{x},\bm{\theta})$
$\displaystyle=-\alpha\cdot{\bm{p}}_{\nu_{t}}(\bm{x},s)^{\top}\Big{(}\nabla_{\bm{\theta}}(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta})\cdot\bm{\theta})-\nabla_{\bm{\theta}}\Delta_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta})\Big{)}.$

Therefore, by Lemma B.4, and the estimate
$\|\bm{Z}_{\nu_{t}}(\bm{x},s)\|_{2}\leq
C_{\bm{Z}}(\|\nu_{t}\|_{\infty}^{2};\alpha)$ from Lemma B.6, we have

$\displaystyle\|\nabla(\nabla_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta})\cdot\bm{\theta})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{\theta}\|_{2}+1)\cdot(C_{\bm{Z}}(\|\nu_{t}\|_{\infty}^{2};\alpha)+1)$
$\displaystyle\|\nabla_{\bm{\theta}}\Delta_{\bm{\theta}}\bm{\sigma}(\bm{Z}_{\nu_{t}}(\bm{x},s),\bm{\theta})\|_{F}$
$\displaystyle\leq
C_{\bm{\sigma}}\cdot(\|\bm{\theta}\|_{2}+1)\cdot(C_{\bm{Z}}(\|\nu_{t}\|_{\infty}^{2};\alpha)^{3}+1).$

Therefore, we can estimate
$\nabla_{\bm{\theta}}J_{t}^{s}(\bm{x},\bm{\theta})$,

$\displaystyle\|\nabla_{\bm{\theta}}J_{t}^{s}(\bm{x},\bm{\theta})\|_{2}$
$\displaystyle\leq
2C_{\bm{\sigma}}\cdot({C_{\bm{Z}}(\|\nu_{t}\|_{\infty}^{2};\alpha)}^{3}+1)(\|\bm{\theta}\|_{2}+1)\|{\bm{p}}_{\nu_{t}}(\bm{x},s)\|_{2}$
$\displaystyle\leq
2C_{\bm{\sigma}}\cdot({C_{\bm{Z}}(\|\nu_{t}\|_{\infty}^{2};\alpha)}^{3}+1)\cdot
C_{\bm{p}}(\|\nu_{t}\|_{\infty}^{2},\|\tau_{t}\|_{2}^{2};\alpha)\cdot(\|\bm{\theta}\|_{2}+1)$
$\displaystyle\leq
C(\|\tau_{t}\|_{2}^{2},\|\nu_{t}\|_{\infty}^{2};\alpha)\cdot(\|\bm{\theta}\|_{2}+1).$

By Lemma B.1 and Lemma B.2

$\displaystyle\mathbb{E}_{\nu_{t}^{s}}J_{t}^{s}(\bm{x},\bm{\theta})-\mathbb{E}_{\nu_{0}^{s}}J_{0}^{s}(\bm{x},\bm{\theta}_{0})$
$\displaystyle\leq
C(\|\nu_{t}\|_{\infty}^{2},\|\tau_{t}\|_{2}^{2};\alpha)\mathcal{W}_{2}(\nu_{t}^{s},\nu_{0}^{s})$
$\displaystyle\leq
C(\|\nu_{t}\|_{\infty}^{2},\|\tau_{t}\|_{2}^{2};\alpha)\sqrt{{\rm
KL}(\nu_{t}^{s}||\nu_{0}^{s})}.$

For $t\in[0,t_{\max})$, by Definition C.4, we have
$\|\nu_{t}\|_{\infty}^{2},\|\tau_{t}\|_{2}^{2}=O(d)$. We have

$\displaystyle\mathbb{E}_{\nu_{t}^{s}}J_{t}^{s}(\bm{x},\bm{\theta})-\mathbb{E}_{\nu_{0}^{s}}J_{0}^{s}(\bm{x},\bm{\theta}_{0})\leq
C(d,\alpha)\sqrt{{\rm KL}(\nu_{t}^{s}||\nu_{0}^{s})}.$

Since $\bm{p}_{\nu_{0}}(\bm{x},s)=\bm{0}$,

$\displaystyle\mathbb{E}_{\nu_{0}^{s}}J_{0}^{s}(\bm{x},\bm{\theta})=\bm{0}.$

Therefore,

$\displaystyle\frac{\partial{\rm KL}(\nu_{t}^{s}||\nu_{0}^{s})}{\partial t}$
$\displaystyle=\beta\cdot\mathbb{E}_{x\sim{\mathcal{D}_{n}}}(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))\mathbb{E}_{\nu_{t}^{s}}J_{t}^{s}(\bm{x},\bm{\theta})$
$\displaystyle=\beta\cdot\mathbb{E}_{x\sim{\mathcal{D}_{n}}}(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))\mathbb{E}_{\nu_{t}^{s}}(J_{t}^{s}(\bm{x},\bm{\theta})-J_{0}^{s}(\bm{x},\bm{\theta}))$
$\displaystyle\leq\beta\cdot C(d,\alpha)\sqrt{{\rm
KL}(\nu_{t}^{s}||\nu_{0}^{s})}\mathbb{E}_{x\sim{\mathcal{D}_{n}}}(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))$
$\displaystyle\leq\beta\cdot C(d,\alpha)\sqrt{{\rm
KL}(\nu_{t}^{s}||\nu_{0}^{s})}\sqrt{\mathbb{E}(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))^{2}}$
$\displaystyle=\beta\cdot C(d,\alpha)\sqrt{{\rm
KL}(\nu_{t}^{s}||\nu_{0}^{s})}\sqrt{\widehat{L}(\tau_{t},\nu_{t})},$

where the last inequality holds owing to the Jesen’s inequality.

By the relation $\mathrm{d}2\sqrt{x}=\mathrm{d}x/\sqrt{x}$,

$\displaystyle\mathrm{d}\left(2\sqrt{{\rm
KL}(\nu_{t}^{s}||\nu_{0}^{s})}\right)\leq\beta\cdot
C(d,\alpha)\sqrt{L(\nu_{t},\tau_{t})}\mathrm{d}t.$

We have for $t\in[0,t_{\max})$,

$\displaystyle\widehat{L}(\tau_{t},\nu_{t})\leq
e^{-\frac{\beta^{2}\Lambda}{2n}t}\widehat{L}(\tau_{0},\nu_{0}).$

Hence,

$\displaystyle 2\sqrt{{\rm KL}(\nu_{t}^{s}\|\nu_{0}^{s})}\leq\beta\cdot
C(d,\alpha)\int_{0}^{t}\sqrt{\widehat{L}(\nu_{t_{0}},\tau_{t_{0}})}\mathrm{d}t_{0}\leq\frac{4C(d,\alpha)}{\Lambda\bar{\beta}}\sqrt{\widehat{L}(\tau_{0},\nu_{0})}\,.$

Since $\tau_{0}(\bm{u})$ is standard normal distribution, we have

$\displaystyle
f_{\tau_{0},\nu_{0}}(\bm{x})=\beta\cdot\bm{a}^{\top}\int_{\mathbb{R}^{k_{\tau}}\times\mathbb{R}^{k_{\tau}}\times\mathbb{R}}\bm{u}^{\top}\bm{\sigma}_{0}(\bm{w}^{\top}\bm{Z}_{\nu_{0}}(\bm{x},1)+b)\mathrm{d}\tau_{0}(\bm{u},\bm{w},b)=0\,,$

and $|y(\bm{x})|\leq 1$ (3.1), we have $\widehat{L}(\tau_{0},\nu_{0})\leq 1$.
Therefore, we obtain

$\displaystyle{\rm KL}(\nu_{t}^{s}\|\nu_{0}^{s})\leq\frac{C_{\rm
KL}(d,\alpha)}{\Lambda^{2}\bar{\beta}^{2}}\,,$

where $C_{\rm KL}$ is a constant dependent only on $d,\alpha$. ∎

######  Lemma C.8.

Assume the PDE (12) has solution $\tau_{t}\in\mathcal{P}^{2}$, and the PDE
(10) has solution $\tau_{t}\in\mathcal{P}^{2}$. Under 3.3, 3.1, 3.2, then for
all $t\in[0,t_{\max})$, the following results hold:

$\displaystyle{\rm
KL}(\tau_{t}\|\tau_{0})\leq\frac{1}{\Lambda^{2}\bar{\beta}^{2}}C_{\rm
KL}(d,\alpha).$

###### Proof of Lemma C.8.

By Gaussian initialization of $\tau_{0}$,
$\log\tau_{0}(\bm{\omega})=-\frac{\|\bm{\omega}\|_{2}^{2}}{2}+C$, and we have
$\nabla_{\bm{\omega}}\frac{\partial{\rm
KL}(\tau_{t}||\tau_{0})}{\partial\tau_{t}}=\nabla_{\bm{\omega}}\log\tau_{t}+\bm{\omega}$.
Therefore, by Lemma 4.6 and Eq. 14, we have

$\displaystyle\frac{\partial{\rm KL}(\tau_{t}||\tau_{0})}{\partial t}$
$\displaystyle=$
$\displaystyle-\beta\cdot\int_{\mathbb{R}^{k_{\tau}}}(\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{t}}(\bm{x},1),\bm{\omega}))\cdot(\nabla_{\bm{\omega}}\log\tau_{t}+\bm{\omega})\mathrm{d}\nu_{t}^{s}(\bm{\theta}).$

Define

$\displaystyle\tilde{\mathbf{u}}_{t}(\bm{\theta})$
$\displaystyle:=\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}[(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})],$

we have

$\displaystyle\frac{\partial{\rm KL}(\tau_{t}||\tau_{0})}{\partial
t}=-\alpha\cdot\int_{\mathbb{R}^{k_{\tau}}}\tau_{t}\tilde{\mathbf{u}}_{t}\cdot(\nabla_{\bm{\omega}}\log\tau_{t}+\bm{\omega})\mathrm{d}\bm{\omega}=-\alpha\cdot\int_{\mathbb{R}^{k_{\tau}}}\tau_{t}[\tilde{\mathbf{u}}_{t}\cdot\bm{\omega}-\nabla_{\bm{\omega}}\cdot\tilde{v}_{t}^{s}]\mathrm{d}\bm{\omega}.$

We also define

$I_{t}(\bm{x},\bm{\omega}):=-\Big{(}\nabla_{\bm{\omega}}h(\bm{Z}_{\nu_{t}}(\bm{x},1),\bm{\omega})\cdot\bm{\omega}-\Delta_{\bm{\omega}}h(\bm{Z}_{\nu_{t}}(\bm{x},1),\bm{\omega})\Big{)}.$

By the definition of $I_{t}(\bm{x},\bm{\omega})$,

$\displaystyle\frac{\partial{\rm KL}(\tau_{t}||\tau_{0})}{\partial
t}=\beta\cdot\mathbb{E}_{\bm{x}\sim\mathcal{D}_{n}}[(f_{\tau_{t},\nu_{t}}(\bm{x})-y(\bm{x}))\mathbb{E}_{\bm{\omega}\sim\tau_{t}}I_{t}(\bm{x},\bm{\omega})].$

Similar to the estimate $J_{t}^{s}$, we have the estimation of $I_{t}$ as

$\displaystyle\mathbb{E}_{\tau_{t}}I_{t}(\bm{x},\bm{\omega})-\mathbb{E}_{\tau_{0}}I_{0}(\bm{x},\bm{\omega})$
$\displaystyle\leq
C(\|\tau_{t}\|_{2}^{2},\|\nu_{t}\|_{\infty}^{2};\alpha)\sqrt{{\rm
KL}(\tau_{t}||\tau_{0})}.$

and we can have, by setting $\bm{\omega}=(a,\bm{w},b)$, we have

$\displaystyle\mathbb{E}_{\tau_{0}}I_{0}(\bm{x},\bm{\omega})$
$\displaystyle=\mathbb{E}_{\tau_{0}}(-\nabla_{\bm{\omega}}h(\bm{x},\bm{\omega})\cdot\bm{\omega}+\Delta_{\bm{\omega}}h(\bm{x},\bm{\omega}))$
$\displaystyle=\mathbb{E}_{(a,\bm{w},b)\sim\mathcal{N}(0,I)}(-a\sigma_{0}(\bm{w}^{\top}\bm{x}+b)-a(\nabla_{\bm{w}}\sigma_{0}(\bm{w}^{\top}\bm{x}+b))\bm{w}-ab\sigma_{0}^{\prime}(\bm{w}^{\top}\bm{x}+b))$
$\displaystyle+\mathbb{E}_{(a,\bm{w},b)\sim\mathcal{N}(0,I)}(\Delta_{\bm{w}}\sigma_{0}(\bm{w}^{\top}\bm{x}+b)+a\sigma_{0}^{\prime\prime}(\bm{w}^{\top}\bm{x}+b))=\mathbf{0}.$

Therefore, we obtain the KL divergence in a similar fashion.

$\displaystyle{\rm KL}(\tau_{t}\|\tau_{0})\leq\frac{C_{\rm
KL}(d,\alpha)}{\Lambda^{2}\bar{\beta}^{2}}.$

∎

######  Lemma C.9 (Lower bound on the KL divergence).

For any $\tau,\tau^{\prime}\in\mathcal{P}^{2}$,
$\nu,\nu^{\prime}\in\mathcal{C}(\mathcal{P}^{2};[0,1])$, if
$\tau^{\prime},\nu^{\prime}$ satisfy the Talagrand inequality
$T(\frac{1}{2})$, (_ref_ Lemma B.2) We have the lower bound for the KL
divergence of $\tau,\tau^{\prime}$ and $\nu,\nu^{\prime}$, such that for
constant $C_{\rm
low}(\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu^{\prime}\|_{\infty}^{2};\alpha)$,

$\displaystyle\sqrt{{\rm KL}(\tau\|\tau^{\prime})}+\sqrt{{\rm
KL}(\nu\|\nu^{\prime})}\geq\frac{\mathbb{E}_{\bm{\omega}\sim\tau}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}h(\bm{Z}_{\nu^{\prime}}(\bm{x},1),\bm{\omega}^{\prime})}{C_{\rm
low}(\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu^{\prime}\|_{\infty}^{2};\alpha)}.$

.

###### Proof of Lemma C.9.

We have the following estimaation

$\displaystyle\mathbb{E}_{\bm{\omega}\sim\tau}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}h(\bm{Z}_{\nu^{\prime}}(\bm{x},1),\bm{\omega}^{\prime})$
$\displaystyle=$
$\displaystyle\underbrace{\left(\mathbb{E}_{\bm{\omega}\sim\tau}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega}^{\prime})\right)}_{\tt(A)}$
$\displaystyle+$
$\displaystyle\underbrace{\left(\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}h(\bm{Z}_{\nu^{\prime}}(\bm{x},1),\bm{\omega}^{\prime})\right)}_{\tt(B)},$

By Lemma B.4, we have

$\displaystyle\|\nabla_{\bm{\omega}}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})\|_{2}\leq
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu}(\bm{x},1)\|_{2}+1)\cdot(\|\bm{\omega}\|_{2}+1),$

and by Lemma B.1, Lemma B.6 and Lemma B.2, we have

$\displaystyle{\tt(A)}\leq$ $\displaystyle
C_{\bm{\sigma}}\cdot(\|\bm{Z}_{\nu}(\bm{x},1)\|_{2}+1)\cdot\max\\{\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2}\\}\cdot\mathcal{W}_{2}(\tau,\tau^{\prime})$
$\displaystyle\leq$ $\displaystyle
C_{\bm{\sigma}}\cdot(C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2};\alpha)+1)\cdot\max\\{\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2}\\}\mathcal{W}_{2}(\tau,\tau^{\prime}).$
$\displaystyle\leq$ $\displaystyle
2C_{\bm{\sigma}}\cdot(C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2};\alpha)+1)\cdot\max\\{\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2}\\}\sqrt{{\rm
KL}(\tau\|\tau^{\prime})}.$

Besides, by Lemma B.5 and Lemma B.2, we have

$\displaystyle{\tt(B)}$
$\displaystyle\leq\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}C_{\bm{\sigma}}(\|\bm{\omega}\|_{2}^{2}+1)\cdot(\|\bm{Z}_{\nu}(\bm{x},1)-\bm{Z}_{\nu^{\prime}}(\bm{x},1)\|_{2})$
$\displaystyle\leq(\|\tau^{\prime}\|_{2}^{2}+1)\cdot
C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu_{1},\nu_{2})$
$\displaystyle\leq 2(\|\tau^{\prime}\|_{2}^{2}+1)\cdot
C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)\cdot\sqrt{{\rm
KL}(\nu\|\nu^{\prime})}.$

We let $C_{\rm low}$ be

$\displaystyle C_{\rm
low}(\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu^{\prime}\|_{\infty}^{2};\alpha)$
$\displaystyle=$
$\displaystyle\min\\{2C_{\bm{\sigma}}\cdot(C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2};\alpha)+1)\cdot\max\\{\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2}\\},2(\|\tau^{\prime}\|_{2}^{2}+1)\cdot
C_{\bm{Z}}(\|\nu_{1}\|_{\infty}^{2},\|\nu_{2}\|_{\infty}^{2};\alpha)\\}$

Therefore, we have

$\displaystyle\sqrt{{\rm KL}(\tau\|\tau^{\prime})}+\sqrt{{\rm
KL}(\nu\|\nu^{\prime})}\geq\frac{\mathbb{E}_{\bm{\omega}\sim\tau}h(\bm{Z}_{\nu}(\bm{x},1),\bm{\omega})-\mathbb{E}_{\bm{\omega}^{\prime}\sim\tau^{\prime}}h(\bm{Z}_{\nu^{\prime}}(\bm{x},1),\bm{\omega}^{\prime})}{C_{\rm
low}(\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu^{\prime}\|_{\infty}^{2};\alpha)}$

Since
$\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu^{\prime}\|_{\infty}^{2}=O(d)$,
we have that the average movement of the KL divergence is on the same order as
the change in output value.

∎

######  Lemma C.10.

Assume the PDE (12) has solution $\tau_{t}\in\mathcal{P}^{2}$, and the PDE
(10) has solution $\tau_{t}\in\mathcal{P}^{2}$. Under 3.3, 3.1, 3.2, then for
all $t\in[0,t_{\max})$, We have the lower bound for the KL divergence of
$\tau_{t}$ and $\nu_{t}$, we have for constant $C_{\rm low}(d;\alpha)$, such
that

$\displaystyle\sqrt{{\rm KL}(\tau_{t}\|\tau_{0})}+\sqrt{{\rm
KL}(\nu_{t}\|\nu_{0})}\geq\frac{\mathbb{E}_{\bm{\omega}\sim\tau_{t}}h(\bm{Z}_{\nu_{t}}(\bm{x},1),\bm{\omega})}{C_{\rm
low}(d;\alpha)}.$

###### Proof of Lemma C.10.

By the definition of $r_{\max}\leq\sqrt{d}$, and the proof of Lemma C.2, we
have $\|\tau_{t}\|_{2}^{2},\|\nu_{t}\|_{\infty}^{2}=O(d)$, and we can directly
obtain $\|\tau_{0}\|_{2}^{2}=d+2,\|\nu_{0}\|_{\infty}^{2}=2d+1$, and
$\mathbb{E}_{\bm{\omega}_{0}\sim\tau_{0}}h(\bm{Z}_{\nu_{0}}(\bm{x},1),\bm{\omega}_{0})=0$.
Besides, Gaussian initialization satisfies $T(\frac{1}{2})$ condition in Lemma
B.2. We have, by Lemma C.9,

$\displaystyle\sqrt{{\rm KL}(\tau_{t}\|\tau_{0})}+\sqrt{{\rm
KL}(\nu_{t}\|\nu_{0})}\geq\frac{\mathbb{E}_{\bm{\omega}\sim\tau_{t}}h(\bm{Z}_{\nu_{t}}(\bm{x},1),\bm{\omega})}{C_{\rm
low}(d;\alpha)},$

where $C_{\rm low}(d;\alpha)$ is a constant depending on $d,\alpha$ derived
from $C_{\rm
low}(\|\tau\|_{2}^{2},\|\tau^{\prime}\|_{2}^{2},\|\nu\|_{\infty}^{2},\|\nu^{\prime}\|_{\infty}^{2};\alpha)$.
∎

######  Lemma C.11.

Under the Assumptions in Lemma C.7, let $\bar{\beta}\geq\frac{4\sqrt{C_{\rm
KL}(d,\alpha)}}{\Lambda r_{\max}}$, we have $t_{\max}=\infty$.

###### Proof of Lemma C.11.

Otherwise, we have the following inequality, for $\forall t<t_{\max}$:

$\displaystyle W_{2}(\nu_{t}^{s},\nu_{0}^{s})\leq 2\sqrt{{\rm
KL}(\nu_{t}^{s}\|\nu_{t}^{0})}\leq\frac{2}{\Lambda\bar{\beta}}\sqrt{C_{\rm
KL}(d,\alpha)},\quad\forall s\in[0,1].$

Therefore,

$\displaystyle\mathcal{W}_{2}(\nu_{t},\nu_{0})\leq\frac{2}{\Lambda\bar{\beta}}\sqrt{C_{\rm
KL}(d,\alpha)}.$

According to the definition of $t_{\max}$, we have
$\mathcal{W}_{2}(\nu_{t},\nu_{0})\leq r_{\max}$. Let

$\displaystyle\bar{\beta}\geq\frac{4\sqrt{C_{\rm KL}(d,\alpha)}}{\Lambda
r_{\max}}$

we have $\mathcal{W}_{2}(\nu_{t},\nu_{0})\leq r_{\max}/2,\forall
t\in[0,t_{\max})$, which contradict to the definition of $t_{\max}$ in
Definition C.4. ∎

###### Proof of Theorem 4.7.

Combine the results of Lemma C.11, Lemma C.8, and Lemma C.7, we prove the
theorem. ∎

###  C.5 Rademacher Complexity

###### Proof of Lemma 4.8.

Let $\gamma$ be a parameter whose value will be determined later in the proof.
Let $\eta_{i},1\leq i\leq n$ be the i.i.d. Rademacher random variables,

$\displaystyle\mathcal{R}_{n}(\mathcal{F}_{\rm KL}(r))$
$\displaystyle=\frac{\beta}{\gamma}\cdot\mathbb{E}_{\eta}\left(\sup_{\tau:{\rm
KL}(\tau\|\tau_{0})\leq r,\nu:{\rm KL}(\nu\|\nu_{0})\leq
r}\mathbb{E}_{\tau}\left(\frac{\gamma}{n}\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})\right)\right)$
$\displaystyle\leq\frac{\beta}{\gamma}\cdot\left(r+\mathbb{E}_{\eta}\sup_{\nu:{\rm
KL}(\nu\|\nu_{0})\leq
r}\log\mathbb{E}_{\tau_{0}}\exp\left(\frac{\gamma}{n}\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})\right)\right)$
$\displaystyle\leq\frac{\beta}{\gamma}\cdot\left(r+\mathbb{E}_{\eta}\log\mathbb{E}_{\tau_{0}}\exp\left(\frac{\gamma}{n}\sup_{\nu:{\rm
KL}(\nu\|\nu_{0})\leq
r}\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})\right)\right)$
$\displaystyle\leq\frac{\beta}{\gamma}\cdot\left(r+\log\mathbb{E}_{\tau_{0}}\mathbb{E}_{\eta}\exp\left(\frac{\gamma}{n}\sup_{\nu:{\rm
KL}(\nu\|\nu_{0})\leq
r}\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})\right)\right),$

where the first inequality is followed by the Donsker-Varadhan representation
of KL-divergence in Lemma B.3. The second inequality follows from the
increasing function $\log(\cdot),\exp(\cdot)$, and
$\sup_{y}\mathbb{E}_{x}f(x,y)\leq\mathbb{E}_{x}\sup_{y}f(x,y)$, for general
variable $x,y$ and function $f$. The third inequality follows from
$\log(\cdot)$’s convexity.

By 3.3, where $\bm{\omega}=(a,\bm{w},b)$, we have

$\displaystyle|h(z_{1},\bm{\omega})-h(z_{2},\bm{\omega})|\leq
C_{1}\cdot\|z_{1}-z_{2}\|_{2}\cdot a\|\bm{w}\|_{2}.$

We further estimate

$\displaystyle\left|\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu}(\bm{x}_{i},1),\bm{\omega})-\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu_{0}}(\bm{x}_{i},1),\bm{\omega})\right|$
$\displaystyle\leq$ $\displaystyle C_{1}n\cdot
a\|\bm{w}\|_{2}\cdot\|\bm{Z}_{\nu}(\bm{x}_{i},1)-\bm{Z}_{\nu_{0}}(\bm{x}_{i},1)\|_{2}$
$\displaystyle\leq$ $\displaystyle C_{1}n\cdot a\|\bm{w}\|_{2}\cdot
C_{\bm{Z}}(\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};\alpha)\cdot\mathcal{W}_{2}(\nu,\nu_{0})$
$\displaystyle\leq$ $\displaystyle C_{1}n\cdot a\|\bm{w}\|_{2}\cdot
C_{\bm{Z}}(\|\nu\|_{\infty}^{2},\|\nu_{0}\|_{\infty}^{2};\alpha)\cdot
2\sqrt{r}$

given ${\rm KL}(\nu\|\nu_{0})\leq r$. Further, we have
$\|\nu_{0}\|_{\infty}^{2}=2d+1$, and we have $\|\nu\|_{\infty}^{2}\leq 6d+2$.

In the following, we use $C_{d}:=2C_{1}\cdot C_{\bm{Z}}(6d+2,2d+1;\alpha)$ to
denote the constant.

$\displaystyle\mathcal{R}_{n}(\mathcal{F}_{\rm
KL}(r))\leq\frac{\beta}{\gamma}\cdot\left(r+\log\mathbb{E}_{\tau_{0}}\mathbb{E}_{\eta}\exp\left(\frac{\gamma}{n}\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu_{0}}(\bm{x}_{i},1),\bm{\omega})+\gamma\cdot
C_{d}\sqrt{r}\cdot a\|\bm{w}\|_{2}\right)\right)$ $\displaystyle=$
$\displaystyle\frac{\beta}{\gamma}\cdot\left(r+\log\mathbb{E}_{\tau_{0}}\exp\left(\gamma\cdot
C_{d}\sqrt{r}\cdot
a\|\bm{w}\|_{2}\right)\mathbb{E}_{\eta}\exp\left(\frac{\gamma}{n}\sum_{i=1}^{n}\eta_{i}h(\bm{Z}_{\nu_{0}}(\bm{x}_{i},1),\bm{\omega})\right)\right)$
$\displaystyle\leq$