# Neural Network-Based Score Estimation in Diffusion Models: Optimization and
Generalization

Yinbin Han, Meisam Razaviyayn & Renyuan Xu
Department of Industrial and Systems Engineering
University of Southern California
{yinbinha, razaviya<EMAIL_ADDRESS>
###### Abstract

Diffusion models have emerged as a powerful tool rivaling GANs in generating
high-quality samples with improved fidelity, flexibility, and robustness. A
key component of these models is to learn the score function through score
matching. Despite empirical success on various tasks, it remains unclear
whether gradient-based algorithms can learn the score function with a provable
accuracy. As a first step toward answering this question, this paper
establishes a mathematical framework for analyzing score estimation using
neural networks trained by gradient descent. Our analysis covers both the
optimization and the generalization aspects of the learning procedure. In
particular, we propose a parametric form to formulate the denoising score-
matching problem as a regression with noisy labels. Compared to the standard
supervised learning setup, the score-matching problem introduces distinct
challenges, including unbounded input, vector-valued output, and an additional
time variable, preventing existing techniques from being applied directly. In
this paper, we show that with a properly designed neural network architecture,
the score function can be accurately approximated by a reproducing kernel
Hilbert space induced by neural tangent kernels. Furthermore, by applying an
early-stopping rule for gradient descent and leveraging certain coupling
arguments between neural network training and kernel regression, we establish
the first generalization error (sample complexity) bounds for learning the
score function despite the presence of noise in the observations. Our analysis
is grounded in a novel parametric form of the neural network and an innovative
connection between score matching and regression analysis, facilitating the
application of advanced statistical and optimization techniques.

##  1 Introduction

Diffusion models excel in diverse generative tasks, spanning image, video, and
audio generation (Song & Ermon, 2019; Dathathri et al., 2019; Song et al.,
2020; Ho et al., 2020), often outperforming their contemporaries, including
GANs, VAEs, normalizing flows, and energy-based models (Goodfellow et al.,
2014; Kingma & Welling, 2013; Rezende & Mohamed, 2015; Zhao et al., 2016).

A typical diffusion model consists of two diffusion processes (Song et al.,
2020; Sohl-Dickstein et al., 2015; Ho et al., 2020): one moving forward in
time and the other moving backward. The forward process transforms a given
data sample into white noise in the limit by gradually injecting noise through
the diffusion term, while the backward process transforms noise to a sample
from the data distribution by sequentially removing the added noise. The
implementation of the backward process depends on the score function, defined
as the gradient of the logarithmic density, at each timestamp of the forward
process. In practice, however, the score function is unknown and one can only
access the true data distribution via finitely many samples. To ensure the
fidelity of the backward process in generating realistic samples, it is
essential to develop efficient methods to estimate the score function using
samples. This estimation is typically achieved through a process known as
score matching, employing powerful nonlinear functional approximations such as
neural networks.

Despite the empirical success, it is theoretically less clear whether a
gradient-based algorithm can train a neural network to learn the score
function. Existing theoretical work (De Bortoli et al., 2021; Chen et al.,
2022a; b; Lee et al., 2023; Chen et al., 2023a; Oko et al., 2023; Mei & Wu,
2023; Li et al., 2023; Chen et al., 2023b; Shah et al., 2023) predominantly
focuses on algorithm-agnostic properties of diffusion models such as score
approximation, score estimation, and distribution recovery, leaving the
theoretical performance of widely-used gradient-based algorithms an open
problem. This paper bridges this gap between theory and practice. Our
contributions are summarized as follows.

#### Our Work and Contributions.

This work investigates the training of a two-layer fully connected neural
network via gradient descent (GD) to learn the score function. First, we
propose a neural network-based parametric form for the score estimator based
on the score decomposition (see Lemma 3.1). This novel design transforms the
score-matching objective into a regression with noisy labels. To show the
trained neural network minimizes the excess risk of this regression problem,
we overcome three main challenges that do not exist in the traditional
supervised learning set-ups: 1) unbounded input, 2) vector-valued output, and
3) an additional time variable. To handle unbounded input, we employ a
truncation argument and control the tail behavior using the properties of
diffusion processes (see Lemma 3.3). Next, we establish a universal
approximation theorem with respect to the score function using the reproducing
kernel Hilbert space (RKHS), induced by the neural tangent kernel (NTK); see
Theorem 3.6. In addition, we leverage the recent NTK-based analysis of neural
networks to show the equivalence between neural network training and kernel
regression (see Theorem 3.9). Consequently, we transform the score matching
into a kernel regression problem. Furthermore, we propose a virtual dataset to
address the issue of target shifting caused by the approximation step. In the
presence of multi-output labels, a vector-valued localized Rademacher
complexity bound is utilized to control the prediction error of two kernel
regressions (see Theorem 3.10). Finally, we employ an early stopping rule for
the kernel regression to minimize the score-matching objective and provide the
generalization result (see Theorem 3.12).

To the best of our knowledge, this is the first work to establish sample
complexity bounds of GD-trained neural networks for score matching.
Specifically, our paper is the first to utilize NTK for establishing
theoretical results for diffusion models. Although the idea of NTK has been
used in many fields, the utilization of existing techniques in the structure
of diffusion models brings about its own significant challenges. Our analysis
is grounded in a novel parametric form of the neural network and an innovative
connection between score matching and regression analysis, facilitating the
application of advanced statistical and optimization techniques. In addition,
the building blocks of our results can be applied to other supervised learning
problems in non-standard forms (such as unbounded input and vector-valued
output), which goes beyond score-matching problems.

#### Related Literature.

Our work is related to three categories of prior work:

First, our framework is closely related to the recent study of diffusion
models. A line of work on this topic provides theoretical guarantees of
diffusion models for recovering data distribution, assuming access to an
accurate score estimator under $L^{2}$ or $L^{\infty}$ norm (De Bortoli et
al., 2021; Chen et al., 2022a; b; Lee et al., 2023; Shah et al., 2023; Li et
al., 2023; Chen et al., 2023b). These results offer only a partial
understanding of diffusion models as the score estimation part is omitted. To
our best knowledge, Chen et al. (2023a) and Oko et al. (2023) are the only
results that provide score estimation guarantees under $L^{2}$ norm, assuming
linear data structure or compactly supported data density. However, their
emphasis is on algorithm-agnostic analysis without evaluation of any specific
algorithms, creating a gap between theory and practical implementation. In
contrast, our work offers the first generalization error (sample complexity)
bounds for GD-trained neural networks.

Second, our techniques relate to the rich literature of deep learning theory.
Inspired by the framework of NTK introduced by Jacot et al. (2018), recently
Du et al. (2018; 2019); Allen-Zhu et al. (2019b); Zou et al. (2020) establish
linear convergence rate of neural networks for fitting random labels. One key
property of GD-trained neural networks is the so-called implicit
regularization of parameters. Namely, the minimizer of over-parameterized
neural networks is close to the random initialization. Combined with uniform
convergence results in statistical learning, this implicit regularization
leads to the generalization property of neural networks in the absence of
label noise (Arora et al., 2019a). However, none of these works delves into
the generalization ability of neural networks when confronted with noisy
labels. Kuzborskij & Szepesvári (2022) is the only work that attempts to study
the GD-trained neural networks with additive noise. To tackle the challenge
posed by the score matching, our approach and, consequently, our theoretical
results differ from the existing literature on deep learning theory for
supervised learning in three key aspects: 1) handling unbounded input, 2)
dealing with vector-valued output, and 3) incorporating an additional time
variable.

Lastly, our work is connected to a body of research focused on early stopping
rules in kernel regression. See Celisse & Wahl (2021) for a comprehensive
overview of this topic. Our work considers a multi-output extension of the
early stopping rule developed in Raskutti et al. (2014), which controls the
complexity of the predictor class based on empirical distribution.

##  2 Preliminaries and Problem Statement

In this section, we introduce the mathematical framework of diffusion models
Song et al. (2020).

#### Forward Process.

The forward process progressively injects noise into the original data
distribution. In the context of data generation, we have the flexibility to
work with any forward diffusion process of our choice. For the sake of
theoretical convenience, we adhere to the standard convention in the
literature (Song & Ermon, 2020; Ho et al., 2020) and focus on the Ornstein-
Ulhenbeck (OU) process, characterized by the following Stochastic Differential
Equation (SDE):

$\displaystyle{\rm d}X_{t}=-\frac{1}{2}g(t)X_{t}{\rm d}t+\sqrt{g(t)}{\rm
d}B_{t},\quad X_{0}\sim p_{0},$ (1)

where $g(t)>0$ is a deterministic weighting function; and
$\left(B_{t}\right)_{t\geq 0}$ is a standard $d$-dimensional Brownian motion.
Here, $p_{0}$ represents the unknown data distribution from which we have
access to only a limited number of samples. Our objective is to generate
additional samples from this distribution. Denoting the distribution of
$X_{t}$ at time $t$ by $p_{t}$, the explicit solution to (1) is given by

$\displaystyle X_{t}=e^{-\int_{0}^{t}\frac{1}{2}g(s){\rm
d}s}X_{0}+e^{-\int_{0}^{t}\frac{1}{2}g(s){\rm
d}s}\int_{0}^{t}e^{\int_{0}^{s}\frac{1}{2}g(u){\rm d}u}\sqrt{g(s)}{\rm
d}B_{s}.$

Consequently, the conditional distribution $X_{t}|X_{0}$ follows a multi-
variate Gaussian distribution $\mathcal{N}(\alpha(t)X_{0},h(t)I_{d})$ with
$\alpha(t)\coloneqq\exp\left(-\int_{0}^{t}\frac{1}{2}g(s){\rm d}s\right)$ and
$h(t)\coloneqq 1-\alpha^{2}(t)$. Furthermore, under mild assumptions, the OU
process converges exponentially to the standard Gaussian distribution. In
practice, the forward process (1) will terminate at a sufficiently large
timestamp $T>0$ such that the distribution $p_{T}$ is close to the standard
Gaussian distribution.

#### Backward Process.

By reversing the forward process in time, we obtain a process
$\bar{X}_{t}:=X_{T-t}$ (well defined under mild assumptions (Haussmann &
Pardoux, 1986)) that transforms white noise into samples from the targeted
data distribution, fulfilling the purpose of generative modeling. To start,
let us first define a backward process associated with (1):

$\displaystyle{\rm d}Y_{t}=\left(\frac{1}{2}g(T-t)Y_{t}+g(T-t)\nabla\log
p_{T-t}(Y_{t})\right){\rm d}t+\sqrt{g(T-t)}{\rm d}\bar{B}_{t},\quad Y_{0}\sim
q_{0}$ (2)

where $(\bar{B}_{t})_{t\geq 0}$ is another $d$-dimensional Brownian motion,
the score function $\nabla\log p_{t}(\cdot)$ is defined as the gradient of log
density of $X_{t}$, and $q_{0}$ is the initial distribution of the backward
process. If the score function is known at each time $t$ and if $q_{0}=p_{T}$,
under mild assumptions, the backward process $(Y_{t})_{0\leq t\leq T}$ has the
same distribution as the time-reserved process $(X_{T-t})_{0\leq t\leq
T}$—see(Föllmer, 2005; Cattiaux et al., 2021; Haussmann & Pardoux, 1986) for
details.

In practice, however, (2) cannot be directly used to generate samples from the
targeted data distribution as both the score function and the distribution
$p_{T}$ are unknown. To address this issue, it is common practice to replace
$p_{T}$ by the standard Gaussian distribution as the initial distribution of
the backward process. Then, we replace the ground-truth score $\nabla\log
p_{t}(x)$ by an estimator $s_{\theta}(x,t)$. The estimator $s_{\theta}$ is
parameterized (and learned) by a neural network. With these modifications, we
obtain an approximation of the backward process, which is practically
implementable:

$\displaystyle{\rm
d}Y_{t}=\left(\frac{1}{2}g(T-t)Y_{t}+g(T-t)s_{\theta}(Y_{t},t)\right){\rm
d}t+\sqrt{g(T-t)}{\rm d}W_{t},\quad Y_{0}\sim\mathcal{N}(0,I_{d}).$ (3)

To generate data using (3), SDE solvers or discrete-time approximation schemes
can be used (Chen et al., 2023a; Ho et al., 2020; Chen et al., 2022b; Song et
al., 2020; Chen et al., 2023b).

#### Score Matching.

To implement the backward process, we need to use samples to estimate the
score function. A natural choice is to minimize the $L^{2}$ loss between the
estimated and actual score:

$\displaystyle\min_{\theta}\;\;\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\lambda(t)\mathbb{E}\left[\left\|s_{\theta}(X_{t},t)-\nabla\log
p_{t}(X_{t})\right\|_{2}^{2}\right]{\rm d}t,$ (4)

where $\lambda(t)$ is the weighting function that captures time inhomogeneity
and $s_{\theta}$ is the estimator of the score function. Here, $T_{0}>0$ is
some small value to prevent the score function from blowing up and to
stabilize the training procedure (Song & Ermon, 2019; Chen et al., 2023a;
Vahdat et al., 2021). A major drawback of the score-matching loss (4) is its
intractability as $\nabla\log p_{t}$ cannot be computed based on the available
samples. Thus, instead of minimizing the loss in (4), one can equivalently
minimize the following denoising score matching as shown by Vincent (2011):

$\displaystyle\min_{\theta}\;\;\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\lambda(t)\mathbb{E}\left[\left\|s_{\theta}(X_{t},t)-\nabla\log
p_{t|0}(X_{t}|X_{0})\right\|_{2}^{2}\right]{\rm d}t.$ (5)

Here, $p_{t|0}(X_{t}|X_{0})$ denotes the conditional probability of $X_{t}$
given $X_{0}$. It is easy to show that the choice of our forward process in
(1) implies

$\displaystyle\nabla\log
p_{t|0}(X_{t}|X_{0})=\frac{\alpha(t)}{h(t)}X_{0}-\frac{X_{t}}{h(t)}.$ (6)

Now, we can plug (6) into (5) and learn the score function estimator. In
practice, however, the score function estimator is parameterized by a neural
network. Next, we discuss such a parameterization.

Algorithm 1 Sample Collection Procedure

1: Input: number of samples $N$ and a small value $T_{0}>0$

2: for $j=1,2,\dots,N$ do

3: Sample $X_{0,j}\sim p_{0}$

4: Sample $t_{j}\sim{\rm Unif}[T_{0},T]$

5: Sample $X_{t_{j}}\sim p_{t_{j}|0}({}\cdot{}|X_{0,j})$

6: end for

7: return $\left\\{\left(t_{j},X_{0,j},X_{t_{j}}\right)\right\\}_{j=1}^{N}$

#### Neural Network-Based Parameterization.

To parametrize the function $s_{\theta}$, we consider a two-layer ReLU neural
network $f_{{\bf W},a}=\left(f^{i}_{{\bf W},a}\right)_{i=1}^{d}$ of the
following form:

$\displaystyle f^{i}_{{\bf W},a}(x,t)$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\sigma(w_{r}^{\top}(x,t-T_{0})).$
(7)

Here, $(x,t)=(x^{1},\dots,x^{d},t)^{\top}\in\mathbb{R}^{d+1}$ is the input
vector, $w_{r}\in\mathbb{R}^{d+1}$ is a weight vector in the first layer,
$a_{r}^{i}\in\mathbb{R}$ is a weight vector in the second layer, and
$\sigma(\cdot)$ is the ReLU activation. The specific bias term $T_{0}$
introduced in the architecture plays an important role in the theoretical
analysis and also offers valuable insights for practical design. For ease of
exposition, we denote ${\bf W}=(w_{1},\dots,w_{m})\in\mathbb{R}^{(d+1)\times
m}$ and $a=[a_{r}^{i}]\in\mathbb{R}^{m\times d}$. We adopt the usual trick in
the over-parameterization literature (Cai et al., 2019; Wang et al., 2019;
Allen-Zhu et al., 2019b) with $a$ fixed throughout the training and only
updating ${\bf W}$. This seemingly shallow architecture poses significant
challenges when analyzing the convergence of gradient-based algorithms due to
its non-convex and non-smooth objective. On the other hand, its ability to
effectively approximate a diverse set of functions makes it a promising
starting point for advancing theoretical developments.

To train the neural network, we need to have samples measuring the “goodness-
of-fit” of the neural network. We use Algorithm 1 to generate $N$ i.i.d. data
samples. In particular, for each $j=1,\dots,N$, we first sample $X_{0,j}$ from
$p_{0}$ and a timestamp $t_{j}$ uniformly over the interval $[T_{0},T]$. Given
$X_{0,j}$ and $t_{j}$, we then sample $X_{t_{j}}$ from the Gaussian
distribution $p_{t_{j}|0}({}\cdot{}|X_{0,j})$. Given the output dataset
$S:=\left\\{(t_{j},X_{0,j},X_{t_{j}})\right\\}_{j=1}^{N}$, we train the neural
network by minimizing a quadratic loss:

$\displaystyle\min_{\bf W}\;\;\widehat{\mathcal{L}}({\bf
W})\coloneqq\frac{1}{2}\sum_{j=1}^{N}\left\|f_{\bf
W}(X_{t_{j}},t_{j})-X_{0,j}\right\|_{2}^{2}.$ (8)

Particularly, we perform the gradience descent (GD) update rule:

$\displaystyle
w_{r}(\tau+1)-w_{r}(\tau)=-\eta\frac{\partial\widehat{\mathcal{L}}(w_{r}(\tau))}{\partial
w_{r}(\tau)}$
$\displaystyle=-\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}(f_{\bf
W}^{i}(X_{t_{j}},t_{j})-X^{i}_{0,j})a_{r}^{i}(X_{t_{j}},t_{j}-T_{0})\mathbb{I}\left\\{w_{r}^{\top}(X_{t_{j}},t_{j}-T_{0})\geq
0\right\\},$ (9)

for $r=1,\ldots,m$. Here, $\eta>0$ is the learning rate. We initialize the
parameter $\bf W$ and $a$ according to the following neural tangent kernel
(NTK) regime (Jacot et al., 2018):

$\displaystyle w_{r}\sim\mathcal{N}(0,I_{d+1}),a_{r}^{i}\sim{\rm
Unif}\left\\{-1,+1\right\\},\quad\forall r\in[m]\text{ and }i\in[d].$

One can show that the training loss (8) is an empirical version of the
denoising score-matching loss defined in (5) under a carefully chosen
$s_{\theta}$. Correspondingly, the finite sample performance of $s_{\theta}$
w.r.t. (5) is referred to as generalization ability. We would like to remark
that the two-layer neural network parameterization has not been explored in
the literature for approximating score functions. While the work Chen et al.
(2023a) considered multi-layer neural networks for score approximation,
generalization, and distribution recovery; our work is complementary to them
as they did not analyze the optimization procedure and no specific learning
algorithm is considered in their work.

#### Neural Tangent Kernels.

For a two-layer ReLU neural network of the form (7), we follow (Jacot et al.,
2018) to introduce an associated neural tangent kernel
$K:\mathbb{R}^{d+1}\times\mathbb{R}^{d+1}\to\mathbb{R}^{d\times d}$ whose
$(i,k)$-th entry is defined as

$\displaystyle K^{ik}\left((x,t),\left(\tilde{x},\tilde{t}\right)\right)$
$\displaystyle\coloneqq\lim_{m\to\infty}\frac{1}{m}z^{\top}\tilde{z}\sum_{r=1}^{m}a_{r}^{i}a_{r}^{k}\mathbb{I}\left\\{w_{r}(0)^{\top}z\geq
0\right\\}\mathbb{I}\left\\{w_{r}(0)^{\top}\tilde{z}\geq 0\right\\}$
$\displaystyle=z^{\top}\tilde{z}\mathbb{E}\left[a_{1}^{i}a_{1}^{k}\mathbb{I}\left\\{w_{1}(0)^{\top}z\geq
0\right\\}\mathbb{I}\left\\{w_{1}(0)^{\top}\tilde{z}\geq 0\right\\}\right],$

where $z=(x,t-T_{0})$ and $\tilde{z}=(\tilde{x},\tilde{t}-T_{0})$. Here, the
expectation is taken over all the randomness of $a_{1}^{i}$, $a_{1}^{k}$ and
$w_{1}(0)$. Similarly, we define a scalar-valued NTK
$\kappa:\mathbb{R}^{d+1}\times\mathbb{R}^{d+1}\to\mathbb{R}$ associated with
each coordinate of the neural network:

$\displaystyle\kappa\Big{(}(x,t),(\tilde{x},\tilde{t})\Big{)}\coloneqq
z^{\top}\tilde{z}\mathbb{E}\left[\mathbb{I}\left\\{w_{1}(0)^{\top}z\geq
0\right\\}\mathbb{I}\left\\{w_{1}(0)^{\top}\tilde{z}\geq 0\right\\}\right],$

