# Spectral evolution and invariance in linear-width neural networks

Zhichao Wang  University of California, San Diego  Andrew Engel  Pacific
Northwest National Laboratory  Anand Sarwate  Rutgers, The State University of
New Jersey  Ioana Dumitriu  University of California, San Diego  Tony Chiang
Pacific Northwest National Laboratory  University of Washington

###### Abstract

We investigate the spectral properties of linear-width feed-forward neural
networks, where the sample size is asymptotically proportional to network
width. Empirically, we show that the weight spectra in this high dimensional
regime are invariant when trained by gradient descent for small constant
learning rates and the changes in both operator and Frobenius norm are
$\Theta(1)$ in the limit. This implies the bulk spectra for both the conjugate
and neural tangent kernels are also invariant. We demonstrate similar
characteristics for models trained with mini-batch (stochastic) gradient
descent with small learning rates and provide a theoretical justification for
this special scenario. When the learning rate is large, we show empirically
that an outlier emerges with its corresponding eigenvector aligned to the
training data structure. We also show that after adaptive gradient training,
where we have a lower test error and feature learning emerges, both the weight
and kernel matrices exhibit heavy tail behavior. Different spectral properties
such as invariant bulk, spike, and heavy-tailed distributions correlate to how
far the kernels deviate from initialization. To understand this phenomenon
better, we focus on a toy model, a two-layer network on synthetic data, which
exhibits different spectral properties for different training strategies.
Analogous phenomena also appear when we train on conventional neural networks
with real-world data. Our results show that monitoring the evolution of the
spectra during training is an important step toward understanding the training
dynamics and feature learning.

Emails<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
##  1 Introduction

Deep learning theory has made insightful connections between the behavior of
neural networks (NNs) and kernel machines through asymptotic analyses of the
so-called kernel regime [Nea95, Wil96, LBN+17, JGH18, BMM18, ADH+19, Yan19].
When the neural network (NN) is _infinitely wide_ , the behavior of NN
coincides with a kernel machine, and the training process as well as the
generalization performance of this ultra-wide NN can be fully described. The
performance of _finite-width_ NNs, however, does not correspond to this
theory, as NNs optimized with gradient-based methods perform better than
infinitely-wide networks in many circumstances [LSP+20, HXAP20, GMMM20, LMZ20,
DM20, GSJW20, RGKZ21, KWLS21]. This gap heavily relies on the task complexity,
data distribution, architecture of the NN and the training strategy [FDP+20].
We consider a more realistic setting, a linear-width regime (LWR), when the
sample size $n$, the input feature dimension $d$, and the width of the hidden
layer $N$ approach infinity at comparable rates. Under the LWR, we aim to
empirically study this theoretical gap in generalization and spectral
properties by training various NNs with different optimization tools [FDP+20,
LM20, JSF+20, TSR22].

For the ultra-wide NN, the change in each neuron is very small throughout the
gradient-based training processes, and so the kernel stays close to its value
at initialization [ZCZG18, DZPS18, DLL+19, BMR21]. There are two kernels
commonly studied in theory: the conjugate kernel (CK) and the neural tangent
kernel (NTK). CK (or the equivalent Gaussian process kernel) is the Gram
matrix of the last hidden layer, which represents random feature regression,
i.e. training only the last layer of the network [Nea95, Wil96, LBN+17, LLC18,
PW17, MM22]; by contrast, the NTK is the Gram matrix of the Jacobian of the NN
with respect to all trainable parameters, which governs the gradient flow of
NN [JGH18, DZPS18, JT20, ADH+19]. In most theoretical results, these kernels
remain fixed throughout training, which leads to a kernel gradient descent
with the initial kernel [JGH18, BM19], whereas in practice the spectra of the
weight matrix, CK and NTK of the NN change while learning the features from
the training data [MM19, MM21, FW20, CHS20, OJMDF21]. In this paper, we
explore the LWR experimentally to explore the following question:

How do the spectra of weight and kernel matrices of the NN evolve during the
training process?

This question is crucial to develop an understanding beyond the kernel regime
and will help us analyze the generalization of the NN in instances when it
performs better than the kernel machine. For this case, the spectral
properties of the trained NN could be entirely different from the initial
kernel [Lon21, BGL+21, SB21]. Also, various spectral properties of weight and
kernel matrices can reveal different features learned by different training
procedures [WHS22]. Understanding the dynamics of the spectral properties may
aid in finding better approaches to training and to tuning hyper-parameters
for NNs. From a theoretical perspective, random matrix theory (RMT) can be
further exploited to study and elucidate the NN training under proportional
limit in high dimensions [LCKS91, PB17, LLC18, PW18, HN20, MM22].

Our main contributions are:

  * •

We find a simple scenario that exhibits different spectral properties for both
weight and kernel matrices through different training procedures. With the
kernel regime as a benchmark, we compare how NN generalizes with different
spectral evolutions.

  * •

The spectra of NNs trained with full batch gradient descent (GD) are globally
invariant, indicating that the NN is still close to a kernel machine; we
provide precise controls of the weight and kernel matrices, and prove global
convergence of GD and the invariance of the limiting spectra in a simple
scenario.

  * •

We observe a transition of the alignment and the emergence of a spike when the
the learning rate exceeds some threshold. The alignment of the spike with a
teacher model confirms that the NN is indeed learning germane features. This
feature learning with large step sizes points to a new theoretical direction
of study for the generalization of NNs.

  * •

The evolution towards heavy-tailed spectra is also discovered by using
adaptive methods and tuning the hyper-parameters, though heavy-tailed spectra
can be antagonistic for generalization in some simple student-teacher tasks.
We also observe the emergence of heavy-tailed spectra for kernel matrices in
more complicated NNs like large language models trained with real-world
datasets.

##  2 Notation and preliminaries

Throughout this paper, $\|\cdot\|$ denotes the $\ell_{2}$ norm for vectors,
$\ell_{2}\to\ell_{2}$ is the operator norm for matrices, while $\|\cdot\|_{F}$
is the Frobenius norm; and $\odot$ represents the Hadamard product between
matrices. $\text{P}_{>1}$ denotes the orthogonal projection onto the nonlinear
part of the function w.r.t. Gaussian measure, and $o_{d,\mathbb{P}}(\cdot)$
represents little-o in probability as $d\to\infty$.

#### Neural Tangent Kernel parameterization.

Consider a $L$-layer fully-connected feedforward NN at initialization without
bias term: for $1\leq\ell\leq L-1$,

${\boldsymbol{h}}_{0}=\frac{{\boldsymbol{x}}}{\sqrt{d}},~{}~{}{\boldsymbol{h}}^{(\ell)}=\frac{1}{\sqrt{n_{\ell}}}\sigma({\boldsymbol{W}}_{\ell}{\boldsymbol{h}}^{(\ell-1)}),$
(1)

and
$f_{{\boldsymbol{\theta}}}({\boldsymbol{x}})={\boldsymbol{v}}^{\top}{\boldsymbol{h}}^{(L-1)},$
where the input vector is ${\boldsymbol{x}}\in\mathbb{R}^{d}$,
${\boldsymbol{W}}_{\ell}\in\mathbb{R}^{n_{\ell}\times n_{\ell-1}}$ is the
weight matrix for the $\ell$-th layer, and
${\boldsymbol{v}}:=[v_{1},\ldots,v_{h}]^{\top}\in\mathbb{R}^{n_{L-1}}$ is the
last-layer weight. For simplicity, let $n_{0}=d$. Denote all trainable
parameters by
${\boldsymbol{\theta}}:=[\text{vec}({\boldsymbol{W}}_{1}),\ldots,\text{vec}({\boldsymbol{W}}_{L-1}),{\boldsymbol{v}}]^{\top}\in\mathbb{R}^{p}$
where each parameter’s initial value is independently sampled from some
distribution. Here $p$ is the total number of parameters. Let the training
dataset be
$({\boldsymbol{X}},{\boldsymbol{y}}):=([{\boldsymbol{x}}_{1},\ldots,{\boldsymbol{x}}_{n}],{\boldsymbol{y}})\in\mathbb{R}^{d\times
n}\times R^{1\times n}$; we can write the output of this NN with respect to
this dataset by
$f_{\boldsymbol{\theta}}({\boldsymbol{X}})=[f_{\boldsymbol{\theta}}({\boldsymbol{x}}_{1}),\ldots,f_{\boldsymbol{\theta}}({\boldsymbol{x}}_{n})]$.
The above parameterization is called the NTK parameterization. For the
purposes of training this neural network, we minimize MSE loss

$L({\boldsymbol{\theta}}):=\frac{1}{2n}\left\|{\boldsymbol{y}}-f_{\boldsymbol{\theta}}({\boldsymbol{X}})\right\|^{2}.$
(2)

We will focus on the NTK parameterization and consider the kernel machine
induced by initial NTK of the NN (see (7)). We aim to seek the cases when the
NN outperforms this kernel during the training process. For this purpose, we
adopt different optimization strategies of training this NN to obtain
different testing performances and spectral properties of trained weights and
empirical kernels.

#### Training process of the neural networks.

NNs are usually trained by gradient-based methods: e.g., full-batch gradient
descent (GD), mini-batch stochastic gradient descent (SGD), Adaptive Gradients
(AdaGrad) and Adam from [KB14]. In general, we can present GD by

${\boldsymbol{\theta}}_{t+1}={\boldsymbol{\theta}}_{t}-\eta\nabla_{\boldsymbol{\theta}}L({\boldsymbol{\theta}}_{t}),$
(3)

where $\eta$ is the learning rate and
$\nabla_{\boldsymbol{\theta}}L({\boldsymbol{\theta}}_{t})$ is the gradient of
the training loss w.r.t. trainable parameters ${\boldsymbol{\theta}}$ at step
$t\geq 0$. We will prove the global convergence of GD in some special cases
ensuring the convergence to some interpolation point in this overparameterized
case. We will also show the hyper-parameters (e.g. learning rate $\eta$)
affect the spectral properties of NNs during training.

#### Conjugate kernel and neural tangent kernel.111In this work, we only
consider empirical conjugate and neural tangent kernels of finite-width NNs.

Below we present the CK and NTK for two-layer NN ($L=2$), denoting $n_{1}=h$
and $n_{0}=d$. The CK is defined as the Gram matrix of the output of the first
hidden layer:

${\boldsymbol{K}}^{\text{CK}}:={\boldsymbol{X}}_{1}^{T}{\boldsymbol{X}}_{1}\in\mathbb{R}^{n\times
n},$ (4)

where
${\boldsymbol{X}}_{1}:=\frac{1}{\sqrt{h}}\sigma\Big{(}{\boldsymbol{W}}{\boldsymbol{X}}/\sqrt{d}\Big{)}$.
In principle, we can view the NN as a function of all training parameters
${\boldsymbol{\theta}}$ and input data ${\boldsymbol{X}}$. The neural tangent
kernel (NTK) is related to the gradient of this neural network function with
respect to ${\boldsymbol{\theta}}$, which is the Gram matrix of the Jacobian
of the neural network function with respect to ${\boldsymbol{\theta}}$:
${\boldsymbol{K}}^{\text{NTK}}:=(\nabla_{\boldsymbol{\theta}}f_{\boldsymbol{\theta}}({\boldsymbol{X}}))^{\top}(\nabla_{\boldsymbol{\theta}}f_{\boldsymbol{\theta}}({\boldsymbol{X}}))$.
Specifically, for a two-layer NN, the NTK can be explicitly written as (14) in
Appendix A. In this paper we are interested in comparing the spectral
distributions for these three matrices (weight, CK, and NTK) at initialization
and at the end of training. Prior random matrix theory (RMT) results can
characterize the limiting eigenvalue distribution of the initial weight matrix
[MP67], the CK matrix [PW17, LLC18], and the NTK matrix [AP20, FW20].

##  3 Additional related work

#### Nonlinear random matrix theory and random feature regression.

The limiting spectrum of CK with random input dataset has been investigated by
[PW17, BP21]; whereas [LLC18, FW20] studied the spectrum of CK with more
general input data. This spectrum is actually a deformed Marčenko–Pastur
distribution [FW20], which becomes a deformed semicircular distribution [WZ21]
when $h\gg n$. The largest eigenvalue of the CK matrix has been studied in
[BP22], and the spectrum of the NTK was analyzed in [AP20, FW20]. As an
application, random feature ridge regression was fully determined by the
limiting spectra of CK or NTK: [LLC18, GLK+20, LCM20, MM22, AP20]. All of
these results belong to LWR.

#### Global convergence of GD for ultra wide NNs.

A recent line of work has shown the global convergences of the learning
dynamics of gradient-based methods in a certain overparameterized regime, e.g.
[DZPS18, DLL+19, OFLS19, OS19, Ngu21, LZB22, PC22, Cha22]. We refer to Table 1
of [PC22] as a summary of these recent results. Most of the theorems in
literature require $h\gg n$, which implies that the NTK is almost static
during training, while [OS19, Ngu21] can consider LWR under some specific
assumptions. Recently, [Cha22] established a new criterion for the convergence
of GD which results in the global convergence of general NNs with finite width
$h$ and $d\geq n$.

#### Beyond NTK regime.

Under the proportional limit, the initial kernel regression can only learn a
linear component of the target [GMMM21]. Thus, it is reasonable to consider
the cases beyond the NTK regime. To this end, [DGA20, HY20] considered the
dynamics of NTK throughout training while [AZLL19, BL20] have shown a second-
order approximation of NTK, outperforming the initial kernel. In addition,
there are many theoretical works analyzing when a NN outperforms the initial
kernels in some specific settings: [LMZ20] proved a two-layer ReLU NN that is
shown to beat any kernel method; [KWLS21] verified a two-layer CNN with some
simple dataset can outperform the initial NTK for image classifications;
[BES+22] showed a NN can escape the kernel regime by only taking one specific
large gradient step; [DLS22] showed a specific gradient-based training can
even learn polynomials with low-dimensional latent representation.

#### Evolution of NTK and alignment in NNs.

The feature learning can be characterized by the evolution of the kernel
during training [FDP+20, OJMDF21, Lon21, ABP21, LHAR22]. Specifically, [Lon21]
studied the hard-margin SVM for “after kernels” which are the CK and NTK
matrices of trained NNs. One of effective ways of depicting how the kernels
evolve during training is to capture the evolution of kernel alignment
[BGL+21, SB21, ABP21, LHAR22]. Kernel alignments between kernels and training
labels essentially reveal how the NN accelerates training [SB21]. Also,
several papers showed that the top eigenfunctions of the kernel align with the
target function learned by the NN [KI20, OJMMF20, OJMDF21]. This becomes an
efficient way of analyzing how NNs learn features through a particular
gradient-based optimization.

#### Large learning rate regime.

As mentioned earlier, the large learning rate may contribute to feature
learning. The benefits of large-learning-rate training have been studied from
different aspects [LWM19, Nak20, BMR22, AVPVF22]. Specifically, [LM20]
observed that training dynamics with large learning rates differs from the
small learning rate regime, where the latter regime exhibits monotone and fast
convergence of training loss but may not generalize well on test data. At the
early phase phase of training, [JSF+20] showed using lower learning rates may
result in finding a region of the loss surface with worse conditioning of
kernel and Hessian matrices. In [Lon21], the after kernels of NNs trained with
larger learning rates generalize better and stay more stable. [LBD+20] raised
a “catapult mechanism”, where gradient descent dynamics converge to flatter
minima for extremely large learning rates. There is a transition as a function
of the learning rate, from the lazy training to the catapult regime. Section
5.2 illustrates a similar transition in our situations.

#### Heavy-tailed phenomenon.

The heavy-tailed phenomenon has appeared in many places in deep learning
theory; [MM19, MM21] observed that many state-of-the-art pre-trained models
obtain heavy-tailed weight spectra. More precisely, these spectra have a “5+1”
phase transition which relates to different degrees of regularization of the
NN. With this heavy-tailed self-regularization theory, [MPM21] further showed
how to distinguish well-trained and poorly trained models by a power-law-based
approximation. [MY21] classified trained weight spectra into three types:
Marčenko–Pastur law, Marčenko–Pastur law with (few) outliers, and heavy-tailed
spectra. We extend this classification to both weight and kernel matrices in
Figure 1.

##  4 Case study for linear-width NNs

In this section, we investigate a two-layer NN with synthetic data. This
setting is promising for future theoretical studies by virtue of RMT. We will
showcase the evolution of its spectral properties over training. When $L=2$,
the NN in (1) is a two-layer NN defined by

$f_{\boldsymbol{\theta}}({\boldsymbol{x}}):=\frac{1}{\sqrt{h}}\sum_{i=1}^{h}v_{i}\sigma({\boldsymbol{w}}_{i}^{\top}{\boldsymbol{x}}/\sqrt{d}).$
(5)

At initialization, we assume that the first hidden-layer
${\boldsymbol{W}}=[{\boldsymbol{w}}_{1},\ldots,{\boldsymbol{w}}_{h}]^{\top}\in\mathbb{R}^{h\times
d}$ is composed of independent standard normal random vectors.

######  Assumption 4.1 (Linear-width regime (LWR)).

