Title: Spectral Evolution and Invariance in Linear-width Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Random matrix theory

Heavy tails

Feature learning

Linear-width neural networks

Spike phase transition
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 spectra of weight in this high dimensional regime are invariant when trained by gradient descent for small constant learning rates; we provide a theoretical justification for this observation and prove the invariance of the bulk spectra for both conjugate and neural tangent kernels. We demonstrate similar characteristics when training with stochastic gradient descent with small learning rates. When the learning rate is large, we exhibit the emergence of an outlier whose corresponding eigenvector is aligned with the training data structure. We also show that after adaptive gradient training, where a lower test error and feature learning emerge, both weight and kernel matrices exhibit heavy tail behavior. Simple examples are provided to explain when heavy tails can have better generalizations. We exhibit different spectral properties such as invariant bulk, spike, and heavy-tailed distribution from a two-layer neural network using different training strategies, and then correlate them to the feature learning. Analogous phenomena also appear when we train conventional neural networks with real-world data. We conclude that monitoring the evolution of the spectra during training is an essential step toward understanding the training dynamics and feature learning.

Spectral Evolution and Invariance in Linear-width Neural Networks

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 spectra of weight in this high dimensional regime are invariant when trained by gradient descent for small constant learning rates; we provide a theoretical justification for this observation and prove the invariance of the bulk spectra for both conjugate and neural tangent kernels. We demonstrate similar characteristics when training with stochastic gradient descent with small learning rates. When the learning rate is large, we exhibit the emergence of an outlier whose corresponding eigenvector is aligned with the training data structure. We also show that after adaptive gradient training, where a lower test error and feature learning emerge, both weight and kernel matrices exhibit heavy tail behavior. Simple examples are provided to explain when heavy tails can have better generalizations. We exhibit different spectral properties such as invariant bulk, spike, and heavy-tailed distribution from a two-layer neural network using different training strategies, and then correlate them to the feature learning. Analogous phenomena also appear when we train conventional neural networks with real-world data. We conclude that monitoring the evolution of the spectra during training is an essential step toward understanding the training dynamics and feature learning.

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 [65, 83, 48, 40, 12, 4, 84]. 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 [47, 38, 34, 51, 25, 32, 76, 42]. This gap heavily relies on the task complexity, data distribution, architecture of the NN and the training strategy [31]. We consider a more realistic

setting, a linear-width regime (LWR), when the sample size n , the input feature dimension d , and the width h of the hidden layer 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.

The ultra-wide NN ( h ≫ n , fixed d ) stays close to the kernel machine induced by initial NN, throughout the gradient-based training processes [88, 28, 27, 11]. 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 training only the last layer of the network [48, 57, 61]; by contrast, the NTK is the Gram matrix of the Jacobian of the NN for all trainable parameters, which governs the gradient flow of NN [40, 28, 4]. In most theoretical results, these kernels remain fixed throughout training, which leads to a kernel gradient descent with the initial kernel [40, 16], whereas in practice the spectra of the weight matrix, CK, and NTK of the NN change while learning the features from the training data [58, 59, 30, 19, 68]. In this paper, under the LWR, we experimentally and theoretically 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 extend our 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 [55, 10, 78]. Also, various spectral properties of weight and kernel matrices can reveal different features learned by different training procedures [82]. Understanding the dynamics of the spectral properties may aid in finding better approaches to training and 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 the proportional limit in high dimensions [45, 72, 57, 74, 37, 61].

Our main findings/contributions are as follows.

Wefind 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 of weight and kernel matrices in NN.

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 prove the global convergence of GD and the invariance of the limiting spectra for both weight and kernel matrices in this scenario.

We observe a phase transition of the alignment and the emergence of a spike outside the bulk of the spectrum when the learning rate exceeds some threshold. The strong alignment of the spike with the teacher model when step sizes are large confirms that the NN is indeed learning germane features from data. This observation justifies the theoretical result of [6] in an ideal two-stage training process.

The evolution towards heavy-tailed spectra is also discovered by using adaptive methods. Our experiments rule out a causal relationship between the occurrence of a heavy-tailed spectrum for the weight matrices and a good generalization. This complements the work of [60, 63, 86] where the authors had observed a strong correlation between the two; while at the same time, we provide simple examples of when heavy-tailed spectra exhibit feature learning and better generalizations.

For more details on how our results fit into existing literature, please see Appendix A.

2 Notation and Preliminaries

Throughout this paper, ∥·∥ denotes the ℓ 2 norm for vectors, ℓ 2 → ℓ 2 is the operator norm for matrices, while ∥ · ∥ F is the Frobenius norm, and ⊙ represents the Hadamard product between matrices. o d, P ( · ) represents little-o in probability as d →∞ .

Neural Tangent Kernel Parameterization. Consider a L -layer fully connected feedforward NN at initialization without bias term: for 1 ≤ ℓ ≤ L -1 ,

and f θ ( x ) = v ⊤ h ( L -1) , where the input vector is x ∈ R d , W ℓ ∈ R n ℓ × n ℓ -1 is the weight matrix for the ℓ -th layer, and v := [ v 1 , . . . , v h ] ⊤ ∈ R n L -1 is the last-layer weight. Let n 0 = d . Denote all trainable parameters by θ := [ vec ( W 1 ) , . . . , vec ( W L -1 ) , v ] ⊤ ∈ R p where each parameter's initial value is independently sampled from some distribution and p is the total number of parameters. Let the training dataset be ( X , y ) := ([ x 1 , . . . , x n ] , y ) ∈ R d × n × R 1 × n ; the output of this NN with respect to this dataset is f θ ( X ) = [ f θ ( x 1 ) , . . . , f θ ( x n )] . We call the above parameterization the NTK parameterization . The loss function for training is a mean squared error (MSE)

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

Training Processes of NNs. NNs are usually trained by gradient-based methods such as full-batch gradient descent (GD), mini-batch stochastic gradient descent (SGD), Adaptive Gradients (AdaGrad), and Adam [43]. We can represent GD by

where η is the learning rate and ∇ θ L ( θ t ) is the gradient of the training loss w.r.t. trainable parameters θ at step t ≥ 0 . We will prove the global convergence of GD in some special (overparameterized) cases ensuring the convergence to a NN that interpolates the data. We will also show the hyperparameters (e.g. learning rate η ) affect the spectral properties of NNs during training.

Conjugate Kernel and Neural Tangent Kernel. 1 When L = 2 , let n 1 = h and n 0 = d be the widths of the output and input layer. The CK is defined as

where X 1 := 1 √ h σ ( WX / √ d ) . We can view the NN as a function of all training parameters θ and input data X . The neural tangent kernel (NTK) is related to the gradient of this neural network function with respect to θ , which is the Gram matrix of the Jacobian of the neural network function with respect to θ , K NTK := ( ∇ θ f θ ( X )) ⊤ ( ∇ θ f θ ( X )) . Specifically, the empirical NTK of two-layer NN can be explicitly written 2 as

In this paper, we are interested in comparing the spectral distributions for these three matrices (weight, CK, and NTK) at initialization and the end of training.

Lazy Training. Lazy training [21] can be viewed as a linear approximation of the NN, i.e. f θ ( x ) ≈ f θ 0 ( x ) + ( θ -θ 0 ) ⊤ ∇ θ f θ 0 ( x ) , defined by minimum-norm interpolation ˆ θ := arg min { ∥ θ -θ 0 ∥ : ( θ -θ 0 ) ⊤ ∇ θ f θ 0 ( X ) = y -f θ 0 ( X ) } . Then, lazy training also represents a kernel machine

where ˆ f ( x ) is the unregularized regression prediction on test data x ∈ R d , the kernel K ( X , X ) is the initial K NTK on training data, and K ( X , x ) = ( ∇ θ f θ 0 ( X )) ⊤ ( ∇ θ f θ 0 ( x )) . The asymptotic performance of ˆ f ( x ) has been analyzed by [1] under the LWR. We view this regime as a benchmark : [21, 11] prove that NN through gradient flow is close to lazy training if h ≫ n ; [38] shows NN can go beyond lazy training under a non-proportional regime.

3 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. A two-layer NN in (1) is defined by

At initialization, we assume that the first hidden-layer W = [ w 1 , . . . , w h ] ⊤ ∈ R h × d is composed of independent standard normal random vectors.

Assumption 3.1 (Linear-width regime (LWR)) . Assume that n d → γ 1 and h d → γ 2 as n →∞ where the aspect ratios γ 1 , γ 2 ∈ (0 , ∞ ) are two fixed constants.

LWR stands as a pivotal setting grounded in high-dimensional statistics [1, 61]. It offers valuable insights especially when addressing real-world datasets. This is in contrast to the infinite-width regime, in which we are already in the asymptotic limit for width at first. Hence, LWR is a better approximation of real-world datasets and practical neural networks compared with the infinite-width regime.

Assumption 3.2 (Activation function) . Suppose that the activation function σ ( x ) is nonlinear and λ σ -Lipschitz with | σ ′ ( x ) | , | σ ′′ ( x ) | ≤ λ σ for all x ∈ R . Moreover, E [ σ ( z )] = 0 for z ∼ N (0 , 1) .

Though the LWR is somewhat impractical, it is still more aligned with deployed models than the infinite-width regime ( h ≫ n , fixed d ). As a kernel machine, the infinite-width NN has been studied extensively [40, 28, 27, 84]. This infinite-width limit is special, however, as NNs may generally evolve beyond the kernel regime and achieve superior performance [31, 55, 10, 78].

Table 1: Four models with the same architecture ( n = 2000 , h = 1500 , d = 1000 , and σ is normalized tanh ), but different choices of initial learning rates and optimizers listed in Table 1. The training label noise σ ε = 0 . 3 and the teacher model is defined by (9) with σ ∗ a normalized softplus and τ = 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.

Assumption 3.3 (Synthetic dataset and teacher model) . Training data is X := [ x 1 , . . . , x n ] ∈ R d × n , where x i i.i.d. ∼ N ( 0 , I d ) . The training labels y = [ y 1 , . . . , y n ] are defined by y i = f ∗ ( x i ) + ε i , for i ∈ [ n ] , where f ∗ : R d → R is the teacher model, and ε i is centered sub-Gaussian noise with variance σ 2 ε .

One of the simplest nonlinear teacher models we can generate is the single-index model, namely f ∗ ( x ) = σ ∗ ( x ⊤ β ) for a fixed vector β with ∥ β ∥ = 1 and nonlinear function σ ∗ ; the hidden feature is simply β ∈ R d . In general, we can consider a multiple-index model

0

1

2

3

4

5

(a) Weight spectra

0.0

0.2

0.4

0.6

0

2

Initial

0

2

Trained

Q-Q plot

Initial

Trained

0

2

4

6

8

10

12

(b) CK spectra

0.0

0.1

0.2

0.3

0.4

0.5

0

5

Initial

0

10

Trained

Q-Q plot

Initial

Trained

0.0

2.5

5.0

7.5

10.0

12.5

15.0

(c) CK spectra

0.00

0.05

0.10

0.15

0.20