where $\mathbb{I}$ denotes the indicator function. From the definition of the
matrix-valued NTK, it is easy to see that $K$ is a diagonal matrix and in
particular,
$K\left((x,t),\left(\tilde{x},\tilde{t}\right)\right)=\kappa((x,t),(\tilde{x},\tilde{t}))I_{d},$
where $I_{d}$ is the $d$-dimensional identity matrix. Moreover, we let
$\mathcal{H}$ be the reproducing Hilbert space (RKHS) induced by the matrix-
valued NTK $K$ and $\mathcal{H}_{1}$ be the RKHS induced by the scalar-valued
NTK $\kappa$ (Jacot et al., 2018; Carmeli et al., 2010). Finally, given a
dataset $S$ and defining $z_{j}=(X_{t_{j}},t_{j}-T_{0})$, the Gram matrix $H$
of the kernel $K$ is defined as a $dN\times dN$ block matrix with

$\displaystyle H\coloneqq\begin{pmatrix}H_{11}&\cdots&H_{1N}\\\
\vdots&\ddots&\vdots\\\ H_{N1}&\cdots&H_{NN}\end{pmatrix},\quad
H^{ik}_{j\ell}\coloneqq
z_{j}^{\top}z_{\ell}\mathbb{E}\left[a_{1}^{i}a_{1}^{k}\mathbb{I}\left\\{z_{j}^{\top}w_{1}(0)\geq
0,z_{\ell}^{\top}w_{1}(0)\geq 0\right\\}\right].$ (10)

##  3 Main Results

This section introduces our main theoretical results. We first propose a
parametric form of $s_{\theta}$ to simplify the score-matching loss in (4).
Next, we show that the empirical version of DSM (5) is indeed equivalent to
the quadratic loss defined in (8). Finally, we provide a decomposition of an
upper bound on the loss function into four terms: a coupling term, a label
mismatch term, a term related to early stopping, and an approximation error.
These terms are carefully analyzed later.

To motivate our parametric form of $s_{\theta}$, we start by the following
decomposition of the score function:

######  Lemma 3.1.

The score function $\nabla\log p_{t}(x)$ admits the following decomposition:

$\nabla\log
p_{t}(x)=\frac{\alpha(t)}{h(t)}\mathbb{E}\left[X_{0}|X_{t}=x\right]-\frac{x}{h(t)}.$
(11)

The proof, which follows the Gaussianity of the transition kernel $p_{t|0}$,
is deferred to the appendix. A similar decomposition has been proved in (Chen
et al., 2023a, Lemma 1) for data with linear structure, and in Li et al.
(2023) for discrete time analysis and the concurrent work (Mei & Wu, 2023).
Compared to the expression of $\nabla\log p_{t|0}(x_{t}|x_{0})$ computed in
(6), we replace $X_{0}$ by $\mathbb{E}\left[X_{0}|X_{t}\right]$ to obtain the
ground-truth score function in (11). Consequently, we call $X_{0}$ the noisy
label and $\mathbb{E}\left[X_{0}|X_{t}\right]$ the true label. We also make
the following assumption on the diffusion models (1).

######  Assumption 3.2.

The target density function $p_{0}$ has a compact support with
$\left\|X_{0}\right\|_{2}\leq D$ almost surely, for some constant $D>0$.

Assumption 3.2 is satisfied in most practical settings, including the
generation of image, video, or audio. This assumption simplifies the
subsequent analysis and can be relaxed to the sub-Gaussian tail assumption.
Next, we propose the parametric form of $s_{\theta}$ and $\lambda(t)$ in the
score-matching loss (4):

$\displaystyle s_{{\bf W},a}(x,t)=\frac{\alpha(t)}{h(t)}\Pi_{D}(f_{{\bf
W},a}(x,t))-\frac{x}{h(t)},\quad\textit{with}\,\,\lambda(t)=\frac{h(t)^{2}}{\alpha(t)^{2}},$

where $\Pi_{D}$ is the projection operator to the $L^{2}$-ball with radius $D$
centered at zero. With the choice of $s_{{\bf W},a}$ and $\lambda(t)$
specified above, the score-matching loss (4) becomes

$\displaystyle\min_{{\bf
W}}\;\;\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{{\bf
W},a}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\right]{\rm d}t,$ (12)

in which we define the target function as
$f_{*}(x,t)\coloneqq\mathbb{E}\left[X_{0}|X_{t}=x\right]$ and the expectation
is taken over $X_{t}$. Given that only ${\bf W}$ is updated during
optimization, in what follows, we omit $a$ in the subscript of the neural
network. Our loss function (12) is also supported by empirical studies (Ho et
al., 2020). In addition, (12) can be viewed as a regression task with noisy
labels. In what follows, we will show that neural networks trained on noisy
labels generalize well w.r.t. (12).

One major challenge in the theoretical analysis, which distinguishes us from
the standard supervised learning problems, is the unboundedness of the input
$X_{t}$ in the objective function. To overcome this challenge, we employ a
truncation argument with a threshold $R$:

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{\bf
W}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\right]{\rm d}t$ $\displaystyle=\;$
$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{\bf
W}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq
R\right\\}\right]{\rm d}t$ (13)
$\displaystyle\qquad+\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{\bf
W}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}>R\right\\}\right]{\rm
d}t.$ (14)

The next lemma controls the tail behavior in (14).

######  Lemma 3.3.

Suppose Assumption 3.2 holds. Then, uniformly over all $\bm{W}$, it holds that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{\bf
W}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}>R\right\\}\right]{\rm
d}t=\mathcal{O}({R}^{d-2}e^{-R^{2}/4}).$

Lemma 3.3 states the term (14) is exponentially small in the threshold $R$.
Thus, it suffices to focus on the loss (13) over the ball with radius $R$.
Inspired by Kuzborskij & Szepesvári (2022) for learning Lipschitz functions,
we upper bound (13) by the following decomposition at each iteration $\tau$:

$\displaystyle\frac{1}{4(T-T_{0})}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}\left(f_{\bf
W(\tau)}(X_{t},t)\right)-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$ $\displaystyle\leq\;$
$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}\left(f_{\bf
W(\tau)}(X_{t},t)\right)-f^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$ (coupling)
$\displaystyle\qquad+\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|f^{K}_{\tau}(X_{t},t)-\tilde{f}^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$ (label mismatch)
$\displaystyle\qquad+\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\tilde{f}^{K}_{\tau}(X_{t},t)-f_{\mathcal{H}}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$ (early stopping)
$\displaystyle\qquad+\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|f_{\mathcal{H}}(X_{t},t)-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t.$ (approximation)

The first term is the coupling error between neural networks $f_{\bf W(\tau)}$
and a function $f_{\tau}^{K}$ defined as:

$\displaystyle
f_{\tau}^{K}(x,t)=\sum_{j=1}^{N}K((X_{t_{j}},t_{j}),(x,t))\gamma_{j}(\tau),\quad\gamma(\tau+1)=\gamma(\tau)-\eta(H\gamma(\tau)-y),$

where $\gamma(0)$ is initialized in (69). The fourth term is the approximation
error of the target function $f_{*}$ by a function $f_{\mathcal{H}}$ in the
RKHS $\mathcal{H}$. These two terms transforms the training of neural networks
into a problem of kernel regression. To learn the function $f_{\mathcal{H}}$,
we define an auxiliary function $\tilde{f}_{\tau}^{K}$ of the same functional
form as ${f}_{\tau}^{K}$, but trained on a different dataset
$\tilde{S}=\\{(t_{j},\tilde{X}_{0,j},X_{t_{j}})\\}_{j=1}^{N}$ with

$\displaystyle\tilde{X}_{0,j}\coloneqq
f_{\mathcal{H}}(X_{t_{j}},t_{j})+\varepsilon_{j},\quad\varepsilon_{j}\coloneqq
X_{0,j}-f_{*}(X_{t_{j}},t_{j}).$

Finally, we control the third term in the above decomposition by the early
stopping rule, which is a classical technique in the statistical learning
literature (Raskutti et al., 2014; Wei et al., 2017).

###  3.1 Approximation

We start by analyzing the approximation term in our decomposition. This
subsection focuses on the approximation error of the target function $f_{*}$
by a function in the RKHS $\mathcal{H}$ induced by the NTK $K$. We start with
a regularity assumption on the coefficient $g(t)$ in the OU process.

######  Assumption 3.4.

The function $g$ is almost everywhere continuous and bounded on $[0,\infty)$.

Assumption 3.4 imposes a minimal requirement to guarantee that both
$\alpha(t)$ and $h(t)$ are well defined at each timestamp $t\geq 0$. In
addition, the boundedness assumption of $g$ is used to establish the Lipschitz
property of the score function with respect to $t$ in the literature (Chen et
al., 2023a; 2022a; 2022b). We also make the following smoothness assumption on
the target function $f_{*}$.

######  Assumption 3.5.

For all $(x,t)\in\mathbb{R}^{d}\times[T_{0},\infty)$, the function
$f_{*}(x,t)$ is $\beta_{x}$-Lipschitz in $x$, i.e.,
$\left|f_{*}(x,t)-f_{*}(x^{\prime},t)\right|_{2}\leq\beta_{x}\left\|x-x^{\prime}\right\|_{2}$.

Assumption 3.5 implies the score function is Lipschitz w.r.t. the input $x$.
This assumption is standard in the literature (Chen et al., 2022a; b; 2023a).
Yet the Lipschitz continuity in Assumption 3.5 is only imposed on the
regression function $f_{*}$, which is a consequence of the score
decomposition. To justify Assumption 3.5, we provide an upper bound of the
Lipschitz constant $\beta_{x}$ in Lemma G.1. The following theorem states a
universal approximation theorem of using RKHS for score functions.

######  Theorem 3.6 (Universal Approximation of Score Function).

Suppose Assumptions 3.2, 3.4 and 3.5 hold. Let $R\geq T-T_{0}$ and
$R_{\mathcal{H}}$ be larger than a constant $c_{1}$ 111The constant $c_{1}$
equals to $C(d+1,0)$ in (Bach, 2017, Proposition 6) that depends only on $d$.
There exists a function $f_{\mathcal{H}}\in\mathcal{H}$ such that
$\left\|f_{\mathcal{H}}\right\|_{\mathcal{H}}^{2}\leq dR_{\mathcal{H}}$ and

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|f_{\mathcal{H}}(X_{t},t)-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t\leq dA^{2}(R_{\mathcal{H}},{R}),$

where $A(R_{\mathcal{H}},R)\coloneqq
c_{1}\Lambda(R)\left(\frac{\sqrt{R_{\mathcal{H}}}}{\Lambda(R)}\right)^{-\frac{2}{d}}\log\left(\frac{\sqrt{R_{\mathcal{H}}}}{\Lambda(R)}\right)$
and $\Lambda(R)=O(\sqrt{d}R^{2})$.

Theorem 3.6 provides an approximation of the target function by the RHKS under
the $L^{2}$ norm. For each given $R$, we can choose $R_{\mathcal{H}}$ large
enough such that $A(R_{\mathcal{H}},R)$ is arbitrarily small. Let us provide a
sketch of the proof of Theorem 3.6. We first construction an auxiliary
function $\tilde{f}_{*}(x,t)\coloneqq f_{*}(x,\left|t\right|+T_{0})$. One can
show $\tilde{f}_{*}$ is Lipschitz continuos in $(x,t)\in\mathbb{R}^{d+1}$.
Then for each coordinate $i$, we apply an approximation result on RKHS for
Lipschitz functions over a $L^{\infty}$-ball (cf. Lemma C.2) to find a
function that approximates $\tilde{f}_{*}^{i}$ well. Since NTK is not a
translation invariant kernel, we need to construct a shifted NTK such that
$f_{\mathcal{H}}^{i}\in\mathcal{H}_{1}$ is close to $f_{*}^{i}$ after
translation. The rest is to show that
$f_{\mathcal{H}}=(f^{i}_{\mathcal{H}})_{i=1}^{d}$ lies in the vector-valued
RKHS $\mathcal{H}$. The complete detailed proof of Theorem 3.6 is deferred to
the appendix.

###  3.2 Coupling

This subsection provides a coupling argument to control the error between the
neural network training and the kernel regression. We make the following
assumption on the dataset $S$:

######  Assumption 3.7.

There exists a $\delta_{1}(\Delta,R)\in[0,1)$ such that $\delta_{1}\rightarrow
0$ when $R\rightarrow\infty$ and $\Delta\to 0$, and we have the following
holds with probability at least  $1-\delta_{1}(\Delta,R)$,

$t_{j}\in[T_{0}+\Delta,T]\text{ and }\left\|X_{t_{j}}\right\|_{2}\leq
R\;\text{ for all sample }j.$

Assumption 3.7, which imposes regularity conditions on the input data
$(t_{j},X_{t_{j}})$,  can be verified by utilizing the tail property of
$X_{t_{j}}$ and the uniform sampling scheme for $t_{j}$; see Lemma G.2 in the
appendix. The next assumption is on the minimum eigenvalue of the Gram matrix
$H$ of the kernel $K$ and is standard in literature (Du et al., 2018; Bartlett
et al., 2021; Nguyen et al., 2021).

######  Assumption 3.8.

There exists a constant $\lambda_{0}(d,N)\geq 1$ such that the smallest
eigenvalue $\lambda_{\min}(H)\geq\lambda_{0}$ with probability at least
$1-\delta_{2}(d,N)$ with $\delta_{2}\to 0$ as $d$ increases, where $N={\rm
Poly}(d)$.

As shown in the literature of deep learning theory (Allen-Zhu et al., 2019a;
Arora et al., 2019a; Liu et al., 2022), the Gram matrix $H$ is a fundamental
quantity that determines the convergence rate of neural network optimizations.
We also remark that Assumption 3.8 is usually satisfied with a sample-
dependent lower bound $\lambda_{0}$; see Lemma G.3 in the appendix for a
justification and see also Nguyen et al. (2021)) for analysis of scalar NTK.
Now we are ready to state our main theorem for the coupling error. Let
$C_{\min}=\Delta$ and $C_{\max}=\sqrt{R^{2}+(T-T_{0})^{2}}$.

######  Theorem 3.9 (Coupling Error).

Suppose Assumptions 3.2, 3.7 and 3.8 hold. If we set
$m=\Omega\left(\frac{(dN)^{6}C_{\max}^{6}}{\lambda_{0}^{10}\delta^{3}C_{\min}^{2}}\right)$,
initialize $w_{r}\sim\mathcal{N}(0,I_{d+1})$ and $a_{r}^{i}\sim{\rm
Unif}\left\\{-1,1\right\\}$ i.i.d., initialize $\gamma(0)$ properly, and set
$\eta=\mathcal{O}\left(\frac{\lambda_{0}}{(dN)^{2}C_{\max}^{4}}\right)$, then
with probability at least $1-\delta$, for all $\tau\geq 0$ and $r=1,\dots,m$
simultaneously, we have

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}\left(f_{\bf
W(\tau)}(X_{t},t)\right)-f^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$ $\displaystyle\leq\frac{4\Delta
D^{2}}{T-T_{0}}+\tilde{\mathcal{O}}\left(\frac{d^{10}N^{9}C_{\max}^{12}}{\sqrt{m}\lambda_{0}^{2}\delta^{4}C_{\min}^{2}}\right).$

The proof is deferred to the appendix. Theorem 3.9 controls the error between
the neural network training and the kernel regression. One can choose $m={\rm
Poly}(d,N,R,\Delta,\lambda_{0},\delta)$ and optimize over $R$ and $\Delta$ to
make the error term small. For each fixed input data sample, (Arora et al.,
2019b, Theorem 3.2) shows that the coupling error is small with high
probability. Our analysis improves this result by showing that the $L^{2}$
coupling error also remains small with high probability. To prove Theorem 3.9,
we first show that the training loss (8) converges with a linear rate (cf.
Theorem D.1). Next, we show that $f_{{\bf W}(\tau)}$ performs similarly as a
linearized function $f_{{\bf\bar{W}}(\tau)}^{\rm lin}$ at each iteration
$\tau$. Finally, we argue that the $L^{2}$ loss between the
$f_{{\bf\bar{W}}(\tau)}^{\rm lin}$ and $f_{\tau}^{K}$ is small because of the
concentration of kernels and a carefully chosen initialization $\gamma(0)$
depending on the neural network initialization.

###  3.3 Label Mismatch

In this subsection, we upper bound the error term induced by the label
mismatch. Recall that $f_{\tau}^{K}$ is trained by kernel regression on the
dataset $S$ while $\tilde{f}_{\tau}^{K}$ is trained on the dataset
$\tilde{S}$. We can control the error induced by the label mismatch in the
following theorem.

######  Theorem 3.10 (Label Mismatch).

Suppose Assumptions 3.7 and 3.8 hold. If we initialize both $f_{0}^{K}$ and
$\tilde{f}_{0}^{K}$ properly, then with probability at least $1-\delta$ it
holds simultaneously for all $\tau$ that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|f_{\tau}^{K}(x,t)-\tilde{f}_{\tau}^{K}(x,t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq
R\right\\}\right]dt\leq
dA(R_{\mathcal{H}},R)+C_{0}\left(\sqrt{dA(R_{\mathcal{H}},R)\Gamma_{\delta}}+\Gamma_{\delta}\right),$

where

$\displaystyle\Gamma_{\delta}$
$\displaystyle\coloneqq\left(2d\left(d\log^{3/2}\left(\frac{eC_{\max}(dN)^{3/2}A(R_{\mathcal{H}},R)}{\lambda_{0}}\right)\frac{A(R_{\mathcal{H}},R)C_{\max}}{\lambda_{0}}\right)+\frac{1}{\sqrt{N}}\right)^{2}$
$\displaystyle\qquad+\frac{d^{2}A^{2}(R_{\mathcal{H}},R)C_{\max}^{2}}{\lambda_{0}^{2}}\left(\log(1/\delta)+\log\left(\log
N\right)\right),$

and $C_{0}$ is a constant defined in (Reeve & Kaban, 2020, Theorem 1).

Theorem 3.10 links the error between $f_{\tau}^{K}$ and $\tilde{f}_{\tau}^{K}$
to the approximation error $A(R_{\mathcal{H}},R)$. The proof of Theorem 3.10
consists of two parts. We first utilize the kernel regression structure to
show that the predictions of $f_{\tau}^{K}$ and $\tilde{f}_{\tau}^{K}$ are
similar over all the samples $(t_{j},X_{t_{j}})$. Next, we apply the vector-
valued localized Rademacher complexity (cf. Lemma E.2) to show that the
performance of these two functions is also close on the population loss. We
defer the proof of Theorem 3.10 to the appendix.

###  3.4 Early Stopping and the final result

Given the function $\tilde{f}_{\tau}^{K}$ trained on the data set
$\tilde{S}=\left\\{\left(t_{j},\tilde{X}_{0,j},X_{t_{j}}\right)\right\\}_{j=1}^{N}$
and the target function $f_{\mathcal{H}}\in\mathcal{H}$ that generates the
virtual label $\tilde{X}_{0,j}$, we transform the score matching problem to a
classical kernel regression problem. The next technical assumption allows us
to reduce the excess risk bound for the early-stopped GD learning in RKHS to
the excess risk bound for learning Lipschitz functions.

######  Assumption 3.11.

Fix any $f_{\mathcal{H}}\in\mathcal{H}$ with
$\left\|f_{\mathcal{H}}\right\|_{\mathcal{H}}^{2}\leq R_{\mathcal{H}}$ and
assume labels are generated as
$\tilde{X}_{0,j}=f_{\mathcal{H}}(X_{t_{j}},t_{j})+\varepsilon_{j}$. Suppose
$\tilde{f}_{\widehat{T}}^{K}$ is obtained by GD-trained kernel regression with
the number of iterations $\widehat{T}$. We assume that there exists $\epsilon$
such that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\tilde{f}^{K}_{\widehat{T}}(X_{t},t)-f_{\mathcal{H}}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t\leq\epsilon(N,\widehat{T}),$

and $\epsilon(N,\widehat{T})\to 0$ as $N\to\infty$.

Here, $\widehat{T}$ is a data-dependent early stopping rule to control the
excess risk of kernel regression. For supervised learning with noisy labels,
early stopping rule for GD is necessary to minimize the excess risk (Hu et
al., 2021; Bartlett & Mendelson, 2002; Li et al., 2020). Assumption 3.11 can
be satisfied by an extension of classical early stopping rules. For the case
of scalar-valued kernel regression, see (Raskutti et al., 2014). Next, we
provide a generalization result for the score estimator:

