Generalized Leverage Score Sampling for Neural Networks

Introduction

In this work, we follow the the approach in [AKM+17] and naturally generalize the result to a broader class of kernels, which is of the form

We summarize our main results and contribution as following:

Generalize the leverage score sampling theory for kernel ridge regression to a broader class of kernels.

Connect the leverage score sampling theory with neural network training.

Theoretically prove the equivalence between training regularized neural network and kernel ridge regression under both random Gaussian initialization and leverage score sampling initialization.

Related work

Given a $m\times n$ matrix $A$. Let $a_{i}^{\top}$ be the $i$-th rows of $A$ and the leverage score of the $i$-th row of $A$ is $\sigma_{i}(A)=a_{i}^{\top}(A^{\top}A)^{\dagger}a_{i}$. A row’s leverage score measures how important it is in composing the row space of $A$. If a row has a component orthogonal to all other rows, its leverage score is $1$. Removing it would decrease the rank of $A$, completely changing its row space. The coherence of $A$ is $\|\sigma(A)\|_{\infty}$. If $A$ has low coherence, no particular row is especially important. If $A$ has high coherence, it contains at least one row whose removal would significantly affect the composition of $A$’s row space.

Leverage score is a fundamental concept in graph problems and numerical linear algebra. There are many works on how to approximate leverage scores [SS11, DMIMW12, CW13, NN13] or more general version of leverages, e.g. Lewis weights [Lew78, BLM89, CP15]. From graph perspective, it has been applied to maximum matching [BLN+20, LSZ20], max-flow [DS08, Mad13, Mad16, LS20b, LS20a], generate random spanning trees [Sch18], and sparsify graphs [SS11]. From matrix perspective, it has been used to give matrix CUR decomposition [BW14, SWZ17, SWZ19] and tensor CURT decomposition [SWZ19]. From optimization perspective, it has been used to approximate the John Ellipsoid [CCLY19], linear programming [LS14, BLSS20, JSWZ20], semi-definite programming [JKL+20], and cutting plane methods [Vai89, LSW15, JLSW20].

Kernel methods can be thought of as instance-based learners: rather than learning some fixed set of parameters corresponding to the features of their inputs, they instead “remember” the $i$-th training example $(x_{i},y_{i})$ and learn for it a corresponding weight $w_{i}$. Prediction for unlabeled inputs, i.e., those not in the training set, is treated by the application of similarity function $\mathsf{K}$, called a kernel, between the unlabeled input $x^{\prime}$ and each of the training inputs $x_{i}$.

There are three lines of works that are closely related to our work. First, our work is highly related to the recent discoveries of the connection between deep learning and kernels [DFS16, Dan17, JGH18, CB18]. Second, our work is closely related to development of connection between leverage score and kernels [RR08, CW17, CMM17, MW17b, MW17a, LTOS18, AKM+17, AKM+19, ACSS20]. Third, our work is related to kernel ridge regression [Bac13, AM15, ZDW15, ACW17, MM17, ZNV+20].

There is a long line of work studying the convergence of neural network with random input assumptions [BG17, Tia17, ZSJ+17, Sol17, LY17, ZSD17, DLT+18, GLM18, BJW19]. For a quite while, it is not known to remove the randomness assumption from the input data points. Recently, there is a large number of work studying the convergence of neural network in the over-parametrization regime [LL18, DZPS19, AZLS19b, AZLS19a, DLL+19, ADH+19b, ADH+19a, SY19, BPSW20]. These results don’t need to assume that input data points are random, and only require some much weaker assumption which is called “data-separable”. Mathematically, it says for any two input data points $x_{i}$ and $x_{j}$, we have $\|x_{i}-x_{j}\|_{2}\geq\delta$. Sufficiently wide neural network requires the width $m$ to be at least $\operatorname{poly}(n,d,L,1/\delta)$, where $n$ is the number of input data points, $d$ is the dimension of input data point, $L$ is the number of layers.

Main results

In this section, we state our results. In Section 3.1, we consider the large-scale kernel ridge regression (KRR) problem. We generalize the Fourier transform result [AKM+17] of accelerating the running time of solving KRR using the tool of leverage score sampling to a broader class of kernels. In Section 3.2, we discuss the interesting application of leverage score sampling for training deep learning models due to the connection between regularized neural nets and kernel ridge regression.

In this section, we generalize the leverage score theory in [AKM+17], which analyzes the number of random features needed to approximate kernel matrix under leverage score sampling regime for the kernel ridge regression task. In the next a few paragraphs, we briefly review the settings of classical kernel ridge regression.

Due to the regularization $\lambda>0$ in this setting, instead of constructing the feature map directly from the distribution $q$, we consider the following ridge leveraged distribution:

The leverage score sampling distribution $q_{\lambda}/s_{\lambda}(K)$ takes the regularization term into consideration and achieves Eq. (2) using the following modified random features vector:

holds with probability at least $1-\delta$.

2 Application in training regularized neural network

Past literature such as [DZPS19],[ADH+19a] have already witnessed the equivalence between training a neural network and solving a kernel regression problem in a broad class of network models. In this work, we first generalize this result to the regularization case, where we connect regularized neural network with kernel ridge regression. Then we apply the above discussed the leverage score sampling theory for KRR to the task of training neural nets.

To illustrate the idea, we consider a simple model two layer neural network with ReLU activation function as in [DZPS19, SY19]Our results directly extends to multi-layer deep neural networks with all layers trained together.

Let $\kappa\in(0,1]$ be a small multiplierTo establish the training equivalence result, we assign $\kappa=1$ back to the normal case. For the training equivalence result, we pick $\kappa>0$ to be a small multiplier only to shrink the initial output of the neural network. The is the same as what is used in [AKM+17].. Let $\lambda\in(0,1)$ be the regularization parameter. We initialize the network as $a_{r}\overset{i.i.d.}{\sim}\operatorname{unif}[\{-1,1\}]$ and $w_{r}(0)\overset{i.i.d.}{\sim}\mathcal{N}(0,I_{d})$. Then we consider solving the following optimization problem using gradient descent:

On the other hand, we consider the following neural tangent kernel ridge regression problem:

We connect the problem Eq. (6) and Eq. (7) by building the following equivalence between their training and test predictors with polynomial widths:

Given any accuracy $\epsilon\in(0,1/10)$ and failure probability $\delta\in(0,1/10)$. Let multiplier $\kappa=1$, number of iterations $T=\widetilde{O}(\frac{1}{\Lambda_{0}})$, network width $m\geq\widetilde{O}(\frac{n^{4}d}{\Lambda_{0}^{4}\epsilon})$ and the regularization parameter $\lambda\leq\widetilde{O}(\frac{1}{\sqrt{m}})$. Then with probability at least $1-\delta$ over the Gaussian random initialization, we have

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