Assume that $\frac{n}{d}\to\gamma_{1}$ and $\frac{h}{d}\to\gamma_{2}$ as
$n\to\infty$ where the aspect ratios $\gamma_{1},\gamma_{2}\in(0,\infty)$ are
two fixed constants.

######  Assumption 4.2 (Activation function).

Suppose that the activation function $\sigma(x)$ is nonlinear and
$\lambda_{\sigma}$-Lipschitz with
$|\sigma^{\prime}(x)|,|\sigma^{\prime\prime}(x)|\leq\lambda_{\sigma}$ for all
$x\in\mathbb{R}$. Moreover, $\mathbb{E}[\sigma(z)]=0$ for
$z\sim\mathcal{N}(0,1)$.

Although the LWR is somewhat impractical, it is still more aligned with
deployed models than the infinite-width regime $(h\gg n,d)$. As a kernel
machine, the NN under infinite-width regime has been studied extensively
[JGH18, DZPS18, DLL+19, Yan19]. This infinite-width limit is a more special
case, however, as NNs may generally evolve beyond the kernel regime and
achieve superior performance [FDP+20, Lon21, BGL+21, SB21].

######  Assumption 4.3 (Synthetic dataset and teacher model).

Suppose that training data is
${\boldsymbol{X}}:=[{\boldsymbol{x}}_{1},\ldots,{\boldsymbol{x}}_{n}]$, where
${\boldsymbol{x}}_{i}$’s are i.i.d. standard Gaussian random vectors in
$\mathbb{R}^{d}$. The training labels ${\boldsymbol{y}}=[y_{1},\ldots,y_{n}]$
are defined by $y_{i}=f^{*}({\boldsymbol{x}}_{i})+\varepsilon_{i},~{}~{}\text{
for }i\in[n],$ where $f^{*}:\mathbb{R}^{d}\to\mathbb{R}$ is the teacher model,
and $\varepsilon_{i}$ is centered sub-Gaussian noise with variance
$\sigma_{\varepsilon}^{2}$.

The student-teacher model can capture how a NN (student) learns the features
from the teacher model and has been theoretically explored by [AP20, GMKZ20,
AS21, MM22, BES+22, DLS22]. Gaussian data is unnecessary but simplifies the
theoretical analysis. One of the simplest nonlinear teacher models we can
generate is the single-index model, namely
$f^{*}({\boldsymbol{x}})=\sigma^{*}({\boldsymbol{x}}^{\top}{\boldsymbol{\beta}})$
for a fixed vector ${\boldsymbol{\beta}}$ with $\|{\boldsymbol{\beta}}\|=1$
and nonlinear function $\sigma^{*}$; the hidden feature is simply
${\boldsymbol{\beta}}\in\mathbb{R}^{d}$. In general, we can also introduce a
hidden manifold model [GMKZ20, GLK+20]:
$f^{*}({\boldsymbol{x}})={\boldsymbol{a}}^{\top}\sigma^{*}({\boldsymbol{W}}_{*}{\boldsymbol{x}}/\sqrt{d})/\sqrt{n_{h}}$
for a fixed signal vector ${\boldsymbol{a}}\in\mathbb{R}^{n_{h}}$, matrix
${\boldsymbol{W}}_{*}\in\mathbb{R}^{n_{h}\times d}$, and nonlinear function
$\sigma^{*}$.

| Optimization | Learning rate $\eta$ |  $R^{2}$ score | Test error | Spectra
---|---|---|---|---|---
Case 1 | GD | 5.0 | 0.63582 | 0.36381 | Invariant Bulk
Case 2 | SGD | 0.1 | 0.60605 | 0.36879 | Invariant Bulk
Case 3 | SGD | 22.0 | 0.76081 | 0.23791 | Bulk+spike
Case 4 | Adam | 0.092 | 0.78829 | 0.21071 | Heavy tail
| Lazy regime |  | 0.68092 | 0.3185 |
Table 1: Four models with the same architecture ($n=2000$, $h=1500$, $d=1000$,
and $\sigma$ is normalized $\tanh$), but different choices of initial learning
rates and optimizers listed in Table 1. The training label noise
$\sigma_{\varepsilon}=0.3$ and the teacher model is defined by (6) with
$\sigma^{*}$ a normalized softplus and $\tau=0.2$. We observe that simply
choosing an optimizer and learning rate can affect the shapes of the final
spectra and the performance of the NN, as measured by $R^{2}$ scores and test
errors.

Figure 1:  Different spectral behaviors in Table 1: (a) The initial and
trained spectra of ${\boldsymbol{W}}$ in Case 1. The spectrum is invariant
based on the Q-Q subplot. (b) The initial and trained spectra of
${\boldsymbol{K}}^{\text{CK}}$ in Case 3. There is an outlier (orange arrow)
in the spectrum after training. (c) The initial and trained spectra of
${\boldsymbol{K}}^{\text{CK}}$ in Case 4. We refer to Appendix A.2 for other
spectra of weight, CK and NTK, where analogous phenomena hold for other
matrices.

We will specifically consider a mixture of single-index and quadratic models
as our teacher model:

$f^{*}({\boldsymbol{x}})=\sigma^{*}({\boldsymbol{x}}^{\top}{\boldsymbol{\beta}})+\frac{\tau}{d}\|{\boldsymbol{x}}\|^{2},$
(6)

for some nonlinear target $\sigma^{*}$, signal ${\boldsymbol{\beta}}$ and
constant $\tau$. Following the above assumptions and constructions, we show
different spectral properties (Figure 1) for this two-layer NN using different
training procedures (Table 1). Figure 1 exhibits three types of spectra after
training: unchanged bulk distribution, bulk with one spike, and heavy tail in
spectra. The performance of the NN in each of these three cases is summarized
in Table 1. Putting things together, we can see close relationships between
the spectra and the generalization of the NN. These different spectral
properties and spectral evolutions actually reveal disparate features learned
via different training strategies.

#### Lazy training.

The advantage of this toy model is that we can easily extract the spectral
behaviors over training and then compare with the kernel machine. We use lazy
training (where the kernel is the initial NTK) as our benchmark to assist us
in determining whether a NN outperforms the associated kernel machine (Table 1
and Figure 10(c)). Lazy training [COB19] can be viewed as a linear
approximation of the NN, i.e.
$f_{{\boldsymbol{\theta}}}({\boldsymbol{x}})\approx
f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{x}})+({\boldsymbol{\theta}}-{\boldsymbol{\theta}}_{0})^{\top}\nabla_{{\boldsymbol{\theta}}}f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{x}})$,
defined by minimum-norm interpolation $\hat{\boldsymbol{\theta}}:=\arg\min$
$\small{\left\\{\|{\boldsymbol{\theta}}-{\boldsymbol{\theta}}_{0}\|:({\boldsymbol{\theta}}-{\boldsymbol{\theta}}_{0})^{\top}\nabla_{{\boldsymbol{\theta}}}f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{X}})={\boldsymbol{y}}-f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{X}})\right\\}}$.
Also, lazy training represents a kernel machine

$\hat{f}({\boldsymbol{x}})=f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{x}})+\left({\boldsymbol{y}}-f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{X}})\right){\boldsymbol{K}}({\boldsymbol{X}},{\boldsymbol{X}})^{-1}{\boldsymbol{K}}({\boldsymbol{X}},{\boldsymbol{x}}),$
(7)

where $\hat{f}({\boldsymbol{x}})$ is the ridgeless regression prediction on
test data ${\boldsymbol{x}}\in\mathbb{R}^{d}$, the kernel
${\boldsymbol{K}}({\boldsymbol{X}},{\boldsymbol{X}})$ is the initial
${\boldsymbol{K}}^{\text{NTK}}$ on training data, and
${\boldsymbol{K}}({\boldsymbol{X}},{\boldsymbol{x}})=(\nabla_{\boldsymbol{\theta}}f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{X}}))^{\top}(\nabla_{\boldsymbol{\theta}}f_{{\boldsymbol{\theta}}_{0}}({\boldsymbol{x}}))$.
The asymptotic performance of $\hat{f}({\boldsymbol{x}})$ has been analyzed by
[AP20] under LWR. We view this regime as a benchmark: [COB19, BMR21] prove
that NN through gradient flow is close to lazy training if $h\gg n$; [HXAP20]
shows NN can go beyond lazy training under a non-proportional regime. Figure
10(c) in Appendix A.2 indicates that, in some cases (green line in Figure
10(c)), NN can become more efficient than lazy training in LWR.

Table 1 compares the test errors and $R^{2}$ scores for different optimization
cases. By tuning the hyper-parameters, we can find specific situations where
NN outperforms the lazy training. From Figures 1(a), 12 and 13 in Appendix
A.2, one can observe the spectral distributions of weight, CK and NTK matrices
remain invariant and static during training in Cases 1$\&$2, which indicates
both cases still belong to lazy regime. This spectral invariance impedes
further feature learning during the training process. The emergence of the
outlier in Figures 1(b) and 14 of Appendix A.2, however, shows the improvement
over lazy training and potential feature learning via the training process,
where the spectra possibly inherit the structures in teacher models (see
Section 5.2). Comparing with Case 2, the results in Case 3 of Table 1 suggest
the importance of the large learning rate regime for training NNs [LWM19,
Nak20, Lon21, BMR22, AVPVF22]. As a remark, our spectral results of Case 3 are
consistent with the observations in [TSR22] through RMT hypothesis testing,
where the majority of trained weight matrix remain random, and the learned
feature may be contained in the largest singular value (outlier) and
associated vector only. From Figures 1(c) and 16 in Appendix A.2, Case 4
further exhibits more spikes and heavy tails in the trained spectra, which
thoroughly goes beyond the realm of initial kernel machine. Notably, this
phenomenon is not unique to Adam since heavy tails also occur when training
with AdaGrad in Figure 19 in Appendix A.4. Though all of these cases have the
same identical initialization, different methods of optimization eventually
lead to various training trajectories and evolutions of the spectra of the
weight and kernel matrices. To acquire feature learning, Cases 3$\&$4 cause
weights to deviate far from initialization.

##  5 Different spectral behaviors in NNs

We further explore the spectral behaviors in different circumstances of Table
1. We aim to clarify how the spectra evolve through different training
processes and how this evolution may affect the NN. Following Figure 1, we
will study the training processes case-by-case: invariant bulk, spikes outside
the bulk, and heavy-tailed distribution with few outliers. Additional
experiments are exhibited in Appendix A.

###  5.1 Invariant bulk distributions

In Figure 1(a), Figures 12 and 13 in Appendix A.2, we observe the bulk
distributions of weight, CK, and NTK matrices in Cases $1\&2$ remain globally
unchanged (invariant) over the training process. Under LWR, this is
empirically verified by Figures 10(a)$\&$(b) in Appendix A.2 for our typical
examples. The difference between initial ${\boldsymbol{W}}_{0}$ and final-step
${\boldsymbol{W}}_{t}$ stays bounded when aspect ratios $\gamma_{1}$ and
$\gamma_{2}$ are fixed. In this section, by investigating the global
convergence of GD, we prove this invariant-bulk phenomenon under certain
assumptions. For detailed proofs, we refer to Appendix B.

For simplicity, we focus on analyzing the training process of the first-hidden
layer with second layer ${\boldsymbol{v}}$ fixed. Denote
$f_{\boldsymbol{\theta}}({\boldsymbol{X}})$ by
$f_{{\boldsymbol{W}}}({\boldsymbol{X}})$ in this case. At any time
$t\in\mathbb{N}$, consider the gradient steps:

${\boldsymbol{W}}_{t+1}={\boldsymbol{W}}_{t}-\eta\nabla_{{\boldsymbol{W}}}L({\boldsymbol{W}}_{t}).$
(8)

In this case, we denote the CK and NTK at gradient step $t\in\mathbb{N}$ by
${\boldsymbol{K}}_{t}^{\text{CK}}:=\frac{1}{h}\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}})^{\top}\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}),$
and
${\boldsymbol{K}}_{t}^{\text{NTK}}:=\frac{1}{d}{\boldsymbol{X}}^{\top}{\boldsymbol{X}}\odot\frac{1}{h}\sigma^{\prime}\left(\frac{1}{\sqrt{d}}{\boldsymbol{W}}_{t}{\boldsymbol{X}}\right)^{\top}\text{diag}({\boldsymbol{v}}_{t})^{2}\sigma^{\prime}\left(\frac{1}{\sqrt{d}}{\boldsymbol{W}}_{t}{\boldsymbol{X}}\right)$
respectively. First, we present an elaborate description of the changes in the
weight, CK and NTK at the early phase of the training (after any finite $t$
steps) as follows.

######  Lemma 5.1 (Early phase).

Under Assumptions 4.1, 4.2 and 4.3, we further assume that
$\left\|{\boldsymbol{v}}\right\|_{\infty}\leq 1$ and $f^{*}$ is a
$\lambda_{\sigma}$-Lipschitz function. Given any fixed $t\in\mathbb{N}$ and
learning rate $\eta=\Theta(1)$, after $t$ gradient steps, the changes
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$,
$\left\|{\boldsymbol{K}}_{t}^{\text{CK}}-{\boldsymbol{K}}_{0}^{\text{CK}}\right\|_{F}$,
and
$\left\|{\boldsymbol{K}}_{t}^{\text{NTK}}-{\boldsymbol{K}}_{0}^{\text{NTK}}\right\|$
are all less than $\frac{C}{n}$, with probability at least
$1-4n\operatorname{exp}\left({-cn}\right)$, for some positive constants
$c,C>0$ which only depend on step $t$ and parameters
$\eta,\gamma_{1},\gamma_{2},\lambda_{\sigma},\sigma_{\varepsilon}$.

With Lemma 5.1, we can further conclude that
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|$,
$\left\|{\boldsymbol{K}}_{t}^{\text{CK}}-{\boldsymbol{K}}_{0}^{\text{CK}}\right\|$,
and
$\left\|{\boldsymbol{K}}_{t}^{\text{NTK}}-{\boldsymbol{K}}_{0}^{\text{NTK}}\right\|$
are asymptotically vanishing for any fixed time $t$. Therefore, all the
eigenvalues/eigenvectors are asymptotically unchanged at the early phase of
the training (see Corollary B.3 in Appendix B). Now we aim to analyze the
spectra at the end of training process (8). In this case, although we are
unable to show the invariance for each eigenvalue, we can verify the
invariance of the limiting bulk distributions for
${\boldsymbol{K}}^{\text{CK}}_{t}$ and ${\boldsymbol{K}}^{\text{CK}}_{t}$ for
all $t$.

By [WZ21, Theorem 2.9], the smallest eigenvalue of
${\boldsymbol{K}}_{0}^{\text{NTK}}$ has an asymptotic lower bound:

$\lambda_{\min}({\boldsymbol{K}}_{0}^{\text{NTK}})\geq\left(a_{\sigma}-\sum_{k=0}^{2}\eta_{k}^{2}\right)(1-o_{d,\mathbb{P}}(1)),$
(9)

where $a_{\sigma}:=\mathbb{E}[\sigma^{\prime}(\xi)^{2}]$ and $\eta_{k}$ is the
$k$-th Hermite coefficient of $\sigma^{\prime}$. Hence, we can claim there
exists some constant $\alpha>0$ only dependent on $\sigma$ such that
$\lambda_{\min}({\boldsymbol{K}}_{0}^{\text{NTK}})\geq 4\alpha^{2}$ with high
probability. Note that $\alpha$ is not vanishing since $\sigma$ is nonlinear.
With this lower bound, we obtain the following global convergence for (8) and
norm control of ${\boldsymbol{W}}_{t}$ as $n/d\to\gamma_{1}$ and
$h/d\to\gamma_{2}$.

######  Theorem 5.2 (Global convergence).

Under the same assumptions of Lemma 5.1, we further assume $v_{i}$’s are
independent and centered random variables in the second layer. For any
$\eta<\min\\{\frac{\alpha^{2}n}{2},\frac{n}{4\lambda_{\sigma}^{2}(1+\sqrt{\gamma_{1}})^{2}}\\}$
and all $t\in\mathbb{N}$, there exists some $\gamma^{*}>0$ such that, when
$\gamma_{2}\geq\gamma^{*}$, the gradient steps (8) will satisfy

$\displaystyle\ell({\boldsymbol{W}}_{t})\leq\left(1-\frac{\eta\alpha^{2}}{2n}\right)^{t}\ell({\boldsymbol{W}}_{0}),$
(10)
$\displaystyle\frac{1}{4}\alpha\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|_{F}+\ell({\boldsymbol{W}}_{t})\leq\ell({\boldsymbol{W}}_{0}),$
(11)
$\displaystyle\sum_{t=0}^{\infty}\left\|{\boldsymbol{W}}_{t+1}-{\boldsymbol{W}}_{t}\right\|_{F}\leq\frac{4\ell({\boldsymbol{W}}_{0})}{\alpha},$
(12)

with high probability, as $n/d\to\gamma_{1}$ and $h/d\to\gamma_{2}$. Here,
$\ell({\boldsymbol{W}}):=\left\|{\boldsymbol{y}}-f_{{\boldsymbol{W}}}({\boldsymbol{X}})\right\|$
corresponds to the training loss for any ${\boldsymbol{W}}$.