######  Theorem 3.12 (Score Estimation and Generalization).

Suppose Assumption 3.2, 3.4, 3.5, 3.7, 3.8 hold and we set $m$ and $\eta$ as
prescribed in Theorem 3.9. Moreover, suppose $\widehat{T}$ satisfies
Assumption 3.11 with corresponding $\epsilon(N,\widehat{T})$. Then for any
large enough $R$ and $R_{\mathcal{H}}$ and small enough $\Delta$, with
probability at least $1-\delta$, it holds that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\lambda(t)\mathbb{E}\left[\left\|s_{{\bf
W}(\widehat{T})}(X_{t},t)-\nabla\log p_{t}(X_{t})\right\|_{2}^{2}\right]{\rm
d}t$
$\displaystyle\leq\mathcal{O}(R^{d-2}e^{-R^{2}/4})+4dA^{2}(R_{\mathcal{H}},R)+\frac{16\Delta
D^{2}}{T-T_{0}}+\tilde{\mathcal{O}}\left(\frac{d^{10}N^{9}C_{\max}^{12}}{\sqrt{m}\lambda_{0}^{2}\delta^{4}C_{\min}^{2}}\right)$
$\displaystyle\qquad+4dA(R_{\mathcal{H}},R)+4C_{0}\left(\sqrt{dA(R_{\mathcal{H}},R)\Gamma_{\delta}}+\Gamma_{\delta}\right)+4\epsilon(N,\widehat{T}),$

where $A(R_{\mathcal{H}},R)$ is defined in Theorem 3.6 and $\Gamma_{\delta}$
is given in Theorem 3.10.

Theorem 3.12 shows that early-stopped neural network $s_{{\bf
W}(\widehat{T})}$ learns the score function $\nabla\log p_{t}$ well in the
$L^{2}$ sense over the interval $[T_{0},T]$. To the best of our knowledge,
this is the first algorithm-based analysis for score estimation with neural
network parameterization. Combined with recent findings in the distribution
recovery property of diffusion models, we are the first to obtain an end-to-
end guarantee with a provably efficient algorithm for diffusion models. The
proof of Theorem 3.12, which relies on Lemma 3.3 and Theorems 3.6, 3.9 and
3.10, can be found in the appendix.

##  4 Conclusion and Discussions

In this paper, we establish the first algorithm-based analysis of neural
network-based score estimation in diffusion models. We demonstrate that GD-
trained overparametrized neural networks can learn the ground truth score
function with a sufficient number of samples when an early stopping rule is
applied. Our work investigates all three aspects of the score estimation task:
approximation, optimization, and generalization. The analytical framework laid
out in this paper sheds light on the understanding of diffusion models and
inspires innovative architecture design.

In addition, our work leaves several interesting questions for future
investigation. For instance, the dimension dependency in our convergence
results remains sub-optimal. To address this, one approach is to consider the
manifold structure of the data distribution, such as the linear subspace
assumption as suggested by Chen et al. (2023a) and Oko et al. (2023). Another
direction is to understand the role of neural network architectures like
U-nets and transformers in the implementation of diffusion models for image
tasks. Finally, an extension of training algorithms to stochastic gradient
descent (SGD) and Adam would be of independent interest.

## References

  * Allen-Zhu et al. (2019a) Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang.  Learning and generalization in overparameterized neural networks, going beyond two layers.  _Advances in neural information processing systems_ , 32, 2019a.
  * Allen-Zhu et al. (2019b) Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In _International conference on machine learning_ , pp. 242–252. PMLR, 2019b.
  * Arora et al. (2019a) 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, 2019a.
  * Arora et al. (2019b) Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  _Advances in neural information processing systems_ , 32, 2019b.
  * Bach (2017) Francis Bach.  Breaking the curse of dimensionality with convex neural networks.  _The Journal of Machine Learning Research_ , 18(1):629–681, 2017.
  * Bartlett & Mendelson (2002) Peter L Bartlett and Shahar Mendelson.  Rademacher and Gaussian complexities: Risk bounds and structural results.  _Journal of Machine Learning Research_ , 3(Nov):463–482, 2002.
  * Bartlett et al. (2021) Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin.  Deep learning: a statistical viewpoint.  _Acta numerica_ , 30:87–201, 2021.
  * Cai et al. (2019) Qi Cai, Zhuoran Yang, Jason D Lee, and Zhaoran Wang.  Neural temporal-difference learning converges to global optima.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Carmeli et al. (2010) Claudio Carmeli, Ernesto De Vito, Alessandro Toigo, and Veronica Umanitá.  Vector valued reproducing kernel hilbert spaces and universality.  _Analysis and Applications_ , 8(01):19–61, 2010.
  * Cattiaux et al. (2021) Patrick Cattiaux, Giovanni Conforti, Ivan Gentil, and Christian Léonard.  Time reversal of diffusion processes under a finite entropy condition.  _arXiv preprint arXiv:2104.07708_ , 2021.
  * Celisse & Wahl (2021) Alain Celisse and Martin Wahl.  Analyzing the discrepancy principle for kernelized spectral filter learning algorithms.  _The Journal of Machine Learning Research_ , 22(1):3498–3556, 2021.
  * Chen et al. (2022a) Hongrui Chen, Holden Lee, and Jianfeng Lu.  Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions.  _arXiv preprint arXiv:2211.01916_ , 2022a.
  * Chen et al. (2023a) Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang.  Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data.  _arXiv preprint arXiv:2302.07194_ , 2023a.
  * Chen et al. (2022b) Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang.  Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions.  _arXiv preprint arXiv:2209.11215_ , 2022b.
  * Chen et al. (2023b) Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim.  The probability flow ode is provably fast.  _arXiv preprint arXiv:2305.11798_ , 2023b.
  * Dathathri et al. (2019) Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu.  Plug and play language models: A simple approach to controlled text generation.  _arXiv preprint arXiv:1912.02164_ , 2019.
  * De Bortoli et al. (2021) Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet.  Diffusion Schrödinger bridge with applications to score-based generative modeling.  _Advances in Neural Information Processing Systems_ , 34:17695–17709, 2021.
  * Du et al. (2019) Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai.  Gradient descent finds global minima of deep neural networks.  In _International conference on machine learning_ , pp. 1675–1685. PMLR, 2019.
  * Du et al. (2018) Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  _arXiv preprint arXiv:1810.02054_ , 2018.
  * Folland (1999) Gerald B Folland.  _Real analysis: modern techniques and their applications_ , volume 40.  John Wiley & Sons, 1999.
  * Föllmer (2005) Hans Föllmer.  An entropy approach to the time reversal of diffusion processes.  In _Stochastic Differential Systems Filtering and Control: Proceedings of the IFIP-WG 7/1 Working Conference Marseille-Luminy, France, March 12–17, 1984_ , pp. 156–163. Springer, 2005.
  * Goodfellow et al. (2014) Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  _Advances in neural information processing systems_ , 27, 2014.
  * Haussmann & Pardoux (1986) Ulrich G Haussmann and Etienne Pardoux.  Time reversal of diffusions.  _The Annals of Probability_ , pp. 1188–1205, 1986.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Hu et al. (2021) Tianyang Hu, Wenjia Wang, Cong Lin, and Guang Cheng.  Regularization matters: A nonparametric perspective on overparametrized neural network.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 829–837. PMLR, 2021.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Kuzborskij & Szepesvári (2022) Ilja Kuzborskij and Csaba Szepesvári.  Learning Lipschitz functions by GD-trained shallow overparameterized ReLU neural networks.  _arXiv preprint arXiv:2212.13848_ , 2022.
  * Lee et al. (2023) Holden Lee, Jianfeng Lu, and Yixin Tan.  Convergence of score-based generative modeling for general data distributions.  In _International Conference on Algorithmic Learning Theory_ , pp. 946–985. PMLR, 2023.
  * Li et al. (2023) Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi.  Towards faster non-asymptotic convergence for diffusion-based generative models.  _arXiv preprint arXiv:2306.09251_ , 2023.
  * Li et al. (2020) Mingchen Li, Mahdi Soltanolkotabi, and Samet Oymak.  Gradient descent with early stopping is provably robust to label noise for overparameterized neural networks.  In _International conference on artificial intelligence and statistics_ , pp. 4313–4324. PMLR, 2020.
  * Liu et al. (2022) Chaoyue Liu, Libin Zhu, and Mikhail Belkin.  Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.  _Applied and Computational Harmonic Analysis_ , 59:85–116, 2022.
  * Mei & Wu (2023) Song Mei and Yuchen Wu.  Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models.  _arXiv preprint arXiv:2309.11420_ , 2023.
  * 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.
  * Oko et al. (2023) Kazusato Oko, Shunta Akiyama, and Taiji Suzuki.  Diffusion models are minimax optimal distribution estimators.  _arXiv preprint arXiv:2303.01861_ , 2023.
  * Qi & Mei (1999) Feng Qi and Jia-Qiang Mei.  Some inequalities of the incomplete Gamma and related functions.  _Zeitschrift für Analysis und ihre Anwendungen_ , 18(3):793–799, 1999.
  * Raskutti et al. (2014) Garvesh Raskutti, Martin J Wainwright, and Bin Yu.  Early stopping and non-parametric regression: an optimal data-dependent stopping rule.  _The Journal of Machine Learning Research_ , 15(1):335–366, 2014.
  * Reeve & Kaban (2020) Henry Reeve and Ata Kaban.  Optimistic bounds for multi-output learning.  In _International Conference on Machine Learning_ , pp. 8030–8040. PMLR, 2020.
  * Rezende & Mohamed (2015) Danilo Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International conference on machine learning_ , pp. 1530–1538. PMLR, 2015.
  * Shah et al. (2023) Kulin Shah, Sitan Chen, and Adam Klivans.  Learning mixtures of Gaussians using the DDPM objective.  _arXiv preprint arXiv:2307.01178_ , 2023.
  * Sohl-Dickstein et al. (2015) Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International conference on machine learning_ , pp. 2256–2265. PMLR, 2015.
  * Song & Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in neural information processing systems_ , 32, 2019.
  * Song & Ermon (2020) Yang Song and Stefano Ermon.  Improved techniques for training score-based generative models.  _Advances in neural information processing systems_ , 33:12438–12448, 2020.
  * Song et al. (2020) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020.
  * Vahdat et al. (2021) Arash Vahdat, Karsten Kreis, and Jan Kautz.  Score-based generative modeling in latent space.  _Advances in Neural Information Processing Systems_ , 34:11287–11302, 2021.
  * Vincent (2011) Pascal Vincent.  A connection between score matching and denoising autoencoders.  _Neural computation_ , 23(7):1661–1674, 2011.
  * Wang et al. (2019) Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang.  Neural policy gradient methods: Global optimality and rates of convergence.  _arXiv preprint arXiv:1909.01150_ , 2019.
  * Wei et al. (2017) Yuting Wei, Fanny Yang, and Martin J Wainwright.  Early stopping for kernel boosting algorithms: A general analysis with localized complexities.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Zhao et al. (2016) Junbo Zhao, Michael Mathieu, and Yann LeCun.  Energy-based generative adversarial network.  _arXiv preprint arXiv:1609.03126_ , 2016.
  * Zou et al. (2020) Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu.  Gradient descent optimizes over-parameterized deep ReLU networks.  _Machine learning_ , 109:467–492, 2020.

##  Appendix A Proof of Lemma 3.1

###### Proof.

Recall that the density function $p_{t}$ can be written as

$\displaystyle p_{t}(x)=\int p_{t|0}(x|x_{0})p_{0}(x_{0}){\rm d}x_{0},$

where the transition kernel $p_{t|0}(x|x_{0})=(2\pi
h(t))^{-d/2}\exp\left(-\frac{1}{2h(t)}\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}\right)$.
Utilizing the dominated convergence theorem leads to

$\displaystyle\nabla\log p_{t}(x)$ $\displaystyle=\frac{\nabla\int
p_{t|0}(x|x_{0})p_{0}(x_{0}){\rm d}x_{0}}{p_{t}(x)}$
$\displaystyle=\frac{\left(2\pi
h(t)\right)^{-d/2}\int-\frac{x-\alpha(t)x_{0}}{h(t)}\exp\left(-\frac{\left\|x-\alpha(t)x_{0}\right\|^{2}}{2h(t)}\right)p_{0}(x_{0}){\rm
d}x_{0}}{p_{t}(x)}$
$\displaystyle=\int-\frac{x-\alpha(t)x_{0}}{h(t)}\cdot\frac{p_{t|0}(x|x_{0})p_{0}(x_{0})}{p_{t}(x)}{\rm
d}x_{0}$ $\displaystyle=\int-\frac{x-\alpha(t)x_{0}}{h(t)}\cdot
p_{0|t}(x_{0}|x){\rm d}x_{0}$
$\displaystyle=\mathbb{E}\left[\left.\frac{\alpha(t)X_{0}-X_{t}}{h(t)}\right|X_{t}=x\right]$
$\displaystyle=\frac{\alpha(t)}{h(t)}\mathbb{E}\left[X_{0}|X_{t}=x\right]-\frac{x}{h(t)},$

which completes the proof. ∎

##  Appendix B Proof of Lemma 3.3

###### Proof.

The proof essentially follows the ideas in Chen et al. (2023a). First, note
that

$\displaystyle p_{t|0}(x_{t}|x_{0})$ $\displaystyle=(2\pi
h(t))^{-d/2}\exp\left(-\frac{1}{2h(t)}\left\|x_{t}-\alpha(t)x_{0}\right\|_{2}^{2}\right)$
$\displaystyle\leq(2\pi
h(t))^{-d/2}\exp\left(-\frac{1}{2h(t)}\left(\frac{1}{2}\left\|x_{t}\right\|_{2}^{2}-\alpha^{2}(t)\left\|x_{0}\right\|_{2}^{2}\right)\right)$
$\displaystyle\leq(2\pi
h(t))^{-d/2}\exp\left(-\frac{1}{2h(t)}\left(\frac{1}{2}\left\|x_{t}\right\|_{2}^{2}-\left\|x_{0}\right\|_{2}^{2}\right)\right).$
(15)

Denote the expectation with respect to the marginal distribution of $X_{0}$ as
$\mathbb{E}_{X_{0}}$. With inequality (15), we can bound

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{\bf
W}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}>R\right\\}\right]{\rm
d}t$
$\displaystyle\leq\frac{4D^{2}}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}_{X_{0}}\left[\int_{\left\|x_{t}\right\|_{2}\geq{R}}p_{t|0}(x_{t}|X_{0}){\rm
d}x_{t}\right]{\rm d}t$
$\displaystyle\leq\frac{4D^{2}}{T-T_{0}}\int_{T_{0}}^{T}(2\pi
h(t))^{-d/2}\mathbb{E}_{X_{0}}\left[\exp\left(\frac{\left\|X_{0}\right\|_{2}^{2}}{2h(t)}\right)\right]\left(\int_{\left\|x_{t}\right\|\geq{R}}\exp\left(-\frac{\left\|x_{t}\right\|^{2}}{4h(t)}\right){\rm
d}x_{t}\right){\rm d}t$
$\displaystyle=O\left(\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\int_{\left\|x_{t}\right\|\geq{R}}\exp\left(-\frac{\left\|x_{t}\right\|^{2}}{4h(t)}\right){\rm
d}x_{t}{\rm d}t\right),$ (16)

where the last step is due to the facts that both $h(t)\in[h(T_{0}),h(T)]$ and
$\left\|X_{0}\right\|\leq D$. We bound the inner integral in (16) by using the
polar coordinate (Folland, 1999, Corollary 2.51):

$\displaystyle\int_{\left\|x_{t}\right\|\geq{R}}\exp\left(-\frac{\left\|x_{t}\right\|^{2}}{4h(t)}\right){\rm
d}x_{t}$
$\displaystyle=\frac{2\pi^{d/2}}{\Gamma(d/2)}\int_{{R}}^{\infty}\exp\left(-\frac{r^{2}}{4h(t)}\right)r^{d-1}{\rm
d}r$
$\displaystyle=\frac{(4h(t))^{d/2}\pi^{d/2}}{\Gamma(d/2)}\int_{{R}^{2}/(4h(t))}^{\infty}\exp\left(-u\right)u^{d/2-1}{\rm
d}u$
$\displaystyle=\frac{2(4h(t))^{d/2}\pi^{d/2}}{d\Gamma(d/2)}\int_{\left({R}^{2}/(4h(t))\right)^{d/2}}^{\infty}\exp\left(-v^{2/d}\right){\rm
d}v$
$\displaystyle\leq\frac{8h(t)\pi^{d/2}}{\Gamma(d/2)}{R}^{d-2}e^{-{R}^{2}/(4h(t))},$

where the last inequality follow from (Qi & Mei, 1999, Equation 10).
Therefore, we conclude that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}(f_{\bf
W}(X_{t},t))-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}>R\right\\}\right]{\rm
d}t$
$\displaystyle=O\left(\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\frac{8h(t)\pi^{d/2}}{\Gamma(d/2)}{R}^{d-2}e^{-{R}^{2}/(4h(t))}{\rm
d}t\right)=O({R}^{d-2}e^{-{R}^{2}/4}).$

∎

##  Appendix C Proof of Theorem 3.6

We first show that $f_{*}(x,t)$ is Lipschitz in $t$ for each fixed $x$ in the
next lemma:

######  Lemma C.1.

Suppose Assumptions 3.2 and 3.4 hold. For each $R>0$, the regression function
$f_{*}$ is $\beta_{t}(R)$-Lipschitz in $t$ for all
$\left\|x\right\|_{\infty}\leq R$ and $t\in[T_{0},\infty)$, i.e.,
$\left|f_{*}(x,t)-f_{*}(x,t^{\prime})\right|_{2}\leq\beta_{t}(R)\left|t-t^{\prime}\right|$,
where $\beta_{t}(R)=O(\sqrt{d}R)$.

###### Proof.

We start with computing the derivative of $f*$ with respect to $t$. The
appication of the dominated convergence theorem implies

$\displaystyle\frac{\partial}{\partial t}f_{*}(x,t)$
$\displaystyle=\frac{\partial}{\partial t}\int x_{0}p_{0|t}(x_{0}|x){\rm
d}x_{0}$ $\displaystyle=\frac{\partial}{\partial
t}\int\frac{x_{0}p_{t|0}(x|x_{0})p_{0}(x_{0})}{\int
p_{t|0}(x|x_{0}^{\prime})p_{0}(x_{0}^{\prime}){\rm d}x_{0}^{\prime}}{\rm
d}x_{0}$ $\displaystyle=\int\frac{x_{0}\frac{\partial}{\partial
t}p_{t|0}(x|x_{0})p_{0}(x_{0})}{\int
p_{t|0}(x|x_{0}^{\prime})p_{0}(x_{0}^{\prime}){\rm d}x_{0}^{\prime}}{\rm
d}x_{0}$
$\displaystyle\qquad-\int\frac{x_{0}p_{t|0}(x|x_{0})p_{0}(x_{0})\int\frac{\partial}{\partial
t}p_{t|0}(x|x_{0}^{\prime\prime})p_{0}(x_{0}^{\prime\prime}){\rm
d}x_{0}^{\prime\prime}}{\left(\int
p_{t|0}(x|x_{0}^{\prime})p_{0}(x_{0}^{\prime}){\rm
d}x_{0}^{\prime}\right)^{2}}{\rm d}x_{0}.$ (17)

To proceed, recall that $X_{t}|X_{0}\sim\mathcal{N}(\alpha(t)X_{0},h(t)I_{d})$
with $\alpha(t)=\exp\left(-\int_{0}^{t}\frac{g(s)}{2}{\rm d}s\right)$ and
$h(t)=1-\alpha^{2}(t)$. We can compute

$\displaystyle\frac{\partial}{\partial t}p_{t|0}(x|x_{0})$
$\displaystyle=\frac{\partial}{\partial t}\left((2\pi
h(t))^{-d/2}\exp\left(-\frac{\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}}{2h(t)}\right)\right)$
$\displaystyle=-\frac{d}{2}(2\pi
h(t))^{-\frac{d}{2}-1}(2\pi)h^{\prime}(t)\exp\left(-\frac{\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}}{2h(t)}\right)$
$\displaystyle\qquad+(2\pi
h(t))^{-d/2}\exp\left(-\frac{\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}}{2h(t)}\right)\left(\frac{2(x-\alpha(t)x_{0})^{\top}x_{0}\alpha^{\prime}(t)}{2h(t)}+\frac{\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}h^{\prime}(t)}{2h^{2}(t)}\right)$
$\displaystyle=\frac{p_{t|0}(x|x_{0})}{2h^{2}(t)}\left(-dh(t)h^{\prime}(t)+2(x-\alpha(t)x_{0})^{\top}x_{0}\alpha^{\prime}(t)h(t)+\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}h^{\prime}(t)\right).$
(18)