0

5

Initial

0

100

Trained

Q-Q plot

Initial

Trained

where β i are some orthogonal unit vectors. We will specifically consider a mixture of single-index and quadratic models as our teacher model in this section:

for some nonlinear target σ ∗ , signal β and constant τ 3 . 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. Putting things together, we can see close relationships between the spectra and the generalization of the NN. These different spectral properties actually reveal disparate features learned via different training strategies.

The advantage of this toy model is that we can easily extract the spectral behaviors over training and then compare them with the kernel machine. We use lazy training defined in (6) as our benchmark to assist us in determining whether a NN outperforms the associated kernel machine. Table 1 compares the test errors and R 2 scores for different optimization cases and the lazy training. By tuning the hyper-parameters, we can find specific situations where NN outperforms the lazy training (see also Figure 10(c) in Appendix B.2).

From Figures 1(a), 12 and 13 in Appendix B.2, one can observe the spectral distributions of the weight, CK and NTK matrices remain invariant and static during training in Cases 1 & 2, which indicates both cases still belong to the 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 B.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 4.2). Comparing with Case 2, Case 3 of Table 1 suggests the importance of the large learning rate regime for training NNs [52, 64, 55, 15, 3]. As a remark, our spectral results of Case 3 are consistent with the observations in [79] through RMT hypothesis testing, where the majority of trained weight matrices 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 B.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 with AdaGrad in Figure 22 in Appendix B.6. Although 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. In the following Section 4, we prove the invariance of the bulk distributions and provide more refined analyses of spikes and heavy tails in terms of feature learning.

4 Different Spectral Behaviors in NNs

We now further explore the spectral behaviors in different cases of Table 1 by clarifying how the spectra evolve through different training processes and how this evolution may affect the NN. Following Figure 1, we study the training processes case-by-case: invariant bulk, spikes outside the bulk, and heavy-tailed distribution. Additional experiments are exhibited in Appendix B.

4.1 Invariant Bulk Distributions

In Figure 1(a) (also Figures 12 and 13 in Appendix B.2), we observe the bulk distributions of weight and kernel matrices in Cases 1&2 remain globally unchanged (invariant) over the training process. under the LWR, this is also empirically verified by Figures 10(a) & (b) in Appendix B.2. In this section, by investigating the global convergence of GD, we prove this invariant-bulk phenomenon under certain assumptions.

For simplicity, we focus on analyzing the training process of the first-hidden layer with the second layer v fixed. Denote f θ ( X ) by f W ( X ) in this case. At any time t ∈ N , consider the gradient steps:

Denote the CK and NTK at gradient step t ∈ N by K CK t := 1 h σ ( W t X ) ⊤ σ ( W t X ) , and K NTK t := 1 d X ⊤ X ⊙ 1 h σ ′ ( 1 √ d W t X ) ⊤ diag ( v t ) 2 σ ′ ( 1 √ d W t X ) 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 4.1 (Early phase) . Under Assumptions 3.1, 3.2 and 3.3, we further assume that ∥ v ∥ ∞ ≤ 1 and f ∗ is a λ σ -Lipschitz function. Given any fixed t ∈ N and learning rate η = Θ(1) , after t gradient steps, the changes 1 √ d ∥ W t -W 0 ∥ F , ∥ ∥ K CK t -K CK 0 ∥ ∥ F , and ∥ ∥ K NTK t -K NTK 0 ∥ ∥ are all less than C n , with probability at least 1 -4 n exp( -cn ) , for some positive constants c, C > 0 which only depend on step t and parameters η, γ 1 , γ 2 , λ σ , σ ε .

Lemma 4.1 shows 1 √ d ∥ W t -W 0 ∥ , ∥ ∥ K CK t -K CK 0 ∥ ∥ , and ∥ ∥ K NTK t -K NTK 0 ∥ ∥ 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 C.3 in Appendix C). Now we aim to analyze the spectra at the end of the training process (10). 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 K CK t and K CK t for all t .

By [81, Theorem 2.9], the smallest eigenvalue of K NTK 0 has an asymptotic lower bound:

where a σ := E [ σ ′ ( ξ ) 2 ] and η k is the k -th Hermite coefficient of σ ′ . Hence, we can claim there exists some constant α > 0 only dependent on σ such that λ min ( K NTK 0 ) ≥ 4 α 2 with high probability. Note that α is not vanishing since σ is nonlinear. With this lower bound, we obtain the following global convergence for (10) and norm control of W t as n/d → γ 1 and h/d → γ 2 .

Theorem 4.2 (Global convergence) . Under the same assumptions of Lemma 4.1, we further assume v i 's are independent and centered random variables in the second layer. For any η < min { α 2 n 2 , n 4 λ 2 σ (1+ √ γ 1 ) 2 } and all t ∈ N , there exists some γ ∗ > 0 such that, when γ 2 ≥ γ ∗ , the gradient steps (10) will satisfy

with high probability, as n/d → γ 1 and h/d → γ 2 . Here, training loss ℓ ( W ) := ∥ y -f W ( X ) ∥ .

We apply the techniques and results by [69, 70] to obtain Theorem 4.2. Notice that, unlike Lemma 4.1, the largest learning rate we can choose is of order Θ( n ) . As a byproduct, the Frobenius norm in (13) implies the following corollary for the invariance of limiting bulk distribution.

Corollary 4.3. Under the same assumptions of Theorem 4.2, for all t ∈ N , with high probability, there exists some constant R > 0 such that the changes 1 √ d ∥ W t -W 0 ∥ F , ∥ ∥ K CK t -K CK 0 ∥ ∥ F , and ∥ ∥ K NTK t -K NTK 0 ∥ ∥ F are all less than R with high probability. This implies the limiting empirical spectra of 1 h W ⊤ t W t , K CK t and K NTK t are the same as the limiting spectra of 1 h W ⊤ 0 W 0 , K CK 0 and K NTK 0 respectively, almost surely as n/d → γ 1 and h/d → γ 2 .

Corollary 4.3 is empirically validated by Figure 29 in Appendix C. In addition, based on Figure 31 in Appendix C, one can further extend Corollary 4.3 to the SGD training process. The total path is O ( √ h ) in (13) and (14), which is negligible compared with the Frobenius norm of initial weight matrix (which is of order Θ( h ) ). Thus, gradient descent iterates (10) remain close to initialization and small perturbation of NTK ensures the smallest eigenvalue (11) of NTK is always lower bounded away from zero. Theorem 4.2, however, does not require that the NTK stays unchanged all the time. Moreover, Corollary 4.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 [11, Theorem 5.4] is not directly applicable to show that a NN is still close to lazy training under the LWR. It requires deeper analysis to claim whether a NN still belongs to the kernel regime or already goes beyond in our case. As we will show in Section 4.2, this also relies on the magnitude of the learning rate for GD/SGD.

4.2 Alignments for Spiked Models

The outliers appear in the spectra of the trained weight, CK, and NTK matrices (Figure 14 in Appendix B.2) when NNs are optimized with large learning rates. The outlier is especially clear for the NTK matrix (Figure 14(b)). Heuristically, this indicates that the NN is learning the feature from the teacher model f ∗ . In Figure 2(a), Figures 15 and 27 in Appendix B, we empirically exhibit these phenomena for W , K CK and K NTK respectively through different training processes.

0.05

0.00

0.05

0.10

Signal in the teacher model

0.10

0.05

0.00

0.05

0.10

Top singular vector

Initial

Trained

(a)

β

vs.

u

1

of weight.

4

6

8

10

12

lr

5

6

7

8

9

10

11

Initial right edge

(b)

λ

(

W

t

W

⊤

t

/d

)

4

6

8

10

12

lr

10

12

14

16

18

Initial right edge

(c)

λ

(

K

CK

t

)

4

6

8

10

12

lr

0.0

0.2

0.4

0.6

0.8

alignment

weight

CK

NTK

(d) Alignments

Spikes of Weight Matrices. The differences between Cases 2&3 empirically validate the benefits of training with large learning rates [52, 64, 55, 15, 3]. Inspired by [6], we consider the alignment between the leading right singular vector u 1 of W t and the signal β in the teacher model defined by (9). For Case 3, a notable alignment appearing in Figure 2(a) after training suggests that W t is capturing the feature β during training. Although this does not ensure NN will entirely beat the optimal kernel lower bound, this alignment reveals a non-negligible feature selection [10] via largestepsize 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 27 in Appendix B.7.

0

20

40

60

80

0.05

0.10

0.15

0.20

0.25

0.30

KTA

Adam

SGD

(a) Evolution of KTA.

0

50

100

150

200

0.30

0.35

0.40

0.45

0.50

cos

i

2

3

4

(b) Alignment evolution.

0

50

100

150

200

250

CK spectra

0.000

0.005

0.010

0.015

0.020

0.025

0.030

0

10

20

0

5

10

15

20

Cauchy

Q-Q plot

Cauchy

(c) Heavy-tailed initial CK.

0

2

4

6

8

10

12

14

Weight spectra

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

1

2

3

Initial

1

2

3

Trained

Q-Q plot

Initial

Trained

(d) Bad example.

Transitions of the Spike 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 η . This indicates a transition of the emergence of the spike outside the bulk distribution. Figure 2, analogously to the well-known BBP transition by Baik, Ben Arous, and Péché in [9] from the RMT community, shows there is a threshold (yellow region) for learning rate: the outliers only appear when η 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 β and the leading eigenvector of W ⊤ t W t /d , and y and the leading eigenvectors of K CK t and K NTK t after training for different η . Similarly with [10], when η 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 W t by [6] for a different scenario 4 .

Spikes of Kernel Matrices. The alignment of the kernel matrix with the training labels y is defined by [22] by Kernel Target Alignment (KTA) as follows: when kernel K is either CK or NTK,

Analogously to [10, 5, 78], 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 v 1 of kernel matrix naturally dominates the alignment with y (Figure 15 in Appendix B.2), which is regarded as a kernel rotation during training in [68]. Notice that this is not the common situation in Cases 1 & 2 of Table 1 (and cf. Figure 11 in Appendix B.2). On the other hand, KTA measures the alignment between y and the full eigenbasis of the kernel. These kernel alignments improve the speed of the convergence of training dynamics but may hurt or boost the generalization of the NNs [68, 78, 10]. 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 a heavy tail in the spectrum is closely related to a better generalization of the NN and more significant feature learning.

4.3 Phenomenon of Heavy-tailed Spectra