We apply the techniques and results by [OFLS19, OS19] to obtain Theorem 5.2.
Notice that, unlike Lemma 5.1, the largest learning rate we can choose is of
order $\Theta(n)$. As a byproduct, Frobenius norm in (11) implies the
following corollary for the invariance of limiting bulk distribution.

######  Corollary 5.3.

Under the same assumptions of Theorem 5.2, for all $t\in\mathbb{N}$, with high
probability, there exists some constant $R>0$ such that the changes
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$,
$\left\|{\boldsymbol{K}}_{t}^{\text{CK}}-{\boldsymbol{K}}_{0}^{\text{CK}}\right\|_{F}$,
and
$\left\|{\boldsymbol{K}}_{t}^{\text{NTK}}-{\boldsymbol{K}}_{0}^{\text{NTK}}\right\|_{F}$
are all less than $R$ with high probability. This implies the limiting
empirical spectra of
$\frac{1}{h}{\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}_{t}^{\text{CK}}$ and ${\boldsymbol{K}}_{t}^{\text{NTK}}$ are
the same as the limiting spectra of
$\frac{1}{h}{\boldsymbol{W}}_{0}^{\top}{\boldsymbol{W}}_{0}$,
${\boldsymbol{K}}_{0}^{\text{CK}}$ and ${\boldsymbol{K}}_{0}^{\text{NTK}}$
respectively, almost surely as $n/d\to\gamma_{1}$ and $h/d\to\gamma_{2}$.

Corollary 5.3 is empirically validated by Figure 25 in Appendix B. In
addition, based on Figure 27 in Appendix B, one can further extend Corollary
5.3 to SGD training process. The total path is $O(\sqrt{h})$ in (11) and (12),
which is negligible compared with the Frobenius norm of initial weight matrix
(which is of order $\Theta(h)$). Thus, gradient descent iterates (8) remain
close to initialization and small perturbation of NTK ensures the smallest
eigenvalue (9) of NTK is always lower bounded away from zero. Theorem 5.2,
however, does not require that the NTK stays unchanged all the time. Moreover,
Corollary 5.3 only shows the invariance of the bulk distribution, while the
emergence of outliers cannot be excluded from this result. Though we have
global convergence in general, we may still move out of the kernel regime.
Global convergence cannot explain when a NN in LWR outperforms the kernel
regime (Figure 10(c)). Notice that [BMR21, Theorem 5.4] is not directly
applicable to show that a NN is still close to lazy training under LWR. It
requires more-refined analysis to claim whether a NN still belongs to kernel
regime or already goes beyond in our case. As we will show in Section 5.2,
this also relies on the magnitude of the learning rate for GD/SGD.

###  5.2 Alignments for spiked models

In Figures 1(b) and 14 in Appendix A.2, outliers appear in the trained weight,
CK and NTK matrices when NNs are optimized with large learning rates.
Particularly, there is a discernible change in the spectrum of NTK in Figure
14(b). Heuristically, the emergence of outliers in the spectra indicates that
the NN is learning the feature from the teacher model $f^{*}$. This has been
observed by [FW20, CHS20, OJMDF21] and rigorously verified by [BES+22] for the
weight matrix after one-step gradient descent with a large learning rate. In
Figure 2(a), Figures 15 and 23 in Appendix A, we empirically exhibit these
phenomena for ${\boldsymbol{W}}$, ${\boldsymbol{K}}^{\text{CK}}$ and
${\boldsymbol{K}}^{\text{NTK}}$ respectively through different training
processes.

(a) ${\boldsymbol{\beta}}$ vs. ${\boldsymbol{u}}_{1}$ of weight.

(b) $\lambda_{\max}({\boldsymbol{W}}_{t}{\boldsymbol{W}}_{t}^{\top}/d)$

(c) $\lambda_{\max}({\boldsymbol{K}}_{t}^{\text{CK}})$

(d) Alignments

Figure 2: (a) Alignment between teacher feature ${\boldsymbol{\beta}}$ and
first principle component (PC) of the trained/initial weight matrices in Case
3 of Table 1. (b)-(d) Transitions of largest eigenvalues
$\lambda_{\max}({\boldsymbol{W}}_{t}{\boldsymbol{W}}_{t}^{\top}/d)$,
$\lambda_{\max}({\boldsymbol{K}}_{t}^{\text{CK}})$ and alignments
($|{\boldsymbol{\beta}}^{\top}{\boldsymbol{u}}_{1}|/\left\|{\boldsymbol{\beta}}\right\|$
and
$|{\boldsymbol{y}}^{\top}{\boldsymbol{v}}_{1}|/\left\|{\boldsymbol{y}}\right\|$
where ${\boldsymbol{u}}_{1}$ is the leading right singular vector of
${\boldsymbol{W}}_{t}$ and ${\boldsymbol{v}}_{1}$ is the leading eigenvector
of either ${\boldsymbol{K}}_{t}^{\text{CK}}$ or
${\boldsymbol{K}}_{t}^{\text{NTK}}$) when increasing the learning rate $\eta$
in SGD training process. In the green region, the largest eigenvalues are
attached to right edges of the bulks (black horizontal lines) and the
alignments are weak; while in the orange region, outliers become apparent and
the alignments become stronger. For different $\eta$, we train the same NN
with the same dataset until the training loss is less than $10^{-5}$.

(a) Evolution of KTA.

(b) Evolution of the metrics.

(c) Heavy-tailed initial CK.

(d) Bad example.

Figure 3: (a) Evolution of KTA of CK defined by (13) with respect to training
labels for Cases 1, 3$\&$4 in Table 1. We normalize the epoch scales
($x$-axis) for better observations. Heavy-tailed phenomena: (b) The evolution
of power $\alpha$, weighted Alpha and Log $\alpha$-norm during the training
process in Case 4 of Table 1. (c) The CK spectra at two initializations for
${\boldsymbol{W}}$: standard Gaussian and Cauchy distributions. (d) Weight
spectra at initial and after SGD training for the first layer of NN. After
training the weight reveals a heavy tail, but the generalization is not as
good as former examples (test loss $1.47504$; $R^{2}$ score is $-0.48$).

#### Spikes of weight matrix.

The differences between Cases $2\&3$ empirically validate the benefits of
training with large learning rates [LWM19, Nak20, Lon21, BMR22, AVPVF22].
Inspired by [BES+22], we consider the alignment between the leading right
singular vector ${\boldsymbol{u}}_{1}$ of ${\boldsymbol{W}}_{t}$ and the
signal ${\boldsymbol{\beta}}$ in the teacher model defined by (6). For Case 3,
a notable alignment appearing in Figure 2(a) after training suggests that
${\boldsymbol{W}}_{t}$ is capturing the feature ${\boldsymbol{\beta}}$ during
training. Although this does not ensure NN will entirely beat the optimal
kernel lower bound, this alignment reveals a non-negligible feature selection
[BGL+21] via large-stepsize training. This dynamical alignment along the task-
relevant direction may further interpret the generalization of the NN. We also
observe similar phenomena for the adaptive optimization in Figure 23 in
Appendix A.5.

#### Transitions of the spikes as a function of learning rate.

From Case 2 to Case 3, we observe the emergence of outliers in the trained
spectra when increasing the learning rate $\eta$. This indicates a transition
of the emergence of the spike outside the bulk distribution. Figure 2,
analogously to BBP transition by [BAP05] in RMT, shows there is a threshold
(yellow region) for learning rate: the outliers only appear when $\eta$
exceeds this threshold. We fix the same NN and dataset for all trials of
training. The flat black lines in Figures 2(b) and (c) are the right edges of
the limiting spectra at initialization. Figure 2(d) records the angles between
${\boldsymbol{\beta}}$ and the leading eigenvector of
${\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}/d$, and ${\boldsymbol{y}}$
and the leading eigenvectors of ${\boldsymbol{K}}_{t}^{\text{CK}}$ and
${\boldsymbol{K}}_{t}^{\text{NTK}}$ after training for different $\eta$.
Similarly with [BGL+21], when $\eta$ is sufficiently large (orange region), we
obtain significant alignments which suggest potential feature learning. These
transitions of leading eigenvalue and eigenvector alignment have been proved
for weight spectrum by [BES+22] for a different scenario.

#### Spikes of kernel matrices.

The alignment of the kernel matrix with the training labels ${\boldsymbol{y}}$
has been defined by [CSTEK01] by Kernel Target Alignment (KTA) as follows:
when kernel ${\boldsymbol{K}}$ is either ${\boldsymbol{K}}^{\text{CK}}$ or
${\boldsymbol{K}}^{\text{NTK}}$,

$\displaystyle\mathrm{KTA}=\frac{\langle{\boldsymbol{K}},{\boldsymbol{y}}^{\top}{\boldsymbol{y}}\rangle}{\left\|{\boldsymbol{K}}\right\|_{F}\left\|{\boldsymbol{y}}\right\|^{2}}.$
(13)

Analogously to [BGL+21, ABP21, SB21], Figure 3(a) depicts the evolution of KTA
of CK in several cases. Based on Figure 2(d), when the spike appears outside
the bulk (Case $3$), its corresponding (leading) eigenvector
${\boldsymbol{v}}_{1}$ of kernel matrix naturally dominates the alignment with
${\boldsymbol{y}}$ (Figure 15 in Appendix A.2), which is regarded as a kernel
rotation during training in [OJMDF21]. Notice that this is not the common
situation in Cases 1$\&$2 of Table 1 (and cf. Figure 11 in Appendix A.2). On
the other hand, KTA measures the alignment between ${\boldsymbol{y}}$ and the
full eigenbasis of the kernel. These kernel alignments provably improve the
speed of the convergence of training dynamics, but may hurt or boost the
generalization of the NNs [OJMDF21, SB21, BGL+21]. Figure 3(a) indicates that
Case 4 with heavy-tailed spectra after training has a larger KTA than the
other cases. In this case, the emergence of heavy tail in the spectrum is
closely related to a better generalization of the NN and more significant
feature learning.

###  5.3 Phenomenon of heavy-tailed spectra

We next analyze the heavy-tailed spectra of weight and kernel matrices in
Figure 1(c). [MM19, MM21] found a strong correlation between the heavy-tailed
spectra of trained state-of-the-art models with better generalization (Figure
7 in Appendix A.1). Heavy-tailed spectra can be viewed as an extreme of
“bulk+spikes”, where a fraction of the eigenvalues move out of the initial
bulk. In RMT, heavy-tailed spectra generally appear when the entries of the
matrix are highly correlated [MM21]. This could heuristically explain heavy-
tailed phenomena in the spectra since the entries of well-trained
${\boldsymbol{W}}_{t}$ should be strongly correlated. We emphasize that the
emergence of heavy tails after training demands a large initial learning rate,
because Figures 20 and 22 in Appendix A.5 indicate that there is only one
outlier appearing after training when using a smaller initial learning rate.
Additionally, Figure 19 in Appendix A.4 reveals similar heavy-tailed spectra
after training with AdaGrad. Unlike [MM19, MM21], we focus on the heavy-tailed
phenomena for both weight and kernel matrices in a simpler model (5), which
opens an important avenue for further theoretical analysis from RMT.

#### Heavy-tailed and Power Law.

Heavy tails are essentially power laws. To measure how “heavy” the spectrum
is, [MM19, MM21, MPM21] provide estimates on the power law and several metrics
of ${\boldsymbol{W}}$. Consider the empirical spectrum of ${\boldsymbol{W}}$
as $\rho(x)\sim x^{-\alpha}$ for large $x$ and some positive constant
$\alpha$. The spectrum with a heavier tail has smaller value of $\alpha$.
Figure 3(b) shows how $\alpha$ evolves through training in Case 4 of Table 1.
As $\alpha$ decreases, a heavy tail in the spectrum of the weight matrix
emerges in Figure 1(c). In Figure 3(b) we introduce two more metrics to show
this evolution: Weighted Alpha $\hat{\alpha}:=\alpha\lambda_{1}$ and Log
$\alpha$-norm $\log\left(\sum_{i=1}^{N}\lambda_{i}^{\alpha}\right)$ where
$\lambda_{N}\leq\ldots\leq\lambda_{1}$ are the eigenvalues of
${\boldsymbol{W}}{\boldsymbol{W}}^{\top}$. Remarkably, Figure 3(b) indicates
the spectra changes dramatically at the early stage of training, which matches
the observation from fine tuning via BERT on real-world data in Figure 5 in
Section 6.

From [MM21, MM19, MPM21], one may be left with the impression that NNs with
heavy tails always generalize well. Unfortunately, this is not case. One can
construct random NNs with heavy-tailed weight in the absence of good
performance at initialization (e.g. Figure 3(c) applies the Cauchy
distribution at initialization). Additionally, Figure 3(d) shows a situation
where the trained weight has perfect power-law tail with a poor
generalization.

(a) GD

(b) SGD

(c) Adam

Figure 4: Different NTK spectra for a small-CNN model on CIFAR-2. Test
accuracy: (a) 79$\%$, (b) 84$\%$, (c) 86.4$\%$.

##  6 Discussions and future directions

We empirically investigated how the spectra of ${\boldsymbol{W}}$,
${\boldsymbol{K}}^{\text{CK}}$ and ${\boldsymbol{K}}^{\text{NTK}}$ evolve
under proportional limit for an idealized student-teacher setting. The
evolution relies on the choice of optimizer and initial learning rate,
resulting in either belonging to or breaking beyond the kernel regime. Our
work implies that understanding the relationship between feature learning and
training processes, requires understanding the evolution of the spectra of
both weight and kernel matrices. In particular, we show that different
training processes affect the eigen-structures of weight and kernel matrices.
Since the evolution is sensitive to feature learning, we can link feature
learning and different training dynamics by studying the spectra of these
matrices. We provide a concise but general model for future theoretical
analysis in deep learning theory and exhibit some hidden spectral properties
of trained NNs which will aid in future theoretical endeavors.

Although it is easier to theoretically analyze and extract spectral properties
of a two-layer NN with synthetic data, we would like to investigate these
spectral properties for general state-of-art architectures, on real-world data
and in more complicated tasks/teacher models. As some preliminary examples, we
present the spectral behaviors for two practical architectures on real-world
datasets. In practice, people mainly focused on the weight matrices in fully-
connected layers and analyzing their spectra. We emphasize, however, the
spectral properties of general kernel matrices induced by the NNs which
contain abundant information of training and generalization [FDP+20, CHS20,
Lon21, ABP21, SB21, WHS22].

First, we show the spectra of ${\boldsymbol{K}}^{\text{NTK}}$ before and after
training for binary classification on CIFAR-2 through small CNNs in Figure 4.
Similarly with Case 1, Figure 4(a) (especially in the Q-Q subplot) manifests
the invariant spectral distribution of NTK through GD training while SGD
exhibits a heavier tail in NTK spectrum in Figure 4(b). This phenomenon is
more evident when trained by Adam in Figure 4(c) with improved accuracy.
Figure 4 suggests that our observations on synthetic data in Section 4 can be
extended to real-world data and on more practical architectures. We note that
there is a lack of the emergence of spikes after training because spikes
already exist in the initial NTK spectrum for this complicated neural
architecture on real-world dataset. Figure 4(a) also indicates that the
spectral invariance of NTK through training will impede the feature learning
and the NN does not generalize well in this training process.

We also investigate the spectral properties on the pre-trained model, BERT
from [DCLT18], with fine-tuning on Sentiment140 dataset of
tweets222https://www.kaggle.com/datasets/kazanova/sentiment140 from [GBH09].
We fine-tune BERT model to train a binary classifier on Sentiment140 and
capture the evolution of CK spectra, rather than the NTK due to the size of
BERT, in Figure 5 (see also Figure 7 in Appendix A.1). A heavy-tailed CK
spectrum with several spikes already exists in this pre-trained model. Unlike
Figure 4 (and every other randomly initialized NN) where the first spike of
NTK becomes larger than at random initialization after training, in Figure
5(a), the leading eigenvalue first decreases and then increases. Moreover,
similarly with Figure 2(d), Figure 5(b) shows that the alignment of the first
eigenvector of the CK and training labels becomes more apparent through fine-
tuning with the highest increase correlating with the leading eigenvalue
decrease. Heuristically, this process seems to unlearn the features in the
pre-trained model and, remarkably, learn new features on the new dataset in
only a few epochs of fine-tuning (For more details of these two experiments,
we refer to Appendix A.1). We believe that the evolutions of the kernel
matrices and some spectral metrics are crucial for understanding feature
learning through fine-tuning [WHS22]. A more comprehensive exploration of the
evolutionary spectral properties of “foundation models” may help shed further
light on these phenomena.

(a) Eigenvalues

(b) Alignments

Figure 5: We use SGD for fine-tuning of the BERT model. The training accuracy
is 95.90$\%$ and test accuracy is $84\%$. (a) The evolution of first and
second eigenvalues of empirical CK during fine-tuning. (b) The alignments of
training labels with first and second eigenvectors of CK during fine-tuning
process respectively.

## Acknowledgement