Since $\alpha^{\prime}(t)=-\alpha(t)g(t)/2$ and
$h^{\prime}(t)=-2\alpha(t)\alpha^{\prime}(t)=\alpha^{2}(t)g(t)$, we can
rewrite (18) as

$\displaystyle\frac{\partial}{\partial t}p_{t|0}(x|x_{0})$
$\displaystyle=\frac{p_{t|0}(x|x_{0})}{2h^{2}(t)}\left(-dh(t)\alpha^{2}(t)g(t)-(x-\alpha(t)x_{0})^{\top}x_{0}\alpha(t)g(t)h(t)+\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}\alpha^{2}(t)g(t)\right)$
$\displaystyle=p_{t|0}(x|x_{0})\frac{\alpha(t)g(t)}{2h^{2}(t)}\left(-dh(t)\alpha(t)-(x-\alpha(t)x_{0})^{\top}x_{0}h(t)+\left\|x-\alpha(t)x_{0}\right\|_{2}^{2}\alpha(t)\right)$
$\displaystyle=p_{t|0}(x|x_{0})\frac{\alpha(t)g(t)}{2h^{2}(t)}\left(-dh(t)\alpha(t)+\alpha(t)\left\|x\right\|_{2}^{2}-(1+\alpha^{2}(t))x_{0}^{\top}x+\alpha(t)\left\|x_{0}\right\|_{2}^{2}\right).$
(19)

Plugging (19) back into (17), we have

$\displaystyle\int\frac{x_{0}\frac{\partial}{\partial
t}p_{t|0}(x|x_{0})p_{0}(x_{0})}{\int
p_{t|0}(x|x_{0}^{\prime})p_{0}(x_{0}^{\prime}){\rm d}x_{0}^{\prime}}{\rm
d}x_{0}$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\mathbb{E}\left[X_{0}\left(-dh(t)\alpha(t)+\alpha(t)\left\|X_{t}\right\|_{2}^{2}-(1+\alpha^{2}(t))X_{0}^{\top}X_{t}+\alpha(t)\left\|X_{0}\right\|_{2}^{2}\right)\bigg{|}X_{t}=x\right]$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\bigg{(}-dh(t)\alpha(t)\mathbb{E}\left[X_{0}|X_{t}=x\right]+\alpha(t)\left\|x\right\|_{2}^{2}\mathbb{E}\left[X_{0}|X_{t}=x\right]$
$\displaystyle\qquad\qquad\qquad\qquad-(1+\alpha^{2}(t))x\mathbb{E}\left[\left\|X_{0}\right\|^{2}_{2}\big{|}X_{t}=x\right]+\alpha(t)\mathbb{E}\left[X_{0}\left\|X_{0}\right\|_{2}^{2}|X_{t}=x\right]\bigg{)},$

and also

$\displaystyle\int\frac{x_{0}p_{t|0}(x|x_{0})p_{0}(x_{0})\int\frac{\partial}{\partial
t}p_{t|0}(x|x_{0}^{\prime\prime})p_{0}(x_{0}^{\prime\prime}){\rm
d}x_{0}^{\prime\prime}}{\left(\int
p_{t|0}(x|x_{0}^{\prime})p_{0}(x_{0}^{\prime}){\rm
d}x_{0}^{\prime}\right)^{2}}{\rm d}x_{0}$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\mathbb{E}\left[X_{0}\mathbb{E}\left[-dh(t)\alpha(t)+\alpha(t)\left\|X_{t}\right\|_{2}^{2}-(1+\alpha^{2}(t))X_{0}^{\top}X_{t}+\alpha(t)\left\|X_{0}\right\|_{2}^{2}\bigg{|}X_{t}\right]\bigg{|}X_{t}=x\right]$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\mathbb{E}\left[-dh(t)\alpha(t)+\alpha(t)\left\|X_{t}\right\|_{2}^{2}-(1+\alpha^{2}(t))X_{0}^{\top}X_{t}+\alpha(t)\left\|X_{0}\right\|_{2}^{2}\bigg{|}X_{t}=x\right]\mathbb{E}\left[X_{0}|X_{t}=x\right]$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\left(-dh(t)\alpha(t)+\alpha(t)\left\|x\right\|_{2}^{2}-(1+\alpha^{2}(t))x^{\top}\mathbb{E}\left[X_{0}|X_{t}=x\right]+\alpha(t)\mathbb{E}\left[\left\|X_{0}\right\|^{2}_{2}|X_{t}=x\right]\right)\mathbb{E}\left[X_{0}|X_{t}=x\right].$

Therefore, we conclude that

$\displaystyle\frac{\partial}{\partial t}f_{*}(x,t)$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\bigg{(}\alpha(t)\mathbb{E}\left[\left\|X_{0}\right\|_{2}^{2}\left(X_{0}-\mathbb{E}\left[X_{0}|X_{t}\right]\right)\big{|}X_{t}=x\right]$
$\displaystyle\qquad\qquad\qquad\qquad-(1+\alpha^{2}(t))x\left(\mathbb{E}\left[\left\|X_{0}\right\|_{2}^{2}|X_{t}=x\right]-\left\|\mathbb{E}\left[X_{0}|X_{t}=x\right]\right\|_{2}^{2}\right)\bigg{)}$
$\displaystyle=\frac{\alpha(t)g(t)}{2h^{2}(t)}\bigg{(}\alpha(t)\mathbb{E}\left[\left\|X_{0}\right\|_{2}^{2}\left(X_{0}-\mathbb{E}\left[X_{0}|X_{t}\right]\right)\big{|}X_{t}=x\right]-(1+\alpha^{2}(t))x{\rm
Cov}(X_{0}|X_{t}=x)\bigg{)}.$

The Pythagorean theorem implies that
$\left\|X_{0}-\mathbb{E}\left[X_{0}|X_{t}\right]\right\|_{2}\leq\left\|X_{0}\right\|_{2}$.
Since $\left\|X_{0}\right\|_{2}\leq D$ by Assumption 3.2, we can apply the
triangle inequality to obtain

$\displaystyle\sup_{t\in[T_{0},\infty)}\sup_{\left\|x\right\|_{\infty}\leq
R}\left\|\frac{\partial}{\partial t}f_{*}(x,t)\right\|_{2}$
$\displaystyle\leq\frac{\alpha(t)g(t)}{2h^{2}(t)}\bigg{[}\alpha(t)\mathbb{E}\left[\left\|X_{0}\right\|_{2}^{2}\left\|X_{0}-\mathbb{E}\left[X_{0}|X_{t}\right]\right\|_{2}\bigg{|}X_{t}=x\right]$
$\displaystyle\qquad\qquad\qquad+(1+\alpha^{2}(t))\left\|x\right\|_{2}\left\|{\rm
Cov}(X_{0}|X_{t}=x)\right\|_{2}\bigg{]}$
$\displaystyle=O(\sqrt{d}R)\eqqcolon\beta_{t}(R),$

where we have used the facts that $\alpha(t)\leq 1$, $h(t)\geq h(T_{0})$ and
$g(t)$ is bounded on $[T_{0},\infty)$. ∎

Next, we define two kernels without the bias term compared to NTKs in Section
2. Let $\tilde{\mathcal{H}}_{1}$ be a real-valued RKHS induced by the scalar-
valued NTK
$\tilde{\kappa}:\mathbb{R}^{d+1}\times\mathbb{R}^{d+1}\to\mathbb{R}$ defined
as

$\displaystyle\tilde{\kappa}(z,\tilde{z})\coloneqq
z^{\top}\tilde{z}\mathbb{E}\left[\mathbb{I}\left\\{w_{1}(0)^{\top}z\geq
0\right\\}\mathbb{I}\left\\{w_{1}(0)^{\top}\tilde{z}\geq 0\right\\}\right].$

Similarly, let $\tilde{\mathcal{H}}$ be a vector-valued RKHS induced by the
matrix-valued NTK
$\tilde{K}:\mathbb{R}^{d+1}\times\mathbb{R}^{d+1}\to\mathbb{R}^{d\times d}$
defined as

$\displaystyle\tilde{K}(z,\tilde{z})=\tilde{\kappa}(z,\tilde{z})I_{d}.$

The next lemma shows the approximation of of a Lipschitz target function over
a ball with radius $R$.

######  Lemma C.2.

(Bach, 2017, Proposition 6) Let $R_{\tilde{\mathcal{H}}_{1}}$ be larger than a
constant $c_{1}$ that depends only on $d$. For any function
$f:\mathbb{R}^{d+1}\to\mathbb{R}$ such that for any
$\left\|z\right\|_{\infty},\left\|z^{\prime}\right\|_{\infty}\leq R$,
$\sup_{\left\|z\right\|_{\infty}\leq R}\left|f(z)\right|\leq\Lambda$ and
$\left|f(z)-f(z^{\prime})\right|\leq\frac{\Lambda}{R}\left\|z-z^{\prime}\right\|_{2}$,
there exists $f_{\tilde{\mathcal{H}}_{1}}\in\tilde{\mathcal{H}}_{1}$ with
$\left\|f_{\mathcal{H}_{1}}\right\|_{\mathcal{H}_{1}}^{2}\leq
R_{\mathcal{H}_{1}}$ and

$\displaystyle\sup_{\left\|z\right\|_{\infty}\leq
R}\left|f(z)-f_{\tilde{\mathcal{H}}_{1}}(z)\right|\leq
A(R_{\tilde{\mathcal{H}}_{1}}),\quad A(R_{\tilde{\mathcal{H}}_{1}})\coloneqq
c_{1}\Lambda\left(\frac{\sqrt{R_{\tilde{\mathcal{H}}_{1}}}}{\Lambda}\right)^{-\frac{2}{d}}\log\left(\frac{\sqrt{R_{\tilde{\mathcal{H}}_{1}}}}{\Lambda}\right).$

Lemma C.2 comes from (Bach, 2017, Proposition 6) for $d+1$, $\alpha=0$ and
$q=\infty$. Now we are prepared to prove that the regression function $f_{*}$
can be approximated by a function in the RKHS $\mathcal{H}$ induced by
$K((x,t),(x^{\prime},t^{\prime}))=\tilde{K}((x,t-T_{0}),(x^{\prime},t^{\prime}-T_{0}))$.

######  Theorem C.3 (Approximation of the Score Function on the Ball).

Suppose Assumptions 3.2, 3.4 and 3.5 hold. Let $R\geq T-T_{0}$ and
$R_{\mathcal{H}}$ be larger than a constant $c_{1}$ that depends only on $d$.
There exists a function $f_{\mathcal{H}}\in\mathcal{H}$ with
$\left\|f_{\mathcal{H}}\right\|_{\mathcal{H}}^{2}\leq dR_{\mathcal{H}}$ and

$\displaystyle\sup_{\left\|x\right\|_{\infty}\leq
R}\sup_{t\in[T_{0},T]}\left\|f_{*}(x,t)-f_{\mathcal{H}}(x,t)\right\|_{\infty}\leq
A(R_{\mathcal{H}},R)\coloneqq
c_{1}\Lambda(R)\left(\frac{\sqrt{R_{\mathcal{H}}}}{\Lambda(R)}\right)^{-\frac{2}{d}}\log\left(\frac{\sqrt{R_{\mathcal{H}}}}{\Lambda(R)}\right),$

where $\Lambda(R)=O(\sqrt{d}R^{2})$.

###### Proof.

We define an auxiliary target function
$\tilde{f}_{*}:\mathbb{R}^{d}\times\mathbb{R}\to\mathbb{R}^{d}$ as
$\tilde{f}_{*}(x,t)\coloneqq f_{*}(x,\left|t\right|+T_{0})$. By Assumption 3.5
and Lemma C.1, the function $f_{*}(x,t)$ is $\beta_{x}$-Lipschitz in $x$ and
$\beta_{t}(R)$-Lipschitz in $t$ for all $\left\|x\right\|_{\infty}\leq R$ and
$t\in[T_{0},\infty)$; so is each coordinate map. Since
$\sup_{\left\|(x,t)\right\|_{\infty}\leq
R}\left\|\tilde{f}_{*}(x,t)\right\|_{2}\leq D$ and for all
$\left\|(x,t)\right\|_{\infty},\left\|(x^{\prime},t^{\prime})\right\|_{\infty}\leq
R$,

$\displaystyle\left\|\tilde{f}_{*}(x,t)-\tilde{f}_{*}(x^{\prime},t^{\prime})\right\|_{2}$
$\displaystyle\leq\left\|\tilde{f}_{*}(x,t)-\tilde{f}_{*}(x^{\prime},t)\right\|_{2}+\left\|\tilde{f_{*}}(x^{\prime},t)-\tilde{f}_{*}(x^{\prime},t^{\prime})\right\|_{2}$
$\displaystyle=\left\|f_{*}(x,\left|t\right|+T_{0})-f_{*}(x^{\prime},\left|t\right|+T_{0})\right\|_{2}+\left\|f_{*}(x^{\prime},\left|t\right|+T_{0})-f_{*}(x^{\prime},\left|t^{\prime}\right|+T_{0})\right\|_{2}$
$\displaystyle\leq\beta_{x}\left\|x-x^{\prime}\right\|_{2}+\beta_{t}(R)\left|\left|t\right|-\left|t^{\prime}\right|\right|$
$\displaystyle\leq\left(\beta_{x}+\beta_{t}(R)\right)\left\|(x,t)-(x^{\prime},t^{\prime})\right\|_{2},$
(20)

one can apply Lemma C.2 by choosing
$\Lambda(R)=\max\left\\{D,R\left\\{\beta_{x}+\beta_{t}(R)\right\\}\right\\}$
to conclude that for each coordinate $i=1,\dots,d$, there exists
$\tilde{f}_{\tilde{\mathcal{H}}_{1}}^{i}\in\tilde{\mathcal{H}}_{1}$ with
$\left\|\tilde{f}^{i}_{\tilde{\mathcal{H}}_{1}}\right\|_{\tilde{\mathcal{H}}_{1}}^{2}\leq
R_{\mathcal{H}}$ such that

$\displaystyle\sup_{\left\|(x,t)\right\|_{\infty}\leq
R}\left|\tilde{f}_{*}^{i}(x,t)-\tilde{f}^{i}_{\tilde{\mathcal{H}}_{1}}(x,t)\right|\leq
A(R_{\mathcal{H}},R)=c_{1}\Lambda(R)\left(\frac{\sqrt{R_{\mathcal{H}}}}{\Lambda(R)}\right)^{-\frac{2}{d}}\log\left(\frac{\sqrt{R_{\mathcal{H}}}}{\Lambda(R)}\right).$

Defining
$f^{i}_{\mathcal{H}}(x,t)\coloneqq\tilde{f}_{\tilde{\mathcal{H}}_{1}}^{i}(x,t-T_{0})$,
we have

$\displaystyle\sup_{\left\|x\right\|_{\infty}\leq
R}\,\,\sup_{t\in[T_{0},R+T_{0}]}\left|f_{*}^{i}(x,t)-f_{\mathcal{H}}^{i}(x,t)\right|\leq
A(R_{\mathcal{H}},R).$

Note that $f^{i}_{\mathcal{H}}:\mathbb{R}^{d+1}\to\mathbb{R}$ lies in the RKHS
induced by the kernel
${\kappa}((x,t),(x^{\prime},t^{\prime}))=\tilde{\kappa}((x,t-T_{0}),(x^{\prime},t^{\prime}-T_{0}))$
and
$\left\|f^{i}_{\mathcal{H}}\right\|_{\tilde{\kappa}}=\left\|\tilde{f}^{i}_{\tilde{\mathcal{H}}_{1}}\right\|_{\tilde{\mathcal{H}}_{1}}$.
We next show that
$f_{\mathcal{H}}=(f^{1}_{\mathcal{H}},\dots,f^{d}_{\mathcal{H}})$ is in the
RKHS induced by $K$. Since each coordinate of $f^{i}_{\mathcal{H}}$ lies in
the RKHS induced by ${\kappa}$, by relabeling data points, without loss of
generality, it suffices to consider

$\displaystyle
f^{i}_{\mathcal{H}}(\cdot)=\sum_{p=1}^{P}\alpha_{p}^{i}{\kappa}((x,t)_{p},\cdot),\quad(x,t)_{p}\in\mathbb{R}^{d+1},\alpha_{p}^{i}\in\mathbb{R}.$

It follows

$\displaystyle
f_{\mathcal{H}}(\cdot)=\sum_{i=1}^{d}{f}_{\mathcal{H}}^{i}(\cdot)\mathbf{e}_{i}=\sum_{i=1}^{d}\left(\sum_{p=1}^{p}\alpha_{p}^{i}\kappa((x,t)_{p},\cdot)\right)\mathbf{e}_{i}=\sum_{p=1}^{P}{K}((x,t)_{p},\cdot)\left(\sum_{i=1}^{d}\alpha_{p}^{i}\mathbf{e}_{i}\right)\in\mathcal{H}.$

Moreover,

$\displaystyle\left\|f_{\mathcal{H}}\right\|_{\mathcal{H}}^{2}$
$\displaystyle=\left\langle\sum_{p=1}^{P}{K}((x,t)_{p},\cdot)\left(\sum_{i=1}^{d}\alpha_{p}^{i}\mathbf{e}_{i}\right),\sum_{q=1}^{P}{K}((x,t)_{q},\cdot)\left(\sum_{k=1}^{d}\alpha_{q}^{k}\mathbf{e}_{k}\right)\right\rangle_{{K}}$
$\displaystyle=\sum_{p,q}\sum_{i,k}\alpha_{p}^{i}\alpha_{q}^{k}\mathbf{e}_{i}^{\top}{K}((x,t)_{p},(x,t)_{q})\mathbf{e}_{k}$
$\displaystyle=\sum_{i=1}^{d}\sum_{p,q}\alpha_{p}^{i}\alpha_{q}^{i}{\kappa}((x,t)_{p},(x,t)_{q})$
$\displaystyle=\sum_{i=1}^{d}\left\langle\sum_{p=1}^{P}\alpha_{p}^{i}{\kappa}((x,t)_{p},\cdot),\sum_{q=1}^{P}\alpha_{q}^{i}{\kappa}((x,t)_{q},\cdot)\right\rangle$
$\displaystyle=\sum_{i=1}^{d}\left\|f^{i}_{\mathcal{H}}\right\|_{{\kappa}}^{2}=\sum_{i=1}^{d}\left\|\tilde{f}^{i}_{\tilde{\mathcal{H}}_{1}}\right\|_{\tilde{\mathcal{H}}_{1}}^{2}\leq
dR_{\mathcal{H}}.$ (21)

Therefore, we have found a function
$f_{\mathcal{H}}:\mathbb{R}^{d+1}\to\mathbb{R}^{d}$ in the RKHS induced by
${K}$ such that $\left\|f_{\mathcal{H}}\right\|_{\mathcal{H}}^{2}\leq
dR_{\mathcal{H}}$ and

$\displaystyle\sup_{\left\|x\right\|_{\infty}\leq
R}\,\,\sup_{t\in[T_{0},T]}\left\|f_{*}(x,t)-f_{\mathcal{H}}(x,t)\right\|_{\infty}\leq
A(R_{\mathcal{H}},R).$

∎

As a corollary of Theorem C.3, we can prove Theorem 3.6.

###### Proof of Theorem 3.6.

For any ${R}\geq T-T_{0}$ and $t\in[T_{0},T]$, we have

$\displaystyle\left\|f_{\mathcal{H}}(X_{t},t)-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}$
$\displaystyle\leq
d\sup_{\left\|x\right\|_{\infty}\leq{R}}\sup_{t\in[T_{0},T]}\left\|f_{\mathcal{H}}(x,t)-f_{*}(x,t)\right\|_{\infty}^{2}\leq
dA^{2}(R_{\mathcal{H}},{R}),$

which implies that