We can further show the equivalence between the test data predictors with the help of the multiplier $\kappa$.

Given any accuracy $\epsilon\in(0,1/10)$ and failure probability $\delta\in(0,1/10)$. Let multiplier $\kappa=\widetilde{O}(\frac{\epsilon\Lambda_{0}}{n})$, number of iterations $T=\widetilde{O}(\frac{1}{\kappa^{2}\Lambda_{0}})$, network width $m\geq\widetilde{O}(\frac{n^{10}d}{\epsilon^{6}\Lambda_{0}^{10}})$ and regularization parameter $\lambda\leq\widetilde{O}(\frac{1}{\sqrt{m}})$. Then with probability at least $1-\delta$ over the Gaussian random initialization, we have

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

2.2 Equivalence II, training with leverage scores

To apply the leverage score theory discussed in Section 3.1, Note the definition of the neural tangent kernel is exactly of the form:

Specifically, given regularization parameter $\lambda>0$, we can define the ridge leverage function with respect to neural tangent kernel $H^{\operatorname{cts}}$ defined in Definition 3.1 as

and corresponding probability density function

We consider training the following reweighed neural network using leverage score initialization:

We show that training this reweighed neural net with leverage score initialization is still equivalence to the neural tangent kernel ridge regression problem (7) as in following theorem:

Given any accuracy $\epsilon\in(0,1)$ and failure probability $\delta\in(0,1/10)$. Let multiplier $\kappa=1$, number of iterations $T=O(\frac{1}{\Lambda_{0}}\log(\frac{1}{\epsilon}))$, network width $m=\operatorname{poly}(\frac{1}{\Lambda_{0}},n,d,\frac{1}{\epsilon},\log(\frac{1}{\delta}))$ and regularization parameter $\lambda=\widetilde{O}(\frac{1}{\sqrt{m}})$. Then with probability at least $1-\delta$ over the random leverage score initialization, we have

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

Overview of techniques

To prove Theorem 3.3, we follow the similar proof framework as Lemma 8 in [AKM+17].

Let $K+\lambda I_{n}=V^{\top}\Sigma^{2}V$ be an eigenvalue decomposition of $K+\lambda I_{n}$. Then conclusion (5) is equivalent to

holds with probability at least $1-\delta$, which can be shown by applying matrix concentration results. Note

Applying matrix concentration Lemma 7 in [AKM+17], we complete the proof.

To establish the equivalence between training neural network with regularization and neural tangent kernel ridge regression, the key observation is that the dynamic kernel during the training is always close to the neural tangent kernel.

Then we can show the gradient flow of training regularized neural net satisfies

where term (13) is the primary term characterizing the linear convergence of $u_{\operatorname{nn}}(t)$ to $t^{*}$, and term (14) is the additive term that can be well controlled if $H(t)$ is sufficiently close to $H^{\operatorname{cts}}$. We argue the closeness of $H(t)\approx H^{\operatorname{cts}}$ as the consequence of the following two observations:

Initialization phase: At the beginning of the training, $H(0)$ can be viewed as approximating the neural tangent kernel $H^{\operatorname{cts}}$ using finite dimensional random features. Note the size of these random features corresponds to the width of the neural network (scale by the data dimension $d$). Therefore, when the neural network is sufficiently wide, it is equivalent to approximate the neural tangent kernel using sufficient high dimensional feature vectors, which ensures $H(0)$ is sufficiently close to $H^{\operatorname{cts}}$.

In the case of leverage score initialization, we further take the regularization into consideration. We use the tool of leverage score to modify the initialization distribution and corresponding network parameter, to give a smaller upper-bound of the width of the nets needed.

Training phase: If the net is sufficiently wide, we can observe the over-parametrization phenomenon such that the weight estimate $W(t)$ at time $t$ will be sufficiently close to its initialization $W(0)$, which implies the dynamic kernel $H(t)$ being sufficiently close to $H(0)$. Due to the fact $H(0)\approx H^{\operatorname{cts}}$ argued in initialization phase, we have $H(t)\approx H^{\operatorname{cts}}$ throughout the algorithm.

Combining both observations, we are able to iteratively show the (nearly) linear convergence property of training the regularized neural net as in following lemma:

$\|w_{r}(0)-w_{r}(t)\|_{2}\leq\epsilon_{W}$, $\forall r\in[m]$

$\|H(0)-H(t)\|_{2}\leq\epsilon_{H}^{\prime}$

$\|u_{\operatorname{nn}}(t)-u^{*}\|_{2}^{2}\leq\max\{\epsilon_{\operatorname{train}}^{2},e^{-(\kappa^{2}\Lambda_{0}+\lambda)t/2}\|u_{\operatorname{nn}}(0)-u^{*}\|_{2}^{2}\}$

Given arbitrary accuracy $\epsilon\in(0,1)$, if we choose $\epsilon_{\operatorname{train}}=\epsilon$, $T=\widetilde{O}(\frac{1}{\kappa^{2}\Lambda_{0}})$ and $m$ sufficiently large in Lemma D.14, then we have $\|u_{\operatorname{nn}}(t)-u^{*}\|_{2}\leq\epsilon$, indicating the equivalence between training neural network with regularization and neural tangent kernel ridge regression for the training data predictions.

To further argue the equivalence for any given test data $x_{\operatorname{test}}$, we observe the similarity between the gradient flows of neural tangent kernel ridge regression $u_{\operatorname{ntk},\operatorname{test}}(t)$ and regularized neural networks $u_{\operatorname{nn},\operatorname{test}}(t)$ as following:

By choosing the multiplier $\kappa>0$ small enough, we can bound the initial difference between these two predictors. Combining with above similarity between gradient flows, we are able to show $|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{ntk},\operatorname{test}}(T)|\geq\epsilon/2$ for appropriate $T>0$. Finally, note the linear convergence property of the gradient of the kernel ridge regression, we can prove $|u_{\operatorname{nn},\operatorname{test}}(T)-u^{*}_{\operatorname{ntk},\operatorname{test}}|\geq\epsilon$.

Using the similar idea, we can also show the equivalence for test data predictors and the case of leverage score initialization. We refer to the Appendix for a detailed proof sketch and rigorous proof.

Conclusion

In this paper, we generalize the leverage score sampling theory for kernel approximation. We discuss the interesting application of connecting leverage score sampling and training regularized neural networks. We present two theoretical results: 1) the equivalence between the regularized neural nets and kernel ridge regression problems under the classical random Gaussian initialization for both training and test predictors; 2) the new equivalence under the leverage score initialization. We believe this work can be the starting point of future study on the use of leverage score sampling in neural network training.

In the appendix, we present our complete results and rigorous proofs. Section A presents some well-known mathematically results that will be used in our proof. Section B discusses our first equivalence result between training regularized neural network and kernel ridge regression. Section C discusses our generalization result of the leverage score sampling theory. Section D discusses our second equivalence result under leverage score initialization and potential benefits compared to the Gaussian initialization. Section E discusses how to extend our results to a broader class of neural network models.