Next, we analyze the heavy-tailed spectra of weight and kernel matrices in Figure 1(c). [58, 59] found a strong correlation between the heavy-tailed spectra of trained state-of-the-art models with better generalization (Figure 7 in Appendix B.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, heavytailed spectra generally appear when the entries of the matrix are highly correlated [59]. This could heuristically explain heavy-tailed phenomena in the spectra since the entries of well-trained W t should be strongly correlated. Unlike [58, 59], we focus on the heavy-tailed phenomena for both weight and kernel matrices in a simpler model (7) and provide a connection between feature learning and heavy-tailed spectra, which opens an important avenue for further theoretical analysis.

0

500

1000 1500 2000 2500 3000 3500 4000

0

200

400

600

800

1000

10 4

Initial

10 2

10 5

Trained

Q-Q plot

Trained

Initial

0

1000

2000

3000

4000

0

200

400

600

800

1000

1200

1400

10 4

Initial

10 2

10 5

Trained

Q-Q plot

Trained

Initial

(b) SGD

0

500

1000 1500 2000 2500 3000 3500 4000

0

200

400

600

800

1000

1200

1400

10 3

10 7

Initial

10 3

10 6

Trained

Q-Q plot

Trained

Initial

(c) Adam

Heavy Tails and Generalization. We emphasize that heavy tails are not sufficient for good generalization, in general, [60, 63]. Figures 3(c)&(d) exhibit NNs with heavy-tailed weights but in the absence of good performance at initialization. In fact, it is the alignments between the features learned from the heavy-tailed part and the features in the teacher model that finally determine the generalization error of NNs.

More precisely, we provide an example of when heavy tails indicate better generalizations. Consider the multiple-index teacher model (8) with k = 5 feature directions β i , and train NNs (7) with GD, SGD, and Adam to get invariant bulk, bulk with one spike and heavy tails, respectively, after training. In Figure 3(b), we present the evolutions of the principle angles θ i between feature subspace U = span { β i } k i =1 and top 100 eigenspace of W ⊤ t W t during different training processes. This eigenspace with respect to the top 100 eigenvalues of W ⊤ t W t corresponds to the heavy-tail part of the spectrum in W ⊤ t W t when training NNs with Adam (solid lines in Figure 3(b)). Comparing with GD and SGD training processes, we observe strong alignments between feature space U and eigenspace w.r.t heavy tails in Adam case in Figure 3(b), which explains why Adam case (NNs with heavy-tailed spectra) generalizes better than the other two cases. For more examples, see Figures 17 and 20 in Appendix B.4. This concludes that NNs with heavy-tailed spectra can generalize better only when the teacher features from data are aligned with the heavy-tailed part of spectra. If the feature dimension in the teacher model is high (i.e. the teacher model is more complicated and intrinsically high-dimensional), then we expect to get a heavy-tailed weight spectrum of well-trained NN where the heavy-tailed part learns all the features in the teacher modes. This example explains why we can use the heavy tails to discriminate well-trained and poorly-trained large models [60, 63, 86].

5 Discussions and Future Directions

We empirically investigated how the spectra of W , K CK , and K NTK evolve under the LWR for an idealized student-teacher setting. 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 eigenstructure of weight and kernel matrices. Since evolution is sensitive to feature learning, we can link feature learning and different training dynamics by studying the spectra of these matrices.

While synthetic data is easier to analyze theoretically, we also investigate these spectral properties on real-world data and more complicated tasks in the following. In practice, people mainly focus on analyzing spectra of the weight matrices in fully connected layers; we choose to also focus on the spectral properties of general kernel matrices induced by the NNs, which contain abundant information [19, 55, 5, 78].

First, we show the spectra of K 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 3 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 datasets. 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 [26], with fine-tuning on Sentiment140 dataset of tweets 5 from [36]. We fine-tune the BERT model for 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 B.1).

0

5

10

15

20

0.0

0.2

0.4

0.6

0.8

1.0

1e6

first

second

(a) Eigenvalues

0

5

10

15

20

0.0

0.2

0.4

0.6

0.8

alignment

first

second

(b) Alignments

A heavy-tailed CK spectrum with several spikes already exists in this pre-trained model. Unlike Figure 4 (and cases in Table 1) 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 to Figure 2(d), our 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 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 (see Figure 7in Appendix B.1). We believe that the evolutions of the kernel matrices and some spectral metrics are crucial for understanding feature learning through fine-tuning [82]. A more comprehensive exploration of the evolutionary spectral properties of 'foundation models' may help shed further light on these phenomena.

Limitations. Although LWR has garnered significant attention in recent years, e.g., [77, 50, 17, 87, 23], we recognize the limitations inherent in LWR. Our LWR is more realistic compared with infinite-width neural networks and is one of the ways to approximate finite but very large neural networks with very large datasets, but there are more sophisticated regimes for NNs. We leave this for future theoretical work. The NTK parameterization is another limitation of this work. We expect to apply our spectral analysis for other parameterizations of NNs with more real-world datasets. See the discussion at the beginning of Appendix B.

References

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.

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.

Maksym Andriushchenko, Aditya Vardhan Varre, Loucas Pillaud-Vivien, and Nicolas Flammarion. Sgd with large step sizes learns sparse features. In International Conference on Machine Learning , pages 903-925. PMLR, 2023.

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.

Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations , 2022.

Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. Highdimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems , 35:37932-37946, 2022.

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.

Zhidong Bai and Jack W Silverstein. Spectral analysis of large dimensional random matrices , volume 20. Springer, 2010.

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.

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.

Peter L Bartlett, Andrea Montanari, and Alexander Rakhlin. Deep learning: a statistical viewpoint. Acta numerica , 30:87-201, 2021.

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.

Lucas Benigni and Sandrine Péché. Eigenvalue distribution of some nonlinear models of random matrices. Electronic Journal of Probability , 26:1-37, 2021.

Lucas Benigni and Sandrine Péché. Largest eigenvalues of the conjugate kernel of single-layered neural networks. arXiv preprint arXiv:2201.04753 , 2022.

Gaspard Beugnot, Julien Mairal, and Alessandro Rudi. On the benefits of large learning rates for kernel methods. In Conference on Learning Theory , pages 254-282. PMLR, 2022.

Alberto Bietti and Julien Mairal. On the inductive bias of neural tangent kernels. In Advances in Neural Information Processing Systems , pages 12873-12884, 2019.

David Bosch, Ashkan Panahi, and Babak Hassibi. Precise asymptotic analysis of deep random feature models. arXiv preprint arXiv:2302.06210 , 2023.

Sourav Chatterjee. Convergence of gradient descent for deep neural networks. arXiv preprint arXiv:2203.16462 , 2022.

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.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. In Advances in neural information processing systems , pages 3036-3046, 2018.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems , 32, 2019.

Nello Cristianini, John Shawe-Taylor, Andre Elisseeff, and Jaz Kandola. On kernel-target alignment. Advances in neural information processing systems , 14, 2001.

Hugo Cui, Florent Krzakala, and Lenka Zdeborova. Bayes-optimal learning of deep random networks of extensive-width. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 6468-6521. PMLR, 23-29 Jul 2023.

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.

Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems , 33:20356-20365, 2020.

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.

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.

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.

Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations , 2020.

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.

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.

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.

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.

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.

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.

Alec Go, Richa Bhayani, and Lei Huang. Twitter sentiment classification using distant supervision. CS224N project report, Stanford , 1(12):2009, 2009.

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.

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.

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.

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.

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.

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.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Dmitry Kopitkov and Vadim Indelman. Neural spectrum alignment: Empirical study. In International Conference on Artificial Neural Networks , pages 168-179. Springer, 2020.

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.

Guillaume Leclerc and Aleksander Madry. The two regimes of deep network training. arXiv preprint arXiv:2002.10376 , 2020.

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.

Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In International Conference on Learning Representations , 2018.

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.

Qianyi Li and Haim Sompolinsky. Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization. Physical Review X , 11(3):031059, 2021.

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.

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.

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.

Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in overparameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis , 2022.

Philip M Long. Properties of the after kernel. arXiv preprint arXiv:2105.10585 , 2021.

Noel Loo, Ramin Hasani, Alexander Amini, and Daniela Rus. Evolution of neural tangent kernels under benign and adversarial training. Advances in Neural Information Processing Systems , 35:11642-11657, 2022.

Cosme Louart, Zhenyu Liao, and Romain Couillet. A random matrix approach to neural networks. The Annals of Applied Probability , 28(2):1190-1248, 2018.

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.

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.

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.

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.

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):E7665E7671, 2018.

Xuran Meng and Jianfeng Yao. Impact of classification difficulty on the weight matrices spectra in deep learning and application to early-stopping. Journal of Machine Learning Research , 24(28):1-40, 2023.

Preetum Nakkiran. Learning rate annealing can provably help generalization, even for convex problems. OPT2020 Workshop , 2020.

Radford M Neal. Bayesian learning for neural networks , volume 118. Springer Science & Business Media, 1995.

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.

Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi, and Pascal Frossard. Neural anisotropy directions. Advances in Neural Information Processing Systems , 33:1789617906, 2020.

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.

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.

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.

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.

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.

Jeffrey Pennington and Pratik Worah. Nonlinear random matrix theory for deep learning. Advances in neural information processing systems , 30, 2017.

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.

Bartlomiej Polaczyk and Jacek Cyranka. Improved overparametrization bounds for global convergence of sgd for shallow neural networks. Transactions on Machine Learning Research , 2022.

Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová. Classifying highdimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. In International Conference on Machine Learning , pages 8936-8947. PMLR, 2021.

Dominik Schröder, Hugo Cui, Daniil Dmitriev, and Bruno Loureiro. Deterministic equivalent and error universality of deep random features learning. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 30285-30320. PMLR, 23-29 Jul 2023.

Haozhe Shan and Blake Bordelon. A theory of neural tangent kernel alignment and its influence on training. arXiv preprint arXiv:2105.14301 , 2021.

Matthias Thamm, Max Staats, and Bernd Rosenow. Random matrix analysis of deep neural network weight matrices. Physical Review E , 106(5):054124, 2022.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

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.

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.

Christopher Williams. Computing with infinite networks. Advances in neural information processing systems , 9, 1996.

Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems , 32, 2019.

Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning , pages 11727-11737. PMLR, 2021.

Yaoqing Yang, Ryan Theisen, Liam Hodgkinson, Joseph E Gonzalez, Kannan Ramchandran, Charles H Martin, and Michael W Mahoney. Evaluating natural language processing models with generalization metrics that do not need access to any training or testing data. arXiv preprint arXiv:2202.02842 , 2022.

Jacob A Zavatone-Veth, William L Tong, and Cengiz Pehlevan. Contrasting random and learned features in deep bayesian linear regression. Physical Review E , 105(6):064118, 2022.

Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes overparameterized deep relu networks. Machine learning , 109:467-492, 2020.

A Additional Related Work

Nonlinear Random Matrix Theory and Random Feature Regression. The limiting spectrum of CK with random input dataset has been investigated by [73, 13]; whereas [57, 30] studied the spectrum of CK with more general input data. This spectrum is actually a deformed Marˇ cenko-Pastur distribution [30], which becomes a deformed semicircular distribution [81] when h ≫ n . The largest eigenvalue of the CK matrix has been studied in [14], and the spectrum of the NTK was analyzed in [1, 30]. As an application, random feature ridge regression was fully determined by the limiting spectra of CK or NTK: [57, 33, 53, 61, 1]. 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. [28, 27, 69, 70, 66, 54, 75, 18]. We refer to Table 1 of [75] as a summary of these recent results. Most of the theorems in the literature require h ≫ n , which implies that the NTK is almost static during training, while [70, 66] can consider LWR under some specific assumptions. Recently, [18] established a new criterion for the convergence of GD which results in the global convergence of general NNs with finite width h and d ≥ n .