$\displaystyle\int_{T_{0}}^{T}\mathbb{E}\left[\left\|f_{\mathcal{H}}(X_{t},t)-f_{*}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t\leq d(T-T_{0})A^{2}(R_{\mathcal{H}},{R}).$

Dividing both sides by $T-T_{0}$ will complete the proof. ∎

##  Appendix D Proof of Theorem 3.9

To prove Theorem 3.9, we first show the linear convergence rate of gradient
descent over the training dataset
$S=\left\\{t_{j},X_{0,j},X_{t_{j}}\right\\}$. Define a Gram matrix
$H(\tau)\in\mathbb{R}^{dN\times dN}$ at each iteration $\tau$ as as block
matrix:

$\displaystyle H(\tau)\coloneqq\begin{pmatrix}H_{11}&\dots&H_{1N}\\\
\vdots&\ddots&\vdots\\\ H_{N1}&\dots&H_{NN}\end{pmatrix},\quad
H^{ik}_{j\ell}(\tau)=\frac{1}{m}z_{j}^{\top}z_{\ell}\sum_{r=1}^{m}a_{r}^{i}a_{r}^{k}\mathbb{I}\left\\{z_{j}^{\top}w_{r}(\tau)\geq
0,z_{\ell}^{\top}w_{r}(\tau)\geq 0\right\\}.$

One can check that $H=\mathbb{E}\left[H(0)\right]$ with expectation taken over
the random initialization. For ease of presentation, recall that we have set
$C_{\min}=\Delta$ and $C_{\max}=\sqrt{R^{2}+(T-T_{0})^{2}}$ so that
$\left\|(X_{t_{j}},t_{j}-T_{0})\right\|_{2}\in[C_{\min},C_{\max}]$ by
Assumption 3.7. Moreover, we denote the activation pattern of neural $w_{r}$
for sample $j$ at iteration $\tau$ as
$\mathbb{I}_{j,r}(\tau)\coloneqq\mathbb{I}\left\\{w_{r}(\tau)^{\top}z_{j}\geq
0\right\\}$. The convergence of GD algorithm is given in the next theorem.

######  Theorem D.1 (Convergence Rate of Gradient Descent).

Suppose Assumptions 3.2, 3.7 and 3.8 hold. If we set
$m=\Omega\left(\frac{(dN)^{6}C_{\max}^{6}}{\lambda_{0}^{10}\delta^{3}C_{\min}^{2}}\right)$
with i.i.d. initialize for $w_{r}\sim\mathcal{N}(0,I_{d+1})$ and
$a_{r}^{i}\sim{\rm Unif}\left\\{-1,1\right\\}$, and we set
$\eta=\mathcal{O}\left(\frac{\lambda_{0}}{(dN)^{2}C_{\max}^{4}}\right)$, then
with probability at least $1-\delta$, for all $\tau\geq 0$ and $r=1,\dots,m$
simultaneously, we have

$\displaystyle\widehat{\mathcal{L}}({\bf
W}(\tau))\leq\left(1-\eta\lambda_{0}\right)^{\tau}\widehat{\mathcal{L}}({\bf
W}(0)),$ (22)

and

$\displaystyle\left\|w_{r}(\tau)-w_{r}(0)\right\|_{2}\leq R_{w}\coloneqq
O\left(\frac{dNC^{2}_{\max}}{\sqrt{m}\lambda_{0}\sqrt{\delta}}\right).$ (23)

###### Proof.

Following the ideas in Du et al. (2018); Arora et al. (2019a), we prove the
convergence of GD by induction. The induction is to show that (22) holds for
all $\tau$. It is straightforward to see the inequality holds for $\tau=0$.
Assuming (22) holds for $0\leq\tau^{\prime}\leq\tau$, we will show it is also
true for $\tau^{\prime}=\tau+1$. Let $u(\tau)={\rm
vec}(u_{1},\dots,u_{N})(\tau)$ and $y={\rm vec}(y_{1},\dots,y_{N})$ with
$u_{j}(\tau)=f_{\bf W(\tau)}(X_{t_{j}},t_{j})$ and $y_{j}=X_{0,j}$. We first
need the following result for all $\tau^{\prime}=0,\dots,\tau+1$:

$\displaystyle\left\|w_{r}(\tau^{\prime})-w_{r}(0)\right\|_{2}$
$\displaystyle=\left\|\eta\sum_{\tau^{\prime\prime}=0}^{\tau^{\prime}-1}\frac{\partial\widehat{\mathcal{L}}({\bf
W}(\tau^{\prime\prime}))}{\partial w_{r}(\tau^{\prime\prime})}\right\|_{2}$
$\displaystyle\leq\eta\sum_{\tau^{\prime\prime}=0}^{\tau^{\prime}-1}\left\|\frac{\partial\widehat{\mathcal{L}}({\bf
W}(\tau^{\prime\prime}))}{\partial w_{r}(\tau^{\prime\prime})}\right\|_{2}$
$\displaystyle\leq\eta
C_{\max}\sum_{\tau^{\prime\prime}=0}^{\tau^{\prime}-1}\frac{\sqrt{dN}\left\|u(\tau^{\prime\prime})-y\right\|_{2}}{\sqrt{m}}$
(24) $\displaystyle\leq\frac{\eta
C_{\max}\sqrt{dN}}{\sqrt{m}}\sum_{\tau^{\prime\prime}=0}^{\tau^{\prime}-1}\left(1-\eta\lambda_{0}\right)^{\tau^{\prime\prime}/2}\left\|u(0)-y\right\|_{2}$
(25) $\displaystyle\leq\frac{\eta
C_{\max}\sqrt{dN}}{\sqrt{m}}\sum_{\tau^{\prime\prime}=0}^{\infty}(1-\eta\lambda_{0}/2)^{\tau^{\prime\prime}}\left\|u(0)-y\right\|_{2}$
(26)
$\displaystyle=\frac{2C_{\max}\sqrt{dN}\left\|u(0)-y\right\|_{2}}{\sqrt{m}\lambda_{0}}.$
(27)

Here, we have an upper bound on gradient (9) to derive (24). Also, (25) and
(26) follow from the induction hypothesis and the fact that $\sqrt{1-x}\leq
1-x/2$. We further bound

$\displaystyle\mathbb{E}\left[\left\|u(0)-y\right\|_{2}^{2}\right]$
$\displaystyle=\sum_{i,j}\mathbb{E}\left[\left|u^{i}_{j}(0)-y_{j}^{i}\right|^{2}\right]$
$\displaystyle=\sum_{i,j}\left[(y_{j}^{i})^{2}-2y_{j}^{i}\mathbb{E}\left[f^{i}({\bf
W},a,(X_{t_{j}},t_{j}))\right]+\mathbb{E}\left[(f^{i})^{2}({\bf
W},a,(X_{t_{j}},t_{j}))\right]\right]$
$\displaystyle\leq\sum_{i,j}\left[(y_{j}^{i})^{2}+C_{\max}^{2}\right]=O(dNC_{\max}^{2}),$

where we have used the facts that $\mathbb{E}\left[f^{i}({\bf
W},a,(X_{t_{j}},t_{j}))\right]=0$, $\mathbb{E}\left[(f^{i})^{2}({\bf
W},a,(X_{t_{j}},t_{j}))\right]\leq C_{\max}^{2}$ and
$\left\|y_{j}\right\|_{2}\leq D$. Thus, the Markov’s inequality yields
$\left\|u(0)-y\right\|_{2}^{2}=O(dNC_{\max}^{2}/\delta)$ with probability at
least $1-\delta$. Therefore, with probability at least $1-\delta$, we have

$\displaystyle\left\|w_{r}(\tau^{\prime})-w_{r}(0)\right\|_{2}\leq
R_{w}\coloneqq
O\left(\frac{dNC^{2}_{\max}}{\sqrt{m}\lambda_{0}\sqrt{\delta}}\right),\quad\forall\tau^{\prime}=0,\dots,\tau+1,r=1,\dots,m.$
(28)

Define index sets

$\displaystyle
S_{j}\coloneqq\left\\{r\in[m]:\mathbb{I}\left\\{A_{j,r}\right\\}=0\right\\},\quad\bar{S}_{j}\coloneqq\left\\{r\in[m]:\mathbb{I}\left\\{A_{j,r}\right\\}\neq
0\right\\},$

where $A_{j,r}\coloneqq\left\\{\left|w_{r}(0)^{\top}z_{j}\right|\leq
R_{w}C_{\max}\right\\}$. Note that

$\displaystyle\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau^{\prime})\neq\mathbb{I}_{j,r}(0)\right\\}\leq\mathbb{I}\left\\{A_{j,r}\right\\}+\mathbb{I}\left\\{\left\|w_{r}(\tau^{\prime})-w_{r}(0)\right\|_{2}>R_{w}\right\\}.$
(29)

To see this, note that if
$\left\|w_{r}(\tau^{\prime})-w_{r}(0)\right\|_{2}\leq R_{w}$ it follows
$\left|w_{r}(\tau^{\prime})^{\top}z_{j}-w_{r}(0)^{\top}z_{j}\right|\leq
R_{w}C_{\max}$. If $w_{r}(0)^{\top}z_{j}>R_{w}C_{\max}$, then
$w_{r}(\tau^{\prime})^{\top}z_{j}>0$. Similarly, if
$w_{r}(0)^{\top}z_{j}<-R_{w}C_{\max}$, then
$w_{r}(\tau^{\prime})^{\top}z_{j}<0$. Hence, we must have
$\mathbb{I}_{j,r}(\tau^{\prime})=\mathbb{I}_{j,r}(0)$. From (28) and (29), we
conclude that with probability at least $1-\delta$, all neurons in $S_{j}$
will not change their activation pattern on $z_{j}$ during optimization, i.e.,

$\displaystyle r\in
S_{j}\implies\mathbb{I}_{j,r}(\tau^{\prime})=\mathbb{I}_{j,r}(0),\quad\forall\tau^{\prime}=0,\dots,\tau+1.$
(30)

With such a partition, we can write the dynamics of $u_{j}^{i}(\tau)$ as

$\displaystyle u_{j}^{i}(\tau+1)-u_{j}^{i}(\tau)$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\left[\sigma(w_{r}(\tau+1)^{\top}z_{j})-\sigma(w_{r}(\tau)^{\top}z_{j})\right]$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r\in
S_{j}}a_{r}^{i}\left[\sigma(w_{r}(\tau+1)^{\top}z_{j})-\sigma(w_{r}(\tau)^{\top}z_{j})\right]$
$\displaystyle\qquad+\frac{1}{\sqrt{m}}\sum_{r\in\bar{S}_{j}}a_{r}^{i}\left[\sigma(w_{r}(\tau+1)^{\top}z_{j})-\sigma(w_{r}(\tau)^{\top}z_{j})\right].$
(31)

By utilizing the condition (30), we bound the first term in (31) as

$\displaystyle\frac{1}{\sqrt{m}}\sum_{r\in
S_{j}}a_{r}^{i}\left[\sigma(w_{r}(\tau+1)^{\top}z_{j})-\sigma(w_{r}(\tau)^{\top}z_{j})\right]$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r\in
S_{j}}a_{r}^{i}\mathbb{I}_{j,r}(\tau)\left(w_{r}(\tau+1)^{\top}z_{j}-w_{r}(\tau)^{\top}z_{j}\right)$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r\in
S_{j}}a_{r}^{i}\mathbb{I}_{j,r}(\tau)\left(-\frac{\eta}{\sqrt{m}}\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})a_{r}^{k}z_{\ell}\mathbb{I}_{\ell,r}(\tau)\right)^{\top}z_{j}$
(32)
$\displaystyle=-\frac{\eta}{m}\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})z_{j}^{\top}z_{\ell}\sum_{r\in
S_{j}}a_{r}^{i}a_{r}^{k}\mathbb{I}_{j,r}(\tau)\mathbb{I}_{\ell,r}(\tau)$
$\displaystyle=-\eta\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})H_{j\ell}^{ik}(\tau)+\epsilon_{j}^{i}(\tau),$
(33)

where we have set
$\epsilon_{j}^{i}(\tau)\coloneqq\frac{\eta}{m}\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})z_{j}^{\top}z_{\ell}\sum_{r\in\bar{S}_{j}}a_{r}^{i}a_{r}^{k}(\tau)\mathbb{I}_{j,r}(\tau)\mathbb{I}_{\ell,r}(\tau)$.
Here, we have used the GD update rule and the definition of
$H_{j\ell}^{ik}(\tau)$ to derive (32) and (33). We can further upper bound the
error term

$\displaystyle\left|\epsilon_{j}^{i}(\tau)\right|\leq\frac{\eta}{m}\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})\left\|z_{j}\right\|_{2}\left\|z_{\ell}\right\|_{2}\left|\bar{S}_{j}\right|\leq\frac{\eta
C_{\max}^{2}\left|\bar{S}_{j}\right|\sqrt{dN}}{m}\left\|u(\tau)-y\right\|_{2}.$
(34)

Next, we denote the second term in (31) by $\bar{\epsilon}_{j}^{i}(\tau)$,
which can be bounded by

$\displaystyle\left|\bar{\epsilon}_{j}^{i}(\tau)\right|$
$\displaystyle=\left|\frac{1}{\sqrt{m}}\sum_{r\in\bar{S}_{j}}a_{r}^{i}\left[\sigma(w_{r}(\tau+1)^{\top}z_{j})-\sigma(w_{r}(\tau)^{\top}z_{j})\right]\right|$
$\displaystyle\leq\frac{1}{\sqrt{m}}\sum_{r\in\bar{S}_{j}}\left|a_{r}^{i}\right|\left|(w_{r}(\tau+1)-w_{r}(\tau))^{\top}z_{j}\right|$
(35)
$\displaystyle\leq\frac{C_{\max}}{\sqrt{m}}\sum_{r\in\bar{S}_{j}}\left\|w_{r}(\tau+1)-w_{r}(\tau)\right\|_{2}$
(36)
$\displaystyle=\frac{C_{\max}}{\sqrt{m}}\sum_{r\in\bar{S}_{j}}\left\|-\frac{\eta}{\sqrt{m}}\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})a_{r}^{k}z_{\ell}\mathbb{I}_{\ell,r}(\tau)\right\|_{2}$
(37) $\displaystyle\leq\frac{\eta
C_{\max}}{m}\sum_{r\in\bar{S}_{j}}\sum_{\ell=1}^{N}\sum_{k=1}^{d}\left|u_{\ell}^{k}(\tau)-y_{\ell}^{k}\right|\left\|z_{\ell}\right\|_{2}$
$\displaystyle\leq\frac{\eta
C^{2}_{\max}\left|\bar{S}_{j}\right|\sqrt{dN}}{m}\left\|u(\tau)-y\right\|_{2},$
(38)

where we have applied the $1$-Lipschitz property of the ReLU activation
function to obtain (35). Also, we have employed the facts
$\left|a_{r}^{i}\right|\leq 1$ and $\left\|z_{j}\right\|_{2}\leq C_{\max}$ in
(36). The GD update rule has been utilized to reach (37). Combining (31), (33)
and (38), we have

$\displaystyle
u_{j}^{i}(\tau+1)-u_{j}^{i}(\tau)=-\eta\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})H_{j\ell}^{ik}(\tau)+\epsilon_{j}^{i}(\tau)+\bar{\epsilon}_{j}^{i}(\tau),$

which can be further written in a compact form through vectorization:

$\displaystyle u(\tau+1)-u(\tau)$ $\displaystyle=-\eta
H(\tau)(u(\tau)-y)+\epsilon(\tau)+\bar{\epsilon}(\tau)$ $\displaystyle=-\eta
H(u(\tau)-y)+\eta(H-H(\tau))(u(\tau)-y)+\epsilon(\tau)+\bar{\epsilon}(\tau),$
(39)

where $\epsilon(\tau)$ and $\bar{\epsilon}(\tau)$ are defined as similar to
$u(\tau)$ by vectorization.

We move on to show $H(\tau)$ is close to $H$ for sufficiently wide neural
networks. First, the Hoeffding’s inequality implies, with probability at least
$1-\delta^{\prime}$, we have

$\displaystyle\left|H_{j\ell}^{ik}(0)-H_{j\ell}^{ik}\right|\leq
C_{\max}^{2}\sqrt{\frac{2\log(2/\delta^{\prime})}{m}}.$

Setting $\delta^{\prime}=\delta/(dN)^{2}$ and applying the union bound, we
obtain

$\displaystyle\left\|H-H(0)\right\|_{F}^{2}=\sum_{i,k,j,\ell}\left|H_{j\ell}^{ik}(0)-H_{j\ell}^{ik}\right|^{2}\leq(dN)^{2}C_{\max}^{4}\cdot\frac{2\log(2(dN)^{2}/\delta)}{m},$
(40)

with probability at least $1-\delta$. Next, note that (29) also implies

$\displaystyle\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau^{\prime})\neq\mathbb{I}_{j,r}(0)\right\\}\leq\sum_{r=1}^{m}\mathbb{I}\left\\{A_{j,r}\right\\}+\mathbb{I}\left\\{\left\|w_{r}(\tau^{\prime})-w_{r}(0)\right\|_{2}>R_{w}\text{
for some }r\right\\}.$

It follows

$\displaystyle\left|H_{j\ell}^{ik}(\tau)-H_{j\ell}^{ik}(0)\right|$
$\displaystyle=\left|\frac{1}{m}z_{j}^{\top}z_{\ell}\sum_{r=1}^{m}a_{r}^{i}a_{r}^{k}\left[\mathbb{I}_{j,r}(\tau)\mathbb{I}_{\ell,r}(\tau)-\mathbb{I}_{j,r}(0)\mathbb{I}_{\ell,r}(0)\right]\right|$
$\displaystyle\leq\frac{C_{\max}^{2}}{m}\sum_{r=1}^{m}\left[\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau)\neq\mathbb{I}_{j,r}(0)\right\\}+\mathbb{I}\left\\{\mathbb{I}_{\ell,r}(\tau)\neq\mathbb{I}_{\ell,r}(0)\right\\}\right]$
$\displaystyle\leq\frac{C_{\max}^{2}}{m}\left(\sum_{r=1}^{m}\left[\mathbb{I}\left\\{A_{j,r}\right\\}+\mathbb{I}\left\\{A_{\ell,r}\right\\}\right]+2\mathbb{I}\left\\{\left\|w_{r}(\tau)-w_{r}(0)\right\|_{2}>R_{w}\text{
for some }r\right\\}\right).$

By taking expectation on both sides and applying (28), we have

$\displaystyle\mathbb{E}\left[\left|H_{j\ell}^{ik}(\tau)-H_{j\ell}^{ik}(0)\right|\right]$
$\displaystyle\leq\frac{C_{\max}^{2}}{m}\sum_{r=1}^{m}\mathbb{E}\left[\mathbb{I}\left\\{A_{j,r}\right\\}+\mathbb{I}\left\\{A_{\ell,r}\right\\}\right]+\frac{2C_{\max}^{2}}{m}\mathbb{E}\left[\mathbb{I}\left\\{\left\|w_{r}(\tau)-w_{r}(0)\right\|_{2}>R_{w}\text{
for some }r\right\\}\right]$
$\displaystyle\leq\frac{4R_{w}C_{\max}^{3}}{\sqrt{2\pi}C_{\min}}+\frac{2C_{\max}^{2}}{m}\delta,$
(41)

where we have used the following anti-concentration inequality for Gaussian
random variables:

$\displaystyle\mathbb{E}\left[\mathbb{I}\left\\{A_{j,r}\right\\}\right]=\mathbb{P}_{z\sim\mathcal{N}(0,\left\|z_{j}\right\|^{2}_{2})}\left(\left|z\right|\leq
R_{w}C_{\max}\right)=\int_{-R_{w}C_{\max}}^{R_{w}C_{\max}}\frac{1}{\sqrt{2\pi\left\|z_{j}\right\|_{2}^{2}}}e^{-z^{2}/2\left\|z_{j}\right\|_{2}^{2}}\leq\frac{2R_{w}C_{\max}}{\sqrt{2\pi}C_{\min}}.$
(42)

Hence, we have

$\displaystyle\mathbb{E}\left[\left\|H(\tau)-H(0)\right\|_{F}\right]\leq\sum_{i,k,j,\ell}\mathbb{E}\left[\left|H_{j\ell}^{ik}(\tau)-H_{j\ell}^{ik}(0)\right|\right]\leq\frac{4(dN)^{2}R_{w}C_{\max}^{3}}{\sqrt{2\pi}C_{\min}}+\frac{2(dN)^{2}C_{\max}^{2}}{m}\delta.$

Finally, from the Markov’s inequality, we know that with probability at least
$1-\delta$ it holds

$\displaystyle\left\|H(\tau)-H(0)\right\|_{F}=\mathcal{O}\left(\frac{(dN)^{3}C_{\max}^{4}}{\sqrt{m}\lambda_{0}\delta^{3/2}C_{\min}}\right).$
(43)

Therefore, combining (40) and (43) leads to

$\displaystyle\left\|H-H(\tau)\right\|_{2}$
$\displaystyle\leq\left\|H-H(0)\right\|_{2}+\left\|H(0)-H(\tau)\right\|_{2}$
$\displaystyle=\mathcal{O}\left(\frac{(dN)\sqrt{\log((dN)^{2}/\delta)}}{\sqrt{m}}\right)+\mathcal{O}\left(\frac{(dN)^{3}C_{\max}^{4}}{\sqrt{m}\lambda_{0}\delta^{3/2}C_{\min}}\right)$
$\displaystyle=\frac{(dN)^{3}C_{\max}^{4}}{\sqrt{m}\lambda_{0}\delta^{3/2}C_{\min}},$
(44)