Appendix

Appendix A Preliminaries

In this section we introduce the probability tools we use in the proof.

We state Chernoff, Hoeffding and Bernstein inequalities.

Let $X_{1},\cdots,X_{n}$ denote $n$ independent bounded variables in $[a_{i},b_{i}]$. Let $X=\sum_{i=1}^{n}X_{i}$, then we have

Let $X_{1},\cdots,X_{n}$ be independent zero-mean random variables. Suppose that $|X_{i}|\leq M$ almost surely, for all $i$. Then, for all positive $t$,

We state three inequalities for Gaussian random variables.

Let $X\sim N(0,\sigma^{2})$, that is, the probability density function of $X$ is given by $\phi(x)=\frac{1}{\sqrt{2\pi\sigma^{2}}}e^{-\frac{x^{2}}{2\sigma^{2}}}$. Then

Let $X\sim\mathcal{N}(\mu,\sigma^{2})$ be a Gaussian random variable with mean $\mu$ and variance $\sigma^{2}$. Then for all $t\geq 0$, we have

Let $X\sim{\cal X}_{k}^{2}$ be a chi-squared distributed random variable with $k$ degrees of freedom. Each one has zero mean and $\sigma^{2}$ variance. Then

We state two inequalities for random matrices.

Let $M_{1}$ and $M_{2}$ be semidefinite upper bounds for the expected squares:

where each $R_{k}$ is an independent copy of $R$. Then, for all $t\geq\sqrt{m/n}+2L/3n$,

A.2 Neural tangent kernel and its properties

Let $\lambda=\lambda_{\min}(H^{\operatorname{cts}})$. If $m=\Omega(\lambda^{-2}n^{2}\log(n/\delta))$, we have

hold with probability at least $1-\delta$.

holds with probability at least $1-n^{2}\cdot\exp(-mR/10)$.

Appendix B Equivalence between sufficiently wide neural net and kernel ridge regression

In this section, we extend the equivalence result in [JGH18, ADH+19a] to the case with regularization term, where they showed the equivalence between a fully-trained infinitely wide/sufficiently wide neural net and the kernel regression solution using the neural tangent kernel (NTK). Specifically, we prove Theorem 3.6 and Theorem 3.7 in this section.

Section B.1 introduces key notations and standard data assumptions. Section B.2 restates and supplements the definitions introduced in the paper. Section B.3 presents several key lemmas about the gradient flow and linear convergence of neural network and kernel ridge regression predictors, which are crucial to the final proof. Section B.4 provides a brief proof sketch. Section B.5 restates the main equivalence Theorem 3.7 and provides a complete proof following the proof sketch. Section B.6 restates and proves Theorem 3.6 by showing it as a by-product of previous proof.

$u^{*}=\lim_{t\rightarrow\infty}u_{\operatorname{ntk}}(t)$ (See Eq. (24))

$u_{\operatorname{test}}^{*}=\lim_{t\rightarrow\infty}u_{\operatorname{ntk},\operatorname{test}}(t)$ (See Eq. (25))

We made the following assumptions: 1. For each $i\in[n]$, we assume $|y_{i}|=O(1)$. 2. $H^{\operatorname{cts}}$ is positive definite, i.e., $\Lambda_{0}:=\lambda_{\min}(H^{\operatorname{cts}})>0$. 3. All the training data and test data have Euclidean norm equal to 1.

B.2 Definitions

To establish the equivalence between neural network and kernel ridge regression, we prove the similarity of their gradient flow and initial predictor. Note kernel ridge regression starts at 0 as initialization, so we hope the initialization of neural network also close to zero. Therefore, using the same technique in [ADH+19a], we apply a small multiplier $\kappa>0$ to both predictors to bound the different of initialization.

We define a two layer neural networks with rectified linear unit (ReLU) activation as the following form

We denote $w_{r}(t),r\in[m]$ as the variable at iteration $t$. We denote

as the test data predictor at iteration $t$.

We define the neural tangent kernel(NTK) and the feature function corresponding to the neural networks $f_{\operatorname{nn}}$ defined in Definition B.2 as following

as the test data predictor at iteration $t$. Note the gradient flow converge the to optimal solution of problem (20) due to the strongly convexity of the problem. We denote

B.3 Gradient, gradient flow, and linear convergence

Denote $L(t)=\frac{1}{2}\|Y-u_{\operatorname{ntk}}(t)\|_{2}^{2}+\frac{1}{2}\lambda\|\beta(t)\|_{2}^{2}$. By the rule of gradient descent, we have

where $\Phi$ is defined in Definition B.4. Thus we have

Note $[u_{\operatorname{ntk}}(t)]_{i}=\kappa f_{\operatorname{ntk}}(\beta(t),x_{i})$ and $[H^{\operatorname{cts}}]_{:,i}=\mathsf{K}_{\operatorname{ntk}}(x_{i},X)$, so writing all the data in a compact form, we have

where the first step follows the chain rule, the second step follows Corollary B.8, the third step uses basic linear algebra, the fourth step follows Eq. (B.3), the fifth step simplifies the expression, and the last step follows the definition of $\Lambda_{0}$. Further, since

where the first step calculates the gradient, and the second step follows from Eq. (B.3). Thus, $e^{2(\kappa^{2}\Lambda_{0}+\lambda)t}\|u_{\operatorname{ntk}}(t)-u^{*}\|_{2}^{2}$ is non-increasing, which implies

Denote $L(t)=\frac{1}{2}\|Y-u_{\operatorname{nn}}(t)\|_{2}^{2}+\frac{1}{2}\lambda\|W(t)\|_{F}^{2}$. By the rule of gradient descent, we have

Also note for ReLU activation $\sigma$, we have

where the first step calculates the derivatives, the second step follows basic linear algebra, the third step follows the property of ReLU activation: $\sigma(l)=l\sigma^{\prime}(l)$, and the last step follows from the definition of $f_{\operatorname{nn}}$. Thus, we have

where the first step follows from chain rule, the second step follows from Eq. (28), the third step follows from the definition of $\mathsf{K}_{t}$ and Eq. (B.3), and the last step rewrites the formula in a compact form. ∎

Note $[u_{\operatorname{nn}}(t)]_{i}=\kappa f_{\operatorname{nn}}(W(t),x_{i})$ and $[H(t))]_{:,i}=\mathsf{K}_{t}(x_{i},X)$, so writing all the data in a compact form, we have

Note by definition, $u_{\operatorname{nn},\operatorname{test}}(t)=\kappa f_{\operatorname{nn}}(W(t),x_{\operatorname{test}})$, so we have