Beyond NTK Regime. Under the proportional limit, the initial kernel regression can only learn a linear component of the target [35]. Thus, it is reasonable to consider the cases beyond the NTK regime. To this end, [29, 39] considered the dynamics of NTK throughout training while [2, 7] 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: [51] proved a two-layer ReLU NN that is shown to beat any kernel method; [42] verified a two-layer CNN with some simple dataset can outperform the initial NTK for image classifications; [6] showed a NN can escape the kernel regime by only taking one specific large gradient step; [24] 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 [31, 68, 55, 5, 56]. Specifically, [55] studied the hard-margin SVM for 'after kernels' which are the CK and NTK matrices of trained NNs. One of the effective ways of depicting how the kernels evolve during training is to capture the evolution of kernel alignment [10, 78, 5, 56]. Kernel alignments between kernels and training labels essentially reveal how the NN accelerates training [78]. Also, several papers showed that the top eigenfunctions of the kernel align with the target function learned by the NN [44, 67, 68]. 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 [52, 64, 15, 3]. Specifically, [46] observed that training dynamics with large learning rates differ 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 of training, [41] showed using lower learning rates may result in finding a region of the loss surface with worse conditioning of kernel and Hessian matrices. In [55], the after kernels of NNs trained with larger learning rates generalize better and stay more stable. [49] 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 lazy training to the catapult regime. Section 4.2 illustrates a similar transition in our situations.

Heavy-tailed Phenomenon. The heavy-tailed phenomenon has appeared in many places in deep learning theory; [58, 59] 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, [60] further showed how to distinguish well-trained and poorly trained models by a power-law-based approximation. [63] classified trained weight spectra into three types: Marˇ cenko-Pastur law, bulk with (few) outliers, and heavy-tailed spectra. We extend this classification to both weight and kernel matrices in Figure 1. Additionally, similarly to the discussion in 4.3, [63] showed that the difficulty of the classification problem is related to the emergence of heavy-tailed spectra in weight matrices. This

heavy-tailed phenomenon can be used to construct metrics for evaluating the generalization of NNs [60, 86], and early stopping of NNs to avoid over-fitting [63].

B Additional Empirical Results

There are different parameterizations for NNs at initialization. The orders of the output of NN are distinct in different cases [21, 25, 85]. This affects the size of stable and non-trivial gradient steps. The distance of trainable parameters from initialization determines whether the NN learns any features from the training data [6, Figure 2]. The performance of networks with different initializations indicates whether the NN belongs to the kernel regime or not [85]. Unlike the NTK parameterization, the mean-field parameterization [62, 20] and maximal update parameterization [85] tend to be feature learning.

For all NNs in our experiments, we apply a normalized and centered nonlinear activation function such that Assumption 3.2 holds ( E [ σ ( z )] = 0 for z ∼ 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 W and v in (7) 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 X . In addition, we consider the training label noise defined in Assumption 3.3 as ε ∼ N (0 , σ 2 ε I n ) .

B.1 Further Discussions on Real-world Data Experiments

For the first experiment in Section 5, we fix the small-CNN architecture, which is similar to the 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 × 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.

0

100

200

300

400

0.50

0.55

0.60

0.65

0.70

0.75

0.80

acc

train

test

0

50

100

150

200

250

300

0.5

0.6

0.7

0.8

0.9

1.0

acc

train

test

(b) SGD

0

20

40

60

80

100

0.5

0.6

0.7

0.8

0.9

1.0

acc

train

test

(c) Adam

In the experiment of the transformer language model in Section 5, we fine-tuned the BERT model with SGD for a binary classification on the Sentiment140 dataset. We apply this transformer and finetuning 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 the 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 to 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 transformer [82].

0

5

10

15

20

25

30

CK spectra

0

20

40

60

80

100

10 0

Initial

10

3

10 2

1-epoch

Q-Q plot

Initial, test acc: 49.3

1-epoch, test acc: 81.875

0

5

10

15

20

25

CK spectra

0

20

40

60

80

100

10 0

Initial

10

3

10 3

2-epoch

Q-Q plot

Initial, test acc: 49.3

2-epoch, test acc: 83.225

0

10

20

30

CK spectra

0

20

40

60

80

100

10 0

Initial

10

3

10 2

12-epoch

Q-Q plot

Initial, test acc: 49.3

12-epoch, test acc: 84.4

B.2 Additional Results for Cases in Table 1

To complement the findings in Figure 1 and Section 3, 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:

at every epoch t through training. The changes in Figures 8(a) and (b) are much smaller than in 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 4.2, in Figure 9, we compute the differences between initial W 0 and final W s in Frobenius norm, operator norm and 2 , ∞ -norm. Empirically, Figure 9 shows

as n →∞ with n/d → γ 1 and N/d → γ 2 , where s is the final time for GD. Here, the entry-wise 2 -∞ matrix norm is defined as

for any matrix M ∈ R N × d with the i -th row m i ∈ R d and 1 ≤ i ≤ N . Notice that

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 grow to keep the NNs residing 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 the same for different n 's. Figures 9(b)-(d) empirically validate Corollary 4.3. Moreover, the observation that 1 √ d ∥ W 0 -W s ∥ F and 1 √ d ∥ W 0 -W s ∥ are Θ(1) may suggest that 1 √ d ( W 0 -W s ) 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 understand the spectral evolution during training. Notice that these norms of the change are different from ultra-wide NN [27, 28]. Similar phenomena can be also observed in Figures 10(a) and (b). Analogous result with different σ and σ ∗ is exhibited in Figure 30 in Appendix C. In addition, Figure 10(c) further investigates the cases when NNs can outperform lazy training as defined by (6). In these experiments, we compare the performances of GD, and SGD with small or large learning rates, and lazy training as n →∞ . Each time, we take 10 trials to average. We observe that SGD with a large learning rate (green line) can asymptotically outperform lazy training.

0

2000

4000

6000

8000

0.0

2.5

5.0

7.5

10.0

12.5

15.0

norms

NTK operator norm

Weight operator norm

(a) Case 1

0

1000

2000

3000

4000

0

5

10

15

norms

NTK operator norm

Weight operator norm

(b) Case 2

0

5

10

15

20

0

50

100

150

norms

NTK operator norm

Weight operator norm

(c) Case 4

=IIÉtr- affettuoso

fgEEggggpEBFBgf

☆

gtEEEEBEfgff

☆

(a) Train/test losses

(b) Change in weight

(c) Change in CK

(d) Change in NTK

2.50

3.0

0.8

1.3

Taining and Test Losses

0.6

1.2

1.1

2.25

2.00

2.5

0.4

0.2

0.2

•risk for initial NTK

十

train

test

R²score

1.0

0.9

0.8

1.50

1.75

1.25

十

十

|CK  CKs]

|CK-CK]

2.0

1.5

十

[NTKo  NTK]2

NTKo NTK]ro

0.7

十 [W-Ws]

0.0

0.1

10002000

0.6

+W-W

+W-w).-

1.00

1.0

750

1000

125015001750200022502500

Sample Dimension n

2750

750

10001250

150017502000225025002750

Sample Dimension n

750

1000

1250150017502000225025002750

Sample Dimension n

750

10001250

150017502000225025002750

Sample Dimension n

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 a spectral point of view, the NTK is invariant/static through training. Based on Figures 1(a) and 12, we can empirically verify Corollary C.3 stated in Appendix C. Globally, the spectra of W , K CK and K NTK are not changing over training as n/d → γ 1 and N/d → γ 2 . The initial spectrum of weight W 0 converges to Marˇ cenko-Pastur law; the initial spectrum of NTK under proportional limit has been studied by [1, 30]. Figure 12(c) demonstrates the global convergence for GDunder the proportional regime, as proved in Theorem 4.2. We can observe this global convergence even for SGD, Case 2 in Table 1, although we do not have proof for it.

Case 2. As a complement, Figure 13 exhibits the spectra of W , K CK and K 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 4.1 can be extended to SGD training with sufficiently small learning rates, which is subject to future work. Analogously to Theorem 4.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 [69, 70]. Once we have the invariant global spectra in Corollary 4.3, we can apply the nonlinear RMT [73, 57, 13, 30] to characterize the limiting spectra under LWR.

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 is 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

1000

2000

3000

4000

Sample Dimension n

0.0

0.2

0.4

0.6

0.8

1.0

|

W

0

Ws

|2

|

W

0

Ws

|

F

|

W

0

Ws

|

2,

(a) Change in weight.

1000

2000

3000

4000

Sample Dimension n

1.0

1.2

1.4

1.6

1.8

2.0

2.2

|

CK

0

CKs

|2

|

CK

0

CKs

|

F

(b) Change in CK.

10 3

sample size

10

1

10 0

test loss

SGD large lr

SGD small lr

Lazy training risks

NTK lower bound

(c) Prediction risks vs.

γ

1

.

0.10

0.05

0.00

0.05

0.10

Signal in the teacher model

0.10

0.05

0.00

0.05

0.10

Top singular vector

Initial

Trained

1

0

1

Training labels

0.05

0.00

0.05

Pc of CK

Initial

Trained

1

0

1

Training labels

0.05

0.00

0.05

Pc of NTK

Initial

Trained

of the NTK shrinks after training (Figure 14(b)), which indicates NNs converge to flatter minima. This resembles the catapult phase in [49] 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 from [49].

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 23 and 25 in Appendix B.6, where we have comparable performances to Case 4 in Table 1.

Heavy tails are essentially power laws. To measure how 'heavy' the spectrum is, [58, 59, 60] provide estimates on the power law of W . Consider the empirical spectrum of W as ρ ( x ) ∼ x -α for large x and some positive constant α . The spectrum with a heavier tail has a smaller value of α . Figure 17(a) shows how α evolves through training in Case 4 of Table 1. As α decreases, a heavy tail in the spectrum of the weight matrix emerges in Figure 1(c). In Figure 17(a) we introduce two more metrics to show this evolution: Weighted Alpha ˆ α := αλ 1 and Log α -norm log ( ∑ N i =1 λ α i )

where λ N ≤ . . . ≤ λ 1 are the eigenvalues of WW ⊤ . Remarkably, Figure 17(a) indicates the spectra change 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 5. These metrics are applied to measure the tails in pre-trained models [60].

B.3 Additional Results for the Emergence of A Spike

As a complement to section 4.2, in Figure 18, we show the training dynamics for SGD training with a larger learning rate in the example of Figure 2(b-d). Here we consider η = 24 which belongs to the orange region in Figure 2(b-d), where the spike and eigenvector alignment emerge. Figure 18 presents the details of the training dynamics of the NN in this case: the largest eigenvalues of CK and

0

1

2

3

4

5

(a) Weight spectra

0.0

0.2

0.4

0.6

0

2

Initial

0

2

Trained

Q-Q plot

Initial

Trained

0

5

10

15

(c) NTK spectra

0.00

0.25