It remains to bound two error terms in (39). From (34) and (38), we know that

$\displaystyle\left\|\epsilon(\tau)+\bar{\epsilon}(\tau)\right\|_{2}$
$\displaystyle\leq\left\|\epsilon(\tau)+\bar{\epsilon}(\tau)\right\|_{1}$
$\displaystyle=\sum_{j=1}^{N}\sum_{i=1}^{d}\left|\epsilon_{j}^{i}(\tau)+\bar{\epsilon}_{j}^{i}\right|(\tau)$
$\displaystyle\leq\sum_{j=1}^{N}\sum_{i=1}^{d}\frac{\eta\left(C_{\max}+C_{\max}^{2}\right)\left|\bar{S}_{j}\right|\sqrt{dN}}{m}\left\|u(\tau)-y\right\|_{2}$
$\displaystyle=\frac{2\eta
C_{\max}^{2}d\sqrt{dN}}{m}\left\|u(\tau)-y\right\|_{2}\sum_{j=1}^{N}\left|\bar{S}_{j}\right|.$
(45)

Furthermore, it follows from (28) and (42) that

$\displaystyle\mathbb{E}\left[\left|\bar{S}_{j}\right|\right]=\mathbb{E}\left[\sum_{r=1}^{m}\mathbb{I}\left\\{A_{j,r}\right\\}\right]=\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}}=\mathcal{O}\left(\frac{\sqrt{m}(dN)C_{\max}^{3}}{\lambda_{0}\sqrt{\delta}C_{\min}}\right).$

Thus, the Markov’s inequality implies
$\sum_{j=1}^{N}\left|\bar{S}_{j}\right|=\mathcal{O}\left(\frac{\sqrt{m}dN^{2}C_{\max}^{3}}{\lambda_{0}\delta^{3/2}C_{\min}}\right)$
with probability at least $1-\delta$.

We need the last result before proving the induction hypothesis. Following the
same argument as in (38), we have

$\displaystyle\left\|u(\tau+1)-u(\tau)\right\|_{2}^{2}$
$\displaystyle\leq\sum_{i,j}\left|u_{j}^{i}(\tau+1)-u_{j}^{i}(\tau)\right|^{2}_{2}$
$\displaystyle\leq(dN)\left(\eta
C_{\max}^{2}\sqrt{dN}\left\|u(\tau)-y\right\|_{2}\right)^{2}$
$\displaystyle=\eta^{2}(dN)^{2}C_{\max}^{4}\left\|u(\tau)-y\right\|_{2}^{2}.$
(46)

With the prediction dynamics (39) and all the estimates (44), (45) and (46),
we can prove the induction hypothesis:

$\displaystyle\left\|u(\tau+1)-y\right\|_{2}^{2}$
$\displaystyle=\left\|u(\tau+1)-u(\tau)+u(\tau)-y\right\|_{2}^{2}$
$\displaystyle=\left\|u(\tau)-y\right\|_{2}^{2}+\left\|u(\tau+1)-u(\tau)\right\|_{2}^{2}+2(u(\tau+1)-u(\tau))^{\top}(u(\tau)-y)$
$\displaystyle=\left\|u(\tau)-y\right\|_{2}^{2}+\left\|u\left(\tau+1\right)-u(\tau)\right\|_{2}^{2}-2\eta(u(\tau)-y)^{\top}H(u(\tau)-y)$
$\displaystyle\qquad+2\eta(u(\tau)-y)^{\top}(H-H(\tau))(u(\tau)-y)+2(\epsilon(\tau)+\bar{\epsilon}(\tau))^{\top}(u(\tau)-y)$
$\displaystyle\leq\left(1-2\eta\lambda_{0}-O\left(\eta^{2}(dN)^{2}C_{\max}^{4}\right)+O\left(\frac{\eta(dN)^{3}C_{\max}^{4}}{\sqrt{m}\lambda_{0}\delta^{3/2}C_{\min}}\right)+O\left(\frac{\eta(dN)^{5/2}C_{\max}^{5}}{\sqrt{m}\lambda_{0}\delta^{3/2}C_{\min}}\right)\right)\left\|u(\tau)-y\right\|_{2}^{2}$
$\displaystyle\leq(1-\eta\lambda_{0})\left\|u(\tau)-y\right\|_{2}^{2},$

where we have used the assumption $\lambda_{0}=\lambda_{\min}(H)>0$ and the
bounds
$m=\Omega\left(\frac{(dN)^{6}C_{\max}^{10}}{\lambda_{0}^{4}\delta^{3}C_{\min}^{2}}\right)$
and $\eta=\mathcal{O}\left(\frac{\lambda_{0}}{(dN)^{2}C_{\max}^{4}}\right)$.
Therefore, we finish the induction and conclude the proof by scaling $\delta$.
∎

To upper bound the coupling term, the non-expansive property of the projection
operator and Assumption 3.2 imply that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\mathbb{E}\left[\left\|\Pi_{D}\left(f_{\bf
W(\tau)}(X_{t},t)\right)-f^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$
$\displaystyle\leq\frac{1}{T-T_{0}}\int_{T_{0}}^{T_{0}+\Delta}\mathbb{E}\left[\left\|\Pi_{D}\left(f_{\bf
W(\tau)}(X_{t},t)\right)-f^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$
$\displaystyle\qquad+\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\mathbb{E}\left[\left\|\Pi_{D}\left(f_{\bf
W(\tau)}(X_{t},t)\right)-f^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t$ $\displaystyle\leq\frac{4\Delta
D^{2}}{T-T_{0}}+\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\mathbb{E}\left[\left\|f_{\bf
W(\tau)}(X_{t},t)-f^{K}_{\tau}(X_{t},t)\right\|_{2}^{2}\mathbb{I}\left\\{\left\|X_{t}\right\|_{2}\leq{R}\right\\}\right]{\rm
d}t.$ (47)

To upper bound the second term in (47), we introduce a linearized neural
network $f_{\bar{\bf W}(\tau)}^{\rm lin}$ updated by

$\displaystyle\bar{w}_{r}\left(\tau+1\right)=\bar{w}_{r}(\tau)-\eta\nabla\widehat{\mathcal{L}}^{{\rm
lin}}(\bar{w}_{r}(\tau)),\quad\widehat{\mathcal{L}}^{{\rm lin}}(\bar{\bf
W})=\frac{1}{2}\sum_{j=1}^{N}\left\|f_{\bar{\bf W}(\tau)}^{\rm
lin}(X_{t_{j}},t_{j})-X_{0,j}\right\|_{2}^{2},$

where $\bar{w}_{r}(0)=w_{r}(0)$ and

$\displaystyle f^{{\rm lin},i}_{\bar{\bf
W}(\tau)}(x,t)\coloneqq\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\bar{w}_{r}(\tau)^{\top}(x,t-T_{0})\mathbb{I}\left\\{{{w}_{r}}(0)^{\top}(x,t-T_{0})\geq
0\right\\}.$

Our next lemma provides the coupling error between $f_{{\bf W}(\tau)}$ and
$f^{{\rm lin}}_{\bar{\bf W}(\tau)}$.

######  Lemma D.2.

Assume the same conditions as in Theorem 3.9. Then with probability at least
$1-\delta$, it holds simultaneously for each $\tau$ that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f_{\bf{W}(\tau)}(x,t)-f^{{\rm lin}}_{\bar{\bf
W}(\tau)}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm
d}t=\mathcal{O}\left(\frac{d(dN)^{9}C_{\max}^{12}}{\sqrt{m}\delta^{4}\lambda_{0}^{2}C_{\min}^{2}}\right).$

###### Proof.

Denote by
$\mathbb{I}_{r}(\tau)\coloneqq\mathbb{I}\left\\{w_{r}(\tau)^{\top}(x,t-T_{0})\geq
0\right\\}$. Note that for each $i=1,\dots,d$ we have

$\displaystyle\left|f_{\bf{W}(\tau)}^{i}(x,t)-f^{{\rm lin},i}_{\bar{\bf
W}(\tau)}(x,t)\right|$
$\displaystyle=\left|\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\sigma\left(w_{r}(\tau)^{\top}(x,t-T_{0})\right)-\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\bar{w}_{r}(\tau)^{\top}(x,t-T_{0})\mathbb{I}_{r}(0)\right|$
$\displaystyle\leq\left|\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\sigma\left(w_{r}(\tau)^{\top}(x,t-T_{0})\right)-\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau)^{\top}(x,t-T_{0})\mathbb{I}_{r}(0)\right|$
$\displaystyle\qquad+\bigg{|}\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau)^{\top}(x,t-T_{0})\mathbb{I}_{r}(0)-\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\bar{w}_{r}(\tau)^{\top}(x,t-T_{0})\mathbb{I}_{r}(0)\bigg{|}$
$\displaystyle=\left|\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau)^{\top}(x,t-T_{0})\left(\mathbb{I}_{r}(\tau)-\mathbb{I}_{r}(0)\right)\right|+\left|\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}(w_{r}(\tau)-\bar{w}_{r}(\tau))^{\top}(x,t-T_{0})\mathbb{I}_{r}(0)\right|$
$\displaystyle\leq\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|\left(w_{r}(\tau)-w_{r}(0)\right)^{\top}(x,t-T_{0})\right|\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}$
$\displaystyle\qquad+\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|\left(w_{r}(\tau)-\bar{w}_{r}(\tau)\right)^{\top}(x,t-T_{0})\right|\mathbb{I}_{r}(0),$
(48)

where we have used the fact that

$\displaystyle\left|a\right|\mathbb{I}\left\\{{\rm sgn}(a)\neq{\rm
sgn}(b)\right\\}\leq\left|a-b\right|\mathbb{I}\left\\{{\rm sgn}(a)\neq{\rm
sgn}(b)\right\\},\quad\forall a,b\in\mathbb{R}.$

Taking square both sides of (48) and apply the Jensen’s inequality, we have

$\displaystyle\left|f_{\bf{W}(\tau)}^{i}(x,t)-f^{{\rm lin},i}_{\bar{\bf
W}(\tau)}(x,t)\right|^{2}$ $\displaystyle\leq
2\sum_{r=1}^{m}\left|\left(w_{r}(\tau)-w_{r}(0)\right)^{\top}(x,t-T_{0})\right|^{2}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}$
$\displaystyle\quad+2\left(\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|\left(w_{r}(\tau)-\bar{w}_{r}(\tau)\right)^{\top}(x,t-T_{0})\right|\mathbb{I}_{r}(0)\right)^{2}$
(49)

We first bound the first term in (49).

Recall that Theorem D.1 implies that with probability at least $1-\delta$, we
have for all $\tau\geq 0$ and $r=1,\dots m$ simultaneously that

$\displaystyle\left\|w_{r}(\tau)-w_{r}(0)\right\|_{2}\leq
R_{w}=O\left(\frac{dNC^{2}_{\max}}{\sqrt{m}\lambda_{0}\sqrt{\delta}}\right).$

With this result, we apply the Cauchy-Schwarz inequality to conclude that with
probability at least $1-\delta$, it simultaneously holds for all
$\left\|x\right\|_{2}\leq R$ and $t\in[T_{0}+\Delta,T]$ that

$\displaystyle\sum_{r=1}^{m}\left|(w_{r}(\tau)-w_{r}(0))^{\top}(x,t-T_{0})\right|^{2}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}$
$\displaystyle\leq\left\|w_{r}(\tau)-w_{r}(0)\right\|_{2}^{2}\left\|(x,t-T_{0})\right\|^{2}_{2}\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}$
$\displaystyle\leq
R_{w}^{2}C_{\max}^{2}\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}.$
(50)

Thus, taking expectation over $X_{t}$ and integration over
$t\in[T_{0}+\Delta,T]$, with probability at least $1-\delta$, we can have

$\displaystyle\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\sum_{r=1}^{m}\left|(w_{r}(\tau)-w_{r}(0))^{\top}(x,t-T_{0})\right|^{2}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}{\rm
d}P_{X_{t}}(x){\rm d}t$ $\displaystyle\leq
R_{w}^{2}C^{2}_{\max}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}{\rm
d}P_{X_{t}}(x){\rm d}t.$ (51)

Next, similar to (29), we have for all $\left\|x\right\|_{2}\leq R$ and
$t\in[T_{0}+\Delta,T]$ that

$\displaystyle\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}\leq\mathbb{I}\left\\{\left|w_{r}(0)^{\top}(x,t-T_{0})\right|\leq
R_{w}C_{\max}\right\\}+\mathbb{I}\left\\{\left\|w_{r}(\tau)-w_{r}(0)\right\|>R_{w}\right\\}.$
(52)

Also, similar to (41), by taking expectation w.r.t.
$\left\\{w_{r}(0)\right\\}_{r=1}^{m}$ in (52), we have for each $(x,t)$ in the
range that

$\displaystyle\mathbb{E}\left[\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}\right]$
$\displaystyle\leq\sum_{r=1}^{m}\mathbb{E}\left[\mathbb{I}\left\\{\left|w_{r}(0)^{\top}(x,t-T_{0})\right|\leq
R_{w}C_{\max}\right\\}\right]+\mathbb{E}\left[\mathbb{I}\left\\{\left\|w_{r}(\tau)-w_{r}(0)\right\|>R_{w}\text{
for some }r\right\\}\right]$
$\displaystyle\leq\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}}+\delta.$ (53)

Now integrating over all $(x,t)$ in the range we get

$\displaystyle\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\int\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}{\rm
d}\mathcal{N}(w_{r}(0)){\rm d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\leq(T-T_{0}-\Delta)\left(\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}}+\delta\right).$

Since $w_{r}(0)$ is independent of $X_{t}$, the Fubini’s theorem and the
Markov inequality implies that with probability at least $1-\delta$ over
random initialization, we can bound the integration in (51) as

$\displaystyle\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\sum_{r=1}^{m}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}{\rm
d}P_{X_{t}}(x){\rm
d}t\leq(T-T_{0}-\Delta)\left(\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}\delta}+1\right).$

Therefore, applying the union bound, with probability at least $1-2\delta$, we
conclude that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\sum_{r=1}^{m}\left|(w_{r}(\tau)-w_{r}(0))^{\top}(x,t-T_{0})\right|^{2}\mathbb{I}\left\\{\mathbb{I}_{r}(\tau)\neq\mathbb{I}_{r}(0)\right\\}{\rm
d}P_{X_{t}}(x){\rm d}t$ $\displaystyle\leq
R^{2}_{w}C^{2}_{\max}\left(\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}\delta}+1\right)\frac{T-T_{0}-\Delta}{T-T_{0}}$
$\displaystyle\leq\frac{2(dN)^{3}C_{\max}^{9}}{\sqrt{2\pi}\sqrt{m}\delta^{5/2}\lambda_{0}^{3}}+\frac{(dN)^{2}C_{\max}^{6}}{m\lambda_{0}^{2}{\delta}}=\mathcal{O}\left(\frac{(dN)^{3}C_{\max}^{9}}{\sqrt{m}\delta^{5/2}\lambda_{0}^{2}}\right).$
(54)

We move on to bound the second term in (48). Note that for all
$\left\|x\right\|_{2}\leq R$ and $t\in[T_{0}+\Delta,T]$, the Cauchy-Schwarz
inequality implies

$\displaystyle\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|(w_{r}(\tau)-\bar{w}_{r}(\tau))^{\top}(x,t-T_{0})\right|\mathbb{I}_{r}(0)$
$\displaystyle\leq\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}\left\|(x,t-T_{0})\right\|_{2}\mathbb{I}_{r}(0)$
$\displaystyle\leq\frac{C_{\max}}{\sqrt{m}}\sum_{r=1}^{m}\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}.$
(55)

Recall the GD updating rule for $w_{r}(\tau)$ and $\bar{w}_{r}(\tau)$:

$\displaystyle w_{r}(\tau+1)$
$\displaystyle=w_{r}(\tau)-\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}(u_{j}^{i}(\tau)-y_{j}^{i})a_{r}^{i}z_{j}\mathbb{I}\left\\{w_{r}(\tau)^{\top}z_{j}\geq
0\right\\},$ $\displaystyle\bar{w}_{r}(\tau+1)$
$\displaystyle=\bar{w}_{r}(\tau)-\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}(u_{j}^{{\rm
lin},i}(\tau)-y_{j}^{i})a_{r}^{i}z_{j}\mathbb{I}\left\\{w_{r}(0)^{\top}z_{j}\geq
0\right\\}.$

Here, we have denoted $u_{j}^{i}(\tau)=f_{{\bf W}(\tau)}^{i}$ and $u_{j}^{{\rm
lin},i}(\tau)=f_{\bar{\bf W}(\tau)}^{{\rm lin},i}$, both evaluated at the
sample $(X_{t_{j}},t_{j})$. Thus, we can write

$\displaystyle w_{r}(\tau+1)-\bar{w}_{r}(\tau+1)$
$\displaystyle=w_{r}(\tau)-\bar{w}_{r}(\tau)-\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}(u_{j}^{i}(\tau)-y_{j}^{i})a_{r}^{i}z_{j}\left(\mathbb{I}_{j,r}(\tau)-\mathbb{I}_{j,r}(0)\right)$
$\displaystyle\quad-\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}\left(u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)\right)a_{r}^{i}z_{j}\mathbb{I}_{j,r}(0).$

Taking norm both sides and apply the Cauchy-Schwarz inequality, we have

$\displaystyle\left\|w_{r}(\tau+1)-\bar{w}_{r}(\tau+1)\right\|_{2}$
$\displaystyle\leq\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}+\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}\left|u_{j}^{i}(\tau)-y_{j}^{i}\right|\left|a_{r}^{i}\right|\left\|z_{j}\right\|_{2}\left|\mathbb{I}_{j,r}(\tau)-\mathbb{I}_{j,r}(0)\right|$
$\displaystyle\qquad+\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}\left|u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)\right|\left|a_{j}^{i}\right|\left\|z_{j}\right\|_{2}\left|\mathbb{I}_{j,r}(0)\right|$
$\displaystyle\leq\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}+\frac{\eta\sqrt{d}C_{\max}}{\sqrt{m}}\sqrt{\sum_{j=1}^{N}\sum_{i=1}^{d}\left(u_{j}^{i}(\tau)-y_{j}^{i}\right)^{2}}\sqrt{\sum_{j=1}^{N}\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau)\neq\mathbb{I}_{j,r}(0)\right\\}}$
$\displaystyle\quad+\frac{\eta\sqrt{d}C_{\max}}{\sqrt{m}}\sqrt{\sum_{j=1}^{N}\sum_{i=1}^{d}\left(u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)\right)^{2}}\sqrt{\sum_{j=1}^{N}\mathbb{I}_{j,r}(0)}.$

Summation over all neurons and apply the Cauchy-Schwarz inequality again, we
can conclude

$\displaystyle\sum_{r=1}^{m}\left\|w_{r}(\tau+1)-\bar{w}_{r}(\tau+1)\right\|_{2}$
$\displaystyle\leq\sum_{r=1}^{m}\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}+\eta\sqrt{d}C_{\max}\sqrt{\sum_{j=1}^{N}\sum_{i=1}^{d}\left(u_{j}^{i}(\tau)-y_{j}^{i}\right)^{2}}\sqrt{\sum_{r=1}^{m}\sum_{j=1}^{N}\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau)\neq\mathbb{I}_{j,r}(0)\right\\}}$
$\displaystyle\quad+\eta\sqrt{d}C_{\max}\sqrt{\sum_{j=1}^{N}\sum_{i=1}^{d}\left(u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)\right)^{2}}\sqrt{\sum_{r=1}^{m}\sum_{j=1}^{N}\mathbb{I}_{j,r}(0)}.$
(56)

Since $w_{r}(0)=\bar{w}_{r}(0)$, telescoping sum over (56) leads to

$\displaystyle\sum_{r=1}^{m}\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}$
$\displaystyle=\eta\sqrt{d}C_{\max}\sum_{s=0}^{\tau-1}\left\|u(s)-y\right\|_{2}\sqrt{\sum_{r=1}^{m}\sum_{j=1}^{N}\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau)\neq\mathbb{I}_{j,r}(0)\right\\}}$
$\displaystyle\quad+\eta\sqrt{d}C_{\max}\sum_{s=0}^{\tau-1}\left\|u(\tau)-u^{\rm
lin}(s)\right\|_{2}\sqrt{\sum_{r=1}^{m}\sum_{j=1}^{N}\mathbb{I}_{j,r}(0)}.$
(57)