where the first step follows the chain rule, the second step follows Corollary B.11, the third step uses basic linear algebra, the fourth step follows Eq. (B.3), the fifth step simplifies the expression, and the last step follows the assumption $\|H(t)-H^{\operatorname{cts}}\|\leq\Lambda_{0}/2$. ∎

B.4 Proof sketch

Our goal is to show with appropriate width of the neural network and appropriate training iterations, the neural network predictor will be sufficiently close to the neural tangent kernel ridge regression predictor for any test data. We follow similar proof framework of Theorem 3.2 in [ADH+19a]. Given any accuracy $\epsilon\in(0,1)$, we divide this proof into following steps:

Firstly, according to the linear convergence property of kernel ridge regression shown in Lemma B.9, we can choose sufficiently large training iterations $T>0$, so that $|u_{\operatorname{test}}^{*}-u_{\operatorname{ntk},\operatorname{test}}(T)|\leq\epsilon/2$, as shown in Lemma B.13.

Once fix training iteration $T$ as in step 1, we bound $|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{ntk},\operatorname{test}}(T)|\leq\epsilon/2$ by showing the following:

Due to the similarity of the the gradient flow of neural network training and neural tangent kernel ridge regression, we can reduce the task of bounding the prediction perturbation at time $T$, i.e., $|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{ntk},\operatorname{test}}(T)|$, back to bounding

the initialization perturbation $|u_{\operatorname{nn},\operatorname{test}}(0)-u_{\operatorname{ntk},\operatorname{test}}(0)|$ and

kernel perturbation $\|H(t)-H^{\operatorname{cts}}\|$, $\|\mathsf{K}_{\operatorname{ntk}}(x_{\operatorname{test}},X)-\mathsf{K}_{t}(x_{\operatorname{test}},X)\|_{2}$, as shown in Lemma B.14.

According to concentration results, we can bound the initialization perturbation $|u_{\operatorname{nn},\operatorname{test}}(0)-u_{\operatorname{ntk},\operatorname{test}}(0)|$ small enough by choosing sufficiently small $\kappa\in(0,1)$, as shown in Lemma B.20.

We characterize the over-parametrization property of the neural network by inductively show that we can bound kernel perturbation $\|H(t)-H^{\operatorname{cts}}\|$, $\|\mathsf{K}_{\operatorname{ntk}}(x_{\operatorname{test}},X)-\mathsf{K}_{t}(x_{\operatorname{test}},X)\|_{2}$ small enough by choosing network width $m>0$ large enough, as shown in Lemma B.21.

Lastly, we combine the results of step 1 and 2 using triangle inequality, to show the equivalence between training neural network with regularization and neural tangent kernel ridge regression, i.e., $|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{test}}^{*}|\leq\epsilon$, as shown in Theorem B.28.

B.5 Equivalence between training net with regularization and kernel ridge regression for test data prediction

In this section, we prove Theorem 3.7 following the proof sketch in Section B.4.

In this section, we give an upper bound for $|u_{\operatorname{ntk},\operatorname{test}}(T)-u_{\operatorname{test}}^{*}|$.

where $\widetilde{O}(\cdot)$ here hides $\operatorname{poly}\log(n/(\epsilon\Lambda_{0}))$.

Due to the linear convergence of kernel ridge regression, i.e.,

where the last step follows from $\beta(0)=0$ and $\|\beta^{*}\|_{2}=\operatorname{poly}(\kappa,n,1/\Lambda_{0})$.

Note $\kappa\in(0,1)$. Thus, by picking $T=\widetilde{O}(\frac{1}{\kappa^{2}\Lambda_{0}})$, we have

where $\widetilde{O}(\cdot)$ here hides $\operatorname{poly}\log(n/(\epsilon\Lambda_{0}))$.∎

The goal of this section is to prove Lemma B.14, which reduces the problem of bounding prediction perturbation to the problem of bounding initialization perturbation and kernel perturbation.

$\|u_{\operatorname{nn}}(0)\|_{2}\leq\sqrt{n}\epsilon_{\operatorname{init}}$ and $|u_{\operatorname{nn},\operatorname{test}}(0)|\leq\epsilon_{\operatorname{init}}$

$\|\mathsf{K}_{\operatorname{ntk}}(x_{\operatorname{test}},X)-\mathsf{K}_{t}(x_{\operatorname{test}},X)\|_{2}\leq\epsilon_{K}$

$\|H(t)-H^{\operatorname{cts}}\|\leq\epsilon_{H}$

Combining results from Lemma B.15, Claim B.16. B.17, B.18, we complete the proof. We have

where the first step follows from Lemma B.15, the second step follows from Claim B.16, B.17 and B.18, and the last step simplifies the expression. ∎

To prove Lemma B.14, we first bound $|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{ntk},\operatorname{test}}(T)|$ by three terms in Lemma B.15, then we bound each term individually in Claim B.16, Claim B.17, and Claim B.18.

Follow the same notation as Lemma B.14, we have

where the first step follows from the definition of integral, the second step follows from the triangle inequality. Note by Corollary B.8, B.11, their gradient flow are given by

where the first step follows from Eq. (32) and Eq. (33), the second step rewrites the formula. Note the term $-\lambda(u_{\operatorname{nn},\operatorname{test}}(t)-u_{\operatorname{ntk},\operatorname{test}}(t))$ will only make

Now let us bound these three terms $A$, $B$ and $C$ one by one. We claim

Note $u_{\operatorname{ntk},\operatorname{test}}(0)=0$, so by assumption we have

where the first step follows from the triangle inequality, the second step follows from Eq. (36), the third step calculates the integration, and the last step follows from the fact $u_{\operatorname{ntk}}(0)=0$. Thus, we have

where the first step follows from Eq. (36), the second step follows from Eq. (B.5.2) and definition of $\epsilon_{K}$. ∎

where the first step follows from the definition of $C$, and the second step follows the Cauchy-Schwartz inequality.

To bound term $\max_{t\in[0,T]}\|u_{\operatorname{ntk}}(t)-u_{\operatorname{nn}}(t)\|_{2}$, notice that for any $t\in[0,T]$, we have

where the first step follows the triangle inequality, and the second step follows the assumption. Further,

where the first step follows from Eq. (B.5.2), the second step follows from Eq. (40), the third step follows from triangle inequality, the fourth step follows from the condition $\|H(\tau)-H^{\operatorname{cts}}\|\leq\epsilon_{H}$ for all $\tau\leq T$ and the triangle inequality, the fifth step follows from the linear convergence of $\|u_{\operatorname{ntk}}(\tau)-u^{*}\|_{2}$ as in Lemma B.9, the sixth step follows the fact $u_{\operatorname{ntk}}(0)=0$, and the last step calculates the maximum. Therefore,