0.50

0.75

1.00

1.25

1.50

0

10

Initial

0

10

Trained

Q-Q plot

Initial

Trained

0

2000

4000

6000

8000

10

6

10

4

10

2

10 0

losses

train

test

1.0

0.5

0.0

0.5

R

2  score

0

1

2

3

4

5

(a) Weight spectra

0.0

0.2

0.4

0.6

0

2

Initial

0

2

Trained

Q-Q plot

Initial

Trained

0.0

2.5

5.0

7.5

10.0

12.5

15.0

(b) NTK spectra

0.0

0.5

1.0

1.5

0

10

Initial

0

10

Trained

Q-Q plot

Initial

Trained

0

2

4

6

8

(c) CK spectra

0.0

0.5

1.0

1.5

2.0

2.5

3.0

0

5

Initial

0

5

Trained

Q-Q plot

Initial

Trained

NTK both increase and the losses first increase and then drop. In Figure 19, we empirically justify that the phase transitions we presented in section 4.2 for SGD can be also extended to full-batch GD cases. We can also observe phase transitions for test losses and R 2 scores when we are gradually increasing the learning rates. Parallel to these, a spike also appears outside the bulk distribution, which corresponds to feature alignments in Figure 19(c)&(f).

B.4 Multiple-index Examples for Heavy-Tailed Spectra in Section 4.3

Figures 17(b) and (c) are additional results for Figure 3(b) in Section 4.3. In this experiment, we consider σ = ReLU, n = 5000 , h = 2500 and d = 1000 for NN (1). Comparing with the teacher model (9) used in Table 1, we employ the multiple-index teacher model (8) with k = 5 and σ ∗ = σ . We trained this student-teacher model using GD ( η = 15 ), SGD ( η = 7 . 25 and batch size 8), and Adam ( η = 0 . 007 and batch size 16) for training this NN, respectively. Similarly with Figure 1, correspondingly, we observe invariant spectrum, bulk with one spike, and heavy tails after training respectively. Heuristically, to learn this f ∗ , the weight W of NN should gradually align with the feature space U spanned by β i 's. Hence, to study feature learning, we can apply principle angles to measure the alignment between W and U . Consider the eigen-decomposition of W ⊤ t W t = ∑ d i =1 λ i v i v ⊤ i with λ 1 ≥ λ 2 ≥ . . . ≥ λ d . Figure 3(b) shows the heavy-tailed part (the eigenspace E := span { v i } 100 i =1 ) is aligned with U after training, which shows how features are learned in the heavy-tailed spectra. Remarkably, the test errors for training processes with SGD and Adam are even smaller than ∥ P > 1 f ∗ ∥ 2 and ∥ P > 2 f ∗ ∥ 2 , where P > 1 denotes the orthogonal projection onto the nonlinear part of the function w.r.t. Gaussian measure. Thus, we experimentally showed that NNs with heavy-tailed spectra can obtain feature learning and generalize better than the other two cases. Another example is exhibited in Figure 20. In this case, k = 5 and there are five leading outlier eigenvalues in the spectrum of the trained weight matrix, along with a heavy-tailed bulk. Interestingly, Figure 20 justifies that the eigenspace of these five leading outliers is strongly aligned with features β i for 1 ≤ i ≤ 5 . This indicates that heavy-tailed spectra with large spikes may have a correlation with feature learning and good generalizations.

⑭

Rahmanamo

MReo

...

Raman

mom

a

as

(a) Spectra of weights.

www.#faaas---@

(b) Spectra of the first part in NTKs.

0

20

40

60

80

0

20

40

60

norms

NTK operator norm

Weight operator norm

(c) Norms of change vs. epochs.

0

20

40

60

80

10

5

10

4

10

3

10

2

10

1

10 0

losses

train

test

1.5

1.0

0.5

0.0

0.5

R

2  score

(a) Losses vs. epochs.

2

1

0

1

2

Training labels

0.075

0.050

0.025

0.000

0.025

0.050

0.075

Pc of CK

Initial

Trained

(b)

y

vs.

v

1

of CK.

2

1

0

1

2

Training labels

0.075

0.050

0.025

0.000

0.025

0.050

0.075

Pc of NTK

Initial

Trained

(b)

y

vs.

v

1

of NTK.

B.5 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 29 in the next section. In section 4.1, Corollary 4.3 shows that under LWR with sufficiently large width h , the limiting spectra of 1 h W ⊤ t W t , K CK t and K NTK t are essentially the same as those of the corresponding initial matrices if we train only the first layer with GD (see (10)). Figure 21 further investigates these phenomena when training only the first layer with Adam. In particular, 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 21(c), we only consider the first component of NTK since the gradient is only taken with respect to 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 C.

B.6 Adaptive Gradients

Inspired by Case 4 in Table 1, we show the spectral performances of adaptive gradient (AdaGrad) in Figures 22 and 23. 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 23.

B.7 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 24 and 25, we first take Adam with large stepsizes for a few steps and then use small-stepsize SGD for convergence. In this scenario, we

iwaBAB

☆

Bk

(a) Spectra of weights.

Mammalogist

(b) Spectra of NTKs.

0

5

10

15

0

1

2

3

4

losses

train

test

2

1

0

R

2  score

(c) Losses vs. epochs.

0.10

0.05

0.00

0.05

0.10

Signal in the teacher model

0.10

0.05

0.00

0.05

0.10

Top singular vector

Initial

Trained

(d) First PC in trained & initial

weights.

2

1

0

1

2

Training labels

0.05

0.00

0.05

0.10

Pc of CK

Initial

Trained

(e) First PC in trained & initial CKs.

2

1

0

1

2

Training labels

0.05

0.00

0.05

Pc of NTK

Initial

Trained

(f) First PC in trained & initial

NTKs.

0

5

10

15

20

25

2.5

5.0

7.5

10.0

12.5

15.0

17.5

metrics

alpha

Weighted Alpha

Log alpha norm

(a) Evolution of metrics.

0

50

100

150

200

0.4

0.5

0.6

0.7

0.8

0.9

cos

i

Adam

1

SGD

1

1

(b) Evolution of PC angles.

0

2

4

6

8

10

12

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

5

10

Initial

5

10

Trained

Q-Q plot

Initial

Trained

(c) Weight spectra.

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 26 and 27, are similar to 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 [67, 68, 78]. These two different spectral properties reveal significant differences between the global minima of these two training processes and different evolutions of the

0

20

40

60

80

100

0

1

2

3

4

5

loss

train losses

test losses

(a) Losses.

0

20

40

60

80

100

10

12

14

16

18

eigen

(

CK

)

(b)

λ

(

CK

)

.

0

20

40

60

80

100

14

15

16

17

18

19

eigen

(

NTK

)

(c)

λ

(

NTK

)

10 1

10 2

learning rates

3×10 1

4×10 1

loss

test losses

(a) Losses vs learning rates.

0

1

2

3

4

5

Weight spectrum

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

1

2

initial

1

2

trained

initial

trained

(b) Weight spectra with a small

learning rate.

0.10

0.05

0.00

0.05

0.10

1st PC for weight

0.05

0.00

0.05

0.10

Feature in techer model

initial

trained

(c) Weight alignments with a small

learning rate.

10 1

10 2

learning rates

0.58

0.60

0.62

0.64

R2 score

R2 score

(d)

R

2

vs learning rates.

0

1

2

3

4

5

6

Weight spectrum

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

1

2

initial

1

2

trained

initial

trained

(e) Weight spectra with a large

learning rate.

0.10

0.05

0.00

0.05

0.10

1st PC for weight

0.05

0.00

0.05

0.10

Feature in techer model

initial

trained

(f) Weight alignments with a large

learning rate.

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.

C Proofs of Results in Section 4.1

C.1 GD Analysis at Early Phase

From (10), the GD process with learning rate η > 0 can be written by

for t ∈ N , where y ∈ R 1 × n . Following [6, Appendix B], in this section we prove the control for gradient step G t . For simplicity, denote f t ( X ) := f θ t ( X ) = 1 √ h v ⊤ σ ( W t X / √ d ) for t ∈ N .

0

10

20

30

40

50

60

70

Weight spectra

0.00

0.02

0.04

0.06

0.08

0.10

0.12

0.14

2.5

5.0

7.5

Initial

2.5

5.0

7.5

Trained

Q-Q plot

Initial

Trained

0

100

200

300

400

500

0.2

0.4

0.6

0.8

1.0

angles

average angle, k

principle angle

second angle

third angle

fourth angle

2

4

6

(a) Weight spectra

0.00

0.05

0.10

0.15

0.20

0.25

0

2

Initial

0

2

Trained

Q-Q plot

Initial

Trained

0

2

4

6

(b) CK spectra

0

1

2

3

0

5

Initial

0

5

Trained

Q-Q plot

Initial

Trained

0.0

2.5

5.0

7.5

10.0

12.5

(c) NTK spectra

0.0

0.5

1.0

1.5

2.0

0

10

Initial

0

10

Trained

Q-Q plot

Initial

Trained

for some constants C, c > 0 only depending on σ ε , λ σ , γ 1 , and γ 2 .

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

Here we use the fact that both W 0 and X are i.i.d. Gaussian random matrices. Then by Assumption 3.1, we conclude that we control σ ( W 0 X / √ d ) . Recall that Assumption 3.3 implies that y = f ∗ ( X ) + ε . Hence, by Lipschitz Gaussian concentration inequality [80, Theorem 5.2.2], each entry of f ∗ ( X ) has independent sub-Gaussian coordinates, whence we can get ∥ f ∗ ( X ) ∥ ≤ C √ n with probability at least 1 -2 ne -cn for some constants c, C > 0 . On the other hand, [ ε ] i = ε i are i.i.d. centered sub-Gaussian noises with variance σ 2 ε . By [80, Theorem 3.1.1], we have

where the constant K is the sub-Gaussian norm defined by K = max i ∥ ε i ∥ ψ 2 . Hence, combining all things together, we obtain the second inequality of this lemma.

0

5

10

15

20

(a) Weight spectra

0.0

0.2

0.4

0.6

0

2

Initial

0

20

Trained

Q-Q plot

Initial

Trained

0

10

20

30

40

(b) CK spectra

0.00

0.05

0.10

0.15

0.20

0.25

0

5

Initial

0

500

Trained

Q-Q plot

Initial

Trained

0

20

40

60

80

100

(c) NTK spectra

0.00

0.02

0.04

0.06

0.08

0.10

0

10

Initial

0

500

Trained

Q-Q plot

Initial

Trained

0.10

0.05

0.00

0.05

0.10

Signal in the teacher model

0.05

0.00

0.05

0.10

Top singular vector

Initial

Trained

(a) First PC in trained & initial

weights.

2

1

0

1

2

Training labels

0.1

0.0

0.1

0.2

0.3

Pc of CK

Initial

Trained

(b) First PC in trained & initial CKs.

2

1

0

1

2

Training labels

0.1

0.0

0.1

0.2

0.3

Pc of NTK

Initial

Trained

(C) First PC in trained & initial

NTKs.

0

5