Z.W., A.E., I.D., and T.C. were partially supported by the Mathematics for
Artificial Reasoning in Science (MARS) initiative via the Laboratory Directed
Research and Development (LDRD) Program at Pacific Northwest National
Laboratory (PNNL). A.S. and T.C. were also partially supported by the
Statistical Inference Generates kNowledge for Artificial Learners (SIGNAL)
program at PNNL. PNNL is a multi-program national laboratory operated for the
U.S. Department of Energy (DOE) by Battelle Memorial Institute under Contract
No. DE-AC05-76RL0-1830.

## References

  * [ABP21] Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan.  Neural networks as kernel learners: The silent alignment effect.  arXiv preprint arXiv:2111.00034, 2021.
  * [ADH+19] 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, 2019.
  * [AP20] Ben Adlam and Jeffrey Pennington.  The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization.  In International Conference on Machine Learning, pages 74–84. PMLR, 2020.
  * [AS21] Shunta Akiyama and Taiji Suzuki.  On learnability via gradient method for two-layer relu neural networks in teacher-student setting.  In International Conference on Machine Learning, pages 152–162. PMLR, 2021.
  * [AVPVF22] Maksym Andriushchenko, Aditya Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion.  Sgd with large step sizes learns sparse features.  arXiv preprint arXiv:2210.05337, 2022.
  * [AZLL19] 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, 2019.
  * [BAP05] Jinho Baik, Gérard Ben Arous, and Sandrine Péché.  Phase transition of the largest eigenvalue for nonnull complex sample covariance matrices.  The Annals of Probability, 33(5):1643–1697, 2005.
  * [BES+22] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang.  High-dimensional asymptotics of feature learning: How one gradient step improves the representation.  arXiv preprint arXiv:2205.01445, 2022.
  * [BGL+21] Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien.  Implicit regularization via neural feature alignment.  In International Conference on Artificial Intelligence and Statistics, pages 2269–2277. PMLR, 2021.
  * [BL20] Yu Bai and Jason D. Lee.  Beyond linearization: On quadratic and higher-order approximation of wide neural networks.  In International Conference on Learning Representations, 2020.
  * [BM19] Alberto Bietti and Julien Mairal.  On the inductive bias of neural tangent kernels.  In Advances in Neural Information Processing Systems, pages 12873–12884, 2019.
  * [BMM18] Mikhail Belkin, Siyuan Ma, and Soumik Mandal.  To understand deep learning we need to understand kernel learning.  In International Conference on Machine Learning, pages 541–549. PMLR, 2018.
  * [BMR21] Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin.  Deep learning: a statistical viewpoint.  Acta numerica, 30:87–201, 2021.
  * [BMR22] Gaspard Beugnot, Julien Mairal, and Alessandro Rudi.  On the benefits of large learning rates for kernel methods.  arXiv preprint arXiv:2202.13733, 2022.
  * [BP21] Lucas Benigni and Sandrine Péché.  Eigenvalue distribution of some nonlinear models of random matrices.  Electronic Journal of Probability, 26:1–37, 2021.
  * [BP22] Lucas Benigni and Sandrine Péché.  Largest eigenvalues of the conjugate kernel of single-layered neural networks.  arXiv preprint arXiv:2201.04753, 2022.
  * [BS10] Zhidong Bai and Jack W Silverstein.  Spectral analysis of large dimensional random matrices, volume 20.  Springer, 2010.
  * [CB18] Lenaic Chizat and Francis Bach.  On the global convergence of gradient descent for over-parameterized models using optimal transport.  In Advances in neural information processing systems, pages 3036–3046, 2018.
  * [Cha22] Sourav Chatterjee.  Convergence of gradient descent for deep neural networks.  arXiv preprint arXiv:2203.16462, 2022.
  * [CHS20] Shuxiao Chen, Hangfeng He, and Weijie Su.  Label-aware neural tangent kernel: Toward better generalization and local elasticity.  Advances in Neural Information Processing Systems, 33, 2020.
  * [COB19] Lenaic Chizat, Edouard Oyallon, and Francis Bach.  On lazy training in differentiable programming.  Advances in Neural Information Processing Systems, 32, 2019.
  * [CSTEK01] Nello Cristianini, John Shawe-Taylor, Andre Elisseeff, and Jaz Kandola.  On kernel-target alignment.  Advances in neural information processing systems, 14, 2001.
  * [DCLT18] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [DGA20] Ethan Dyer and Guy Gur-Ari.  Asymptotics of wide networks from feynman diagrams.  In International Conference on Learning Representations, 2020.
  * [DLL+19] 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, pages 1675–1685. PMLR, 2019.
  * [DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi.  Neural networks can learn representations with gradient descent.  In Conference on Learning Theory, pages 5413–5452. PMLR, 2022.
  * [DM20] Amit Daniely and Eran Malach.  Learning parities with neural networks.  Advances in Neural Information Processing Systems, 33:20356–20365, 2020.
  * [DZPS18] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In International Conference on Learning Representations, 2018.
  * [FDP+20] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli.  Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel.  Advances in Neural Information Processing Systems, 33:5850–5861, 2020.
  * [FW20] Zhou Fan and Zhichao Wang.  Spectra of the conjugate kernel and neural tangent kernel for linear-width neural networks.  Advances in neural information processing systems, 33, 2020.
  * [GBH09] Alec Go, Richa Bhayani, and Lei Huang.  Twitter sentiment classification using distant supervision.  CS224N project report, Stanford, 1(12):2009, 2009.
  * [GLK+20] Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc Mézard, and Lenka Zdeborová.  Generalisation error in learning with random features and the hidden manifold model.  In International Conference on Machine Learning, pages 3452–3462. PMLR, 2020.
  * [GMKZ20] Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová.  Modeling the influence of data structure on learning in neural networks: The hidden manifold model.  Physical Review X, 10(4):041044, 2020.
  * [GMMM20] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  When do neural networks outperform kernel methods?  Advances in Neural Information Processing Systems, 33:14820–14830, 2020.
  * [GMMM21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari.  Linearized two-layers neural networks in high dimension.  The Annals of Statistics, 49(2):1029–1054, 2021.
  * [GSJW20] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart.  Disentangling feature and lazy training in deep neural networks.  Journal of Statistical Mechanics: Theory and Experiment, 2020(11):113301, 2020.
  * [HN20] Boris Hanin and Mihai Nica.  Products of many large random matrices and gradients in deep neural networks.  Communications in Mathematical Physics, 376(1):287–322, 2020.
  * [HXAP20] Wei Hu, Lechao Xiao, Ben Adlam, and Jeffrey Pennington.  The surprising simplicity of the early-time learning dynamics of neural networks.  Advances in Neural Information Processing Systems, 33:17116–17128, 2020.
  * [HY20] Jiaoyang Huang and Horng-Tzer Yau.  Dynamics of deep neural networks and neural tangent hierarchy.  In International Conference on Machine Learning, pages 4542–4551. PMLR, 2020.
  * [JGH18] 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.
  * [JSF+20] Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras.  The break-even point on optimization trajectories of deep neural networks.  In International Conference on Learning Representations, 2020.
  * [JT20] Ziwei Ji and Matus Telgarsky.  Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks.  In International Conference on Learning Representations, 2020.
  * [KB14] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [KI20] Dmitry Kopitkov and Vadim Indelman.  Neural spectrum alignment: Empirical study.  In International Conference on Artificial Neural Networks, pages 168–179. Springer, 2020.
  * [KWLS21] Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh.  Local signal adaptivity: Provable feature learning in neural networks beyond kernels.  Advances in Neural Information Processing Systems, 34, 2021.
  * [LBD+20] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari.  The large learning rate phase of deep learning: the catapult mechanism.  arXiv preprint arXiv:2003.02218, 2020.
  * [LBN+17] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein.  Deep neural networks as gaussian processes.  arXiv preprint arXiv:1711.00165, 2017.
  * [LCKS91] Yann Le Cun, Ido Kanter, and Sara A Solla.  Eigenvalues of covariance matrices: Application to neural-network learning.  Physical Review Letters, 66(18):2396, 1991.
  * [LCM20] Zhenyu Liao, Romain Couillet, and Michael W Mahoney.  A random matrix analysis of random fourier features: beyond the gaussian kernel, a precise phase transition, and the corresponding double descent.  Advances in Neural Information Processing Systems, 33:13939–13950, 2020.
  * [LHAR22] Noel Loo, Ramin Hasani, Alexander Amini, and Daniela Rus.  Evolution of neural tangent kernels under benign and adversarial training.  arXiv preprint arXiv:2210.12030, 2022.
  * [LLC18] Cosme Louart, Zhenyu Liao, and Romain Couillet.  A random matrix approach to neural networks.  The Annals of Applied Probability, 28(2):1190–1248, 2018.
  * [LM20] Guillaume Leclerc and Aleksander Madry.  The two regimes of deep network training.  arXiv preprint arXiv:2002.10376, 2020.
  * [LMZ20] Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang.  Learning over-parametrized two-layer neural networks beyond ntk.  In Conference on learning theory, pages 2613–2682. PMLR, 2020.
  * [Lon21] Philip M Long.  Properties of the after kernel.  arXiv preprint arXiv:2105.10585, 2021.
  * [LSP+20] Jaehoon Lee, Samuel Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein.  Finite versus infinite neural networks: an empirical study.  Advances in Neural Information Processing Systems, 33:15156–15172, 2020.
  * [LWM19] Yuanzhi Li, Colin Wei, and Tengyu Ma.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  In Advances in Neural Information Processing Systems, pages 11674–11685, 2019.
  * [LZB22] 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, 2022.
  * [MM19] Michael Mahoney and Charles Martin.  Traditional and heavy tailed self regularization in neural network models.  In International Conference on Machine Learning, pages 4284–4293. PMLR, 2019.
  * [MM21] Charles H Martin and Michael W Mahoney.  Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning.  Journal of Machine Learning Research, 22(165):1–73, 2021.
  * [MM22] Song Mei and Andrea Montanari.  The generalization error of random features regression: Precise asymptotics and the double descent curve.  Communications on Pure and Applied Mathematics, 75(4):667–766, 2022\.
  * [MMN18] Song Mei, Andrea Montanari, and Phan-Minh Nguyen.  A mean field view of the landscape of two-layer neural networks.  Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
  * [MP67] Vladimir A Marčenko and Leonid Andreevich Pastur.  Distribution of eigenvalues for some sets of random matrices.  Mathematics of the USSR-Sbornik, 1(4):457, 1967.
  * [MPM21] Charles H Martin, Tongsu Serena Peng, and Michael W Mahoney.  Predicting trends in the quality of state-of-the-art neural networks without access to training or testing data.  Nature Communications, 12(1):1–13, 2021.
  * [MY21] Xuran Meng and Jianfeng Yao.  Implicit data-driven regularization in deep neural networks under sgd.  arXiv preprint arXiv:2111.13331, 2021.
  * [Nak20] Preetum Nakkiran.  Learning rate annealing can provably help generalization, even for convex problems.  arXiv preprint arXiv:2005.07360, 2020.
  * [Nea95] Radford M Neal.  Bayesian learning for neural networks, volume 118.  Springer Science & Business Media, 1995.
  * [Ngu21] Quynh Nguyen.  On the proof of global convergence of gradient descent for deep relu networks with linear widths.  In International Conference on Machine Learning, pages 8056–8062. PMLR, 2021.
  * [OFLS19] Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi.  Generalization guarantees for neural networks via harnessing the low-rank structure of the jacobian.  arXiv preprint arXiv:1906.05392, 2019.
  * [OJMDF21] Guillermo Ortiz-Jiménez, Seyed-Mohsen Moosavi-Dezfooli, and Pascal Frossard.  What can linearized neural networks actually say about generalization?  Advances in Neural Information Processing Systems, 34, 2021.
  * [OJMMF20] Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi, and Pascal Frossard.  Neural anisotropy directions.  Advances in Neural Information Processing Systems, 33:17896–17906, 2020.
  * [OS19] Samet Oymak and Mahdi Soltanolkotabi.  Overparameterized nonlinear learning: Gradient descent takes the shortest path?  In International Conference on Machine Learning, pages 4951–4960. PMLR, 2019.
  * [OS20] Samet Oymak and Mahdi Soltanolkotabi.  Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks.  IEEE Journal on Selected Areas in Information Theory, 1(1):84–105, 2020.
  * [PB17] Jeffrey Pennington and Yasaman Bahri.  Geometry of neural network loss surfaces via random matrix theory.  In International Conference on Machine Learning, pages 2798–2806. PMLR, 2017.
  * [PC22] Bartlomiej Polaczyk and Jacek Cyranka.  Improved overparametrization bounds for global convergence of stochastic gradient descent for shallow neural networks.  arXiv preprint arXiv:2201.12052, 2022.
  * [PW17] Jeffrey Pennington and Pratik Worah.  Nonlinear random matrix theory for deep learning.  Advances in neural information processing systems, 30, 2017.
  * [PW18] Jeffrey Pennington and Pratik Worah.  The spectrum of the fisher information matrix of a single-hidden-layer neural network.  Advances in neural information processing systems, 31, 2018.
  * [RGKZ21] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová.  Classifying high-dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed.  In International Conference on Machine Learning, pages 8936–8947. PMLR, 2021.
  * [SB21] Haozhe Shan and Blake Bordelon.  A theory of neural tangent kernel alignment and its influence on training.  arXiv e-prints, pages arXiv–2105, 2021.
  * [TSR22] Matthias Thamm, Max Staats, and Bernd Rosenow.  Random matrix analysis of deep neural network weight matrices.  arXiv preprint arXiv:2203.14661, 2022.
  * [Ver18] Roman Vershynin.  High-dimensional probability: An introduction with applications in data science, volume 47.  Cambridge university press, 2018.
  * [WHS22] Alexander Wei, Wei Hu, and Jacob Steinhardt.  More than a toy: Random matrix models predict how real-world neural representations generalize.  In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 23549–23588. PMLR, 17–23 Jul 2022.
  * [Wil96] Christopher Williams.  Computing with infinite networks.  Advances in neural information processing systems, 9, 1996.
  * [WZ21] Zhichao Wang and Yizhe Zhu.  Deformed semicircle law and concentration of nonlinear random matrices for ultra-wide neural networks.  arXiv preprint arXiv:2109.09304, 2021.
  * [Yan19] Greg Yang.  Wide feedforward or recurrent neural networks of any architecture are gaussian processes.  Advances in Neural Information Processing Systems, 32, 2019.
  * [YH20] Greg Yang and Edward J Hu.  Feature learning in infinite-width neural networks.  arXiv preprint arXiv:2011.14522, 2020.
  * [ZCZG18] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu.  Stochastic gradient descent optimizes over-parameterized deep relu networks.  arXiv preprint arXiv:1811.08888, 2018.

##  Appendix A Additional empirical results

There are different parameterizations for NNs at initialization. The orders of
the output of NN are distinct in different cases [COB19, DM20, YH20]. This
affects size of stable and non-trivial gradient step. The distance of
trainable parameters from initialization determines whether the NN learns any
features from the training data [BES+22, Figure 2]. The performance of
networks with different initializations indicate whether the NN belongs to
kernel regime or not [YH20]. Unlike the NTK parameterization, the mean-field
parameterization [MMN18, CB18] and maximal update parameterization [YH20] tend
to be feature learning.

For all NNs in our experiments, we apply a normalized and centered nonlinear
activation function such that Assumption 4.2 holds ($\mathbb{E}[\sigma(z)]=0$
for $z\sim\mathcal{N}(0,1)$) because we can exclude a large but trivial spike
in the initial spectra of kernel matrices. In all architectures of NNs we
considered, we remove the bias term of each layer and apply the NTK
parameterization (1) with standard Gaussian initialization. Specifically, all
entries of ${\boldsymbol{W}}$ and ${\boldsymbol{v}}$ in (5) at initialization
are i.i.d. standard Gaussian random variables. For all experiments on
synthetic datasets, we use standard Gaussian random matrices to generate the
training data ${\boldsymbol{X}}$. In addition, we consider the training label
noise defined in Assumption 4.3 as
$\boldsymbol{\varepsilon}\sim\mathcal{N}(0,\sigma_{\varepsilon}^{2}\mathbf{I}_{n})$.
The empirical NTK of two-layer NN can be explicitly written333Here we train
both layers, so we have two parts in the NTK expression. If we only train the
first-hidden layer, we can simply remove the second CK part. For general
formula of the empirical NTK, see [HY20, FW20]. as

$\displaystyle{\boldsymbol{K}}^{\text{NTK}}=\frac{1}{d}{\boldsymbol{X}}^{\top}{\boldsymbol{X}}\odot\frac{1}{h}\sigma^{\prime}\left(\frac{1}{\sqrt{d}}{\boldsymbol{W}}{\boldsymbol{X}}\right)^{\top}\text{diag}({\boldsymbol{v}})^{2}\sigma^{\prime}\left(\frac{1}{\sqrt{d}}{\boldsymbol{W}}{\boldsymbol{X}}\right)+{\boldsymbol{K}}^{\text{CK}}.$
(14)

In the following, we further introduce more empirical results for practical
NNs in Section 6 and two-layer NNs with Gaussian dataset in Section 4.

###  A.1 Further discussion on real-world data experiments

For the first experiment in Section 6, we fix the small-CNN architecture,
which is similar to VGG model, and CIFAR-2 dataset, and vary the methods of
optimization of training. Corresponding to Figure 4, the training and test
accuracy histories for three cases are shown in Figure 6. Here, in Figure
6(a), we used GD with learning rate $8\times 10^{-3}$; we used SGD with
learning rate $10^{-3}$, batch size 32 and momentum $0.2$ in Figure 6(b);
Figure 6(c) employs the same learning rate and batch size as 6(b) but employs
Adam optimization.

(a) GD

(b) SGD

(c) Adam

Figure 6: Training/test accuracy v.s. epochs for small-CNN model on CIFAR-2
with different optimizers.

In the experiment of transformer language model in Section 6, we fine-tuned
the BERT model with SGD for a binary classification on Sentiment140 dataset.
We apply this transformer and fine-tuning to sentiment analysis for social
media data. For fine-tuning, the learning rate is 0.003, the batch size is 64
and the momentum is 0.8. The purpose of this experiment is to extract the
spectral properties of pre-trained models and the evolution of the CK spectra
over fine-tuning. Combining Figure 5, the following Figure 7 exhibits the
evolution of CK spectrum during fine-tuning. Similarly with Case 4 in Table 1,
the CK spectrum of this pre-trained model (red histogram in Figure 7)
possesses a heavy-tailed distribution, which suggests this transformer has
received adequate training. From Figure 7(a) to 7(c), we observe the bulk
distribution first shrinks then extends during fine-tuning. This is similar
with the evolution of the first eigenvalue of CK in Figure 5(a). Accompanied
by this spectra evolution, there is a rapid transformation of the features
through fine-tuning, linking the features in the pre-trained model with
features in the new dataset. We expect that further spectral analysis will
elucidate the feature learning in this kind of transformers [WHS22].

Figure 7: The spectra of CK of the BERT model on Sentiment140 dataset at epoch
1, 2 and 12.

###  A.2 Additional results for cases in Table 1

To complement the findings in Figure 1 and Section 4, we now present
additional results on synthetic data and two-layer NNs. In this section, we
will always use the same architecture and dataset as the typical examples in
Table 1.

#### Norms of the change.

Based on Figure 1, the trajectories of the weight and kernel matrices are
quite different among all cases in Table 1. Hence, for all cases in Table 1,
we record the changes in the weight and NTK matrices in both Frobenius norm
and operator norm in Figure 8:

$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|_{F},~{}\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|,~{}\left\|{\boldsymbol{K}}^{\text{NTK}}_{0}-{\boldsymbol{K}}^{\text{NTK}}_{t}\right\|_{F},~{}\text{and
}\left\|{\boldsymbol{K}}^{\text{NTK}}_{0}-{\boldsymbol{K}}^{\text{NTK}}_{t}\right\|,$

at every epoch $t$ through training. The changes in Figures 8(a) and (b) are
much smaller than the last case. Figure 8(c) has significant changes in both
norms after training, which is consistent with the heavy-tailed phenomenon in
Figure 1(c). The global optima of the last case is far from the
initialization.

Following the settings of Theorem 5.2, in Figure 9, we compute the differences
between initial ${\boldsymbol{W}}_{0}$ and final ${\boldsymbol{W}}_{s}$ in
Frobenius norm, operator norm and $2,\infty$-norm. Empirically, Figure 9 shows

$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{s}\right\|_{F},~{}\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{s}\right\|,~{}\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{s}\right\|_{2,\infty}=\Theta(1)$
(15)

as $n\to\infty$ with $n/d\to\gamma_{1}$ and $N/d\to\gamma_{2}$, where $s$ is
the final time for GD. Here, the entry-wise $2$-$\infty$ matrix norm is
defined as

$\|{\boldsymbol{M}}\|_{2,\infty}:=\max_{1\leq i\leq N}\|\textbf{m}_{i}\|,$

for any matrix ${\boldsymbol{M}}\in\mathbb{R}^{N\times d}$ with the $i$-th row
$\textbf{m}_{i}\in\mathbb{R}^{d}$ and $1\leq i\leq N$. Notice that

$\|{\boldsymbol{M}}\|_{2,\infty}\leq\|{\boldsymbol{M}}\|\leq\|{\boldsymbol{M}}\|_{F}.$
(16)

Similar observations for CK and NTK in both Frobenius norm and operator norm
are also apparent in Figure 9, which empirically verifies the invariance of
the spectra after training. Here, we fix the aspect ratios and let $n$ be
growing to keep the NNs reside in LWR. For different $n$’s, we repeat the
experiments 10 times for average. In each experiment, we train the NN until it
converges. As shown in Figure 9(a), the test losses are almost same for
different $n$’s. Figures 9(b)-(d) empirically validate Corollary 5.3.
Moreover, the observation that
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{s}\right\|_{F}$
and
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{s}\right\|$
are $\Theta(1)$ may suggest that
$\frac{1}{\sqrt{d}}\big{(}{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{s}\big{)}$ is
a low-rank perturbation. That is, training in LWR may be transferring some
low-rank structures to the weight spectrum. This low-rank perturbation can
help us better understanding the spectral evolution during training. Notice
that these norms of the change are different from ultra wide NN [DLL+19,
DZPS18]. Similar phenomena can be also observed in Figures 10(a) and (b).
Analogous result with different $\sigma$ and $\sigma^{*}$ is exhibited in
Figure 26 in Appendix B. In addition, Figure 10(c) further investigates the
cases when NNs can outperform lazy training as defined by (7). In these
experiments, we compare the performances of GD, SGD with small or large
learning rates, and lazy training as $n\to\infty$. Each time, we take 10
trials to average. We observe that SGD with a large learning rate (green line)
can asymptotically outperform lazy training.

(a) Case 1

(b) Case 2

(c) Case 4

Figure 8: The evolution of the changes in operator/Frobenius norms of the
weight/CK/NTK matrices through different training processes. Each case
corresponds to case in Table 1. Case 3 is exhibited in Figure 14(c) below.

(a) Train/test loss (b) Change in weight (c) Change in CK (d) Change in NTK

Figure 9: Performances of NNs and changes in different norms for weight and
kernels, when $d/n=0.6$ and $N/n=1.2$ are fixed as $n$ is growing. The
activation is normalized $\tanh$ and the teacher model is
$f^{*}({\boldsymbol{x}})=\sigma^{*}({\boldsymbol{\beta}}^{\top}{\boldsymbol{x}})$
where $\sigma^{*}$ is normalized softplus. We average over 10 trials in each
case. All these curves are almost flat, which indicates these values are not
growing with $\gamma_{1}$ and $\gamma_{2}$. Here, in the second figure from
left, we normalized all weights ${\boldsymbol{W}}$ with $\frac{1}{\sqrt{d}}$
to observe (15).

(a) Change in weight.

(b) Change in CK.

(c) Prediction risks vs. $\gamma_{1}$.

Figure 10: (a) Change between initial ${\boldsymbol{W}}_{0}$ and final step
${\boldsymbol{W}}_{s}$ in operator norm, Frobenius norm and $(2,\infty)$-norm
when $d/n=0.5,N/n=0.8$ are fixed as $n\to\infty$. We train NNs by SGD with
$\eta=2.5$ for 15 trials to average. (b) Change of
${\boldsymbol{K}}^{\text{CK}}$ in operator norm and Frobenius norm. (c)
Prediction risks for lazy training defined by (7), GD with $\eta=\Theta(1)$
(red), SGD with $\eta=\Theta(1)$ (blue dot) and $\eta\propto\gamma_{1}$
(green), as $\gamma_{1}\to\infty$ and $\gamma_{2}=2.5$. The black dashed line
stands for the kernel lower bound given by the nonlinear part of the teacher
model.

#### Case 1.

Comparing with Figure 15, Figure 11 shows no alignments with training data in
GD training. This corresponds to the performances in Table 1. The performance
of Case 1 is not as good as the prediction risks in Figure 15, since Figure 11
suggests that no feature learning appears after GD training. Gradient descent
requires the weights to converge to some global minima close to
initialization, thereby offering no guarantees for lower generalization
errors. Next, Figure 12 further presents more results on GD training and
indicates more evidence of kernel regime in Case 1. This shows that, from
spectral point of view, the NTK is invariant/static through training. Based on
Figures 1(a) and 12, we can empirically verify Corollary B.3 stated in
Appendix B. Globally, the spectra of ${\boldsymbol{W}}$,
${\boldsymbol{K}}^{\text{CK}}$ and ${\boldsymbol{K}}^{\text{NTK}}$ are not
changing over training as $n/d\to\gamma_{1}$ and $N/d\to\gamma_{2}$. The
initial spectrum of weight ${\boldsymbol{W}}_{0}$ converges to Marčenko–Pastur
law; the initial spectrum of NTK under proportional limit has been studied by
[AP20, FW20]. Figure 12(c) demonstrates the global convergence for GD under
proportional regime, as proved in Theorem 5.2. We can observe this global
convergence even for SGD, Case 2 in Table 1, although we do not have a proof
for it.

Figure 11: Alignment with leading PCs of trained weight, CK and NTK matrices
in Case 1 of Table 1.

Figure 12: Performances of Case 1 in Table 1: (a) The initial and trained
spectra of first-hidden layer ${\boldsymbol{W}}$. (c) The initial and trained
spectra of empirical NTK matrix defined by (14). Q-Q subplot shows these two
spectra are almost the same. Training and test losses vs. epochs for GD
(right).

#### Case 2.

As a complement, Figure 13 exhibits the spectra of ${\boldsymbol{W}}$,
${\boldsymbol{K}}^{\text{CK}}$ and ${\boldsymbol{K}}^{\text{NTK}}$ for Case 2
in Table 1. The phenomena are similar to Case 1. This observation provides
evidence that all results and conjectures in Section 5.1 can be extended to
SGD training with sufficiently small learning rates, which is subject to
future work. Analogously to Theorem 5.2, we conjecture that the global
convergence when training both layers of NN with SGD still holds in this
proportional limit. The proof strategy for global convergence in this case can
again follow [OFLS19, OS19]. Once we have the invariant global spectra in
Corollary 5.3, we can apply the nonlinear RMT [PW17, LLC18, BP21, FW20] to
characterize the limiting spectra under LWR.

Figure 13: Spectral properties for Case 2 in Table 1: (a) The initial and
trained spectra of first-hidden layer ${\boldsymbol{W}}$. (b) The initial and
trained spectra of empirical NTK defined by (14). (c) The initial and trained
spectra of empirical CK defined by (4).

#### Case 3.

Next, in Figures 14 and 15, we present spectral properties for Case 3 in Table
1, where a spike detaches from the bulk after large-step-size training. Notice
that Figures 1(b) and 14(a) imply that the bulk spectra for weight and CK
remain unchanged over training despite the emergence of spikes. This not true
for NTK by observing Figures 14(b) and (c). The Frobenius norm of NTK changes
significantly during training and is not $O(1)$ anymore; the spectra of the
first component of the NTK shrinks after training (Figure 14(b)), which
indicates NNs converge to flatter minima. This resembles the catapult phase in
[LBD+20] for extremely large learning rates. Figure 15(a) shows the
convergence rate for SGD in Case 3. Empirically, we observe that the training
loss will not monotonically decrease when using a larger learning rate than
Case 3, which may be analogous to catapult phases in [LBD+20].

(a) Weight spectra.

(b) Spectra of first part in NTK.

(c) Norms of change vs. epochs.

Figure 14: Additional performance for Case 3 in Table 1: (a) The initial and
trained weight spectra. Notice that there is one outlier after training, while
the bulk remains invariant. This is analogous to the behavior of CK spectra in
Figure 1(b). (b) The spectra of the first part in (14) at initialization and
after training. Orange arrow points out the outlier of the spectrum. (c) The
changes $\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|$,
$\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$,
$\left\|{\boldsymbol{K}}^{\text{NTK}}_{t}-{\boldsymbol{K}}^{\text{NTK}}_{0}\right\|$
and
$\left\|{\boldsymbol{K}}^{\text{NTK}}_{t}-{\boldsymbol{K}}^{\text{NTK}}_{0}\right\|_{F}$
at each epoch $t$ throughout the training process.

(a) Losses vs. epochs.

(b) ${\boldsymbol{y}}$ vs. ${\boldsymbol{v}}_{1}$ of CK.

(b) ${\boldsymbol{y}}$ vs. ${\boldsymbol{v}}_{1}$ of NTK.

Figure 15: (a) Training/test losses and $R^{2}$ scores at each epoch of
training in Case 3. (b) Alignment between training labels ${\boldsymbol{y}}$
and first PC of trained/initial CK. (c) Alignment between training labels
${\boldsymbol{y}}$ and the first PC of trained/initial NTK. We use the same
setting in Case 3 of Table 1. There are strong alignments between kernels and
training labels as stated in Section 5.2.

#### Case 4.

Additional results for Case 4 in Table 1 are shown in Figure 16. Unlike the
strong alignments in Case 3 (Figure 15), Figures 16 (e)-(f) do not exhibit
strong alignments for the leading singular vectors or eigenvectors. This may
be due to the heavy tails present after training, with the other large spikes
detaching from the bulks are also important for generalization. A similar
phenomenon can be also seen in Figures 19 and 21 in Appendix A.4, where we
have comparable performances to Case 4 in Table 1.

(a) Weight spectra.

(b) Spectra of NTK.

(c) Losses vs. epochs.

(d) First PC in trained weight.

(e) First PC in trained CK.

(f) First PC in trained NTK.

Figure 16: Additional performance for Case 4 in Table 1: (a) The initial and
trained weight spectra. Notice that there are several outliers after training,
while the bulk has a heavier tail. (b) The spectra of the NTK (14) at
initialization and after training. (c) The test/training losses and $R^{2}$
score (green line) at each epoch $t$ throughout training process. (d)
Alignment between the leading PC of weight matrix and the signal
${\boldsymbol{\beta}}$ in the teacher model before (red) and after (blue)
training. (e) Alignment between the leading PC of the CK matrix and the
training labels ${\boldsymbol{y}}$ before/after training. (f) Alignment
between the leading PC of the NTK matrix and ${\boldsymbol{y}}$ before/after
training.

###  A.3 Training only the first hidden layer

We now present additional results when only training the first layer of NN
with Adam. This result resembles Figure 25 in the next section. In section
5.1, Corollary 5.3 shows that under LWR with sufficiently large width $h$, the
limiting spectra of
$\frac{1}{h}{\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}_{t}^{\text{CK}}$ and ${\boldsymbol{K}}_{t}^{\text{NTK}}$ are
essentially the same as those of the corresponding initial matrices if we
train only the first layer with GD (see (8)). Figure 17 further investigates
these phenomena when training only the first layer with Adam. Specially, the
Q-Q plots show the invariant spectra of weight, CK and NTK matrices even when
the training loss is approaching zero. Here, in Figure 17(c), we only consider
the first component of NTK since the gradient is only taken with respect to
${\boldsymbol{W}}_{t}$. Another observation is that the smallest eigenvalues
of the initial and trained NTK are both bounded away from zero. This is
crucial for the proof of the global convergence as shown in Appendix B.

Figure 17: Additional performance for Adam with small learning rate
$\eta=0.001$, where $n=2000,h=3000,d=1000$, $\sigma_{\varepsilon}=0.3$ and
batch size is 100. In this simulation, we only train the first hidden layer
${\boldsymbol{W}}_{t}$. Activation $\sigma$ is normalized softplus and target
is normalized $\tanh$. The final test loss is 0.36219 and $R^{2}$ score is
around $63.70\%$.

###  A.4  Adaptive Gradients

Inspired by Case 4 in Table 1, we show the spectral performances of adaptive
gradient (AdaGrad) in Figures 19 and 19. The performance of this method
matches Case 4 in Table 1, where we can also easily observe heavy tails and
detaching spikes after training, especially in Q-Q subplots. This suggests
that adaptive optimization is more likely to yield heavy-tailed distributions
in trained NNs. Besides, analogously to Figure 16, there is no strong
alignment in the single leading PC of the weight or kernel matrices after
training in Figure 19.

Figure 18: Additional performance for AdaGrad with learning rate $\eta=0.5$,
where $n=2000,h=1500,d=1000$, $\sigma_{\varepsilon}=0.3$ and small batch size
is 8. Activation $\sigma$ is normalized softplus and target is normalized
$\tanh$. The final test loss is 0.23555 and $R^{2}$ score is around $0.76249$.
The black lines in the Q-Q subplots are the line of $y=x$.

(a) First PC in trained weight.

(b) First PC in trained CK.

(C) First PC in trained NTK.

Figure 19: Alignment between the leading PC of the weight/kernel matrices and
the signal ${\boldsymbol{\beta}}$ or the training labels ${\boldsymbol{y}}$
before/after training for experiment in Figure 19. Analogously to Case 4 in
Appendix A.2, there is no strong alignment in the leading component of
weight/kernel matrices.

###  A.5 Different global minima and alignments

In this section, to distinguish the different alignments in Case 3$\&$4 of
Table 1, we introduce the following two simulations with slightly different
optimizers to get quite different spectra and alignments among leading PCs
after training. In the first experiment, Figures 20 and 21, we first take Adam
with large stepsizes for a few steps and then use small-stepsize SGD for
convergence. In this scenario, we can get heavy-tailed distributions after
training, and the phenomena are essentially the same as Case 4 of Table 1.
There is no strong alignment for the first leading eigenvector, while useful
features may be learned by a few top eigenvectors in the heavy-tailed spectra
after training. In the second experiment, we directly apply Adam with a small
initial learning rate for training. In contrast, the results of this case,
presented in Figures 22 and 23, are similar with Case 3 in Table 1. The test
loss and $R^{2}$ score are close to previous examples, though slightly worse.
As explained in the previous section, because there is only one spike
appearing outside the bulk after training in the second case, the leading PC
is highly aligned with the training dataset structure after training and the
feature learning mainly stems from the outliers in this situation. This
interprets the strongly anisotropic structures in trained spectra of NNs in
the second case [OJMMF20, OJMDF21, SB21]. These two different spectral
properties reveal significant differences between the global minima of these
two training processes, and different evolutions of the spectra in NNs.
Remarkably, based on the different spectral behaviors in trained weight and
kernel matrices, these two experiments exhibit disparate features learned by
distinct training procedures. Hence, analyzing the spectral properties in
trained kernel matrices is beneficial for clarifying what features our NNs
have learned during the training processes.

Figure 20: Additional performance for Adam with learning rate $\eta=0.09$ and
$4$ epochs, then SGD with learning rate $\eta=5\times 10^{-4}$ and $100$
epochs, where $n=2000,h=1500,d=1000$, $\sigma_{\varepsilon}=0.3$ and the batch
size is 32. We train the NN until the training loss is less than $10^{-10}$.
The activation $\sigma$ is normalized softplus and the target is normalized
$\tanh$. The final test loss is 0.22511 and $R^{2}$ score is around $0.77462$.

(a) First PC in trained weight.

(b) First PC in trained CK.

(c) First PC in trained NTK.

Figure 21: Alignment between the leading PC of the weight/kernel matrices and
the signal ${\boldsymbol{\beta}}$ or the training labels ${\boldsymbol{y}}$
before/after training for the experiment in Figure 20. This is analogous to
Case 4 in Appendix A.2.

Figure 22: Additional performance for Adam with learning rate $\eta=0.002$ and
$700$ epochs, where $n=2000,h=1500,d=1000$, $\sigma_{\varepsilon}=0.3$ and
batch size is 64. The activation $\sigma$ is normalized softplus and the
target is normalized $\tanh$. The final test loss is 0.23954 and $R^{2}$ score
is around $0.76027$. The orange arrows show the positions of the outliers.
Spectra behaviors in this case differ from Figure 20.

(a) First PC in trained weight.

(b) First PC in trained CK.

(c) First PC in trained NTK.

Figure 23: Alignment between the leading PC of the weight/kernel matrices and
the signal ${\boldsymbol{\beta}}$ or the training labels ${\boldsymbol{y}}$
before/after training for the experiment in Figure 22. We can observe strong
alignments in this case comparing with Figure 21 because of outliers in above
Figure 22. These kernel alignments induce anisotropic structures in the kernel
matrices during training [SB21].

##  Appendix B Proofs of theorems in Section 5.1

###  B.1 GD analysis at early phase

From (8), the GD process with learning rate $\eta>0$ can be written by

$\displaystyle{\boldsymbol{W}}_{t+1}=$
$\displaystyle{\boldsymbol{W}}_{t}+\eta\cdot{\boldsymbol{G}}_{t},\text{ where
}$ (17) $\displaystyle{\boldsymbol{G}}_{t}=$
$\displaystyle\frac{1}{n\sqrt{dh}}\left[\left({\boldsymbol{v}}\left({\boldsymbol{y}}-\frac{1}{\sqrt{h}}{\boldsymbol{v}}^{\top}\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right)\right)\odot\sigma^{\prime}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right]{\boldsymbol{X}}^{\top},$
(18)

for $t\in\mathbb{N}$, where ${\boldsymbol{y}}\in\mathbb{R}^{1\times n}$.
Following [BES+22, Appendix B], in this section we prove the control for
gradient step ${\boldsymbol{G}}_{t}$. For simplicity, denote
$f_{t}({\boldsymbol{X}}):=f_{{\boldsymbol{\theta}}_{t}}({\boldsymbol{X}})=\frac{1}{\sqrt{h}}{\boldsymbol{v}}^{\top}\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})$
for $t\in\mathbb{N}$.

######  Lemma B.1.

Under the same assumptions as in Lemma 5.1, we have

$\displaystyle\mathbb{P}\left({\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|\geq
C\sqrt{n}}\right)\leq~{}$ $\displaystyle 2e^{-cn},$
$\displaystyle\mathbb{P}\left({\left\|{\boldsymbol{y}}\right\|\geq
C\sqrt{n}}\right)\leq~{}$ $\displaystyle 2e^{-cn},$

for some constants $C,c>0$ only depending on $\sigma_{\varepsilon}$,
$\lambda_{\sigma}$, $\gamma_{1}$, and $\gamma_{2}$.

###### Proof.

Due to [FW20, Lemma D.4.], we can directly obtain that

$\mathbb{P}\left({\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|\geq
C^{\prime}(\sqrt{n}+\sqrt{h})\sqrt{\frac{h}{d}}}\right)\leq 2e^{-cn}.$

Here we use the fact that both ${\boldsymbol{W}}_{0}$ and ${\boldsymbol{X}}$
are i.i.d. Gaussian random matrices. Then by Assumption 4.1, we conclude that
we control $\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})$. Recall
that Assumption 4.3 implies that
${\boldsymbol{y}}=f^{*}({\boldsymbol{X}})+\boldsymbol{\varepsilon}$. Hence, by
Lipschitz Gaussian concentration inequality [Ver18, Theorem 5.2.2], each entry
of $f^{*}({\boldsymbol{X}})$ has independent sub-Gaussian coordinates, whence
we can get $\|f^{*}({\boldsymbol{X}})\|\leq C\sqrt{n}$ with probability at
least $1-2ne^{-cn}$ for some constants $c,C>0$. On the other hand,
$[\boldsymbol{\varepsilon}]_{i}=\varepsilon_{i}$ are i.i.d. centered sub-
Gaussian noises with variance $\sigma_{\varepsilon}^{2}$. By [Ver18, Theorem
3.1.1], we have