Given final accuracy $\epsilon$, to ensure $|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{ntk},\operatorname{test}}(T)|\leq\epsilon$, we need to choose $\kappa>0$ small enough to make $\epsilon_{\operatorname{init}}=O(\epsilon)$ and choose width $m>0$ large enough to make $\epsilon_{H}$ and $\epsilon_{\operatorname{test}}$ both $O(\epsilon)$. And we discuss these two tasks one by one in the following sections.

B.5.3 Upper bounding initialization perturbation

In this section, we bound $\epsilon_{\operatorname{init}}$ to our wanted accuracy $\epsilon$ by picking $\kappa$ large enough. We prove Lemma B.20.

Further, given any accuracy $\epsilon\in(0,1)$, if $\kappa=\widetilde{O}(\epsilon(\Lambda_{0}+\lambda)/n)$, let $\epsilon_{\operatorname{init}}=\epsilon(\Lambda_{0}+\lambda)/n$, we have

hold with probability $1-\delta$, where $\widetilde{O}(\cdot)$ hides the $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

Since $w_{r}(0)\sim\mathcal{N}(0,I_{d})$, so $w_{r}(0)^{\top}x_{\operatorname{test}}\sim N(0,\|x_{\operatorname{test}}\|_{2})$. Note $\|x_{\operatorname{test}}\|_{2}\leq 1$, by Gaussian tail bounds Lemma A.5, we have with probability $1-\delta/(2m)$:

Since $u_{\operatorname{nn},\operatorname{test}}(0)=\frac{1}{\sqrt{n}}\sum Z_{r}$, by combining Eq. (42), (43) and union bound over all $r\in[m]$, we have with probability $1-\delta$:

Further, note $[u_{\operatorname{nn}}(0)]_{i}=\kappa f_{\operatorname{nn}}(W(0),x_{i})$ and $u_{\operatorname{nn},\operatorname{test}}(0)=\kappa f_{\operatorname{nn}}(W(0),x_{\operatorname{test}})$. Thus, by choosing $\kappa=\widetilde{O}(\epsilon(\Lambda_{0}+\lambda)/n)$, taking the union bound over all training and test data, we have

hold with probability $1-\delta$, where $\widetilde{O}(\cdot)$ hides the $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$. ∎

B.5.4 Upper bounding kernel perturbation

In this section, we try to bound kernel perturbation by induction. We want to prove Lemma B.21, which also helps to show the equivalence for training data prediction as shown in Section B.6.

1. $\|w_{r}(0)-w_{r}(t)\|_{2}\leq\epsilon_{W}$, $\forall r\in[m]$

2. $\|H(0)-H(t)\|_{2}\leq\epsilon_{H}^{\prime}$

4. $\|\mathsf{K}_{0}(x_{\operatorname{test}},X)-\mathsf{K}_{t}(x_{\operatorname{test}},X)\|_{2}\leq\epsilon_{K}^{\prime}$

Further, $\epsilon_{W}\leq\widetilde{O}(\frac{\epsilon\lambda_{0}^{2}}{n^{2}})$, $\epsilon_{H}^{\prime}\leq\widetilde{O}(\frac{\epsilon\lambda_{0}^{2}}{n})$ and $\epsilon_{K}^{\prime}\leq\widetilde{O}(\frac{\epsilon\lambda_{0}^{2}}{n^{1.5}})$. Here $\widetilde{O}(\cdot)$ hides the $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

We first state some concentration results for the random initialization that can help us prove the lemma.

By lemma A.6, with probability at least $1-\delta$,

holds with probability at least $1-\delta$. Note by definition,

holds for any training data $x_{i}$. By Hoeffding inequality, we have for any $t>0$,

Setting $t=(\frac{2}{m}\log{(2n/\delta)})^{1/2}$, we can apply union bound on all training data $x_{i}$ to get with probability at least $1-\delta$, for all $i\in[n]$,

holds with probability at least $1-\delta$. Using union bound over above three events, we finish the proof. ∎

Now conditioning on Eq. (44), (45), (46) holds, We show all the four conclusions in Lemma B.21 holds using induction.

Note the base case when $t=0$ trivially holds. Now assuming Lemma B.21 holds before time $t\in[0,T]$, we argue that it also holds at time $t$. To do so, Lemmas B.23, B.24, B.25 argue these conclusions one by one.

where the first step follows triangle inequality, the second step follows Eq. (B.5.4), and the last step follows the definition of $\epsilon_{W}$ as Eq. (48). ∎

holds with probability $1-n^{2}\cdot\exp{(-m\epsilon_{W}/10)}$.

Directly applying Lemma A.10, we finish the proof. ∎

Fix $\epsilon_{H}^{\prime}>0$ independent of $t$. If for all $\tau<t$

holds for all $\tau<t$. Denote $\epsilon_{H}=\epsilon_{H}^{\prime}+4n(\log{(n/\delta)}/m)^{1/2}$, we have $\|H(\tau)-H^{\operatorname{cts}}\|\leq\epsilon_{H}\leq\Lambda_{0}/2$, which satisfies the condition of Lemma B.12. Thus, for any $\tau<t$, we have

where the first step follows from Lemma B.12, the second step follows from Eq. (B.5.4). Now let us discuss two cases:

Case 1. If for all $\tau<t$, $\|u_{\operatorname{nn}}(\tau)-u^{*}\|_{2}\geq\epsilon_{\operatorname{train}}$ always holds, we want to argue that

Note by assumption (51) and (52), we have

holds for any $\tau<t$. Thus, plugging into (B.5.4),

Case 2. If there exist $\overline{\tau}<t$, such that $\|u_{\operatorname{nn}}(\overline{\tau})-u^{*}\|_{2}<\epsilon_{\operatorname{train}}$, we want to argue that $\|u_{\operatorname{nn}}(t)-u^{*}\|_{2}<\epsilon_{\operatorname{train}}$. Note by assumption (51) and (52), we have

holds for $\tau=\overline{\tau}$, which implies $e^{(\kappa^{2}\Lambda_{0}+\lambda)\tau}(\|u_{\operatorname{nn}}(\tau)-u^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2})$ is non-increasing at $\tau=\overline{\tau}$. Since $\|u_{\operatorname{nn}}(\overline{\tau})-u^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2}<0$, by induction, $e^{(\kappa^{2}\Lambda_{0}+\lambda)\tau}(\|u_{\operatorname{nn}}(\tau)-u^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2})$ being non-increasing and $\|u_{\operatorname{nn}}(\tau)-u^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2}<0$ holds for all $\overline{\tau}\leq\tau\leq t$, which implies

Fix $\epsilon_{W}\in(0,1)$ independent of $t$. If $\forall r\in[m]$, we have

holds with probability at least $1-n\cdot\exp{(-m\epsilon_{W}/10)}$.

Recall the definition of $\mathsf{K}_{0}$ and $\mathsf{K}_{t}$

Fix $i\in[n]$, by Bernstein inequality (Lemma A.3), we have for any $t>0$,