10

15

(a) Weight spectra

0.0

0.2

0.4

0.6

0

2

Initial

0

10

Trained

Q-Q plot

Initial

Trained

0

5

10

15

20

25

30

(b) NTK spectra

0.0

0.1

0.2

0.3

0.4

0.5

0

10

Initial

0

200

Trained

Q-Q plot

Initial

Trained

0

5

10

15

20

25

30

(c) CK spectra

0.00

0.05

0.10

0.15

0.20

0

5

Initial

0

200

Trained

Q-Q plot

Initial

Trained

Lemma C.2. Under the assumptions of Lemma 4.1, given any fixed t ∈ N and learning rate η = Θ(1) , the weight matrix after t gradient steps W t defined in (18) satisfies

for some positive constants c, C > 0 only depending on t, η, σ ε , λ σ , γ 1 and γ 2 .

Proof. Denote σ ⊥ ( x ) = σ ( x ) -µ 1 x which is the nonlinear part of σ and µ 1 = E [ zσ ( z )] . Thus, E [ σ ⊥ ( z ) z ] = 0 for z ∼ N (0 , 1) . Based on this, we can further decompose the gradient G t into

0.10

0.05

0.00

0.05

Signal in the teacher model

0.10

0.05

0.00

0.05

0.10

Top singular vector

Initial

Trained

(a) First PC in trained weight.

2

1

0

1

2

Training labels

0.10

0.05

0.00

0.05

0.10

Pc of CK

Initial

Trained

(b) First PC in trained CK.

2

1

0

1

2

Training labels

0.10

0.05

0.00

0.05

0.10

Pc of NTK

Initial

Trained

(c) First PC in trained NTK.

0

2

4

6

8

10

(a) Weight spectra

0.0

0.2

0.4

0.6

0

2

Initial

0

2

Trained

Q-Q plot

Initial

Trained

0

5

10

15

20

(b) NTK spectra

0.0

0.1

0.2

0.3

0.4

0.5

0

10

Initial

0

20

Trained

Q-Q plot

Initial

Trained

0

5

10

15

(c) CK spectra

0.00

0.05

0.10

0.15

0.20

0

5

Initial

0

10

Trained

Q-Q plot

Initial

Trained

At first, consider t = 0 and bound the spectral norm of W 1 . By assumption, we know ∥ v ∥ ≤ √ h . Due to Corollary 7.3.3 in [80], we have

Therefore, by (21), we can control A 0 and B 0 separately. Notice that, as a rank-one matrix,

Hence, by Lemma C.1 and (22), one can easily claim that ∥ A 0 ∥ ≤ C/ √ n with probability at least 1 -e -cn for some constants c, C > 0 . On the other hand, since v ( y -f t ( X )) is rank-one and σ ′ ⊥ = σ ′ -µ 1 with | σ ′ ( x ) | ≤ λ σ , we can similarly obtain

As A 0 , we can apply Lemma C.1 and (22) again to conclude (20) for t = 1 .

For general t , we apply induction. We assume that after the t -th gradient step with η = Θ(1) , Eq. (20) holds for some constants C, c > 0 . Following [6, Lemma 16], we now show that the similar high-probability statement also holds for W t +1 (for some different constants c ′ , C ′ ). Firstly,

0.10

0.05

0.00

0.05

0.10

Signal in the teacher model

0.10

0.05

0.00

0.05

0.10

Top singular vector

Initial

Trained

(a) First PC in trained & initial

weights.

2

1

0

1

2

Training labels

0.075

0.050

0.025

0.000

0.025

0.050

0.075

Pc of CK

Initial

Trained

(b) First PC in trained & initial CKs.

2

1

0

1

2

Training labels

0.075

0.050

0.025

0.000

0.025

0.050

0.075

Pc of NTK

Initial

Trained

(c) First PC in trained & initial

NTKs.

following the same argument as [71, Setion 6.6.1], we know that

Note that ∥ W t -W 0 ∥ F = O (1 / √ n ) with high probability by the induction hypothesis. Hence, by LemmaC.1 and (22), we have ∥ f t ( X ) ∥ ≤ C √ n with high probability. Indeed, the difference between f t ( X ) and f 0 ( X ) is significantly negligible comparing with the initial value f 0 ( X ) . Similarly with A 0 , A t satisfies

Analogously for B t , we have

Thus, Lemma C.1, (22), and (23) ensure that

for constants c ′ , C ′ > 0 . Since ∥ W t +1 -W 0 ∥ F ≤ ∥ W t -W 0 ∥ F + η ∥ ∥ A t ∥ ∥ F + η ∥ ∥ B t ∥ ∥ F , by induction hypothesis, we can conclude that (20) 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 (17), we can also deduce the following norm bounds:

Lemma C.2 and the above bounds are empirically verified by Figure 28(a) for t = 3 . Not only upper bounds, this simulation also shows that at early phase ∥ W t -W 0 ∥ , ∥ W t -W 0 ∥ F , and ∥ W t -W 0 ∥ 2 , ∞ are all of the same Θ(1 / √ n ) order.

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

for some t -dependent constant C > 0 , any x ∼ N (0 , I ) and any finite time t . In other words, when η = Θ(1) , the change of the output of the NN at the early phase (i.e. t = Θ(1) ) is negligible and its order is O ( 1 √ n ) .

C.2 Proof of Lemma 4.1

In this section, we complete the proof of Lemma 4.1. We first mention the empirical validation of Lemma 4.1 in Figure 28. Notice that the changes in Frobenius norm for W and K CK are exactly Θ(1 / √ n ) and Θ(1 /n ) , respectively. The operator norm of K NTK matches with Lemma 4.1, while the Frobenius norm of the change decays slower than the rate Θ(1 /n ) . Additionally, in the simulation, we use v ∼ N (0 , I ) , which indicates that our assumption for v in Lemma 4.1 can be weakened.

10 3

2×10 3

3×10 3

Sample Dimension n

10

2

10

1

Weight differences

(1/

n

)

|

W

0

Ws

|2

|

W

0

Ws

|

F

|

W

0

Ws

|

2,

(a) Early phase for

W

t

.

10 3

2×10 3

3×10 3

Sample Dimension n

10

2

10

1

CK differences

(1/

n

)

|

CK

0

CKs

|2

|

CK

0

CKs

|

F

(b) Early phase for

K

CK

.

10 3

2×10 3

3×10 3

Sample Dimension n

10

2

10

1

NTK differences

(1/

n

)

|

NTK

0

NTKs

|2

|

NTK

0

NTKs

|

F

(c) Early phase for

K

NTK

.

Proof of Lemma 4.1. Lemma C.2 directly validates the control of 1 √ d ∥ W t -W 0 ∥ F . By virtue of this result, we now present estimates for CK and NTK. Based on [71, Section 6.6.1], we have

We apply the mean value theorem to obtain this inequality. Recall the operator norm bound for Gaussian random matrix X in (22). We know ∥ X / √ d ∥ ≲ 1 + √ γ 1 with high probability as n/d → γ 1 . Hence, with the help of Lemma C.2, we can claim

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

with probability at least 1 -exp( -cn ) . Therefore, we can control the change in the CK matrix in the Frobenius norm by the following inequalities:

Therefore, by (24), (25) and Lemma C.1, we can claim that there exist constants c, C > 0 such that with probability at least 1 -exp( -cn ) , ∥ ∥ ∥ K CK t -K CK 0 ∥ ∥ ∥ 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 exactly the CK, it suffices to consider the change of the first part of the empirical NTK. Recall that

Following the notation in [71], we denote J ( W t ) := [ J ( w t 1 ) , . . . , J ( w t N )] ∈ R n × hd with J ( w i ) := v i √ h diag ( σ ′ ( X ⊤ w i / √ d )) X ⊤ / √ d ∈ R n × d . Hence, K t = J ( W t ) J ( W t ) ⊤ and

By [71, Lemma 6.6], we know ∥J ( W 0 ) ∥ 2 = ∥ ∥ K NTK 0 ∥ ∥ is upper bounded by some constant C > 0 with high probability. Then, we apply the inequalities from Lemma 6.5 of [71] to obtain

where the last inequality is due to the mean value theorem, the uniform bound on σ ′′ , and the assumption on the second layer v . Notice that Gaussian random vectors satisfy

as n/d → γ 1 and h/d → γ 2 . Thus, with (22) and Lemma C.2, we obtain

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

Corollary C.3. For any fixed t ∈ N , i ∈ [ d ] and k ∈ [ n ] , denote λ t i , ν t k and µ t k the i -th, and k -th eigenvalues of 1 h W ⊤ t W t , K CK t and K NTK t , respectively. Then, under the assumptions of Lemma 4.1, we have

almost surely in LWR. Consequently, the eigenvalues of 1 h W ⊤ t W t , K CK t and K NTK t are the same as corresponding the eigenvalues of initial 1 h W ⊤ 0 W 0 , K CK 0 and K NTK 0 , respectively.

This corollary is a direct outcome of Weyl's inequality from Theorem A.46 in [8]. Consequently, this corollary concludes that for any fixed t ≥ 0 , almost surely, the limiting spectra of 1 h W ⊤ t W t , K CK t and K NTK t are the same as those of 1 h W ⊤ 0 W 0 , K CK 0 and K NTK 0 in LWR. This corollary claims that not only does 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 K NTK t has the same lower bound as K NTK 0 in the early phase of training.

C.3 Global Convergence for GD Under LWR

In this section, we study the final stage of (10) as training loss is approaching zero and prove Theorem 4.2. Figure 29 shows that the spectra are unchanged globally, even after training in this case. In Corollary 4.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 v ∼ N (0 , I ) , which is more general than our assumption on v in Theorem 4.2.

0

2

4

6

(a) Weight spectra

0.0

0.1

0.2

0.3

0.4

0

2

Initial

0

2

Trained

Q-Q plot

Initial

Trained

0

2

4

6

8

(b) CK spectra

0.0

0.5

1.0

1.5

2.0

2.5

0

5

Initial

0

5

Trained

Q-Q plot

Initial

Trained

0

5

10

15

(c) NTK spectra

0.00

0.25

0.50

0.75

1.00

1.25

1.50

0

10

Initial

0

10

Trained

Q-Q plot

Initial

Trained

Proof of Theorem 4.2. Recall the Jacobian matrix J ( W ) defined in the proof of Lemma 4.1, and the definition of α based on (11) in Section 3. Denote the event

By (22), (28) and Theorem 2.9 of [81], we have P ( A ) ≥ 1 -2 e -cn -2 ne -cn -n -7 / 3 for some constant c > 0 and all large n in LWR. In the following, conditionally on event A , we will apply Theorem 6.10 of [71] to obtain the global convergence. Conditionally on A , Lemma 6.6 of [71] implies

for any W . Define β = 2 λ σ (1 + √ γ 1 ) . Moreover, in terms of (27), we can verify the Lipschitz property for the Jacobian matrix as follows: conditionally on A ,