Theorem D.1 implies that with probability at least $1-\delta$,

$\displaystyle\left\|u(\tau)-y\right\|_{2}^{2}\leq(1-\eta\lambda_{0})^{\tau}\left\|u(0)-y\right\|_{2}^{2}=(1-\eta\lambda_{0})^{\tau}\mathcal{O}\left(\frac{dNC_{\max}^{2}}{\delta}\right).$
(58)

Moreover, (53) leads to

$\displaystyle\mathbb{E}\left[\sum_{r=1}^{m}\sum_{j=1}^{N}\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau)\neq\mathbb{I}_{j,r}(0)\right\\}\right]\leq
N\left(\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}}+\delta\right).$

The Markov inequality implies with probability at least $1-\delta$, we have

$\displaystyle{\sum_{r=1}^{m}\sum_{j=1}^{N}\mathbb{I}\left\\{\mathbb{I}_{j,r}(\tau)\neq\mathbb{I}_{j,r}(0)\right\\}}\leq
N\left(\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}\delta}+1\right)=\mathcal{O}\left(\frac{dN^{2}\sqrt{m}C_{\max}^{3}}{\lambda_{0}C_{\min}\delta^{3/2}}\right).$
(59)

It remains to bound $\left\|u(\tau)-u^{\rm lin}(\tau)\right\|_{2}$ with high
probability. From the definitions of $u(\tau)$ and $u^{\rm lin}(\tau)$, we
have

$\displaystyle u_{j}^{i}(\tau+1)-u^{{\rm lin},i}_{j}(\tau+1)$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\sigma\left(w_{r}(\tau+1)^{\top}z_{j}\right)-\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\bar{w}_{r}(\tau+1)^{\top}z_{j}\mathbb{I}_{j,r}(0)$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau+1)^{\top}z_{j}\mathbb{I}_{j,r}(\tau)+\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau+1)^{\top}z_{j}\left(\mathbb{I}_{j,r}(\tau+1)-\mathbb{I}_{j,r}(\tau)\right)$
$\displaystyle\quad-\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\bar{w}_{r}(\tau+1)^{\top}z_{j}\mathbb{I}_{j,r}(0)$
$\displaystyle=\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\left(w_{r}(\tau)-\eta\frac{\partial\widehat{\mathcal{L}}(\bf
W(\tau))}{\partial
w_{r}(\tau)}\right)^{\top}z_{j}\mathbb{I}_{j,r}(\tau)+\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau+1)^{\top}z_{j}\left(\mathbb{I}_{j,r}(\tau+1)-\mathbb{I}_{j,r}(\tau)\right)$
$\displaystyle\quad-\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\left(\bar{w}_{r}(\tau)-\eta\frac{\partial\widehat{\mathcal{L}}^{\rm
lin}(\bar{\bf
W}(\tau))}{\partial\bar{w}_{r}(\tau)}\right)^{\top}z_{j}\mathbb{I}_{j,r}(0)$
$\displaystyle=u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)+\frac{\eta}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}\left(\frac{\partial\widehat{\mathcal{L}}^{\rm
lin}(\bar{\bf
W}(\tau))}{\partial\bar{w}_{r}(\tau)}\mathbb{I}_{j,r}(0)-\frac{\partial\widehat{\mathcal{L}}(\bf
W(\tau))}{\partial w_{r}(\tau)}\mathbb{I}_{j,r}(\tau)\right)^{\top}z_{j}$
$\displaystyle\quad+\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau+1)^{\top}z_{j}\left(\mathbb{I}_{j,r}(\tau+1)-\mathbb{I}_{j,r}(\tau)\right)$
$\displaystyle=u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)+\eta\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{{\rm
lin},k}(\tau)-y_{\ell}^{k})H_{j\ell}^{ik}(0)-\eta\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})H_{j\ell}^{ik}(\tau)$
$\displaystyle\quad+\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau+1)^{\top}z_{j}\left(\mathbb{I}_{j,r}(\tau+1)-\mathbb{I}_{j,r}(\tau)\right)$
$\displaystyle=u_{j}^{i}(\tau)-u_{j}^{{\rm
lin},i}(\tau)+\eta\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{{\rm
lin},k}(\tau)-u_{\ell}^{k})H_{j\ell}^{ik}(0)-\eta\sum_{\ell=1}^{N}\sum_{k=1}^{d}(u_{\ell}^{k}(\tau)-y_{\ell}^{k})(H_{j\ell}^{ik}(0)-H_{j\ell}^{ik}(\tau))$
$\displaystyle\quad+\frac{1}{\sqrt{m}}\sum_{r=1}^{m}a_{r}^{i}w_{r}(\tau+1)^{\top}z_{j}\left(\mathbb{I}_{j,r}(\tau+1)-\mathbb{I}_{j,r}(\tau)\right).$

Define a block matrix ${\bf Z}(\tau)$ such that its $(i,j)$-th row is

$\displaystyle\left({\bf
Z}^{i}_{j}\right)^{\top}(\tau)\coloneqq\frac{1}{\sqrt{m}}\left[a_{1}^{i}z_{j}^{\top}\mathbb{I}_{j,1}(\tau),\dots,a_{m}^{i}z_{j}^{\top}\mathbb{I}_{j,m}(\tau)\right].$

By vectorization, we rewrite the above equation in a compact form:

$\displaystyle u(\tau+1)-u^{\rm lin}(\tau+1)$ $\displaystyle=u(\tau)-u^{\rm
lin}(\tau)+\eta H(0)(u^{\rm lin}(\tau)-u(\tau))-\eta(H(0)-H(\tau))(u(\tau)-y)$
$\displaystyle\quad+({\bf Z}(\tau+1)-{\bf Z}(\tau)){\rm vec}(\bf W)(\tau+1)$
$\displaystyle=\left(I_{dN}-\eta H(0)\right)(u(\tau)-u^{\rm
lin}(\tau))-\eta\underbrace{(H(0)-H(\tau))(u(\tau)-y)}_{\eqqcolon\xi(\tau)}$
$\displaystyle\quad+\underbrace{({\bf Z}(\tau+1)-{\bf Z}(\tau)){\rm vec}(\bf
W)(\tau+1)}_{\eqqcolon\bar{\xi}(\tau)}.$ (60)

Unrolling the recursion (60) and noting that $u(0)=u^{\rm lin}(0)$, we can
have

$\displaystyle u(\tau)-u^{\rm lin}(\tau)=\sum_{s=0}^{\tau-1}(I_{dN}-\eta
H(0))^{\tau-1-s}\left(-\eta\xi(s)+\bar{\xi}(s)\right).$

The summation should be understood as $0$ when $\tau=0$. Taking norm both
sides and apply the Cauchy-Schwarz inequality and the triangle inequality, we
get

$\displaystyle\left\|u(\tau)-u^{\rm lin}(\tau)\right\|_{2}$
$\displaystyle\leq\sum_{s=0}^{\tau-1}\left\|(I_{dN}-\eta
H(0))^{\tau-1-s}\right\|_{2}\left(\eta\left\|\xi(s)\right\|_{2}+\left\|\bar{\xi}(s)\right\|_{2}\right)$
$\displaystyle\leq\sum_{s=0}^{\tau-1}(1-\eta\lambda_{0})^{\tau-1-s}\left(\eta\left\|\xi(s)\right\|_{2}+\left\|\bar{\xi}(s)\right\|_{2}\right).$
(61)

Here, we have applied Assumption 3.8 and Weyl’s inequality to show that
$\lambda_{\min}(H(0))\geq\lambda_{0}/2$ with probability at least $1-\delta$
(Du et al., 2018, Lemma 3.2). We turn to bound $\left\|\xi(s)\right\|_{2}$ and
$\left\|\bar{\xi}(s)\right\|_{2}$, respectively. Note that (43) and (58) imply
that with probability at least $1-2\delta$,

$\displaystyle\left\|\xi(s)\right\|_{2}$
$\displaystyle\leq\left\|H(0)-H(s)\right\|_{2}\left\|u(s)-y\right\|_{2}$
$\displaystyle=\mathcal{O}\left(\frac{(dN)^{3}C_{\max}^{4}}{\sqrt{m}\lambda_{0}\delta^{3/2}C_{\min}}(1-\eta\lambda_{0})^{s/2}\sqrt{\frac{dN}{\delta}}\right)$
$\displaystyle\leq\mathcal{O}\left(\frac{(dN)^{7/2}C_{\max}^{4}}{\sqrt{m}\lambda_{0}\delta^{2}C_{\min}}(1-\eta\lambda_{0})^{s/2}\right).$
(62)

Next, to upper bound $\left\|\bar{\xi}(s)\right\|_{2}$, note that for each
$(i,j)$-entry we have

$\displaystyle\left|\bar{\xi}_{j}^{i}(s)\right|$
$\displaystyle\leq\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|a_{r}^{i}\right|\left|w_{r}(s+1)^{\top}z_{j}\right|\left|\mathbb{I}_{j,r}(s+1)-\mathbb{I}_{j,r}(s)\right|$
$\displaystyle\leq\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|w_{r}(s+1)^{\top}z_{j}-w_{r}(s)^{\top}z_{j}\right|\left|\mathbb{I}_{j,r}(s+1)-\mathbb{I}_{j,r}(s)\right|$
$\displaystyle\leq\frac{C_{\max}}{\sqrt{m}}\sum_{r=1}^{m}\left\|w_{r}(s+1)-w_{r}(s)\right\|_{2}\left|\mathbb{I}_{j,r}(s+1)-\mathbb{I}_{j,r}(s)\right|.$
(63)

To proceed, we apply the GD-updating rule to have

$\displaystyle\left\|w_{r}(s+1)-w_{r}(s)\right\|_{2}$
$\displaystyle\leq\left\|\frac{\eta}{\sqrt{m}}\sum_{j=1}^{N}\sum_{i=1}^{d}(u_{j}^{i}(s)-y_{j}^{i})a_{r}^{i}z_{j}\mathbb{I}_{j,r}(s)\right\|_{2}$
$\displaystyle\leq\frac{\eta
C_{\max}}{\sqrt{m}}\left\|u(s)-y\right\|_{1}\leq\frac{\eta\sqrt{dN}C_{\max}}{\sqrt{m}}\left\|u(s)-y\right\|_{2}.$
(64)

Plugging (64) into (63), we have with probability at least $1-3\delta$,

$\displaystyle\left|\bar{\xi}_{j}^{i}(s)\right|$
$\displaystyle\leq\frac{\eta\sqrt{dN}C_{\max}^{2}}{m}\left\|u(s)-y\right\|_{2}\sum_{r=1}^{m}\left|\mathbb{I}_{j,r}(s+1)-\mathbb{I}_{j,r}(s)\right|$
$\displaystyle\leq\frac{\eta\sqrt{dN}C_{\max}^{2}}{m}\left\|u(s)-y\right\|_{2}\left(\sum_{r=1}^{m}\left|\mathbb{I}_{j,r}(s+1)-\mathbb{I}_{j,r}(0)\right|+\sum_{r=1}^{m}\left|\mathbb{I}_{j,r}(s)-\mathbb{I}_{j,r}(0)\right|\right)$
$\displaystyle=\mathcal{O}\left(\frac{\eta\sqrt{dN}C_{\max}^{3}}{m}(1-\eta\lambda_{0})^{s/2}\sqrt{\frac{dN}{\delta}}\left(\frac{2mR_{w}C_{\max}}{\sqrt{2\pi}C_{\min}\delta^{2}}+1\right)\right)$
$\displaystyle=\mathcal{O}\left(\frac{\eta(dN)^{2}C_{\max}^{5}}{\sqrt{m}\lambda_{0}\delta^{2}C_{\min}}(1-\eta\lambda_{0})^{s/2}\right).$

Thus, we can have with probability at least $1-3\delta$,

$\displaystyle\left\|\bar{\xi}(s)\right\|_{2}\leq\left\|\bar{\xi}(s)\right\|_{1}=\sum_{j=1}^{N}\sum_{i=1}^{d}\left|\bar{\xi}_{j}^{i}(s)\right|=\mathcal{O}\left(\frac{\eta(dN)^{3}C_{\max}^{5}}{\sqrt{m}\lambda_{0}\delta^{2}C_{\min}}(1-\eta\lambda_{0})^{s/2}\right).$
(65)

Noting that

$\displaystyle\sum_{s=0}^{\tau-1}(1-\eta\lambda_{0})^{\tau-1-\frac{s}{2}}$
$\displaystyle=(1-\eta\lambda_{0})^{\frac{\tau-1}{2}}\sum_{s=0}^{\tau-1}(1-\eta\lambda_{0})^{\frac{\tau-1}{2}-\frac{s}{2}}$
$\displaystyle\leq(1-\eta\lambda_{0})^{\frac{\tau-1}{2}}\frac{1}{1-\sqrt{1-\eta\lambda_{0}}}$
$\displaystyle\leq\frac{2(1-\eta\lambda_{0})^{\frac{\tau-1}{2}}}{\eta\lambda_{0}}.$

Therefore, we can conclude that with probability at least $1-5\delta$ that

$\displaystyle\left\|u(\tau)-u^{\rm
lin}(\tau)\right\|_{2}=\mathcal{O}\left(\frac{(dN)^{7/2}C_{\max}^{5}}{\sqrt{m}\lambda_{0}^{2}\delta^{2}C_{\min}}(1-\eta\lambda_{0})^{\frac{\tau-1}{2}}\right).$
(66)

Now, substitution (58), (59) and (66) back into (57), we have with probability
at least $1-7\delta$ that

$\displaystyle\sum_{r=1}^{m}\left\|w_{r}(\tau)-\bar{w}_{r}(\tau)\right\|_{2}$
$\displaystyle\lesssim\eta\sqrt{d}C_{\max}\sum_{s=0}^{\tau-1}(1-\eta\lambda_{0})^{\frac{s}{2}}\sqrt{\frac{dN}{\delta}}\frac{\sqrt{d}Nm^{1/4}C^{2}_{\max}}{\sqrt{\lambda_{0}}\sqrt{C_{\min}}\delta^{3/4}}$
$\displaystyle\quad+\eta\sqrt{d}C_{\max}\sum_{s=1}^{\tau-1}\frac{(dN)^{7/2}C_{\max}^{5}}{\sqrt{m}\lambda_{0}^{2}\delta^{2}C_{\min}}(1-\eta\lambda_{0})^{\frac{s-1}{2}}\sqrt{mN}$
$\displaystyle\lesssim\frac{(dN)^{3/2}m^{1/4}C^{3}_{\max}}{\lambda_{0}^{3/2}\delta^{5/4}\sqrt{C_{\min}}}+\frac{(dN)^{9/2}C_{\max}^{6}}{\lambda_{0}^{3}\delta^{2}C_{\min}}$
$\displaystyle\lesssim\frac{(dN)^{9/2}m^{1/4}C_{\max}^{6}}{\lambda_{0}^{3/2}\delta^{2}C_{\min}}.$
(67)

Since (67) holds with high probability independent of given $(x,t)$, we know
that with probability at least $1-7\delta$, (55) can be upper bounded
simultaneously over all $\left\|x\right\|_{2}\leq R$ and
$t\in[T_{0}+\Delta,T]$ that

$\displaystyle\frac{1}{\sqrt{m}}\sum_{r=1}^{m}\left|(w_{r}(\tau)-\bar{w}_{r}(\tau))^{\top}(x,t-T_{0})\right|\mathbb{I}_{r}(0)$
$\displaystyle\lesssim\frac{(dN)^{9/2}C_{\max}^{6}}{m^{1/4}\lambda_{0}^{3/2}\delta^{2}C_{\min}}.$
(68)

Integrating over (48) and combining (54) and (68), we have with probability at
least $1-9\delta$ that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left|f_{\bf{W}(\tau)}^{i}(x,t)-f^{{\rm lin},i}_{\bar{\bf
W}(\tau)}(x,t)\right|^{2}{\rm d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\lesssim\frac{(dN)^{3}C_{\max}^{6}}{\sqrt{m}\delta^{5/2}\lambda_{0}^{2}}+\left(\frac{(dN)^{9/2}C_{\max}^{6}}{m^{1/4}\lambda_{0}^{3/2}\delta^{2}C_{\min}}\right)^{2}\lesssim\frac{(dN)^{9}C_{\max}^{12}}{\sqrt{m}\delta^{4}\lambda_{0}^{2}C_{\min}^{2}}.$

As a consequence, with probability at least $1-9\delta$, we have

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f_{\bf{W}(\tau)}(x,t)-f^{{\rm lin}}_{\bar{\bf
W}(\tau)}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm
d}t=\mathcal{O}\left(\frac{d(dN)^{9}C_{\max}^{12}}{\sqrt{m}\delta^{4}\lambda_{0}^{2}C_{\min}^{2}}\right).$

∎

Next, we control the coupling error between the linearized neural network
$f^{{\rm lin}}_{\bar{\bf W}(\tau)}$ an the function $f_{\tau}^{K}$ in the next
lemma. Recall the updating rule of $\gamma(\tau)$ is given by

$\displaystyle\gamma(\tau+1)=\gamma(\tau)-\eta(H\gamma(\tau)-y),\quad\gamma(0)=H^{-1}u(0).$
(69)

Consequently, multiplying both sides of the updating rule by $H$ leads to

$\displaystyle u^{K}(\tau+1)=u^{K}(\tau)-\eta H(u^{K}(\tau)-y),\quad
u^{K}(0)=u(0).$

The updating rule of $\gamma$ can be regarded as a GD updating rule under an
alternative coordinate system. Let $\omega=\sqrt{H}\gamma$ and define the
training objective

$\displaystyle\widehat{\mathcal{L}}^{K}(\omega)=\frac{1}{2}\left\|u^{K}-y\right\|_{2}^{2}=\frac{1}{2}\left\|\sqrt{H}\omega-y\right\|_{2}^{2}.$

Here we have used the fact that $u^{K}=H\gamma=\sqrt{H}\omega$. Thus, the GD
updating rule of $\omega$ is

$\displaystyle\omega(\tau+1)=\omega(\tau)-\eta\sqrt{H}\left(u^{K}(\tau)-y\right).$
(70)

Multiplying both sides of equation 70 by $\sqrt{H^{-1}}$, we have the same
updating rule of $\gamma(\tau)$.

######  Lemma D.3.

Assume the same conditions as in Theorem 3.9, if we initialize
$\gamma(0)=\bar{\gamma}(0)=H(0)^{-1}u(0)$, then it holds with probability at
least $1-\delta$ that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f^{\rm lin}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm
d}t=\tilde{\mathcal{O}}\left(\frac{d^{5}N^{4}C_{\max}^{8}}{m\delta^{2}\lambda_{0}^{2}}\right).$

###### Proof.

Note that the gradient of the training loss is

$\displaystyle\frac{\partial\widehat{\mathcal{L}}^{\rm lin}(\bar{\bf
W})}{\partial{\rm vec}(\bar{\bf W})}=\frac{\partial}{\partial{\rm
vec}(\bar{\bf W})}\frac{1}{2}\left\|u^{\rm lin}-y\right\|_{2}^{2}={\bf
Z}(0)^{\top}(u^{\rm lin}-y).$

We first show that at $\tau=0$, there is a vector
$\bar{\gamma}(0)\in\mathbb{R}^{dN}$ such that ${\rm vec}(\bar{\bf W})(0)={\bf
Z}(0)$. Note that our choice implies $\bar{\gamma}(0)=\gamma(0)=\left({\bf
Z}(0){\bf Z}(0)^{\top}\right)^{-1}u^{\rm lin}(0)$. Let ${\bf Z}(0)=U\Sigma
V^{\top}$ be the corresponding singular value decomposition. Since ${\bf
Z}(0)$ has full row rank, we can write the diagonal entries of $\Sigma$ as
$\sigma_{1}\geq\dots\geq\sigma_{dN}>0$. Noting that $u^{\rm lin}(0)={\bf
Z}(0){\rm vec}(\bar{\bf W}(0))$,

$\displaystyle{\bf Z}(0)^{\top}\bar{\gamma}(0)$ $\displaystyle={\bf
Z}(0)^{\top}\left({\bf Z}(0){\bf Z}(0)^{\top}\right)^{-1}u^{\rm lin}(0)$
$\displaystyle=V\Sigma^{\top}U^{\top}(U\Sigma
V^{\top}V\Sigma^{\top}U^{\top})^{-1}U\Sigma V^{\top}{\rm vec}(\bar{\bf W}(0))$
$\displaystyle=V\Sigma^{\top}U^{\top}(U{\rm
diag}(\sigma_{1}^{-2},\dots,\sigma_{dN}^{-2})U^{\top})U\Sigma V^{\top}{\rm
vec}(\bar{\bf W}(0))$ $\displaystyle=V\begin{pmatrix}I_{dN}&0\\\
0&0\end{pmatrix}V^{\top}{\rm vec}(\bar{\bf W}(0))$
$\displaystyle=V\begin{pmatrix}I_{dN}&0\\\
0&0\end{pmatrix}\begin{pmatrix}I_{dN}&0\\\ 0&0\end{pmatrix}V^{\top}{\rm
vec}(\bar{\bf W}(0))={\rm vec}(\bar{\bf W}(0)).$