Note by definition $\epsilon_{K}^{\prime}=2\sqrt{n}\epsilon_{W}$, so we finish the proof. ∎

Now we summarize all the conditions need to be satisfied so that the induction works as in Table 1.

where the first step follows from the definition of $u^{*}$, the second step follows from Cauchy-Schwartz inequality, the third step follows from $Y=O(\sqrt{n})$, and the last step follows from the choice of the parameters.

Further, with probability $1-\delta$, we have

Thus, we have $\|u_{\operatorname{nn}}(0)-u^{*}\|_{2}=\widetilde{O}(\sqrt{n})$.

Now, by direct calculation, we have all the induction conditions satisfied with high probability. Note the failure probability only comes from Lemma B.20, B.22, B.24, B.26, which only depend on the initialization. By union bound over these failure events, we have all four conclusions in Lemma B.21 holds with high probability, which completes the proof.

where $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

By Lemma B.20, we can choose $\epsilon_{\operatorname{init}}=\epsilon(\Lambda_{0})/n$.

where the first step follows from triangle inequality, the second step follows from Lemma B.22, and the last step follows from Lemma B.21. Thus, we can choose $\epsilon_{K}=\frac{\epsilon\Lambda_{0}^{2}}{n^{1.5}}$.

where the first step follows from triangle inequality, the second step follows from Lemma B.22, and the last step follows from Lemma B.21. Thus, we can choose $\epsilon_{H}=\frac{\epsilon\Lambda_{0}^{2}}{n}$.

B.5.6 Main result for test data prediction equivalence

In this section, we restate and prove Theorem 3.7.

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

It follows from combining results of bounding $\|u_{\operatorname{nn},\operatorname{test}}(T)-u_{\operatorname{ntk},\operatorname{test}}(T)\|_{2}\leq\epsilon/2$ as shown in Lemma B.27 and $\|u_{\operatorname{ntk},\operatorname{test}}(T)-u_{\operatorname{test}}^{*}\|_{2}\leq\epsilon/2$ as shown in Lemma B.13 using triangle inequality. ∎

B.6 Equivalence between training net with regularization and kernel ridge regression for training data prediction

In this section, we restate and proof Theorem 3.6.

Note the proof of equivalence results for the test data in previous sections automatically gives us an equivalence results of the prediction for training data. Specifically, the third conclusion in Lemma B.21 characterizes the training prediction $u_{\operatorname{nn}}(t)$ throughout the training process. Thus, we have the following theorem characterize the equivalence between training net with regularization and kernel ridge regression for the training data.

Given any accuracy $\epsilon\in(0,1/10)$ and failure probability $\delta\in(0,1/10)$, if $\kappa=1$, $T=\widetilde{O}(\frac{1}{\Lambda_{0}})$, network width $m\geq\widetilde{O}(\frac{n^{4}d}{\lambda_{0}^{4}\epsilon})$ and regularization parameter $\lambda\leq\widetilde{O}(\frac{1}{\sqrt{m}})$, then with probability at least $1-\delta$ over the random initialization, we have

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

Let $\kappa=1$, $T=\widetilde{O}(\frac{1}{\Lambda_{0}})$, $\lambda\leq\widetilde{O}(\frac{1}{\sqrt{m}})$, $m\geq\widetilde{O}(\frac{n^{4}d}{\lambda_{0}^{4}\epsilon})$ and $\epsilon_{\operatorname{train}}=\epsilon$ in Lemma B.21. We can see all the conditions in Table 1 hold. Thus, the third conclusion in Lemma B.21 holds. So with probability $1-\delta$, we have

where the first step follows from Lemma B.21, the second step follows from $\kappa=1$ and $\epsilon_{\operatorname{train}}=\epsilon$, the last step follows from $T=\widetilde{O}(\frac{1}{\Lambda_{0}})$ and $\|u_{\operatorname{nn}}(0)-u^{*}\|_{2}^{2}\leq n$ with high probability. ∎

Appendix C Generalization result of leverage score sampling for approximating kernels

In this section, we generalize the result of Lemma 8 in [AKM+17] for a more broad class of kernels and feature vectors. Specifically, we prove Theorem 3.3.

Section C.1 introduces the related kernel and random features, we also restate Definition 3.1 and 3.2 for leverage score sampling and random features in this section. Section C.2 restates and proves our main result Theorem 3.3.

If $w_{1},\cdots,w_{m}$ are drawn according to $p(\cdot)$, then

If $w$ are drawn according to $p(\cdot)$, then

If $w_{1},\cdots,w_{m}$ are drawn according to $p(\cdot)$, then

If $w_{1},\cdots,w_{m}$ are drawn according to $q(\cdot)$, then

If $w$ are drawn according to $q(\cdot)$, then

If $w_{1},\cdots,w_{m}$ are drawn according to $q(\cdot)$, then

Let $q_{\lambda}(w)$ denote the leverage score defined in Definition C.4. Let $s_{\lambda}(K)$ denote the statistical dimension defined in Definition C.5. We define the leverage score sampling distribution as

C.2 Main result

holds with probability at least $1-\delta$.

To prove the theorem, we follow the same proof framework as Lemma 8 in [AKM+17].

Let $K+\lambda I_{n}=V^{\top}\Sigma^{2}V$ be an eigenvalue decomposition of $K+\lambda I_{n}$. Note that Eq. (59) is equivalent to

Multiplying $\Sigma^{-1}V$ on the left and $V^{\top}\Sigma^{-1}$ on the left for both sides of Eq. (60), it suffices to show that

holds with probability at least $1-\delta$. Let

where the first step follows from the definition of $Y_{r}$, the second step follows the linearity of expectation, the third step calculations the expectation, and the last step follows Eq. (57).

where the first step follows from the definition of $Y_{r}$, the second step follows from basic linear algebra, and the last step follows from Eq. (58).

holds with probability at least $1-\delta$.

where the first step follows from $\|A\|\leq\operatorname{tr}[A]$ for any positive semidefinite matrix, the second step follows $\operatorname{tr}[AB]=\operatorname{tr}[BA]$, the third step follows from the definition of $V,\Sigma$, the fourth step follows from the definition of leverage score $q_{\lambda}(\cdot)$ as defined in Definition C.4, and the last step follows from the condition $\widetilde{q}_{\lambda}(w)\geq q_{\lambda}(w)$.

where the first step follows from the definition of $Y_{r}$, the second step follows from the definition of $V,\Sigma$, the third step follows from $\|A\|\leq\operatorname{tr}[A]$ for any positive semidefinite matrix, the fourth step follows from the definition of leverage score $q_{\lambda}(\cdot)$ as defined in Definition C.4, the fifth step follows from the definition of $Y_{r}$, the sixth step follows from the definition of $\overline{q}_{\lambda}(\cdot)$, and the last step follows from the condition $\widetilde{q}_{\lambda}(w)\geq q_{\lambda}(w)$.

