SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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čenko-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].
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  $(\mathbb{E}[\sigma (z)] = 0$  for  $z\sim \mathcal{N}(0,1))$  because we can exclude a large but trivial spike in the initial spectra of kernel matrices. In all architectures of NNs we considered, we remove the bias term of each layer and apply the NTK parameterization (1) with standard Gaussian initialization. Specifically, all entries of  $\pmb{W}$  and  $\pmb{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  $\pmb{X}$ . In addition, we consider the training label noise defined in Assumption 3.3 as  $\varepsilon \sim \mathcal{N}(0,\sigma_{\varepsilon}^{2}\mathbf{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 \times 10^{-3}$ ; we used SGD with learning rate  $10^{-3}$ , batch size 32 and momentum 0.2 in Figure 6(b); Figure 6(c) employs the same learning rate and batch size as 6(b) but employs Adam optimization.
(b) SGD
Figure 6: Training/test accuracy v.s. epochs for small-CNN model on CIFAR-2 with different optimizers.
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 fine-tuning to sentiment analysis for social media data. For fine-tuning, the learning rate is 0.003, the batch size is 64 and the momentum is 0.8. The purpose of this experiment is to extract the spectral properties of pre-trained models and the evolution of the CK spectra over fine-tuning. Combining Figure 5, the following Figure 7 exhibits the evolution of 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].
Figure 7: The spectra of CK of the BERT model on Sentiment140 dataset at epoch 1, 2 and 12.
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:
$$
\frac {1}{\sqrt {d}} \left\| \boldsymbol {W} _ {0} - \boldsymbol {W} _ {t} \right\| _ {F}, \frac {1}{\sqrt {d}} \left\| \boldsymbol {W} _ {0} - \boldsymbol {W} _ {t} \right\|, \left\| \boldsymbol {K} _ {0} ^ {\mathrm {N T K}} - \boldsymbol {K} _ {t} ^ {\mathrm {N T K}} \right\| _ {F}, \text {a n d} \left\| \boldsymbol {K} _ {0} ^ {\mathrm {N T K}} - \boldsymbol {K} _ {t} ^ {\mathrm {N T K}} \right\|,
$$
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
$$
\frac {1}{\sqrt {d}} \left\| \boldsymbol {W} _ {0} - \boldsymbol {W} _ {s} \right\| _ {F}, \frac {1}{\sqrt {d}} \left\| \boldsymbol {W} _ {0} - \boldsymbol {W} _ {s} \right\|, \frac {1}{\sqrt {d}} \left\| \boldsymbol {W} _ {0} - \boldsymbol {W} _ {s} \right\| _ {2, \infty} = \Theta (1) \tag {16}
$$
as  $n\to \infty$  with  $n / d\rightarrow \gamma_{1}$  and  $N / d\rightarrow \gamma_{2}$ , where  $s$  is the final time for GD. Here, the entry-wise  $2 - \infty$  matrix norm is defined as
$$
\| M \| _ {2, \infty} := \max  _ {1 \leq i \leq N} \| \mathbf {m} _ {i} \|,
$$
for any matrix  $M \in \mathbb{R}^{N \times d}$  with the  $i$ -th row  $\mathbf{m}_i \in \mathbb{R}^d$  and  $1 \leq i \leq N$ . Notice that
$$
\| M \| _ {2, \infty} \leq \| M \| \leq \| M \| _ {F}. \tag {17}
$$
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  $\frac{1}{\sqrt{d}} \| W_0 - W_s \|_F$  and  $\frac{1}{\sqrt{d}} \| W_0 - W_s \|$  are  $\Theta(1)$  may suggest that  $\frac{1}{\sqrt{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 results with different  $\sigma$  and  $\sigma^*$  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 \to \infty$ . Each time, we take 10 trials to average. We observe that SGD with a large learning rate (green line) can asymptotically outperform lazy training.
(c) Case 4
Figure 8: The evolution of the changes in operator/Frobenius norms of the weight/CK/NTK matrices through different training processes. Each case corresponds to the case in Table 1. Case 3 is exhibited in Figure 14(c) below.
(a) Train/test losses
Figure 9: Performances of NNs and changes in different norms for weight and kernels, when  $d / n = 0.6$  and  $N / n = 1.2$  are fixed as  $n$  is growing. The activation is normalized tanh and the teacher model is  $f^{*}(\pmb {x}) = \sigma^{*}(\pmb{\beta}^{\top}\pmb {x})$  where  $\sigma^{*}$  is a normalized softplus. We average over 10 trials in each case. All these curves are almost flat, which indicates these values are not growing with  $\gamma_{1}$  and  $\gamma_{2}$ . Here, in the second figure from the left, we normalized all weights  $\mathbf{W}$  with  $\frac{1}{\sqrt{d}}$  to observe (16).
(b) Change in weight
(c) Change in CK
(d) Change in NTK
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^{\mathrm{CK}}$  and  $K^{\mathrm{NTK}}$  are not changing over training as  $n / d \to \gamma_1$  and  $N / d \to \gamma_2$ . The initial spectrum of weight  $W_0$  converges to Marčenko-Pastur law; the initial spectrum of NTK under proportional limit has been studied by [1, 30]. Figure 12(c) demonstrates the global convergence for GD under 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^{\mathrm{CK}}$  and  $K^{\mathrm{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
(a) Change in weight.
(b) Change in CK.
(c) Prediction risks vs.  $\gamma_{1}$ .
Figure 11: Alignment with leading PCs of trained weight, CK and NTK matrices in Case 1 of Table 1.
Figure 10: (a) Change between initial  $\mathbf{W}_0$  and final step  $\mathbf{W}_s$  in operator norm, Frobenius norm and  $(2,\infty)$ -norm when  $d / n = 0.5$ ,  $N / n = 0.8$  are fixed as  $n\to \infty$ . We train NNs by SGD with  $\eta = 2.5$  for 15 trials to average. (b) Change of  $K^{\mathrm{CK}}$  in operator norm and Frobenius norm. (c) Prediction risks for lazy training defined by (6), GD with  $\eta = \Theta (1)$  (red), SGD with  $\eta = \Theta (1)$  (blue dot) and  $\eta \propto \gamma_{1}$  (green), as  $\gamma_{1}\rightarrow \infty$  and  $\gamma_{2} = 2.5$ . The black dashed line stands for the kernel lower bound given by the nonlinear part of the teacher model.
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  $\rho(x) \sim x^{-\alpha}$  for large  $x$  and some positive constant  $\alpha$ . The spectrum with a heavier tail has a smaller value of  $\alpha$ . Figure 17(a) shows how  $\alpha$  evolves through training in Case 4 of Table 1. As  $\alpha$  decreases, a heavy tail in the spectrum of the weight matrix emerges in Figure 1(c). In Figure 17(a) we introduce two more metrics to show this evolution: Weighted Alpha  $\hat{\alpha} \coloneqq \alpha \lambda_1$  and Log  $\alpha$ -norm  $\log \left( \sum_{i=1}^{N} \lambda_i^\alpha \right)$ .
where  $\lambda_N \leq \ldots \leq \lambda_1$  are the eigenvalues of  $\mathbf{W}\mathbf{W}^\top$ . Remarkably, Figure 17(a) indicates the spectral 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  $\eta = 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
Figure 12: Performances of Case 1 in Table 1: (a) The initial and trained spectra of the first-hidden layer  $\mathbf{W}$ . (c) The initial and trained spectra of empirical NTK matrix defined by (5). Q-Q subplot shows these two spectra are almost the same. Training and test losses vs. epochs for GD (right).
Figure 13: Spectral properties for Case 2 in Table 1: (a) The initial and trained spectra of the first-hidden layer  $W$ . (b) The initial and trained spectra of empirical NTK are defined by (5). (c) The initial and trained spectra of empirical CK defined by (4).
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  $\sigma = \mathrm{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  $\sigma^{*} = \sigma$ . We trained this student-teacher model using GD ( $\eta = 15$ ), SGD ( $\eta = 7.25$  and batch size 8), and Adam ( $\eta = 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  $\mathbf{W}$  of NN should gradually align with the feature space  $U$  spanned by  $\beta_{i}$ 's. Hence, to study feature learning, we can apply principle angles to measure the alignment between  $\mathbf{W}$  and  $U$ . Consider the eigen-decomposition of  $\mathbf{W}_{t}^{\top}\mathbf{W}_{t} = \sum_{i=1}^{d}\lambda_{i}\mathbf{v}_{i}\mathbf{v}_{i}^{\top}$  with  $\lambda_{1} \geq \lambda_{2} \geq \ldots \geq \lambda_{d}$ . Figure 3(b) shows the heavy-tailed part (the eigenspace  $E := \operatorname{span}\{\pmb{v}_{i}\}_{i=1}^{100}$ ) 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  $\| \mathrm{P}_{>1}f^{*}\|^{2}$  and  $\| \mathrm{P}_{>2}f^{*}\|^{2}$ , where  $\mathrm{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  $\beta_{i}$  for  $1 \leq i \leq 5$ . This indicates that heavy-tailed spectra with large spikes may have a correlation with feature learning and good generalizations.
(a) Spectra of weights.
(b) Spectra of the first part in NTKs.
(c) Norms of change vs. epochs.
Figure 14: Additional performance for Case 3 in Table 1: (a) The initial and trained weight spectra. Notice that there is one outlier after training, while the bulk remains invariant. This is analogous to the behavior of CK spectra in Figure 1(b). (b) The spectra of the first part in (5) at initialization and after training. The orange arrow points out the outlier of the spectrum. (c) The changes  $\| \pmb{W}_t - \pmb{W}_0\|$ ,  $\| \pmb{W}_t - \pmb{W}_0\|_F$ ,  $\| \pmb{K}_t^{\mathrm{NTK}} - \pmb{K}_0^{\mathrm{NTK}}\|$  and  $\left\| \pmb{K}_t^{\mathrm{NTK}} - \pmb{K}_0^{\mathrm{NTK}}\right\|_F$  at each epoch  $t$  throughout the training process.
(a) Losses vs. epochs.
Figure 15: (a) Training/test losses and  $R^2$  scores at each epoch of training in Case 3. (b) Alignment between training labels  $y$  and first PC of trained/initial CK. (c) Alignment between training labels  $y$  and the first PC of trained/initial NTK. We use the same setting in Case 3 of Table 1. There are strong alignments between kernels and training labels as stated in Section 4.2.
(b)  $\pmb{y}$  vs.  $\pmb{v}_1$  of CK.
(b)  $\pmb{y}$  vs.  $\pmb{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  $\frac{1}{h} W_t^\top W_t$ ,  $K_t^{\mathrm{CK}}$  and  $K_t^{\mathrm{NTK}}$  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 step sizes for a few steps and then use small-stepsize SGD for convergence. In this scenario, we
(a) Spectra of weights.
(d) First PC in trained & initial weights.
(b) Spectra of NTKs.
(e) First PC in trained & initial CKs.
(c) Losses vs. epochs.
(f) First PC in trained & initial NTKs.