for any ˜ W , W ∈ R h × d . Therefore, conditionally on A , J ( W ) is a L -Lipschitz function with respect to W where L := 2 β √ h . To complete the proof, it suffices to investigate the smallest singular value of J ( W ) when W is in the vicinity of W 0 . Recall ℓ ( W ) = ∥ y -f W ( X ) ∥ . Notice that for any unit vector u ∈ R n , we have u ⊤ f W 0 ( X ) = 1 √ h ∑ h i =1 v i σ ( w ⊤ i X / √ d ) u , where w ⊤ i is the i -th row of W 0 for i ∈ [ N ] . Consider event B := {∥ ∥ ∥ σ ( W 0 X / √ d ) ∥ ∥ ∥ ≤ C √ n } for some universal constant C > 0 . Lemma C.1 proves P ( B ) ≥ 1 -2 e -cn . By the assumption of v , we know each entry v i is a sub-Gaussian random variable with a sub-Gaussian norm at most 1. Then, according to Hoeffding's inequality, conditionally on the event B , we have

for every t ≥ 0 and some constant c > 0 . Let t = 2 √ n . Considering an 1 4 -net N of the unit sphere S n -1 , we can get

for some constant c ′ > 0 . Hence, based on Lemma C.1 and (31), we can obtain ℓ ( W 0 ) ≤ C 0 √ n with high probability for some universal constant C 0 > 0 . Let us denote this event as C : { ℓ ( W 0 ) ≤ C 0 √ n } . Define R := 4 ℓ ( W 0 ) /α . For any W in a ball of radius R centered at W 0 , we have ∥ W 0 -W ∥ F ≤ R and ∥J ( W ) -J ( W 0 ) ∥ ≤ LR, conditionally on event A . Thus, by (30), on event A∩C , the smallest singular value σ min ( J ( W )) of the Jacobian matrix J ( W ) can be bounded by

for some universal constant C > 0 and sufficiently large n, d, h . Notice that here constants C, β , and α do not rely on γ 2 . Therefore, there exists a sufficiently large γ ∗ > 0 such that for all γ 2 ≥ γ ∗ , we have 2 α -8 Cβ α √ γ 1 γ 2 ≥ α . In other words, when h is sufficiently large but still in the same order as n and d , for all ∥ W -W 0 ∥ F ≤ R , we have σ min ( J ( W )) ≥ α conditionally on C ∩ A . Combining with (29) and (30), conditionally on C ∩ A , all the assumptions of Theorem 6.10 by [71] are satisfied when ∥ W -W 0 ∥ F ≤ R . Therefore, when the learning rate η n ≤ 1 β 2 min { 1 , 4 α LR } , we can get (12)-(14) for all t ∈ N , conditionally on C ∩ A . Both events A and C occur with high probability and only depend on initialization W 0 , X and y . Hence we complete the proof of this theorem. Notice that since γ 2 ≥ γ ∗ is sufficiently large, 4 α LR ≥ α 2 2 Cβ √ γ 2 γ 1 > 1 . Therefore, it suffices to require η ≤ n/β 2 to conclude that (12), (13) and (14) hold with high probability. This completes the proof. Moreover, (14) further shows that for all t ∈ N ,

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

for some constant C > 0 only depending on γ , γ , σ , σ and σ

d

1 2 ε ∗ .

-BffFpSEgE÷-t

(a) Changes for

1

√

W

.

(b) Changes for CK.

(c) Changes for NTK.

As a corollary, (14) controls the deviation of the final step weight from the initial weight. This is empirically shown in Figure 30(a), which shows that 1 √ d ∥ W t -W 0 ∥ , 1 √ d ∥ W t -W 0 ∥ F , and 1 √ d ∥ W t -W 0 ∥ 2 , ∞ are Θ(1) when trainable parameters are convergent. This implies that the final W t is still close to the initial weight W 0 , even after training. Consequently, with this observation, we can prove Corollary 4.3 in the following.

Proof of Corollary 4.3. Based on (32), we know 1 √ d ∥ W o -W t ∥ F ≤ 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 4.1. For CK, we employ Lemma C.1 and (25) to get

_fEEfEgfIfoI-

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

Let us denote w t i ∈ R 1 × d as the i -th row of W t , and x j as the j -th column of X . Additionally, by Assumption 3.2, we know that

for any x, y ∈ R . For NTK, by modifying (27), one can deduce that

where ( i ) is because of [80, Exercise 6.3.3] and the assumption on v , ( ii ) is due to (22), ( iii ) is due to the definition of Frobenius norm and (33), and ( iv ) is due to [80, Exercise 6.3.3] and (22). As a result, from (26), we can finally conclude that ∥ ∥ ∥ K CK t -K CK 0 ∥ ∥ ∥ F = O d, P (1) as n/d → γ 1 and h/d → γ 2 .

As for the limiting spectra of weight and kernel matrices, since we know that

we can automatically apply Corollary A.41 of [8]. This directly implies that the limiting empirical spectra of 1 h W ⊤ t W t , K CK t and K NTK t are the same as the limiting spectra of 1 h W ⊤ 0 W 0 , K CK 0 and K NTK 0 , respectively, as n/d → γ 1 and h/d → γ 2 (see Figure 29).

Further Simulations for Changes in Norms. The simulation of Figure 30 empirically coincides with the norm bounds in Theorem 4.2 for different norms. Because of (17), it suffices to only consider the Frobenius norm of the change for each matrix. As a remark, Theorem 4.2 requires γ 2 to be larger than some threshold γ ∗ to ensure norms of the change throughout training. However, Figure 30 indicates Theorem 4.2 still holds even when γ ∗ is small i.e. the width h is not very large. Here in Figure 30, γ 2 = 1 2 γ 1 = 0 . 6 . Figure 31 suggests that similar results to Theorem 4.2 and Corollary 4.3 hold for SGD when training the first layer. This is also akin to Figure 21. Moreover, we further conjecture that similar results to Theorem 4.2 and Corollary 4.3 will hold even when training both layers (3). Denote the second layer at step t by v t . Then, indicated by Figure 32, 1 √ d ∥ W t -W 0 ∥ , 1 √ d ∥ W t -W 0 ∥ F , 1 √ d ∥ W t -W 0 ∥ 2 , ∞ , ∥ ∥ K CK t -K CK 0 ∥ ∥ , ∥ ∥ K CK t -K CK 0 ∥ ∥ F , ∥ ∥ K NTK t -K NTK 0 ∥ ∥ , and ∥ v t -v 0 ∥ / ∥ v 0 ∥ are all Θ(1) in LWR. Since ∥ v 0 ∥ = O ( √ h ) , we observe that ∥ v t -v 0 ∥ = Θ( √ h ) in this case. Meanwhile, unlike Corollary 4.3, Figure 32(c) cannot verify that ∥ ∥ K NTK t -K NTK 0 ∥ ∥ F is still upper bounded by a constant. One possible explanation is that in this case the change in the second layer v t is much more significant than in the first-layer weight W t , hence the NTK matrix may change a lot in this general case.

Similarly, Figure 33 shows norms of the change when training the NN with SGD for both layers. Here we use the same batch size 128 and learning rate η = 1 . 0 for all experiments when varying the dimension n but fixing the aspect ratios γ 1 and γ 2 . All the observations are similar to the results of Corollary 4.3 except the Frobenius norm of the change for NTK. Figure 33(c) indicates that ∥ ∥ K NTK t -K NTK 0 ∥ ∥ F will not be Θ(1) anymore, which fluctuates more in this case.

2000

4000

6000

8000

10000

Sample Dimension n

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

|

W

0

Ws

|

|

W

0

Ws

|

F

|

W

0

Ws

|

2,

(a) Changes for

1

√

W

.

2000

4000

6000

8000

10000

Sample Dimension n

1.0

1.2

1.4

1.6

1.8

2.0

2.2

2.4

|

CK

0

CKs

|

|

CK

0

CKs

|

F

(b) Changes for CK.

2000

4000

6000

8000

10000

Sample Dimension n

0.000

0.025

0.050

0.075

0.100

0.125

0.150

0.175

|

NTK

0

NTKs

|

|

v

0

vs

|/|

v

0|

|

NTK

0

NTKs

|

Fro

(c) Changes for NTK.

d

Bakr

MB

c-

Meran

AMAR

dharmas onscreen

lmaooooo

M

☒

O

(a) Changes for

1

√

W

.

(b) Changes for CK.

(c) Changes for NTK.

d

A

.

(a) Changes for

1

√

W

.

(b) Changes for CK.

(c) Changes for NTK.

d

_MtMkMawMBmmMmBoB

_

☆

Sdgtf0TffIfI- oEfgfÉTMM÷kfffr_