Thus, let $\lambda_{1}\geq\lambda_{2}\geq\cdots\geq\lambda_{n}$ be the eigenvalues of $K$, we have

where the first step follows from Lemma A.8, the second step follows from the definition of $D$ and $s_{\lambda}(K)=\operatorname{tr}[(K+\lambda I_{n})^{-1}K]=\sum_{i}^{n}\lambda_{i}/(\lambda_{i}+\lambda)=\operatorname{tr}[D]$, the third step follows the condition $\lambda\in(0,\|K\|)$, the fourth step follows the condition $\epsilon\in(0,1/2)$, and the last step follows from the bound on $m$. ∎

In this section, we connected the neural network theory with the leverage score sampling theory by showing a new equivalence result between training reweighed neural network with regularization under leverage score initialization and corresponding neural tangent kernel ridge regression. Specifically, we prove Theorem D.21. Due to the similarity of the results to Section B, we present this section in the same framework.

Section D.1 introduces new notations and states the standard data assumptions again. Section D.2 restates and supplements the definitions in the paper. Section D.3 presents the key lemmas about the leverage score initialization and related properties, which are crucial to the proof. Section D.4 provides a brief proof sketch. Section D.5 restates and proves the main result Theorem D.21 following the proof sketch.

Here, we list the locations where definitions and theorems in the paper are restated. Definition 3.8 is restated in Definition D.2. Theorem 3.9 is restated in Theorem D.21.

$\overline{u}^{*}=\lim_{t\rightarrow\infty}\overline{u}_{\operatorname{ntk}}(t)$ (See Eq. (70))

$\overline{u}_{\operatorname{test}}^{*}=\lim_{t\rightarrow\infty}\overline{u}_{\operatorname{ntk},\operatorname{test}}(t)$ (See Eq. (71))

We made the following assumptions: 1. For each $i\in[n]$, we assume $|y_{i}|=O(1)$. 2. $H^{\operatorname{cts}}$ is positive definite, i.e., $\Lambda_{0}:=\lambda_{\min}(H^{\operatorname{cts}})>0$. 3. All the training data and test data have Euclidean norm equal to 1.

D.2 Definitions

as the test data predictor at iteration $t$.

as the test data predictor at iteration $t$. Note the gradient flow converge the to optimal solution of problem (66) due to the strongly convexity of the problem. We denote

D.3 Leverage score sampling, gradient flow, and linear convergence

Recall in the main body we connect the leverage score sampling theory and convergence theory of the neural network training by observing

and corresponding probability density function

where $c_{1}=O(\frac{1}{n})$ and $c_{2}=O(\frac{1}{\Lambda_{0}})$.

where the first step follows from $H^{\operatorname{cts}}\succeq\Lambda_{0}I_{n}$, the second step follows from the definition of $\Phi$, the third step follows from the linearity of trace operator, the fourth step follows from $\operatorname{tr}(AB)=\operatorname{tr}(BA)$, the fifth step follows from the definition of $\phi$, and the last step follows from $\|x_{i}\|_{2}=1$ and $\sigma^{\prime}(\cdot)\leq 1$.

Similarly, note $H^{\operatorname{cts}}\preceq nI_{n}$, we have

Let $\overline{H}(t)$, $H^{\operatorname{cts}}$ be the kernel defined as in Definition D.3 and Definition B.4. Let $\Lambda_{0}>0$ be defined as in Definition B.4. Let $p(\cdot)$ denotes the probability density function for Gaussian $\mathcal{N}(0,I_{d})$. Let $q(\cdot)$ denotes the leverage sampling distribution with respect to $p(\cdot)$, $\overline{H}(0)$ and $\lambda$ defined in Definition C.6. Let $\Delta\in(0,1/4)$. Then we have

By choosing $m\geq\widetilde{O}(\Delta^{-2}s_{\lambda}(H^{\operatorname{cts}})$, with probability at least $1-\delta$,

Further, if $\lambda\leq\Lambda_{0}$, we have with probability at least $1-\delta$,

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(s_{\lambda}(\overline{H}(0))/\delta)$.

where the first step follows from the definition of $\overline{H}(0)$, the second step follows the definition of $\overline{\Phi}$, the third step calculates the expectation, and the last step follows the definition of $H^{\operatorname{cts}}$.

Also, by applying Theorem C.7 directly, we have with probability at least $1-\delta$,

if $m\geq\widetilde{O}(\Delta^{-2}s_{\lambda}(\overline{H}(0))$.

Denote $L(t)=\frac{1}{2}\|Y-\overline{u}_{\operatorname{ntk}}(t)\|_{2}^{2}+\frac{1}{2}\lambda\|\beta(t)\|_{2}^{2}$. By the rule of gradient descent, we have

Note $[\overline{u}_{\operatorname{ntk}}(t)]_{i}=\kappa\overline{f}_{\operatorname{ntk}}(\beta(t),x_{i})$ and $[\overline{H}(0)]_{:,i}=\overline{\mathsf{K}}_{0}(x_{i},X)$, so writing all the data in a compact form, we have

where the first step follows the chain rule, the second step follows Corollary D.8, the third step uses basic linear algebra, the fourth step follows Eq. (D.3), the fifth step simplifies the expression, and the last step follows from Lemma D.6. Further, since

where the first step calculates the gradient, and the second step follows from Eq. (D.3). Thus, $e^{(\kappa^{2}\Lambda_{0}+\lambda)t}\|\overline{u}_{\operatorname{ntk}}(t)-\overline{u}^{*}\|_{2}^{2}$ is non-increasing, which implies

Denote $L(t)=\frac{1}{2}\|Y-\overline{u}_{\operatorname{nn}}(t)\|_{2}^{2}+\frac{1}{2}\lambda\|W(t)\|_{F}^{2}$. By the rule of gradient descent, we have

Also note for ReLU activation $\sigma$, we have

where the first step calculates the derivatives, the second step follows basic linear algebra, the third step follows the property of ReLU activation: $\sigma(l)=l\sigma^{\prime}(l)$, and the last step follows from the definition of $\overline{f}_{\operatorname{nn}}$. Thus, we have

where the first step follows from chain rule, the second step follows from Eq. (75), the third step follows from the definition of $\overline{\mathsf{K}}_{t}$ and Eq. (D.3), and the last step rewrites the formula in a compact form. ∎

Note $[\overline{u}_{\operatorname{nn}}(t)]_{i}=\kappa\overline{f}_{\operatorname{nn}}(W(t),x_{i})$ and $[\overline{H}(t))]_{:,i}=\overline{\mathsf{K}}_{t}(x_{i},X)$, so writing all the data in a compact form, we have

Note by definition, $\overline{u}_{\operatorname{nn},\operatorname{test}}(t)=\kappa\overline{f}_{\operatorname{nn}}(W(t),x_{\operatorname{test}})$, so we have

where the first step follows the chain rule, the second step follows Corollary D.11, the third step uses basic linear algebra, the fourth step follows Eq. (D.3), the fifth step simplifies the expression, and the last step follows the assumption $\|\overline{H}(t)-\overline{H}(0))\|\leq\Lambda_{0}/4$ and the fact $\|\overline{H}(0))\|\leq\Lambda_{0}/2$. ∎