$\mathbb{P}\left({\|\boldsymbol{\varepsilon}\|\leq
2\sigma_{\varepsilon}\sqrt{n}}\right)\geq
1-2\operatorname{exp}\left({-\frac{cn}{K^{4}}}\right),$

where the constant $K$ is the sub-Gaussian norm defined by
$K=\max_{i}\|\varepsilon_{i}\|_{\psi_{2}}$. Hence, combining all things
together, we obtain the second inequality of this lemma.

∎

######  Lemma B.2.

Under the assumptions of Lemma 5.1, given any fixed $t\in\mathbb{N}$ and
learning rate $\eta=\Theta(1)$, the weight matrix after $t$ gradient steps
${\boldsymbol{W}}_{t}$ defined in (17) satisfies

$\displaystyle\mathbb{P}\left({\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}\geq\frac{C}{\sqrt{n}}}\right)$
$\displaystyle\leq\operatorname{exp}\left({-cn}\right),$ (19)

for some positive constants $c,C>0$ only depending on $t,\eta,$
$\sigma_{\varepsilon}$, $\lambda_{\sigma}$, $\gamma_{1}$ and $\gamma_{2}$.

###### Proof.

Denote $\sigma_{\perp}(x)=\sigma(x)-\mu_{1}x$ which is the nonlinear part of
$\sigma$ and $\mu_{1}=\mathbb{E}[z\sigma(z)]$. Thus,
$\mathbb{E}[\sigma_{\perp}(z)z]=0$ for $z\sim\mathcal{N}(0,1)$. Based on this,
we can further decompose the gradient ${\boldsymbol{G}}_{t}$ into