tfgg÷Tfgg÷r

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies gradient descent training in single-hidden-layer neural networks in the linear-width regime (i.e., that in which the input dimension, hidden layer width, and number of training datapoints together tend proportionally to infinity). Its primary result is that the bulk spectra of the conjugate and tangent kernels do not change over training in this regime, but at large learning rates an outlier eigenvalue can emerge.\n\nSTRENGTHS:\nThe topic of how kernels evolve outside of the lazy regime is timely, and the RMT-inspired setup and results presented should be of interest to the community. The theoretical results appear correct, though I've not checked the proofs line-by-line.\n\nWEAKNESSES:\nThe manuscript presents a series of interesting observations, but to my taste it sacrifices depth for the sake of breadth. The results on heavy-tailed spectra are thought-provoking, but I think focusing on the emergence of spikes in the spectra would make for a stronger manuscript. If the authors could prove the existence of a BBP-like transition, they would have a strong paper. The current comment on Lines 247-248 comparing to [6] is not very satisfying. At the very least, a more detailed empirical investigation of this phenomenon would result in a more convincing story. \n\nIn this vein, the experiments are not very systematic, and as a result cannot clearly disentangle which changes in the setup are required to produce a given phenomenon. For instance, Table 1 contains various settings of batch size and learning rate, but nowhere are these parameters swept jointly over some reasonable range. I have in mind something like the parameters sweeps for learning single-index models performed in recent work by Atanasov et al., \"The Onset of Variance-Limited Behavior for Networks in the Lazy and Rich Regimes\" (ICLR 2023; note also that this relevant recent work is not cited). The goal here would be to establish under precisely what conditions a spike emerges. As it stands, the authors can make only vague statements. \n\nThe transition to using Adam for the heavy-tailed weight experiments is even more drastic: is it possible to observe this phenomenon using a non-adaptive optimizer? I think dissecting this phenomenon is more properly the subject of a separate paper, as it is currently not well-integrated with the other results. The discussion in Lines 286-291 is also rather heuristic, and I don't think the authors provided sufficient evidence for their claim that heavy-tailed spectra naturally emerge in feature-learning networks trained on complex tasks. Moreover, how does \"[t]his example explain[] why we can use the heavy tails to discriminate well-trained and poorly-trained large models?\" This requires elaboration, and more empirical evidence.\n\nQUESTIONS:\n- In the Introduction, it would be nice to mention recent work on Gaussian equivalence for deep random feature models in the linear-width regime by Schröder et al., \"Deterministic equivalent and error universality of deep random features learning\" (ICML 2023) and Bosch et al., \"Precise Asymptotic Analysis of Deep Random Feature Models\" (2023). \n\n- The linear-width regime for deep Bayesian neural networks has recently been the subject of study in the community of researchers using tools from statistical physics to study deep learning, see e.g., Li and Somplinsky, \"Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization\" (PRX 2021), Zavatone-Veth et al., \"Contrasting random and learned features in deep Bayesian linear regression (PRE 2022), and Cui et al., \"Optimal learning of deep random networks of extensive-width\" (ICML 2023). I think it would be useful to at least comment upon the Gaussian equivalence results of Cui et al. as a point of reference for how different inference procedures behave in the extensive-width regime. \n\n- In Figure 2 (b-d) and elsewhere, it would be better to label the abscissa with $\\eta$ or \"learning rate\" rather than \"lr,\" for the sake of consistency with the text.\n\n- In Figure 1 and elsewhere, the histograms should be plotted using a log-scaled ordinate. As it stands, the outlier in 1b is nearly invisible. \n\n- Why use the single-index + quadratic model in eq. 8 rather than a (simpler) single-index model? \n\n- The BERT experiments in Figure 5 strike me as mostly decorative, as the setting is not commensurate with the earlier portions of the paper. What insights would the authors argue can be transferred from their MLP experiments to this setting?\n\nLIMITATIONS:\nThe paper does not provide adequate discussion either of the relevance of the linear-width setting for practice or of the limitations of their experiments.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper examines the spectrum of the conjugate kernel and neural tangent kernel of real networks before and after training. Based on analyzing the spectrum in some experiments, several observations about how the spectum evolves during training are made. In particular, there seems to be different phases depending on the optimizer used and learning rate which effect how the spectrum changes from initialization. Understanding this spectrum-evolution phenomenon may lead to a better understanding of how neural networks train, and shed light on important questions like understanding differences between ADAM vs GD vs SGD.\n\nSTRENGTHS:\n* The review of existing literature is really well done (there is even a more detailed review in the appendices). This puts the current work in context quite nicely and makes this paper a very accessible place to start learning about these spectral questions.\n* The experimental plots are nicely done and clearly showcase the behaviour being discussed (the Q-Q plots also help a lot with the understanding). There are also numerous other experiments in the appendix. This again makes this paper a great resource for these kinds of experiments. \n* I am not an expert on this area, but to me the way the different phases of spectral evolution are laid out (i.e. what they call \"Invariant Bluk\" vs \"Bulk+Spike\" vs \"Heavy tail\") is new to me. (I have heard of this outlier thing happening before for the Hessian...I am not sure how related that is). To me this paper looks like the first step towards building a phase diagram of how choices about optimizer/trianing effect the spectrum which would be a really useful thing to have.\n* I found the conclusion about learning rate (e.g. Figure 2) the most impressive conclusion here, and to me this may be the most actionable of all the various phenomenon observed here. This is a nice conclusion to draw and could pave the way for more results in this direction.\n\nWEAKNESSES:\n* One might complain that the theory developed here is a bit weak, in that the theoretical results are all inequalities which are (probably) not sharp. However, since the paper is mainly empirical I actually think these theoretical results are reasonable here.\n* I am not an expert on this area, so I was not able to tell how much of the results here were expected/known already. Hopefully at least one other reviewer is able to give that context. \n* It is not clear how robust the experimental results are: in the given plots the phenomenon seem to be happening as described, but maybe there could be some extra discussion about \"edge cases\" between the various phases and discussion about when things break down (e.g. when N is too small, or there are not enough examples etc)\n\nQUESTIONS:\n* Is there any possibility to describe how the outlier eigenvalues are evolving in time during training? Particularly when there is only one outlier eigenvalue, this would be an interesing bit of theory to try and understand.\n* Is there any relationship between the spectrum of the CK/NTK and the spectrum of the Hessian?\n\nLIMITATIONS:\n* One limitation that always comes up in these limiting type things is the question of how large real neural networks have to be for the theory to actually work. Some discussion or experiments specifically addressing this (e.g. showing the error in the predictions as a function of network size) could help explicitly address this.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper carries out an empirical study of the spectral properties of finite-width neural network kernels after training, and compares them to their infinite-width counterparts to assess if there is an improvement. They study these changes with respect to various hyperparameters such as learning rate and optimizer. The results suggest that for large learning rates, where the linear dynamics is not a valid approximation, they show that feature learning emerges by inspecting the distribution of the final spectra. They identify two different spectral property that distinguish feature learning from lazy regime: emergence of low-rank spikes in the spectrum and heavy-tailness. They explore when these properties emerge and provide theoretical arguments for these observations.\n\nSTRENGTHS:\nThe paper is well-written and easy to follow. The experiments are complimented well with theoretical arguments, and the discussion is sound. The Supplementary Material also includes nicely organized discussion of previous works and additional experiments and discussion.\n\nIn summary, they introduce the well-known fact that the small-learning rates make a NN with NTK parameterization stuck around its initialization and derive theoretical results to show that the bulk of the spectrum remain invariant under lazy training. In my opinion, the main contribution is that they find two different ways feature learning can happen: emergence of low-rank outliers in the spectrum with large learning rates with GD optimizers and heavy-tailed spectrum with adaptive optimizers.\n\nFurthermore, for GD/SGD, they make a theoretical connection for how low-rank outliers emerge with a phase transition in RMT (ref. [9] in main paper.)\n\nFinally, the relation of heavy-tailed spectrum to generalization has been discussed.\n\nWEAKNESSES:\n1. While the paper brings valuable insights, it is highly restricted by the 2-layer neural networks and synthetic dataset with a particular target. It would be useful to perform experiments on simple real datasets such as MNIST to show that spikes related to number of classes also emerge in the spectrum. For example, in neural collapse phenomenon, one can see that weight matrices of last layer acquire low rank structure that can be tracked to the number of classes.\n\n2. Such experiments can also shed light on how full-batch vs. stochastic GD differ in terms of the after-training spectrum. Currently, I do not think there are any insights related to how SGD helps feature learning since GD experiment does not have a large LR analogue.\n\n3. Is there an intuitive explanation how adaptive optimizers cause heavy-tailed spectrum in CK? More discussion around this phenomenon could be very helpful.\n\nQUESTIONS:\n1. It is hard to parse Figure 1, first panel shows weight spectrum for GD with LR 5, second panel shows CK spectrum for SGD with LR 22 and last one shows CK spectrum for ADAM. Something like CK spectrum for Case 2, 3, 4 would be more helpful.\n\n2. In Figure 4, is there a reason why GD does not show the same properties as SGD? With large enough LR, could GD also show the same heavy-tailness? Also, it is not clear how adaptive optimizer qualitatively differs from SGD.\n\nLIMITATIONS:\nA limitations section is missing.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

comments: SUMMARY:
This paper analyzes the spectral properties of feedforward neural networks under a student-teacher setting. A linear-width regime is considered, where the sample size and network width grow comparably with the input feature dimension. The paper shows that the spectra of weight and kernel matrices are invariant under training, and that outliers occur in large step size regimes. Heavy-tailed spectra are also discussed with their relation to generalization. Evidence from both synthetic and real-world datasets is shown.

STRENGTHS:
- Proposed LWR as a novel perspective for analyzing feedforward networks, and presented an interesting analysis of NN spectral invariance, potentially providing insights into NN feature learning.
- Evidence from both synthetic and real-world datasets enhance the argument.
- The presentation of the paper is clear.

WEAKNESSES:
- The proposed LWR feels under-explained. Please see "Questions". Otherwise I think this is a solid papar.

QUESTIONS:
- In the proposed LWR, how should one understand that the input feature dimension $d$ goes to infinity? This feels somewhat counterintuitive, compared with e.g. the infinite-width regime.
- To what extent does the observed spectral invariance depend on the LWR? I would assume this does not occur under the infinite-width regime.

LIMITATIONS:
n/a, theoretical work
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nConsiderable attention in deep learning theory has been given to the analysis of neural networks in the kernel regime, which neural networks approach as they approach infinite width. While this analysis has led to insights, there is an important gap between the theorized generalization performance of neural networks in this infinite width regime and the practical performance of finite-width neural networks, which typically perform better. \n\nThis paper studies simple, two-layer neural networks in the more realistic *linear-width regime* (LWR), in which sample size, input feature dimension, and the width of the layer approach infinity at comparable rates. In several analyses, the authors target the following question: *How do the spectra of the NN's weight and kernel matrices evolve during the training process?* The authors compare to the kernel regime as a benchmark and analyze changes in spectra that occur through training with different gradient descent algorithms . \n\nFor their first simple experiment (Section 3), they find that the spectra of the weight, conjugate kernel (CK) and neural tangent kernel (NTK) matrices remain invariant during training for gradient descent (GD) and stochastic gradient descent (SGD) with a low learning rate. This is consistent with the Lazy Training (LT), and indicates that the network is not outperforming a kernel machine. However, for SGD with a large learning rate, a spike in the spectral distribution emerges after training, which indicates improvement over lazy training. They note that this spike is consistent with previously published results that demonstrate that learned features may be indicated by an outlier in the spectra with a large singular value.\n\nThey classify the spectral distributions into three categories: invariant bulk, spikes outside the bulk, and heavy-tailed distributions. They prove that the invariant bulk phenomenon is expected under certain assumptions. They then empirically analyze the spike phenomenon, and demonstrate that as the learning rate increases, there is a threshold at which the spike emerges, which corresponds with greater alignment and suggests feature learning. They then provide an explanation for the relationship between heavy tailed spectra and better generalization performance. Lastly, they perform some analyses on CNNs trained on more a more natural dataset, CIFAR-2, as well as BERT with fine-tuning on Sentiment140.\n\nSTRENGTHS:\nThe authors tackle an important question, which is to understand how the structure of neural networks change through the training process. To get a handle on it, they analyze the spectra of the weight and kernel matrices of simple neural networks on controlled datasets. They find some interesting phenomena that relate the spectral properties to learning rate and generalization performance.\n\nWEAKNESSES:\nThe paper would benefit from more effort put towards making clear the meaning and significance of the experiments and results. Why look at the spectra? What is the significance of the three classes of spectral distributions that are observed? (The answer to the latter question is in the paper, but it is somewhat buried. This should be put front and center and made very clear). \n\nThe figures and captions in the paper could be clearer. Oftentimes, variables are used as labels. The reader forgets which variables correspond to what. It would be better to use English names, or more descriptive properties. \n\nThe specific primary findings should be stated more clearly. The statements provided in the introduction are general, and don't provide much insight into the conclusions. It takes work for the reader to tease out these conclusions.\n\nQUESTIONS:\n- What can we understand about neural networks, generally speaking, by looking at the spectra of their weight matrices and kernels? Why use this approach, rather than analyzing other properties of the weights and network structure?\n- How do the results found here advance our theoretical understanding of how deep networks learn?\n- How can these results be used in practice?\n\nLIMITATIONS:\nThe authors note the limitation of the FWR. They do not state limitations of the spectral approach to analyzing networks.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}