D.4 Proof sketch

We introduce a new kernel ridge regression problem with respect to $\overline{H}(0)$ to decouple the prediction perturbation resulted from initialization phase and training phase. Specifically, given any accuracy $\epsilon\in(0,1)$, we divide this proof into following steps:

Firstly, we bound the prediction perturbation resulted from initialization phase $\|u^{*}-\overline{u}^{*}\|_{2}\leq\epsilon/2$ by applying the leverage score sampling theory, as shown in Lemma D.13.

Then we use the similar idea as section B to bound the prediction perturbation resulted from training phase $\|\overline{u}_{\operatorname{nn}}(T)-\overline{u}^{*}\|_{2}\leq\epsilon/2$ by showing the over-parametrization and convergence property of neural network inductively, as shown in Lemma D.14 and Corollary D.20.

Lastly, we combine the results of step 1 and 2 using triangle inequality to show $\|\overline{u}_{\operatorname{nn}}(T)-u^{*}\|_{2}\leq\epsilon$, as shown in Theorem D.21.

D.5 Main result

In this section, we prove Theorem D.21 following the above proof sketch.

with probability at least $1-\delta$. Particularly, given arbitrary $\epsilon\in(0,1)$, if $m\geq\widetilde{O}(\frac{n}{\epsilon\Lambda_{0}})$ and $\lambda\leq\widetilde{O}(\frac{1}{\sqrt{m}})$, we have

Here $\widetilde{O}(\cdot)$ hides $\operatorname{poly}\log(s_{\lambda}(H^{\operatorname{cts}})/\delta)$.

By Lemma D.6, if $m\geq\widetilde{O}(\Delta^{-2}s_{\lambda}(H^{\operatorname{cts}})$, we have

where the first step follows from Cauchy-Schwartz inequality, the second step follows from Eq. (78), and the last step follows from the definition of $\Lambda_{0}$ and $\|Y\|_{2}=O(\sqrt{n})$. ∎

1. $\|w_{r}(0)-w_{r}(t)\|_{2}\leq\epsilon_{W}$, $\forall r\in[m]$

2. $\|\overline{H}(0)-\overline{H}(t)\|_{2}\leq\epsilon_{H}^{\prime}$

Here $\widetilde{O}(\cdot)$ hides the $\operatorname{poly}\log(n/(\epsilon\delta\Lambda_{0}))$.

We first state the following concentration result for the random initialization that can help us prove the lemma.

hold for all $r\in[m]$, where $c_{2}=O(n)$.

By lemma A.6, if $w_{r}(0)\sim\mathcal{N}(0,I_{n})$, then with probability at least $1-\delta$,

Now conditioning on Eq. (44), (45), (46) holds, We show all the four conclusions in Lemma B.21 holds using induction.

which are independent of $t$. Here $\alpha_{w,0}$ are defined in Eq. (79).

Note the base case when $t=0$ trivially holds. Now assuming Lemma D.14 holds before time $t\in[0,T]$, we argue that it also holds at time $t$. To do so, Lemmas D.16, D.17, D.19 argue these conclusions one by one.

where the first step follows triangle inequality, the second step follows Eq. (D.5.2), and the last step follows the definition of $\epsilon_{W}$ as Eq. (80).

holds with probability $1-n^{2}\cdot\exp{(-m\epsilon_{W}c_{1}/10)}$, where $c_{1}=O(1/n)$.

Directly applying Lemma D.18, we finish the proof. ∎

holds with probability at least $1-n^{2}\cdot\exp(-mRc_{1}/10)$, where $c_{1}=O(1/n)$.

where the last step follows from for each $r,i,j$, we define

We consider $i,j$ are fixed. We simplify $s_{r,i,j}$ to $s_{r}$.

Then $s_{r}$ is a random variable that only depends on $\widetilde{w}_{r}$. Since $\{\widetilde{w}_{r}\}_{r=1}^{m}$ are independent, $\{s_{r}\}_{r=1}^{m}$ are also mutually independent.

where the last step follows from the anti-concentration inequality of Gaussian (Lemma A.4).

If $\neg A_{i,r}$ and $\neg A_{j,r}$ happen, then

where the last step follows from Lemma D.5 and $c_{1}=O(n)$. We also have $|s_{r}|\leq 1/c_{1}$. So we can apply Bernstein inequality (Lemma A.3) to get for all $t>0$,

Fix $\epsilon_{H}^{\prime}>0$ independent of $t$. If for all $\tau<t$

Note $\|\overline{H}(\tau)-\overline{H}(0)\|\leq\epsilon_{H}\leq\Lambda_{0}/4$. By Lemma D.12, for any $\tau<t$, we have

where the first step follows from Lemma D.12, the second step follows from definition of $\epsilon_{H}^{\prime}$.

Case 1. If for all $\tau<t$, $\|u_{\operatorname{nn}}(\tau)-u^{*}\|_{2}\geq\epsilon_{\operatorname{train}}$ always holds, we want to argue that

holds for any $\tau<t$. Thus, plugging into (D.5.2),

Case 2. If there exist $\overline{\tau}<t$, such that $\|\overline{u}_{\operatorname{nn}}(\overline{\tau})-\overline{u}^{*}\|_{2}<\epsilon_{\operatorname{train}}$, we want to argue that $\|\overline{u}_{\operatorname{nn}}(t)-\overline{u}^{*}\|_{2}<\epsilon_{\operatorname{train}}$. Note by assumption (84), we have

holds for $\tau=\overline{\tau}$, which implies $e^{(\kappa^{2}\Lambda_{0}+\lambda)\tau/2}(\|\overline{u}_{\operatorname{nn}}(\tau)-\overline{u}^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2})$ is non-increasing at $\tau=\overline{\tau}$. Since $\|\overline{u}_{\operatorname{nn}}(\overline{\tau})-\overline{u}^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2}<0$, by induction, $e^{(\kappa^{2}\Lambda_{0}+\lambda)\tau/2}(\|\overline{u}_{\operatorname{nn}}(\tau)-\overline{u}^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2})$ being non-increasing and $\|\overline{u}_{\operatorname{nn}}(\tau)-\overline{u}^{*}\|_{2}^{2}-\epsilon_{\operatorname{train}}^{2}<0$ holds for all $\overline{\tau}\leq\tau\leq t$, which implies