SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $a_{\sigma} \coloneqq \mathbb{E}[\sigma'(\xi)^2]$  and  $\eta_k$  is the  $k$ -th Hermite coefficient of  $\sigma'$ . Hence, we can claim there exists some constant  $\alpha > 0$  only dependent on  $\sigma$  such that  $\lambda_{\min}(\pmb{K}_0^{\mathrm{NTK}}) \geq 4\alpha^2$  with high probability. Note that  $\alpha$  is not vanishing since  $\sigma$  is nonlinear. With this lower bound, we obtain the following global convergence for (10) and norm control of  $\pmb{W}_t$  as  $n / d \to \gamma_1$  and  $h / d \to \gamma_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  $\eta < \min \left\{\frac{\alpha^{2}n}{2}, \frac{n}{4\lambda_{\sigma}^{2}(1 + \sqrt{\gamma_{1}})^{2}}\right\}$  and all  $t \in \mathbb{N}$ , there exists some  $\gamma^{*} > 0$  such that, when  $\gamma_{2} \geq \gamma^{*}$ , the gradient steps (10) will satisfy
$$
\ell \left(\boldsymbol {W} _ {t}\right) \leq \left(1 - \frac {\eta \alpha^ {2}}{2 n}\right) ^ {t} \ell \left(\boldsymbol {W} _ {0}\right), \tag {12}
$$
$$
\frac {1}{4} \alpha \| \boldsymbol {W} _ {0} - \boldsymbol {W} _ {t} \| _ {F} + \ell (\boldsymbol {W} _ {t}) \leq \ell (\boldsymbol {W} _ {0}), \tag {13}
$$
$$
\sum_ {t = 0} ^ {\infty} \left\| \boldsymbol {W} _ {t + 1} - \boldsymbol {W} _ {t} \right\| _ {F} \leq \frac {4 \ell \left(\boldsymbol {W} _ {0}\right)}{\alpha}, \tag {14}
$$
with high probability, as  $n / d\to \gamma_{1}$  and  $h / d\rightarrow \gamma_{2}$ . Here, training loss  $\ell (\pmb {W})\coloneqq \| \pmb {y} - f_{\pmb{W}}(\pmb {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  $\Theta(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 \in \mathbb{N}$ , with high probability, there exists some constant  $R > 0$  such that the changes  $\frac{1}{\sqrt{d}}\| \pmb{W}_t - \pmb{W}_0\|_F$ ,  $\| \pmb{K}_t^{CK} - \pmb{K}_0^{CK}\|_F$ , and  $\left\| \pmb{K}_t^{NTK} - \pmb{K}_0^{NTK}\right\|_F$  are all less than  $R$  with high probability. This implies the limiting empirical spectra of  $\frac{1}{h}\pmb{W}_t^\top \pmb{W}_t$ ,  $\pmb{K}_t^{CK}$  and  $\pmb{K}_t^{NTK}$  are the same as the limiting spectra of  $\frac{1}{h}\pmb{W}_0^\top \pmb{W}_0$ ,  $\pmb{K}_0^{CK}$  and  $\pmb{K}_0^{NTK}$  respectively, almost surely as  $n / d \to \gamma_1$  and  $h / d \to \gamma_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(\sqrt{h})$  in (13) and (14), which is negligible compared with the Frobenius norm of initial weight matrix (which is of order  $\Theta(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^{\mathrm{CK}}$  and  $K^{\mathrm{NTK}}$  respectively through different training processes.
(a)  $\beta$  vs.  $\pmb{u}_{1}$  of weight.
(b)  $\lambda_{\max}(\pmb{W}_t\pmb{W}_t^\top /d)$
(c)  $\lambda_{\max}(\pmb{K}_t^{\mathrm{CK}})$
Figure 2: (a) Alignment between teacher feature  $\beta$  and first PC of the trained/initial weights in Case 3 of Table 1. (b)-(d) Transitions of  $\lambda_{\mathrm{max}}(\pmb{W}_t\pmb{W}_t^\top / d)$ ,  $\lambda_{\mathrm{max}}(\pmb{K}_t^{\mathrm{CK}})$  and alignments  $(|\beta^\top \pmb{u}_1| / \| \beta\|$  and  $|\pmb{y}^\top \pmb{v}_1| / \| \pmb{y}\|$  where  $\pmb{u}_1$  and  $\pmb{v}_1$  are the first singular vectors of  $\pmb{W}_t$  and either  $\pmb{K}_t^{\mathrm{CK}}$  or  $\pmb{K}_t^{\mathrm{NTK}}$ , respectively) when increasing the learning rate  $\eta$  while training the NN with SGD. In the green region, the largest eigenvalues are attached to the bulk (black horizontal lines) and the alignments are weak; in the orange one, outliers become apparent and the alignments become stronger. For different  $\eta$ , we train the same NN with the same dataset until the training loss is less than  $10^{-5}$ . Here, "lr" in the  $x$ -axis represents varying learning rates.
(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  $\pmb{u}_{1}$  of  $W_{t}$  and the signal  $\beta$  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  $\beta$  during training. Although this does not ensure NN will entirely beat the optimal kernel lower bound, this alignment reveals a non-negligible feature selection [10] via large-stepsize training. This dynamical alignment along the task-relevant direction may further interpret the generalization of the NN. We also observe similar phenomena for the adaptive optimization in Figure 27 in Appendix B.7.
(a) Evolution of KTA.
Figure 3: (a) Evolution of KTA of CK defined by (15) with respect to training labels for Cases 1, 3&4 in Table 1. We normalize the epoch scales ( $x$ -axis) for better observations. Heavy-tailed phenomena: (b) Evolutions of PC angles  $\theta_{i}$  between feature subspace  $U$  of (8) and top 100 eigenspace of  $\boldsymbol{W}_{t}^{\top}\boldsymbol{W}_{t}$  during training with Adam (solid line), SGD (dashed line) and GD (dash-dot). For the first PC  $\theta_{1}$ , see Figure 17 in Appendix B.4. (c) The CK spectra at two initializations for  $\boldsymbol{W}$ : standard Gaussian and Cauchy distributions. (d) Weight spectra at initial and after SGD training. After training the weight reveals a heavy tail, but generalizes not as well as former examples (test loss 1.47504;  $R^2$  score  $-0.48$ ).
(b) Alignment evolution.
(c) Heavy-tailed initial CK.
(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  $\eta$ . 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  $\eta$  exceeds this threshold. We fix the same NN and dataset for all trials of training. The flat black lines in Figures 2(b) and (c) are the right edges of the limiting spectra at initialization. Figure 2(d) records the angles between  $\beta$  and the leading eigenvector of  $\boldsymbol{W}_t^\top \boldsymbol{W}_t / d$ , and  $\boldsymbol{y}$  and the leading eigenvectors of  $\boldsymbol{K}_t^{\mathrm{CK}}$  and  $\boldsymbol{K}_t^{\mathrm{NTK}}$  after training for different  $\eta$ . Similarly with [10], when  $\eta$  is sufficiently large (orange region), we obtain significant alignments which suggest potential feature learning. These transitions of leading eigenvalue and eigenvector alignment have been proved for  $\boldsymbol{W}_t$  by [6] for a different scenario<sup>4</sup>.
Spikes of Kernel Matrices. The alignment of the kernel matrix with the training labels  $\pmb{y}$  is defined by [22] by Kernel Target Alignment (KTA) as follows: when kernel  $\pmb{K}$  is either CK or NTK,
$$
\mathrm {K T A} = \frac {\left\langle \boldsymbol {K} , \boldsymbol {y} ^ {\top} \boldsymbol {y} \right\rangle}{\left\| \boldsymbol {K} \right\| _ {F} \left\| \boldsymbol {y} \right\| ^ {2}}. \tag {15}
$$
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  $\pmb{v}_{1}$  of kernel matrix naturally dominates the alignment with  $\pmb{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  $\pmb{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, heavy-tailed 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.
<sup>4</sup>We apply NTK parameterization for our neural networks and train both layers until convergence, while [6] considers the mean-field initialization and early stage of training dynamics of GD for the first layer.
(a) GD
(b) SGD
Figure 4: Different NTK spectra for a small-CNN model on CIFAR-2. The subplots are Q-Q plots for the comparison between initial and trained spectra. Test accuracies: (a)  $79\%$ , (b)  $84\%$ , (c)  $86.4\%$ .
(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  $\beta_{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  $\theta_{i}$  between feature subspace  $U = \mathrm{span}\{\beta_{i}\}_{i=1}^{k}$  and top 100 eigenspace of  $W_{t}^{\top}W_{t}$  during different training processes. This eigenspace with respect to the top 100 eigenvalues of  $W_{t}^{\top}W_{t}$  corresponds to the heavy-tail part of the spectrum in  $W_{t}^{\top}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^{\mathrm{CK}}$ , and  $K^{\mathrm{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^{\mathrm{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<sup>5</sup> 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).
(a) Eigenvalues
(b) Alignments
Figure 5: We use SGD for fine-tuning the BERT model. The training accuracy is  $95.90\%$  and the test accuracy is  $84\%$ . (a) The evolution of first and second eigenvalues of empirical CK during fine-tuning. (b) The alignments of training labels with first and second eigenvectors of CK during fine-tuning. See Figure 7 for the spectra of CK at different epochs.
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 7 in 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.
Z.W., A.E., I.D., and T.C. were partially supported by the Mathematics for Artificial Reasoning in Science (MARS) initiative via the Laboratory Directed Research and Development (LDRD) Program at Pacific Northwest National Laboratory (PNNL). A.S. and T.C. were also partially supported by the Statistical Inference Generates kKnowledge for Artificial Learners (SIGNAL) program at PNNL. PNNL is a multi-program national laboratory operated for the U.S. Department of Energy (DOE) by Battelle Memorial Institute under Contract No. DE-AC05-76RL0-1830. Z.W. would like to thank Denny Wu and Libin Zhu for their valuable suggestions and comments.
<sup>5</sup>https://www.kaggle.com/datasets/kazanova/sentiment140
[54] 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.
[55] Philip M Long. Properties of the after kernel. arXiv preprint arXiv:2105.10585, 2021.
[56] 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.
[57] Cosme Louart, Zhenyu Liao, and Romain Couillet. A random matrix approach to neural networks. The Annals of Applied Probability, 28(2):1190-1248, 2018.
[58] 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.
[59] 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.
[60] 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.
[61] 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.
[62] Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
[63] 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.
[64] Preetum Nakkiran. Learning rate annealing can provably help generalization, even for convex problems. OPT2020 Workshop, 2020.
[65] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 1995.
[66] 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.
[67] Guillermo Ortiz-Jiménez, Apostolos Modas, Seyed-Mohsen Moosavi, and Pascal Frossard. Neural anisotropy directions. Advances in Neural Information Processing Systems, 33:17896-17906, 2020.
[68] 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.
[69] 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.
[70] 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.
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 Marcenko-Pastur distribution [30], which becomes a deformed semicircular distribution [81] when  $h \gg 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 \gg 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 \geq 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.