It follows for each $\tau$, there is a vector
$\bar{\gamma}(\tau)\in\mathbb{R}^{dN}$ such that

$\displaystyle{\rm vec}(\bar{\bf W}(\tau))$ $\displaystyle={\rm vec}(\bar{\bf
W}(\tau-1))-\eta{\bf Z}(0)^{\top}(u^{\rm lin}(\tau-1)-y)={\bf
Z}(0)^{\top}\bar{\gamma}(\tau).$

Define a matrix ${\bf Z}(x,t)\in\mathbb{R}^{d\times m(d+1)}$ such that its
$i$-th row is

$\displaystyle\left({\bf
Z}^{i}(x,t)\right)^{\top}\coloneqq\frac{1}{\sqrt{m}}\left[a_{1}^{i}(x,t-T_{0})^{\top}\mathbb{I}_{1}(0),\dots,a_{m}^{i}(x,t-T_{0})^{\top}\mathbb{I}_{m}(0)\right].$

Next, one can rewrite

$\displaystyle f^{\rm lin}_{\bar{\bf W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)$
$\displaystyle={\bf Z}(x,t){\rm vec}(\bar{\bf
W}(\tau))-\sum_{j=1}^{N}K((X_{t_{j}},t_{j}),(x,t))\gamma_{j}(\tau)$
$\displaystyle={\bf Z}(x,t){\bf
Z}(0)^{\top}\bar{\gamma}(\tau)-\sum_{j=1}^{N}K((X_{t_{j}},t_{j}),(x,t))\gamma_{j}(\tau)$
$\displaystyle={\bf Z}(x,t){\bf
Z}(0)^{\top}\bar{\gamma}(\tau)-\hat{K}(x,t)\gamma(\tau)$ $\displaystyle={\bf
Z}(x,t){\bf
Z}(0)^{\top}\left(\bar{\gamma}(\tau)-\gamma(\tau)\right)-\left({\bf
Z}(x,t){\bf Z}(0)^{\top}-\hat{K}(x,t)\right)\gamma(\tau),$ (71)

where we have defined

$\displaystyle\hat{K}(x,t)\coloneqq[K((X_{t_{1}},t_{1}),(x,t)),\dots,K((X_{t_{N}},t_{N}),(x,t))],\quad\gamma(\tau)\coloneqq[\gamma_{1}^{\top}(\tau),\dots,\gamma_{N}^{\top}(\tau)]^{\top}.$

Taking square of both sides of (71), we can get

$\displaystyle\left\|f^{\rm lin}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}$ $\displaystyle\leq
2\left\|{\bf Z}(x,t){\bf
Z}(0)^{\top}\left(\bar{\gamma}(\tau)-\gamma(\tau)\right)\right\|_{2}^{2}+2\left\|\left({\bf
Z}(x,t){\bf Z}(0)^{\top}-\hat{K}(x,t)\right)\gamma(\tau)\right\|_{2}^{2}$
$\displaystyle\leq 2\left\|{\bf Z}(x,t){\bf
Z}(0)^{\top}\right\|_{2}^{2}\left\|{\bar{\gamma}(\tau)-\gamma(\tau)}\right\|_{2}^{2}+2\left\|{{\bf
Z}(x,t){\bf
Z}(0)^{\top}-\hat{K}(x,t)}\right\|_{2}^{2}\left\|\gamma(\tau)\right\|_{2}^{2}.$

Since $H(0)={\bf Z}(0){\bf Z}(0)^{\top}$ and the Gram matrix of $K$ is $H$, we
have

$\displaystyle u^{\rm lin}(\tau)-u^{K}(\tau)$
$\displaystyle=H(0)\bar{\gamma}(\tau)-H\gamma(\tau)$
$\displaystyle=H(0)(\bar{\gamma}(\tau)-\gamma(\tau))+(H(0)-H)\gamma(\tau).$

We first upper bound $\left\|u^{\rm lin}(\tau)-u^{K}(\tau)\right\|_{2}$. The
GD updating rules imply

$\displaystyle u^{\rm lin}(\tau+1)$ $\displaystyle=u^{\rm lin}(\tau)-\eta
H(0)(u^{\rm lin}(\tau)-y),$ $\displaystyle u^{K}(\tau+1)$
$\displaystyle=u^{K}(\tau)-\eta H(u^{K}(\tau)-y),$

with $u^{\rm lin}(0)=u^{K}(0)=u(0)$. It follows

$\displaystyle u^{\rm lin}(\tau+1)-u^{K}(\tau+1)$ $\displaystyle=u^{\rm
lin}(\tau)-u^{K}(\tau)-\eta(H-H(0))(u^{K}(\tau)-y)$ $\displaystyle\qquad-\eta
H(0)(u^{\rm lin}(\tau)-u^{K}(\tau))$ $\displaystyle=(I_{dN}-\eta H(0))(u^{\rm
lin}(\tau)-u^{K}(\tau))-\eta(H-H(0))(u^{K}(\tau)-y).$ (72)

Unrolling (72), we have

$\displaystyle u^{\rm lin}(\tau)-u^{K}(\tau)$ $\displaystyle=(I_{dN}-\eta
H(0))^{\tau}(u^{\rm lin}(0)-u^{K}(0))$
$\displaystyle\qquad-\eta\sum_{s=0}^{\tau-1}(I_{dN}-\eta
H(0))^{\tau-1-s}(H-H(0))(u^{K}(s)-y)$
$\displaystyle=-\eta\sum_{s=0}^{\tau-1}(I_{dN}-\eta
H(0))^{\tau-1-s}(H-H(0))(u^{K}(s)-y).$

Taking norm of both sides, we have

$\displaystyle\left\|u^{\rm lin}(\tau)-u^{K}(\tau)\right\|_{2}$
$\displaystyle\leq\eta\left\|H-H(0)\right\|_{2}\sum_{s=0}^{\tau-1}\left\|I_{dN}-\eta
H(0)\right\|_{2}^{\tau-1-s}\left\|u^{K}(s)-y\right\|_{2}$
$\displaystyle\leq\eta\left\|H-H(0)\right\|_{2}\sum_{s=0}^{\tau-1}\left(1-\frac{\eta\lambda_{0}}{2}\right)^{\tau-1-s}\left\|u^{K}(s)-y\right\|_{2}$
$\displaystyle\leq\eta\left\|H-H(0)\right\|_{2}\max_{0\leq
s\leq\tau-1}\left\|u^{K}(s)-y\right\|_{2}\sum_{s=0}^{\tau-1}\left(1-\frac{\eta\lambda_{0}}{2}\right)^{\tau-1-s}.$

Note that with probability at least $1-\delta$,

$\displaystyle\max_{0\leq s\leq\tau-1}\left\|u^{K}(s)-y\right\|_{2}$
$\displaystyle=\left\|u^{K}(0)-y\right\|_{2}=\left\|u(0)-y\right\|_{2}=\mathcal{O}\left(\frac{\sqrt{dN}C_{\max}}{\sqrt{\delta}}\right).$
(73)

With (73), we can obtain that with probability at least $1-2\delta$, it holds

$\displaystyle\left\|u^{\rm lin}(\tau)-u^{K}(\tau)\right\|_{2}$
$\displaystyle\leq\eta\mathcal{O}\left(\frac{dNC_{\max}\sqrt{\log((dN)^{2}/\delta)}}{\sqrt{m}}\right)\mathcal{O}\left(\frac{\sqrt{dN}C_{\max}}{\sqrt{\delta}}\right)\frac{2}{\eta\lambda_{0}}$
$\displaystyle=\tilde{\mathcal{O}}\left(\frac{(dN)^{3/2}(C_{\max})^{2}}{\sqrt{m}\lambda_{0}\delta}\right).$

It remains to bound $\left\|\gamma(\tau)\right\|_{2}$. The GD updating rule
leads to

$\displaystyle\gamma(\tau+1)=\gamma(\tau)-\eta(H\gamma(\tau)-y)=(I_{dN}-\eta
H)\gamma(\tau)+\eta y.$

Unrolling the recursive formula, we can have

$\displaystyle\gamma(\tau)=(I_{dN}-\eta
H)^{\tau}\gamma(0)+\eta\sum_{s=0}^{\tau-1}(I_{dN}-\eta H)^{s}y.$

Taking norm both sides, we have

$\displaystyle\left\|\gamma(\tau)\right\|_{2}\leq\left\|I_{dN}-\eta
H\right\|^{\tau}_{2}\left\|\gamma(0)\right\|_{2}+\eta\left\|\sum_{s=0}^{\tau-1}(I_{dN}-\eta
H)^{s}\right\|_{2}\left\|y\right\|_{2}.$

Note that

$\displaystyle\sum_{s=0}^{\tau-1}(I_{dN}-\eta H)^{s}=(I_{dN}-(I_{dN}-\eta
H)^{\tau})(\eta H)^{-1}\preceq\eta^{-1}H^{-1},$

where we have chosen $\eta$ small enough so that $I_{dN}-\eta H$ is positive
definite. Therefore, with probability at least $1-\mathcal{O}(\delta)$, we
have

$\displaystyle\left\|\gamma(\tau)\right\|_{2}$
$\displaystyle\leq\left\|H^{-1}\right\|_{2}\left\|u(0)\right\|_{2}+\left\|H^{-1}\right\|_{2}\left\|y\right\|_{2}=\mathcal{O}\left(\frac{\sqrt{dN}C_{\max}}{\lambda_{0}\sqrt{\delta}}\right).$

Finally, we have

$\displaystyle\frac{\lambda_{0}}{2}\left\|\bar{\gamma}(\tau)-\gamma(\tau)\right\|_{2}\leq\tilde{\mathcal{O}}\left(\frac{(dN)^{3/2}C_{\max}^{2}}{\sqrt{m}\lambda_{0}\delta}\right)+\tilde{\mathcal{O}}\left(\frac{dNC_{\max}}{\sqrt{m}}\right)\mathcal{O}\left(\frac{\sqrt{dN}C_{\max}}{\lambda_{0}\sqrt{\delta}}\right)=\tilde{\mathcal{O}}\left(\frac{(dN)^{3/2}C_{\max}^{2}}{\sqrt{m}\lambda_{0}\delta}\right).$

With all these preparations, we can bound for all $\left\|x\right\|_{2}\leq R$
and $t\in[T_{0}+\Delta,T]$,

$\displaystyle\left\|f^{\rm lin}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}$ $\displaystyle\leq
2\left\|{\bf Z}(x,t){\bf
Z}(0)^{\top}\right\|_{2}^{2}\tilde{\mathcal{O}}\left(\frac{(dN)^{3}C_{\max}^{4}}{m\lambda_{0}^{4}\delta^{2}}\right)$
$\displaystyle\qquad+2\left\|{{\bf Z}(x,t){\bf
Z}(0)^{\top}-\hat{K}(x,t)}\right\|_{2}^{2}\mathcal{O}\left(\frac{{dN}C^{2}_{\max}}{\lambda^{2}_{0}{\delta}}\right)$
(74)

Since $\left\|(x,t-T_{0})\right\|_{2}\leq C_{\max}$, we have

$\displaystyle\left\|{\bf Z}(x,t)\right\|_{2}^{2}\leq\sum_{i=1}^{d}\left\|{\bf
Z}^{i}(x,t)\right\|_{2}^{2}=\sum_{i=1}^{d}\sum_{r=1}^{m}\left\|\frac{1}{\sqrt{m}}a_{r}^{i}(x^{\top},t-T_{0})\mathbb{I}_{r}(0)\right\|_{2}^{2}\leq
dC_{\max}^{2}.$

Also, we have

$\displaystyle\left\|{\bf
Z}(0)\right\|_{2}^{2}\leq\sum_{i=1}^{d}\sum_{j=1}^{N}\left\|{\bf
Z}_{j}^{i}(0)\right\|_{2}^{2}=\sum_{i=1}^{d}\sum_{j=1}^{N}\sum_{r=1}^{m}\left\|\frac{1}{\sqrt{m}}a_{r}^{i}z_{j}^{\top}\mathbb{I}_{j,r}(0)\right\|_{2}^{2}\leq
dNC_{\max}^{2}.$

Now integration over $\left\|x\right\|_{2}\leq R$ and $t\in[T_{0}+\Delta,T]$
of (74) yields

$\displaystyle\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f^{\rm lin}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\leq\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}2d^{2}NC_{\max}^{4}\tilde{\mathcal{O}}\left(\frac{(dN)^{3}C_{\max}^{4}}{m\lambda_{0}^{4}\delta^{2}}\right){\rm
d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\qquad+2\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|{{\bf Z}(x,t){\bf
Z}(0)^{\top}-\hat{K}(x,t)}\right\|_{2}^{2}\mathcal{O}\left(\frac{{dN}C^{2}_{\max}}{\lambda^{2}_{0}{\delta}}\right){\rm
d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\leq\mathcal{O}\left(\frac{{dN}C^{2}_{\max}}{\lambda^{2}_{0}{\delta}}\right)\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|{{\bf Z}(x,t){\bf Z}(0)^{\top}-\hat{K}(x,t)}\right\|_{2}^{2}{\rm
d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\qquad+\tilde{\mathcal{O}}\left(\frac{d^{5}N^{4}C_{\max}^{8}}{{m}\lambda_{0}^{4}\delta^{2}}\right)(T-T_{0}-\Delta).$

Note that for each $i,k,j$, we can write

$\displaystyle\left({\bf Z}(x,t){\bf
Z}(0)^{\top}\right)_{j}^{ik}=\frac{1}{m}\sum_{r=1}^{m}a_{r}^{i}a_{r}^{k}(X_{t_{j}},t_{j}-T_{0})^{\top}(x,t-T_{0})\mathbb{I}_{j,r}(0)\mathbb{I}_{r}(0).$

as a sum of independent random variable bounded by $C_{\max}^{2}/m$ when
$\left\|x\right\|_{2}\leq R$ and $t\in[T_{0}+\Delta,T]$. Taking expectation
over the initialization, we have

$\displaystyle\mathbb{E}\left[\left|{\left({\bf Z}(x,t){\bf
Z}(0)^{\top}\right)_{j}^{jk}-\hat{K}_{j}^{ik}(x,t)}\right|_{2}^{2}\right]={\rm
Var}\left({\left({\bf Z}(x,t){\bf
Z}(0)^{\top}\right)_{j}^{jk}}\right)=\mathcal{O}\left(\frac{C_{\max}^{4}}{m}\right).$

Integration over all $x$ and $t$ gives us

$\displaystyle\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\mathbb{E}\left[\left|{\left({\bf Z}(x,t){\bf
Z}(0)^{\top}\right)_{j}^{jk}-\hat{K}_{j}^{ik}(x,t)}\right|_{2}^{2}\right]{\rm
d}P_{X_{t}}(x){\rm d}t$
$\displaystyle=\mathcal{O}\left(\frac{C_{\max}^{4}}{m}\right)(T-T_{0}-\Delta).$

The Fubini’s theorem and the Markov inequality implies with probability at
least $1-\delta/(d^{2}N)$, we have

$\displaystyle\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left|{\left({\bf Z}(x,t){\bf
Z}(0)^{\top}\right)_{j}^{jk}-\hat{K}_{j}^{ik}(x,t)}\right|_{2}^{2}{\rm
d}P_{X_{t}}(x){\rm
d}t\leq\mathcal{O}\left(\frac{C_{\max}^{4}d^{2}N}{m\delta}(T-T_{0}-\Delta)\right).$

Therefore, with probability at least $1-\mathcal{O}(\delta)$, we have

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f^{\rm lin}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\leq\mathcal{O}\left(\frac{{dN}C^{2}_{\max}}{\lambda^{2}_{0}{\delta}}\right)\mathcal{O}\left(\frac{C_{\max}^{4}d^{4}N^{2}}{m\delta}\right)+\tilde{\mathcal{O}}\left(\frac{d^{5}N^{4}C_{\max}^{8}}{{m}\lambda_{0}^{4}\delta^{2}}\right)$
$\displaystyle=\tilde{\mathcal{O}}\left(\frac{d^{5}N^{4}C_{\max}^{8}}{m\delta^{2}\lambda_{0}^{2}}\right).$

∎

Now we are ready to prove Theorem 3.9.

###### Proof of Theorem 3.9.

Note that

$\displaystyle\left\|f_{\bf{W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}\leq
2\left\|f_{\bf{W}(\tau)}(x,t)-f^{{\rm lin}}_{\bar{\bf
W}(\tau)}(x,t)\right\|_{2}^{2}+2\left\|f^{{\rm lin}}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}.$

Lemma D.2 and D.3 imply that with probability at least $1-\delta$, it holds
simultaneously over all $\tau\geq 0$ that

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f_{\bf{W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}{\rm
d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\leq\frac{2}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f_{\bf{W}(\tau)}(x,t)-f^{{\rm lin}}_{\bar{\bf
W}(\tau)}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\qquad+\frac{2}{T-T_{0}}\int_{T_{0}+\Delta}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f^{{\rm lin}}_{\bar{\bf
W}(\tau)}(x,t)-f_{\tau}^{K}(x,t)\right\|_{2}^{2}{\rm d}P_{X_{t}}(x){\rm d}t$
$\displaystyle\leq\mathcal{O}\left(\frac{d(dN)^{9}C_{\max}^{12}}{\sqrt{m}\delta^{4}\lambda_{0}^{2}C_{\min}^{2}}\right)+\tilde{\mathcal{O}}\left(\frac{d^{5}N^{4}C_{\max}^{8}}{m\delta^{2}\lambda_{0}^{2}}\right)$
$\displaystyle=\tilde{O}\left(\frac{d^{10}N^{9}C_{\max}^{12}}{\sqrt{m}\lambda_{0}^{2}\delta^{4}C_{\min}^{2}}\right).$

This finishes the proof. ∎

##  Appendix E Proof of Theorem 3.10

In this section, we prove Theorem 3.10. Our target is to bound

$\displaystyle\frac{1}{T-T_{0}}\int_{T_{0}}^{T}\int_{\left\|x\right\|_{2}\leq
R}\left\|f_{\tau}^{K}(x,t)-\tilde{f}_{\tau}^{K}(x,t)\right\|_{2}^{2}{\rm
d}P_{X_{t}}{\rm d}t.$

Here, $f_{\tau}^{K}$ and $\tilde{f}_{\tau}^{K}$ are trained with labels
$X_{0,j}$ and $\tilde{X}_{0,j}$, respectively. We first bound the performance
of these two kernel regressions on training samples. With the same spirit as
in the proof of Theorem D.1, let $u^{K}(\tau)$ and $\tilde{u}^{K}(\tau)$ be
the prediction of $f_{\tau}^{K}$ and $\tilde{f}_{\tau}^{K}$ on the samples,
respectively. The following lemma provides the label mismatch error on the
training samples.

######  Lemma E.1.

Assume the same conditions as in Theorem C.3 and suppose Assumption 3.7 holds.
If we set $\eta$ small enough and initialize $f_{0}^{K}$ and
$\tilde{f}_{0}^{K}$ with the same parameters $H(0)^{-1}u(0)$, then we can
upper bound

$\displaystyle\left\|u^{K}(\tau)-\tilde{u}^{K}(\tau)\right\|_{2}^{2}\leq
dNA(R_{\mathcal{H}},R)^{2}.$

###### Proof.

Note that the GD updating rule leads to

$\displaystyle u^{K}(\tau+1)$ $\displaystyle=u^{K}(\tau)-\eta
H(u^{K}(\tau)-y)$ $\displaystyle=(I_{dN}-\eta H)u^{K}(\tau)+\eta Hy$
$\displaystyle=(I_{dN}-\eta
H)^{\tau+1}u^{K}(0)+\eta\sum_{s=0}^{\tau}(I_{dN}-\eta H)^{s}Hy$
$\displaystyle=(I_{dN}-\eta H)^{\tau+1}u^{K}(0)+(I_{dN}-(I_{dN}-\eta
H)^{\tau+1})y.$

Similary, for $\tilde{u}^{K}(\tau)$, we have

$\displaystyle\tilde{u}^{K}(\tau+1)=(I_{dN}-\eta