$\displaystyle{\boldsymbol{G}}_{t}=\underbrace{\frac{\mu_{1}}{n\sqrt{dh}}{\boldsymbol{v}}\left({\boldsymbol{y}}-f_{t}({\boldsymbol{X}})\right){\boldsymbol{X}}^{\top}}_{{\boldsymbol{A}}^{t}}+\underbrace{\frac{1}{n\sqrt{dh}}\left({\boldsymbol{v}}\left({\boldsymbol{y}}-f_{t}({\boldsymbol{X}})\right)\odot\sigma^{\prime}_{\perp}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right){\boldsymbol{X}}^{\top}}_{{\boldsymbol{B}}^{t}}.$
(20)

At first, consider $t=0$ and bound the spectral norm of
${\boldsymbol{W}}_{1}$. By assumption, we know
$\|{\boldsymbol{v}}\|\leq\sqrt{h}$. Due to Corollary 7.3.3 in [Ver18], we have

$\mathbb{P}\left({\frac{1}{\sqrt{d}}\|{\boldsymbol{X}}\|\geq
2\left(1+\sqrt{\frac{n}{d}}\right)}\right)\leq
2\operatorname{exp}\left({-cn}\right).$ (21)

Therefore, by (20), we can control ${\boldsymbol{A}}^{0}$ and
${\boldsymbol{B}}^{0}$ separately. Notice that, as a rank-one matrix,

$\displaystyle\left\|{\boldsymbol{A}}^{0}\right\|=\left\|{\boldsymbol{A}}^{0}\right\|_{F}$
$\displaystyle\leq\frac{\mu_{1}}{\sqrt{n}}\frac{\left\|{\boldsymbol{X}}\right\|}{\sqrt{d}}\frac{1}{\sqrt{n}}(\left\|{\boldsymbol{y}}\right\|+\left\|f_{0}({\boldsymbol{X}})\right\|)\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}$
$\displaystyle\leq\frac{\mu_{1}}{\sqrt{n}}\frac{\left\|{\boldsymbol{X}}\right\|}{\sqrt{d}}\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}\frac{1}{\sqrt{n}}\left(\left\|{\boldsymbol{y}}\right\|+\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|\right).$

Hence, by Lemma B.1 and (21), one can easily claim that
$\|{\boldsymbol{A}}^{0}\|\leq C/\sqrt{n}$ with probability at least
$1-e^{-cn}$ for some constants $c,C>0$. On the other hand, since
${\boldsymbol{v}}\left({\boldsymbol{y}}-f_{t}({\boldsymbol{X}})\right)$ is
rank-one and $\sigma_{\perp}^{\prime}=\sigma^{\prime}-\mu_{1}$ with
$|\sigma^{\prime}(x)|\leq\lambda_{\sigma}$, we can similarly obtain

$\displaystyle\left\|{\boldsymbol{B}}^{0}\right\|_{F}$
$\displaystyle\leq\frac{1}{n\sqrt{dh}}\left\|{\boldsymbol{v}}\left({\boldsymbol{y}}-f_{t}({\boldsymbol{X}})\right)\odot\sigma^{\prime}_{\perp}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right\|_{F}\|{\boldsymbol{X}}\|$
$\displaystyle\leq\frac{1}{n\sqrt{hd}}\left\|{\boldsymbol{X}}\right\|\left(\left\|{\boldsymbol{y}}\right\|+\left\|f_{0}({\boldsymbol{X}})\right\|\right)\left\|{\boldsymbol{v}}\right\|\max_{i,j}\left|\sigma^{\prime}_{\perp}({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right|_{i,j}$
$\displaystyle\leq\frac{\mu_{1}+\lambda_{\sigma}}{\sqrt{n}}\frac{\left\|{\boldsymbol{X}}\right\|}{\sqrt{d}}\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}\frac{1}{\sqrt{n}}\left(\left\|{\boldsymbol{y}}\right\|+\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|\right).$

As ${\boldsymbol{A}}^{0}$, we can apply Lemma B.1 and (21) to conclude (19)
for $t=1$.

For general $t$, we apply induction. We assume that after the $t$-th gradient
step with $\eta=\Theta(1)$, Eq. (19) holds for some constants $C,c>0$.
Following [BES+22, Lemma 16], we now show that the similar high-probability
statement also holds for ${\boldsymbol{W}}_{t+1}$ (for some different
constants $c^{\prime},C^{\prime}$). Firstly, following the same argument as
[OS20, Setion 6.6.1], we know that

$\displaystyle\left\|f_{t}({\boldsymbol{X}})\right\|\leq$
$\displaystyle\left\|f_{0}({\boldsymbol{X}})\right\|+\left\|f_{t}({\boldsymbol{X}})-f_{0}({\boldsymbol{X}})\right\|$
$\displaystyle\leq$
$\displaystyle\left\|f_{0}({\boldsymbol{X}})\right\|+\frac{\lambda_{\sigma}}{\sqrt{h}}\left\|{\boldsymbol{v}}\right\|\frac{\left\|{\boldsymbol{X}}\right\|}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}.$
(22)

Note that
$\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}=O(1/\sqrt{n})$
with high probability by the induction hypothesis. Hence, by Lemma B.1 and
(21), we have $\left\|f_{t}({\boldsymbol{X}})\right\|\leq C\sqrt{n}$ with high
probability. Indeed, the difference between $f_{t}({\boldsymbol{X}})$ and
$f_{0}({\boldsymbol{X}})$ is significantly negligible comparing with the
initial value $f_{0}({\boldsymbol{X}})$. Similarly with
${\boldsymbol{A}}_{0}$, ${\boldsymbol{A}}^{t}$ satisfies

$\displaystyle\left\|{\boldsymbol{A}}^{t}\right\|=\left\|{\boldsymbol{A}}^{t}\right\|_{F}$
$\displaystyle\leq\frac{\mu_{1}}{\sqrt{n}}\frac{\left\|{\boldsymbol{X}}\right\|}{\sqrt{d}}\frac{1}{\sqrt{n}}(\left\|{\boldsymbol{y}}\right\|+\left\|f_{t}({\boldsymbol{X}})\right\|)\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}.$

Analogously for ${\boldsymbol{B}}^{t}$, we have

$\displaystyle\left\|{\boldsymbol{B}}^{t}\right\|_{F}$
$\displaystyle\leq\frac{\mu_{1}+\lambda_{\sigma}}{\sqrt{n}}\frac{\left\|{\boldsymbol{X}}\right\|}{\sqrt{d}}\frac{\left\|{\boldsymbol{v}}\right\|}{\sqrt{h}}\frac{1}{\sqrt{n}}(\left\|{\boldsymbol{y}}\right\|+\left\|f_{t}({\boldsymbol{X}})\right\|).$

Thus, Lemma B.1, (21), and (22) ensure that

$\displaystyle\mathbb{P}\left({\left\|{\boldsymbol{A}}^{t}\right\|_{F}\geq\frac{C^{\prime}}{\sqrt{n}}}\right)\leq\operatorname{exp}\left({-c^{\prime}n}\right),~{}\mathbb{P}\left({\left\|{\boldsymbol{B}}^{t}\right\|_{F}\geq\frac{C^{\prime}}{\sqrt{n}}}\right)\leq\operatorname{exp}\left({-c^{\prime}n}\right),$

for constants $c^{\prime},C^{\prime}>0$. Since

$\left\|{\boldsymbol{W}}_{t+1}-{\boldsymbol{W}}_{0}\right\|_{F}\leq\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}+\eta\left\|{\boldsymbol{A}}^{t}\right\|_{F}+\eta\left\|{\boldsymbol{B}}^{t}\right\|_{F},$

by induction hypothesis, we can conclude that (19) holds for the $(t+1)$-th
step with some constants $C,c>0$, which are different from the constants at
the $t$-th step.

∎

As a corollary, by (16), we can also deduce the following norm bounds:

$\mathbb{P}\left({\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|\geq\frac{C}{\sqrt{n}}}\right)\leq\operatorname{exp}\left({-cn}\right),~{}~{}~{}\mathbb{P}\left({\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{2,\infty}\geq\frac{C}{\sqrt{n}}}\right)\leq\operatorname{exp}\left({-cn}\right).$

Lemma B.2 and the above bounds are empirically verified by Figure 24(a) for
$t=3$. Not only upper bounds, this simulation bu also shows that at early
phase $\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|$,
$\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$, and
$\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{2,\infty}$ are all
of the same $\Theta(1/\sqrt{n})$ order.

As a remark, from the bound of the second term of (22), we can deduce that the
change of the output of the NN satisfies

$|f_{t}({\boldsymbol{x}})-f_{0}({\boldsymbol{x}})|\leq\frac{C}{\sqrt{n}},$

for some $t$-dependent constant $C>0$, any
${\boldsymbol{x}}\sim\mathcal{N}(0,\mathbf{I})$ and any finite time $t$. In
other words, when $\eta=\Theta(1)$, the change of the output of the NN at the
early phase (i.e. $t=\Theta(1)$) is negligible and its order is
$O(\frac{1}{\sqrt{n}})$.

###  B.2 Proof of Lemma 5.1

In this section, we complete the proof of Lemma 5.1. We first mention the
empirical validation of Lemma 5.1 in Figure 24. Notice that the changes in
Frobenius norm for ${\boldsymbol{W}}$ and ${\boldsymbol{K}}^{\text{CK}}$ are
exactly $\Theta(1/\sqrt{n})$ and $\Theta(1/n)$, respectively. The operator
norm of ${\boldsymbol{K}}^{\text{NTK}}$ matches with Lemma 5.1, while the
Frobenius norm of the change decays slower than the rate $\Theta(1/n)$.
Additionally, in the simulation, we use
${\boldsymbol{v}}\sim\mathcal{N}(0,\mathbf{I})$, which indicates that our
assumption for ${\boldsymbol{v}}$ in Lemma 5.1 can be weakened.

(a) Early phase for ${\boldsymbol{W}}_{t}$.

(b) Early phase for ${\boldsymbol{K}}^{\text{CK}}$.

(c) Early phase for ${\boldsymbol{K}}^{\text{NTK}}$.

Figure 24: Empirical validations for Lemma 5.1 and Lemma B.2 at $t=3$. Here
$\sigma_{\varepsilon}=0.2$, activation $\sigma$ is a normalized ReLU and the
target function $\sigma^{*}$ is normalized $\tanh$. Fix $d/n=0.6$ and
$N/n=1.2$ as $n$ is increasing. At each dimension, we take 25 trials to
average. (a) Norms of the changes for
${\boldsymbol{W}}_{3}-{\boldsymbol{W}}_{0}$. (b) Norms of the changes for
${\boldsymbol{K}}^{\text{CK}}_{3}-{\boldsymbol{K}}^{\text{CK}}_{0}$. (c) Norms
of the changes for
${\boldsymbol{K}}^{\text{NTK}}_{3}-{\boldsymbol{K}}^{\text{NTK}}_{0}$.

###### Proof of Lemma 5.1.

Lemma B.2 directly validates the control of
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$.
By virtue of this result, we now present estimates for CK and NTK. Based on
[OS20, Section 6.6.1], we have

$\displaystyle\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right\|\leq\frac{\lambda_{\sigma}}{\sqrt{d}}\|{\boldsymbol{X}}\|\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\|_{F}.$
(23)

We apply the mean value theorem to obtain this inequality. Recall the operator
norm bound for Gaussian random matrix ${\boldsymbol{X}}$ in (21). We know
$\|{\boldsymbol{X}}/\sqrt{d}\|\lesssim 1+\sqrt{\gamma_{1}}$ with high
probability as $n/d\to\gamma_{1}$. Hence, with the help of Lemma B.2, we can
claim

$\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right\|\leq
C\lambda_{\sigma}(1+\sqrt{\gamma_{1}})/\sqrt{n},$

with probability at least $1-\operatorname{exp}\left({-cn}\right)$, for any
fixed finite $t\in[n]$. Similarly, we can control the change in Frobenius norm
as follows:

$\displaystyle\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})\right\|_{F}^{2}\leq~{}\frac{\lambda_{\sigma}^{2}}{d}\left\|{\boldsymbol{X}}\right\|^{2}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|_{F}^{2}\leq
C\lambda_{\sigma}^{2}(1+\sqrt{\gamma_{1}})^{2}/n,$ (24)

with probability at least $1-\operatorname{exp}\left({-cn}\right)$. Therefore,
we can control the change in CK matrix in Frobenius norm by the following
inequalities:

$\displaystyle\left\|{\boldsymbol{K}}^{\text{CK}}_{t}-{\boldsymbol{K}}^{CK}_{0}\right\|_{F}$
$\displaystyle\leq~{}$
$\displaystyle\frac{1}{h}\left\|\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})^{\top}\left(\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right)\right\|_{F}$
$\displaystyle~{}~{}~{}+\frac{1}{h}\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})^{\top}\left(\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right)\right\|_{F}$
$\displaystyle\leq~{}$
$\displaystyle\frac{1}{h}\left(\left\|\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|+\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d}\right\|\right)\cdot\left\|\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|_{F}$
$\displaystyle+\frac{1}{h}\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|\cdot\left\|\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|_{F}.$

Therefore, by (23), (24) and Lemma B.1, we can claim that there exist
constants $c,C>0$ such that with probability at least
$1-\operatorname{exp}\left({-cn}\right)$,
$\left\|{\boldsymbol{K}}^{\text{CK}}_{t}-{\boldsymbol{K}}^{CK}_{0}\right\|_{F}$
is upper bounded by $C/n$ in the LWR.

Now we consider the change in the NTK matrix during training. Since the
empirical NTK can be decomposed into two parts, one of which is exact the CK,
it suffices to consider the change of the first part of the empirical NTK.
Recall that

${\boldsymbol{K}}_{t}:=\frac{1}{d}{\boldsymbol{X}}^{\top}{\boldsymbol{X}}\odot\frac{1}{h}\sigma^{\prime}\left(\frac{1}{\sqrt{d}}{\boldsymbol{W}}_{t}{\boldsymbol{X}}\right)^{\top}\text{diag}({\boldsymbol{v}}_{t})^{2}\sigma^{\prime}\left(\frac{1}{\sqrt{d}}{\boldsymbol{W}}_{t}{\boldsymbol{X}}\right).$

Following the notation in [OS20], we denote
$\mathcal{J}({\boldsymbol{W}}_{t}):=[\mathcal{J}({\boldsymbol{w}}_{1}^{t}),\ldots,\mathcal{J}({\boldsymbol{w}}_{N}^{t})]\in\mathbb{R}^{n\times
hd}$ with
$\mathcal{J}({\boldsymbol{w}}_{i}):=\frac{v_{i}}{\sqrt{h}}\text{diag}(\sigma^{\prime}({\boldsymbol{X}}^{\top}{\boldsymbol{w}}_{i}/\sqrt{d})){\boldsymbol{X}}^{\top}/\sqrt{d}\in\mathbb{R}^{n\times
d}$. Hence,
${\boldsymbol{K}}_{t}=\mathcal{J}({\boldsymbol{W}}_{t})\mathcal{J}({\boldsymbol{W}}_{t})^{\top}$
and

$\displaystyle\left\|{\boldsymbol{K}}_{t}-{\boldsymbol{K}}_{0}\right\|=~{}$
$\displaystyle\left\|\mathcal{J}({\boldsymbol{W}}_{t})\mathcal{J}({\boldsymbol{W}}_{t})^{\top}-\mathcal{J}({\boldsymbol{W}}_{0})\mathcal{J}({\boldsymbol{W}}_{0})^{\top}\right\|$
$\displaystyle\leq~{}$ $\displaystyle
2\left\|\mathcal{J}({\boldsymbol{W}}_{0})\right\|\left\|\mathcal{J}({\boldsymbol{W}}_{t})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|+\left\|\mathcal{J}({\boldsymbol{W}}_{t})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|^{2}.$
(25)

By [OS20, Lemma 6.6], we know
$\left\|\mathcal{J}({\boldsymbol{W}}_{0})\right\|^{2}=\left\|{\boldsymbol{K}}^{\text{NTK}}_{0}\right\|$
is upper bounded by some constant $C>0$ with high probability. Then, we apply
the inequalities from Lemma 6.5 of [OS20] to obtain

$\displaystyle\left\|\mathcal{J}({\boldsymbol{W}}_{t})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|^{2}$
$\displaystyle=~{}$
$\displaystyle\left\|\left(\left(\sigma^{\prime}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma^{\prime}({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right)^{\top}\frac{\text{diag}({\boldsymbol{v}})^{2}}{h}\left(\sigma^{\prime}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma^{\prime}({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right)\right)\odot\left(\frac{{\boldsymbol{X}}^{\top}{\boldsymbol{X}}}{d}\right)\right\|$
$\displaystyle\leq~{}$
$\displaystyle\left\|\left(\sigma^{\prime}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma^{\prime}({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right)^{\top}\frac{\text{diag}({\boldsymbol{v}})}{\sqrt{h}}\right\|^{2}\left(\max_{i\in[n]}\left\|{\boldsymbol{x}}_{i}/\sqrt{d}\right\|^{2}\right)$
$\displaystyle\leq~{}$
$\displaystyle\frac{1}{h}\|{\boldsymbol{v}}\|^{2}_{\infty}\left\|\sigma^{\prime}({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma^{\prime}({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|^{2}\left(\max_{i\in[n]}\left\|{\boldsymbol{x}}_{i}/\sqrt{d}\right\|^{2}\right)$
$\displaystyle\leq~{}$
$\displaystyle\frac{\lambda_{\sigma}^{2}}{h}\frac{\|{\boldsymbol{X}}\|^{2}}{d}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|^{2}_{F}\left(\max_{i\in[n]}\left\|{\boldsymbol{x}}_{i}/\sqrt{d}\right\|^{2}\right),$
(26)

where the last inequality is due to the mean value theorem, the uniform bound
on $\sigma^{\prime\prime}$, and the assumption on the second layer
${\boldsymbol{v}}$. Notice that Gaussian random vectors satisfy

$\displaystyle\mathbb{P}\left({\max_{i\in[n]}\frac{1}{d}\|{\boldsymbol{x}}_{i}\|^{2}\geq
2}\right)\leq$ $\displaystyle 2ne^{-cn},$ (27)

as $n/d\to\gamma_{1}$ and $h/d\to\gamma_{2}$. Thus, with (21) and Lemma B.2,
we obtain

$\mathbb{P}\left({\left\|\mathcal{J}({\boldsymbol{W}}_{t})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|\geq\frac{C\lambda_{\sigma}(1+\gamma_{1})}{n}}\right)\leq
4ne^{-cn},$

where constant $C$ relies on the number of steps $t$. Hence, by (25), we can
finally bound in norm the difference between the initial and the trained NTK
matrices at early phase ($t$ is finite).

∎

######  Corollary B.3.

For any fixed $t\in\mathbb{N}$, $i\in[d]$ and $k\in[n]$, denote
$\lambda^{t}_{i}$, $\nu_{k}^{t}$ and $\mu_{k}^{t}$ the $i$-th, and $k$-th
eigenvalues of $\frac{1}{h}{\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}_{t}^{\text{CK}}$ and ${\boldsymbol{K}}_{t}^{\text{NTK}}$,
respectively. Then, under the assumptions of Lemma 5.1, we have

$|\lambda^{t}_{i}-\lambda^{0}_{i}|,~{}|\nu_{k}^{t}-\nu_{k}^{0}|,~{}|\mu_{k}^{t}-\mu_{k}^{0}|\to
0,$

almost surely in LWR. Consequently, the eigenvalues of
$\frac{1}{h}{\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}_{t}^{\text{CK}}$ and ${\boldsymbol{K}}_{t}^{\text{NTK}}$ are
the same as corresponding the eigenvalues of initial
$\frac{1}{h}{\boldsymbol{W}}_{0}^{\top}{\boldsymbol{W}}_{0}$,
${\boldsymbol{K}}_{0}^{\text{CK}}$ and ${\boldsymbol{K}}_{0}^{\text{NTK}}$,
respectively.

This corollary is a direct outcome of Weyl’s inequality [BS10, Theorem A.46].
Consequently, this corollary concludes that for any fixed $t\geq 0$, almost
surely, the limiting spectra of
$\frac{1}{h}{\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}_{t}^{\text{CK}}$ and ${\boldsymbol{K}}_{t}^{\text{NTK}}$ are
the same as those of
$\frac{1}{h}{\boldsymbol{W}}_{0}^{\top}{\boldsymbol{W}}_{0}$,
${\boldsymbol{K}}_{0}^{\text{CK}}$ and ${\boldsymbol{K}}_{0}^{\text{NTK}}$ in
LWR. This corollary claims that not only the bulk of distributions stay
identical to the initialization, but also that any eigenvalues stay the same
as at the initialization. This shows that the smallest eigenvalue of
${\boldsymbol{K}}_{t}^{\text{NTK}}$ has the same lower bound as
${\boldsymbol{K}}_{0}^{\text{NTK}}$ in the early phase of training.

###  B.3 Global convergence for GD

In this section, we study the final stage of (8) as training loss is
approaching zero and prove Theorem 5.2. Figure 25 shows that the spectra are
unchanged globally, even after training in this case. In Corollary 5.3, we
confirm this observation for the weight, CK and NTK matrices via Frobenius
norm control. In the simulation, the second layer is initialized as
${\boldsymbol{v}}\sim\mathcal{N}(0,\mathbf{I})$, which is more general than
our assumption on ${\boldsymbol{v}}$ in Theorem 5.2.

Figure 25: The initial and trained spectra (until training loss is less than
$10^{-5}$) when using GD only for the first layer ($h=3000,n=2000,d=1000$):
(a) Weight spectra. (b) ${\boldsymbol{K}}^{\text{CK}}$ spectra. (c)
${\boldsymbol{K}}^{\text{NTK}}$ spectra. Final $R^{2}$ score is 0.55964 and
test loss is 0.44724. The activation is a normalized ReLU, and the target is
Sigmoid.

###### Proof of Theorem 5.2.

Recall the Jacobian matrix $\mathcal{J}({\boldsymbol{W}})$ defined in the
proof of Lemma 5.1, and the definition of $\alpha$ based on (9) in Section 4.
Denote the event

${\mathcal{A}}:=\left\\{\|{\boldsymbol{X}}\|\leq
2\left(1+\sqrt{\gamma_{1}}\right)\sqrt{d},~{}\max_{i\in[n]}\|{\boldsymbol{x}}_{i}\|^{2}\leq
2d,\sigma_{\min}(\mathcal{J}({\boldsymbol{W}}_{0}))\geq 2\alpha\right\\}.$

By (21), (27) and Theorem 2.9 of [WZ21], we have
$\mathbb{P}\left({{\mathcal{A}}}\right)\geq 1-2e^{-cn}-2ne^{-cn}-n^{-7/3}$ for
some constant $c>0$ and all large $n$ in LWR. In the following, conditionally
on event ${\mathcal{A}}$, we will apply Theorem 6.10 of [OS20] to obtain the
global convergence. Conditionally on ${\mathcal{A}}$, Lemma 6.6 of [OS20]
implies

$\displaystyle\left\|\mathcal{J}({\boldsymbol{W}})\right\|\leq\lambda_{\sigma}\left\|{\boldsymbol{v}}\right\|_{\infty}\left\|{\boldsymbol{X}}/\sqrt{d}\right\|\leq
2\lambda_{\sigma}(1+\sqrt{\gamma_{1}}),$ (28)

for any ${\boldsymbol{W}}$. Define
$\beta=2\lambda_{\sigma}(1+\sqrt{\gamma_{1}})$. Moreover, in terms of (26), we
can verify the Lipschitz property for the Jacobian matrix as follows:
conditionally on ${\mathcal{A}}$,

$\displaystyle\left\|\mathcal{J}(\tilde{{\boldsymbol{W}}})-\mathcal{J}({\boldsymbol{W}})\right\|\leq\frac{2\beta}{\sqrt{h}}\left\|\tilde{{\boldsymbol{W}}}-{\boldsymbol{W}}\right\|_{F},$
(29)

for any $\tilde{{\boldsymbol{W}}},{\boldsymbol{W}}\in\mathbb{R}^{h\times d}$.
Therefore, conditionally on ${\mathcal{A}}$, $\mathcal{J}({\boldsymbol{W}})$
is a $L$-Lipschitz function with respect to ${\boldsymbol{W}}$ where
$L:=\frac{2\beta}{\sqrt{h}}$. To complete the proof, it suffices to
investigate the smallest singular value of $\mathcal{J}({\boldsymbol{W}})$
when ${\boldsymbol{W}}$ is in vicinity of ${\boldsymbol{W}}_{0}$. Recall
$\ell({\boldsymbol{W}})=\left\|{\boldsymbol{y}}-f_{{\boldsymbol{W}}}({\boldsymbol{X}})\right\|$.
Notice that for any unit vector ${\boldsymbol{u}}\in\mathbb{R}^{n}$, we have
${\boldsymbol{u}}^{\top}f_{{\boldsymbol{W}}_{0}}({\boldsymbol{X}})=\frac{1}{\sqrt{h}}\sum_{i=1}^{h}v_{i}\sigma({\boldsymbol{w}}_{i}^{\top}{\boldsymbol{X}}/\sqrt{d}){\boldsymbol{u}}$,
where ${\boldsymbol{w}}_{i}^{\top}$ is the $i$-th row of
${\boldsymbol{W}}_{0}$ for $i\in[N]$. Consider event
$\mathcal{B}:=\left\\{\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|\leq
C\sqrt{n}\right\\}$ for some universal constant $C>0$. Lemma B.1 proves
$\mathbb{P}\left({\mathcal{B}}\right)\geq 1-2e^{-cn}$. By the assumption of
${\boldsymbol{v}}$, we know each entry $v_{i}$ is a sub-Gaussian random
variable with sub-Gaussian norm at most 1. Then, according to Hoeffding’s
inequality, conditionally on the event $\mathcal{B}$, we have

$\mathbb{P}\left({\left|\frac{1}{\sqrt{h}}\sum_{i=1}^{h}v_{i}\sigma({\boldsymbol{w}}_{i}^{\top}{\boldsymbol{X}}/\sqrt{d}){\boldsymbol{u}}\right|\geq
t}\right)\leq 2\operatorname{exp}\left({-ct^{2}}\right),$

for every $t\geq 0$ and some constant $c>0$. Let $t=2\sqrt{n}$. Considering an
$\frac{1}{4}$-net $\mathcal{N}$ of the unit sphere $\mathbb{S}^{n-1}$, we can
get

$\mathbb{P}\left({\left\|f_{{\boldsymbol{W}}_{0}}({\boldsymbol{X}})\right\|\geq\sqrt{n}}\right)\leq\mathbb{P}\left({2\max_{{\boldsymbol{u}}\in\mathcal{N}}\left|{\boldsymbol{u}}^{\top}f_{{\boldsymbol{W}}_{0}}({\boldsymbol{X}})\right|\geq\sqrt{n}}\right)\leq
9^{n}2\operatorname{exp}\left({-cn}\right)\leq 2e^{-c^{\prime}n},$ (30)

for some constant $c^{\prime}>0$. Hence, based on Lemma B.1 and (30), we can
obtain $\ell({\boldsymbol{W}}_{0})\leq C_{0}\sqrt{n}$ with high probability
for some universal constant $C_{0}>0$. Let us denote this event as
${\mathcal{C}}:\left\\{\ell({\boldsymbol{W}}_{0})\leq C_{0}\sqrt{n}\right\\}$.
Define $R:=4\ell({\boldsymbol{W}}_{0})/\alpha$. For any ${\boldsymbol{W}}$ in
a ball of radius $R$ centered at ${\boldsymbol{W}}_{0}$, we have
$\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}\right\|_{F}\leq R$ and
$\left\|\mathcal{J}({\boldsymbol{W}})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|\leq
LR,$ conditionally on event ${\mathcal{A}}$. Thus, by (29), on event
${\mathcal{A}}\cap{\mathcal{C}}$, the smallest singular value
$\sigma_{\min}(\mathcal{J}({\boldsymbol{W}}))$ of the Jacobian matrix
$\mathcal{J}({\boldsymbol{W}})$ can be bounded by

$\displaystyle\sigma_{\min}(\mathcal{J}({\boldsymbol{W}}))\geq~{}$
$\displaystyle\sigma_{\min}(\mathcal{J}({\boldsymbol{W}}_{0}))-\left\|\mathcal{J}({\boldsymbol{W}})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|$
$\displaystyle\geq~{}$ $\displaystyle 2\alpha-LR\geq
2\alpha-\frac{8\beta}{\alpha}\frac{\ell({\boldsymbol{W}}_{0})}{\sqrt{h}}\geq
2\alpha-\frac{8C\beta}{\alpha}\sqrt{\frac{\gamma_{1}}{\gamma_{2}}},$

for some universal constant $C>0$ and sufficiently large $n,d,h$. Notice that
here constants $C,\beta$, and $\alpha$ do not rely on $\gamma_{2}$. Therefore,
there exists a sufficiently large $\gamma^{*}>0$ such that for all
$\gamma_{2}\geq\gamma^{*}$, we have
$2\alpha-\frac{8C\beta}{\alpha}\sqrt{\frac{\gamma_{1}}{\gamma_{2}}}\geq\alpha$.
In other words, when $h$ is sufficiently large but still in the same order as
$n$ and $d$, for all
$\left\|{\boldsymbol{W}}-{\boldsymbol{W}}_{0}\right\|_{F}\leq R$, we have
$\sigma_{\min}(\mathcal{J}({\boldsymbol{W}}))\geq\alpha$ conditionally on
${\mathcal{C}}\cap{\mathcal{A}}$. Combining with (28) and (29), conditionally
on ${\mathcal{C}}\cap{\mathcal{A}}$, all the assumptions of Theorem 6.10 by
[OS20] are satisfied when
$\left\|{\boldsymbol{W}}-{\boldsymbol{W}}_{0}\right\|_{F}\leq R$. Therefore,
when the learning rate
$\frac{\eta}{n}\leq\frac{1}{\beta^{2}}\min\left\\{1,\frac{4\alpha}{LR}\right\\}$,
we can get (10)-(12) for all $t\in\mathbb{N}$, conditionally on
${\mathcal{C}}\cap{\mathcal{A}}$. Both events ${\mathcal{A}}$ and
${\mathcal{C}}$ occur with high probability and only depend on initialization
${\boldsymbol{W}}_{0}$, ${\boldsymbol{X}}$ and ${\boldsymbol{y}}$. Hence we
complete the proof of this theorem. Notice that since
$\gamma_{2}\geq\gamma^{*}$ is sufficiently large,
$\frac{4\alpha}{LR}\geq\frac{\alpha^{2}}{2C\beta}\sqrt{\frac{\gamma_{2}}{\gamma_{1}}}>1$.
Therefore, it suffices to require $\eta\leq n/\beta^{2}$ to conclude that
(10), (11) and (12) hold with high probability. This completes the proof.
Moreover, (12) further shows that for all $t\in\mathbb{N}$,

$\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|_{F}\leq R\leq
C\sqrt{n}+o_{d,\mathbb{P}}(1),$ (31)

where we again apply Lemma B.1 in the following way:

$\displaystyle\ell({\boldsymbol{W}}_{0})\leq C\sqrt{n}+o_{d,\mathbb{P}}(1),$

for some constant $C>0$ only depending on
$\gamma_{1},\gamma_{2},\sigma_{\varepsilon},\sigma$ and $\sigma^{*}$. ∎

(a) Changes for $\frac{1}{\sqrt{d}}{\boldsymbol{W}}$.

(b) Changes for CK.

(c) Changes for NTK.

Figure 26: Measuring the change for the weight, CK and NTK matrices when
training NN with (8). We fix $d/n=1.2$ and $h/n=0.6$ when $n$ is increasing.
Here, $\sigma$ is normalized ReLU and the target is normalized $\tanh$. The
largest $n=6400$ and the learning rate $\eta=5.0$ for all training processes.
We train each neural network until the training losses approaches zero. Each
experiment repeats 4 times. In (a), we consider the changes
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|$,$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$,
and
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{2,\infty}$.

As corollary, Eq. (12) controls the deviation of the final-step weight from
the initial weight. This is empirically shown in Figure 26(a), which shows
that
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|$,$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$,
and
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{2,\infty}$
are $\Theta(1)$ when trainable parameters are convergent. This implies that
the final ${\boldsymbol{W}}_{t}$ is still close to the initial weight
${\boldsymbol{W}}_{0}$, even after training. Consequently, with this
observation, we can prove Corollary 5.3 in the following.

###### Proof of Corollary 5.3.

Based on (31), we know
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{o}-{\boldsymbol{W}}_{t}\right\|_{F}\leq
C_{0}$ holds with high probability for some universal constant $C_{0}>0$.
Conditionally on this event, we can then estimate changes in CK and NTK after
training. The method is analogous to Lemma 5.1. For CK, we employ Lemma B.1
and (24) to get

$\displaystyle\left\|{\boldsymbol{K}}^{\text{CK}}_{t}-{\boldsymbol{K}}^{CK}_{0}\right\|_{F}$
$\displaystyle\leq~{}$
$\displaystyle\frac{2}{h}\left(\left\|\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|+\left\|\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d}\right\|\right)\cdot\left\|\sigma({\boldsymbol{W}}_{t}{\boldsymbol{X}}/\sqrt{d})-\sigma({\boldsymbol{W}}_{0}{\boldsymbol{X}}/\sqrt{d})\right\|_{F}$
$\displaystyle\lesssim~{}$
$\displaystyle\frac{2\lambda_{\sigma}^{2}(1+\sqrt{\gamma_{1}})^{2}}{h}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|_{F}^{2}+\frac{2C\sqrt{n}\lambda_{\sigma}(1+\sqrt{\gamma_{1}})}{h}\left\|{\boldsymbol{W}}_{0}-{\boldsymbol{W}}_{t}\right\|_{F}$
$\displaystyle\lesssim~{}$
$\displaystyle\frac{2\lambda_{\sigma}(1+\sqrt{\gamma_{1}})C_{0}}{\gamma_{2}^{2}}\left(\lambda_{\sigma}(1+\sqrt{\gamma_{1}})C_{0}+C\sqrt{\gamma_{1}}\right)=O_{d,\mathbb{P}}(1).$

Hence, this shows control of the change for CK matrix after training,
comparing with the initial CK.

Let us denote ${\boldsymbol{w}}_{i}^{t}\in\mathbb{R}^{1\times d}$ as the
$i$-th row of ${\boldsymbol{W}}_{t}$, and ${\boldsymbol{x}}_{j}$ as the $j$-th
column of ${\boldsymbol{X}}$. Additionally, by Assumption 4.2, we know that

$|\sigma^{\prime}(x)-\sigma^{\prime}(y)|\leq\lambda_{\sigma}|x-y|$ (32)

for any $x,y\in\mathbb{R}$. For NTK, by modifying (26), one can deduce that

$\displaystyle\left\|\mathcal{J}({\boldsymbol{W}}_{t})-\mathcal{J}({\boldsymbol{W}}_{0})\right\|_{F}^{2}=\sum_{i=1}^{h}\left\|\mathcal{J}({\boldsymbol{w}}_{i}^{t})-\mathcal{J}({\boldsymbol{w}}_{i}^{0})\right\|^{2}_{F}$
$\displaystyle\overset{(i)}{\leq}$
$\displaystyle\frac{1}{h}\sum_{i=1}^{h}\left\|\text{diag}\left(\sigma^{\prime}({\boldsymbol{w}}_{i}^{t}{\boldsymbol{X}}/\sqrt{d}-\sigma^{\prime}({\boldsymbol{w}}_{i}^{0}{\boldsymbol{X}}/\sqrt{d})\right)\right\|^{2}_{F}\left\|\frac{{\boldsymbol{X}}}{\sqrt{d}}\right\|^{2}$
$\displaystyle\overset{(ii)}{\leq}$
$\displaystyle\frac{(1+\sqrt{\gamma_{1}})^{2}}{h}\sum_{i=1}^{h}\left\|\text{diag}\left(\sigma^{\prime}({\boldsymbol{w}}_{i}^{t}{\boldsymbol{X}}/\sqrt{d}-\sigma^{\prime}({\boldsymbol{w}}_{i}^{0}{\boldsymbol{X}}/\sqrt{d})\right)\right\|^{2}_{F}+o_{d,\mathbb{P}}(1)$
$\displaystyle\overset{(iii)}{\leq}$
$\displaystyle\frac{\lambda^{2}_{\sigma}(1+\sqrt{\gamma_{1}})^{2}}{h}\sum_{i=1}^{h}\sum_{j=1}^{n}\left(\frac{1}{\sqrt{d}}({\boldsymbol{w}}_{i}^{t}-{\boldsymbol{w}}_{i}^{0}){\boldsymbol{x}}_{j}\right)+o_{d,\mathbb{P}}(1)$
$\displaystyle\overset{(iv)}{\leq}$
$\displaystyle\frac{\lambda^{2}_{\sigma}(1+\sqrt{\gamma_{1}})^{4}}{h}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}^{2}+o_{d,\mathbb{P}}(1)\leq\frac{\lambda^{2}_{\sigma}(1+\sqrt{\gamma_{1}})^{4}C_{0}^{2}}{\gamma_{2}}+o_{d,\mathbb{P}}(1),$

where $(i)$ is because of [Ver18, Exercise 6.3.3] and the assumption on
${\boldsymbol{v}}$, $(ii)$ is due to (21), $(iii)$ is due to the definition of
Frobenius norm and (32), and $(iv)$ is due to [Ver18, Exercise 6.3.3] and
(21). As a result, from (25), we can finally conclude that
$\left\|{\boldsymbol{K}}^{\text{CK}}_{t}-{\boldsymbol{K}}^{CK}_{0}\right\|_{F}=O_{d,\mathbb{P}}(1)$
as $n/d\to\gamma_{1}$ and $h/d\to\gamma_{2}$.

As for the limiting spectra, since we know

$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F},~{}\left\|{\boldsymbol{K}}_{t}^{\text{CK}}-{\boldsymbol{K}}_{0}^{\text{CK}}\right\|_{F},~{}\left\|{\boldsymbol{K}}_{t}^{\text{NTK}}-{\boldsymbol{K}}_{0}^{\text{NTK}}\right\|_{F}=O_{d,\mathbb{P}}(1),$

we can automatically apply [BS10, Corollary A.41]. This directly implies that
the limiting empirical spectra of
$\frac{1}{h}{\boldsymbol{W}}_{t}^{\top}{\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}_{t}^{\text{CK}}$ and ${\boldsymbol{K}}_{t}^{\text{NTK}}$ are
the same as the limiting spectra of
$\frac{1}{h}{\boldsymbol{W}}_{0}^{\top}{\boldsymbol{W}}_{0}$,
${\boldsymbol{K}}_{0}^{\text{CK}}$ and ${\boldsymbol{K}}_{0}^{\text{NTK}}$,
respectively, as $n/d\to\gamma_{1}$ and $h/d\to\gamma_{2}$ (see Figure 25).

∎

#### Further simulations for the norm of changes.

The simulation of Figure 26 empirically coincides with the norm bounds in
Theorem 5.2 for different norms. Because of (16), it suffices to only consider
the Frobenius norm of the change for each matrix. As a remark, Theorem 5.2
requires $\gamma_{2}$ to be larger than some threshold $\gamma^{*}$ to ensure
norms of the change throughout training. However, Figure 26 indicates Theorem
5.2 still holds even when $\gamma^{*}$ is small i.e. the width $h$ is not very
large. Here in Figure 26, $\gamma_{2}=\frac{1}{2}\gamma_{1}=0.6$. Figure 27
suggests that similar results to Theorem 5.2 and Corollary 5.3 hold for SGD
when training the first layer. This is also akin to Figure 17. Moreover, we
further conjecture that similar results to Theorem 5.2 and Corollary 5.3 will
hold even when training both layers (3). Denote the second layer at step $t$
by ${\boldsymbol{v}}_{t}$. Then, indicated by Figure 28,
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|$,$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{F}$,
$\frac{1}{\sqrt{d}}\left\|{\boldsymbol{W}}_{t}-{\boldsymbol{W}}_{0}\right\|_{2,\infty}$,
$\left\|{\boldsymbol{K}}^{\text{CK}}_{t}-{\boldsymbol{K}}^{\text{CK}}_{0}\right\|$,
$\left\|{\boldsymbol{K}}^{\text{CK}}_{t}-{\boldsymbol{K}}^{\text{CK}}_{0}\right\|_{F}$,
$\left\|{\boldsymbol{K}}^{\text{NTK}}_{t}-{\boldsymbol{K}}^{\text{NTK}}_{0}\right\|$,
and
$\left\|{\boldsymbol{v}}_{t}-{\boldsymbol{v}}_{0}\right\|/\left\|{\boldsymbol{v}}_{0}\right\|$
are all $\Theta(1)$ in LWR. Since
$\left\|{\boldsymbol{v}}_{0}\right\|=O(\sqrt{h})$, we observe that
$\left\|{\boldsymbol{v}}_{t}-{\boldsymbol{v}}_{0}\right\|=\Theta(\sqrt{h})$ in
this case. Meanwhile, unlike Corollary 5.3, Figure 28(c) cannot verify that
$\left\|{\boldsymbol{K}}^{\text{NTK}}_{t}-{\boldsymbol{K}}^{\text{NTK}}_{0}\right\|_{F}$
is still upper bounded by a constant. One possible explanation is that in this
case the change in the second layer ${\boldsymbol{v}}_{t}$ is much more
significant than in the first-layer weight ${\boldsymbol{W}}_{t}$, hence the
NTK matrix may change a lot in this general case.

Similarly, Figure 29 shows norms of the change when training NN with SGD for
both layers. Here we use the same batch size 128 and learning rate $\eta=1.0$
for all experiments when varying the dimension $n$ but fixing the aspect
ratios $\gamma_{1}$ and $\gamma_{2}$. All the observations are similar with
the results of Corollary 5.3 except the Frobenius norm of the change for NTK.
Figure 29(c) indicates that
$\left\|{\boldsymbol{K}}^{\text{NTK}}_{t}-{\boldsymbol{K}}^{\text{NTK}}_{0}\right\|_{F}$
will not be $\Theta(1)$ anymore, which fluctuates more in this case.

(a) Changes for $\frac{1}{\sqrt{d}}{\boldsymbol{W}}$.

(b) Changes for CK.

(c) Changes for NTK.

Figure 27: Measuring the change for weight ${\boldsymbol{W}}_{t}$,
${\boldsymbol{K}}^{\text{CK}}$, and ${\boldsymbol{K}}^{\text{NTK}}$ matrices
when training NN with (8) for the first layer ${\boldsymbol{W}}_{t}$ using SGD
with batch size 200. We fix $d/n=0.6$ and $h/n=1.2$ when $n$ is increasing.
Here $\sigma$ is normalized softplus and the target is normalized $\tanh$. The
largest $n$ is $11000$, $\sigma_{\varepsilon}=0.3$ and the learning rate is
$\eta=3.6$ for all training processes. We train each neural network until the
training losses approach zero. Each experiment repeats 15 times.

(a) Changes for $\frac{1}{\sqrt{d}}{\boldsymbol{W}}$.

(b) Changes for CK.

(c) Changes for NTK.

Figure 28: Measuring the change for ${\boldsymbol{W}}_{t}$,