SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {I} _ {t} = \left\{\boldsymbol {\theta} ^ {l} \mid \lambda_ {\max } \left(\boldsymbol {\Theta} _ {t} ^ {l}\right) \geq \lambda_ {\alpha} \right\}, \mathcal {N} _ {t} = \left\{\boldsymbol {\theta} ^ {l} \mid \lambda_ {\max } \left(\boldsymbol {\Theta} _ {t} ^ {l}\right) <   \lambda_ {\alpha} \right\}. \tag {5}
$$
As a result, the computation cost of back propagation is reduced.
We empirically observe that  $\alpha \in (0,1)$  is relatively stable over the training process for a certain model and dataset. Therefore,  $\alpha$  can be empirically assigned at the early stage of training after warming up, or automatically learned by a meta network.
Temporal Policy. The previous empirical analysis shows that during training, after the eigenvalues of a module's mNTK enter the nuisance region, further training this module can lead to poor generalization. Therefore, to avoid overfitting, we terminate the training of this module according to the temporal variation of the principal eigenvalue; we call this the temporal policy.
For each module  $\theta^l$ , its principal eigenvalue after warming up is recorded as the initial value  $\lambda_{\max}(\Theta_0^l)$ . During each training episode  $t$ , the temporal variation of principal eigenvalue is defined as  $\Delta_t^l = |\lambda_{\max}(\Theta_t^l) - \lambda_{\max}(\Theta_0^l)|$ . The information modules will be early stopped if the difference in the temporal variation between two subsequent episodes relative to the initial variation is smaller
than a pre-defined threshold  $\beta$ :
$$
\frac {\left| \Delta_ {t} ^ {l} - \Delta_ {t - 1} ^ {l} \right|}{\Delta_ {1} ^ {l}} <   \beta . \tag {6}
$$
$\beta$  is empirically bounded by eigenvalue variation within the nuisance region.
Optimization Analysis using NTK. Neural Tangent Kernel (NTK) (Jacot et al., 2018), which calculates the Gram matrix of Jacobian, is known as a powerful tool to analyze convergence and generalization properties (Arora et al., 2019). Modern neural networks are generally over-parameterized with positive definite NTKs, meaning that theoretically the gradient flow always converges to the equilibrium where the training loss is zero (Chizat & Bach, 2018; Bombari et al., 2022). A line of works study the training dynamics of over-parameterized models (Li et al., 2020; Liu et al., 2022) based on NTK, mainly under the assumption of two-layer infinite-width networks. Many papers (Xiao et al., 2020) study the spectrum of the NTK, and find in particular the largest eigenvalue dominates the training regime (Jacot et al.,
2018; Bowman & Montufar, 2022). Empirical NTK, developed by many practical tools (Novak et al., 2022; Engel et al., 2022), succeeds in deep learning applications such as neural architecture search (Xu et al., 2021; Chen et al., 2021) and network pruning (Chen et al., 2023; Wang et al., 2023). Our work takes a deeper look into the modules of an over-parameterized model, demonstrating module-level training asynchronizaton of the components by the spatio-temporal distribution of the modular NTK. To the best of our knowledge, this is the first work that proposes to analyze the modular NTK variation for optimization.
Adaptive Training Approaches. Multirate training (Vlaar & Leimkuhler, 2022) is an emerging related technique that partitions neural network parameters into two parts, in which parameters in the slow part are updated less frequently. Our work analyzes the fine-grained training dynamics of modules and proposes a theoretical-inspired method to dynamically update parameters. Besides, there is a wide range of methods that can save computational resources, such as network pruning (Lee et al., 2018a; Rachwan et al., 2022) and dynamic sparse training (Liu et al., 2020; Jiang et al., 2022). These works tend to disable parameters during both the forward and backward propagation processes, while our work only sparsifies the gradients during the backward propagation process. Hence, these techniques are orthogonal to our method, and the proposed mNTK can be employed for evaluating the importance of parameters in pruning-based methods. We also integrate our method with pruning, and the results are presented in Appendix.
This section presents experimental results on various model architectures. The models under consideration are both over-parameterized and highly structured, including BERT (Devlin et al., 2018)
with multi-head self-attention (MHSA), Switch-Transformer (Fedus et al., 2022) with both MHSA and mixture-of-experts (MoE) (Shazeer et al., 2017b) and VGG (Simonyan & Zisserman, 2014) with convolutional filters.
Evaluation Measure. We evaluate the effectiveness of our algorithm for model optimization in terms of training efficiency and generalization. The training efficiency is measured by the validation or test error of the models after a certain number of floating-point operations (FLOPs), which serves as a lower bound for the execution time (Justus et al., 2018). The generalization performance is measured by the converged test error. Additionally, we calculate the total number of FLOPs used by the models until they achieve the converged test error.
Approximation of NTK. It is hard to directly compute the empirical NTK for an over-parameterized model, due to the computational and memory cost when calculating the Jacobian matrix  $J \in \mathbb{R}^{nk \times m}$ . Therefore, two alternatives are widely used to approximate the empirical NTK: (1) sampling a subset of training data instead of using the entire set; (2) computing the gradient by sum-of-logits instead of all output units, as shown in (Mohamadi & Sutherland, 2022).
Implementation Details. In practice, we evaluate the empirical NTK by sampling N training data (N is dynamically adjusted based on the available GPU memory in the experimental setup). To expedite the NTK calculation, we leverage parallelism by utilizing the implementation of Engel et al.  $(2022)^{2}$  which enables us to collect different data samples across multiple GPUs. The FLOPs is computed as the total sum of the forward and backward passes, where the backward FLOPs is approximately estimated as twice the forward pass (Rasley et al., 2020), and we measure it using the DeepSpeed Profiler<sup>3</sup>. All experiments are conducted on  $8 \times$  NVIDIA GeForce RTX 3090 GPUs. For further experimental details, please refer to the Appendix.
5.1 Main Result
Results of BERT. Firstly, we experiment on BERT (Devlin et al., 2018), which stacks 12 Transformer layers with 12 attention heads in each layer. Following the basic setup of Liu et al. (2019), we train BERT from scratch by masked language modeling (MLM) task on WikiText-2 (Merit et al., 2016). We compare the proposed MAT method with directly training a BERT model (BERT), using multiple learning rates (Vlaar & Leimkuhler, 2022) to train BERT (Multirate), and randomly selecting some heads for training (BERT-Rand). In this experiment, MAT applies  $\alpha = 0.1$ ,  $\beta = 10^{-3}$ .
Table 1: Results of BERT on WikiText-2. FLOPs is measured per GPU without embedding. Computation refers to the FLOPs model used until achieving best test loss. Best results are in boldface.
Method	Validation Loss (Log PPL) @ 10 PFLOPs	@ 15 PFLOPs	@ 20 PFLOPs	Test Loss (Log PPL) @ Convergence	Computation (PFLOPs)
BERT	5.39 ± 0.15	4.75 ± 0.06	4.48 ± 0.03	4.41 ± 0.05	28.70
BERT-Rand	5.65 ± 0.18	5.11 ± 0.16	4.96 ± 0.11	4.82 ± 0.06	21.53
Multirate	4.93 ± 0.10	4.57 ± 0.03	4.52 ± 0.02	4.48 ± 0.03	19.46
MAT	4.46 ± 0.04	4.41 ± 0.02	4.35 ± 0.02	4.27 ± 0.03	16.50
Table 1 shows the performance of all methods used in training BERT on WikiText-2. MAT outperforms all the baselines in training efficiency. In particular, MAT achieves the almost identical validation loss of 4.46 using 10 PFLOPs (1 PFLOPs = 10 $^{15}$  FLOPs) of computation whereas the vanilla BERT achieves 4.48 using 20 PFLOPs. Our method saves half of the computation by performing a lighter back propagation pass only for appropriate subsets of the modules. Compared to MAT, the performance degradation of BERT-Rand confirms the effectiveness of using  $\lambda_{\mathrm{max}}$  as the criterion for selecting important modules. While Multirate shows the potential for efficient training, it comes at the cost of sacrificing performance. In contrast, MAT achieves better test performance with a reduction of about  $15\%$  in test perplexity, and also saves  $41\%$  of computational resources.
Results of Switch-Transformer. To further evaluate the scalability of our method, we apply MAT to the feed-forward network (FFN) of transformer layer. Switch-Transformer (Fedus et al., 2021) is a representative implementation of Mixture-of-Expert (MoE) architecture (Shazeer et al., 2017a),
$^{2}$ https://github.com/pnnl/torchntk
which has shown excellent performance in NLP tasks recently. MoE (Shazeer et al., 2017b) architects the FFN into a structured one, which replicates the FFN as experts and dynamically assigns tokens to each expert. In this experiment, we compare the performance of Switch-Transformers using vanilla, Multirate and Switch-Rand training methods on WikiText-103 (Merit et al., 2016). Specifically, MAT regards both experts and attention heads as the modules with  $\alpha_{h} = 0.1$ ,  $\alpha_{e} = 0.2$ ,  $\beta_{h} = \beta_{e} = 10^{-3}$ .
Table 2: Results of Switch-Transformer on WikiText-103. FLOPs is measured per GPU without embedding. Best results are in boldface.
Method	Validation Loss (Log PPL)	Test Loss (Log PPL) @ Convergence	Computation (PFLOPs)
@ 200 PFLOPs	@ 400 PFLOPs	@ 600 PFLOPs
Switch	3.16 ± 0.12	2.31 ± 0.05	1.93 ± 0.03	1.74 ± 0.02	1155.2
Switch-Rand	2.92 ± 0.15	2.15 ± 0.08	2.01 ± 0.07	1.93 ± 0.05	837.5
Multirate	2.67 ± 0.09	2.02 ± 0.08	1.83 ± 0.05	1.77 ± 0.04	740.4
MAT	2.34 ± 0.06	1.92 ± 0.03	1.69 ± 0.02	1.68 ± 0.02	614.8
Table 2 shows that the performance of Switch-Transformer trained on a large dataset with a large number of parameters. It demonstrates that MAT significantly reduces the computation required for training. Switch-Rand and Multirate also improve training efficiency to some extent. Compared to them, our method precisely identifies the appropriate heads and experts trained with features that are easy to learn and generalize, reducing validation perplexity by  $52.3\%$ ,  $25.2\%$ , and  $18.1\%$  ( $25.9\%$ ,  $16.9\%$  and  $12.4\%$  for log perplexity) within 200, 400, and 600 PFLOPs of computation, respectively. Compared to the vanilla model, our method achieves  $10.4\%$  improved test perplexity. These results demonstrate the effectiveness of MAT in larger models and datasets.
Results of VGG. To further validate the versatility of MAT, we deploy it on computer vision tasks. We take classic convolutional network VGG16 as an example, which is over-parameterized for the CIFAR-10 dataset. MAT splits convolutional filters into two parts of modules by  $\alpha = 0.2$ ,  $\beta = 10^{-6}$ . Experimental results listed in Table 3 show that our method helps VGG16 converge faster (47.3%) with better test accuracy compared with the vanilla model.
Table 3: Results of VGG16 on CIFAR-10. Best results are in boldface.
Method	Computation until Train Acc = n (PFLOPs)	Test Accuracy (Top-1, %) @ Convergence	Computation (PFLOPs)
@ n = 95%	@ n = 99%
VGG16	8.85	12.06	93.77 ± 0.08	17.14
VGG16-Rand	9.96	13.35	92.71 ± 0.08	18.84
Multirate	7.19	9.75	93.43 ± 0.14	13.35
MAT	5.31	7.21	93.86 ± 0.05	9.03
5.2 Module-level Training Analysis
This experiment analyzes the histogram of the number of training epochs of all the attention heads in BERT. As shown in Figure 5, more than half of the heads were trained with backpropagation in only  $20\%$  of the epochs, and the average training epoch per head was  $34\%$  of the entire training process. This verifies that the sparse backpropagation induced by MAT enhances training efficiency. Furthermore, approximately  $20\%$  of the heads were never updated, indicating the sparse activation of over-parameterized structured models. This finding suggests the possibility of further reducing the computation of stable modules during the forward pass.
Figure 5: Histogram of epochs where the heads are trained using back-propagation.
We have analyzed the modular-level and temporal training characteristic of structured overparameterized models through a new measure of modular neural tangent kernel (mNTK). Empirical and theoretical analysis demonstrates the relationship between optimization effectiveness and mNTK principal eigenvalues. Based on this finding, we designed a novel training strategy termed Modular Adaptive Training (MAT), which uses a modularly and temporally adaptive dynamic threshold to select partial modules for gradient back propagation. Experimental results show that MAT reduces training cost and increases test accuracy compared to existing training scheme. This work seeks to improve training efficiency by sparsifying the gradient update. Besides pruning, this work can be combined with other techniques to further improve training efficiency. We leave it for future work.
This work was supported by National Natural Science Foundation of China under Grant No. 62090025, National Key R&D Program of China under Grant No. 2022YFB4400400 and China Postdoctoral Science Foundation No. 2022M720767.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019.
Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli. Memorization and optimization in deep neural networks with minimum over-parameterization. In Advances in Neural Information Processing Systems, 2022.
Benjamin Bowman and Guido F Montufar. Spectral bias outside the training set for deep networks in the kernel regime. Advances in Neural Information Processing Systems, 35:30362-30377, 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search onImagenet in fourgpu hours: A theoretically inspired perspective. In International Conference on Learning Representations, 2021.
Yixuan Chen, Yubin Shi, Mingzhi Dong, Xiaochen Yang, Dongsheng Li, Yujiang Wang, Robert Dick, Qin Lv, Yingying Zhao, Fan Yang, et al. Over-parameterized model optimization with polyak- $\{\backslash L\}$  ojasiewicz condition. In The Eleventh International Conference on Learning Representations, 2023.
Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming.(2018). arXiv preprint arXiv:1812.07956, 2018.
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.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
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.
Andrew Engel, Zhichao Wang, Anand D Sarwate, Sutanay Choudhury, and Tony Chiang. Torchntk: A library for calculation of neural tangent kernels of pytorch models. arXiv preprint arXiv:2205.12372, 2022.
Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances in neural information processing systems, 27, 2014.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontonon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
Chiyuan Zhang, Samy Bengio, and Yoram Singer. Are all layers created equal? arXiv preprint arXiv:1902.01996, 2019.
A Trainability and Generalization
A.1 Trainability
Following the previous work (Jacot et al., 2018) training neural networks in function space instead of parameter space, we analyze the trainability of an over-parameterized model by investigating the evolution of its predictions. Specifically, after a step of gradient descent
$$
\boldsymbol {\theta} _ {t + 1} = \boldsymbol {\theta} _ {t} - \eta \nabla_ {\boldsymbol {\theta} _ {t}} \mathcal {L}, \tag {7}
$$
the updated predictions can be approximated using the first-order Taylor expansion as
$$
\begin{array}{l} f (\mathcal {X}; \theta_ {t + 1}) = f (\mathcal {X}; \theta_ {t} - \eta \nabla_ {\theta_ {t}} \mathcal {L}) \\ \approx f \left(\mathcal {X}; \boldsymbol {\theta} _ {t}\right) - \eta \nabla_ {\boldsymbol {\theta} _ {t}} f \left(\mathcal {X}; \boldsymbol {\theta} _ {t}\right) ^ {\top} \nabla_ {\boldsymbol {\theta} _ {t}} \mathcal {L} \tag {8} \\ = f (\mathcal {X}; \boldsymbol {\theta} _ {t}) - \eta \nabla_ {\boldsymbol {\theta} _ {t}} f (\mathcal {X}; \boldsymbol {\theta} _ {t}) ^ {\top} \nabla_ {\boldsymbol {\theta} _ {t}} f (\mathcal {X}; \boldsymbol {\theta} _ {t}) \nabla_ {\mathcal {Z}} \mathcal {L} \\ = f (\mathcal {X}; \boldsymbol {\theta} _ {t}) - \eta \Theta_ {t} (\mathcal {X}, \mathcal {X}) \nabla_ {\mathcal {Z}} \mathcal {L}. \\ \end{array}
$$
Motivated by Arora et al. (2019), we rewrite the equation into the form of norm and eigenpairs:
$$
\begin{array}{l} \left\| f (\mathcal {X}; \boldsymbol {\theta} _ {t + 1}) - f (\mathcal {X}; \boldsymbol {\theta} _ {t}) \right\| _ {2} ^ {2} \approx \left\| \eta \boldsymbol {\Theta} _ {t} (\mathcal {X}, \mathcal {X}) \nabla_ {\mathcal {Z}} \mathcal {L} \right\| _ {2} ^ {2} \\ = \eta^ {2} (\Theta_ {t} (\mathcal {X}, \mathcal {X}) \nabla_ {\mathcal {Z}} \mathcal {L}) ^ {\top} (\Theta_ {t} (\mathcal {X}, \mathcal {X}) \nabla_ {\mathcal {Z}} \mathcal {L}) \\ = \eta^ {2} \nabla_ {\mathcal {Z}} \mathcal {L} ^ {\top} \Theta_ {t} (\mathcal {X}, \mathcal {X}) \Theta_ {t} (\mathcal {X}, \mathcal {X}) \nabla_ {\mathcal {Z}} \mathcal {L} \\ = \eta^ {2} \nabla_ {\mathcal {Z}} \mathcal {L} ^ {\top} \left(\sum_ {i} ^ {n k} \lambda_ {i} \mathbf {u} _ {i} \mathbf {u} _ {i} ^ {\top}\right) \left(\sum_ {i} ^ {n k} \lambda_ {i} \mathbf {u} _ {i} \mathbf {u} _ {i} ^ {\top}\right) \nabla_ {\mathcal {Z}} \mathcal {L} \\ \stackrel {(i)} {=} \eta^ {2} \nabla_ {\mathcal {Z}} \mathcal {L} ^ {\top} \left(\sum_ {i} ^ {n k} \lambda_ {i} ^ {2} \mathbf {u} _ {i} \mathbf {u} _ {i} ^ {\top} \mathbf {u} _ {i} \mathbf {u} _ {i} ^ {\top}\right) \nabla_ {\mathcal {Z}} \mathcal {L} \tag {9} \\ \stackrel {(i i)} {=} \eta^ {2} \nabla_ {\mathcal {Z}} \mathcal {L} ^ {\top} \sum_ {i} ^ {n k} \lambda_ {i} ^ {2} \mathbf {u} _ {i} \mathbf {u} _ {i} ^ {\top} \nabla_ {\mathcal {Z}} \mathcal {L} \\ = \eta^ {2} \sum_ {i} ^ {n k} \lambda_ {i} ^ {2} \left(\mathbf {u} _ {i} ^ {\top} \nabla_ {\mathcal {Z}} \mathcal {L}\right) ^ {\top} \mathbf {u} _ {i} ^ {\top} \nabla_ {\mathcal {Z}} \mathcal {L} \\ = \eta^ {2} \sum_ {i} ^ {n k} \lambda_ {i} ^ {2} (\mathbf {u} _ {i} ^ {\top} \nabla_ {\mathcal {Z}} \mathcal {L}) ^ {2}, \\ \end{array}
$$
where  $(i)$  follows  $\mathbf{u}_i^\top \mathbf{u}_j = 0, \forall i \neq j$  and  $(ii)$  follows  $\mathbf{u}_i^\top \mathbf{u}_i = \| \mathbf{u}_i\|^2 = 1$ . Then we can derive Eq. 9 into:
$$
\left\| f \left(\mathcal {X}; \boldsymbol {\theta} _ {t + 1}\right) - f \left(\mathcal {X}; \boldsymbol {\theta} _ {t}\right) \right\| _ {2} = \sqrt {\sum_ {i = 1} ^ {n k} \left(\eta \lambda_ {i} \mathbf {u} _ {i} ^ {\top} \nabla_ {\mathcal {Z}} \mathcal {L}\right) ^ {2}}. \tag {10}
$$
As we can see, at every step of the gradient descent, the model learns the target function faster along the eigen-directions corresponding to the larger eigenvalues.
Further, for the loss function assumed by squared error loss, we characterize the loss reduction by the following directional derivative (Wang et al., 2020):
$$
\Delta \mathcal {L} (\boldsymbol {\theta}) = \lim  _ {\epsilon \rightarrow 0} \frac {\mathcal {L} (\boldsymbol {\theta} + \epsilon \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta})) - \mathcal {L} (\boldsymbol {\theta})}{\epsilon}, \tag {11}
$$
using Taylor's first-order approximation  $\mathcal{L}(\pmb {\theta} + \epsilon \nabla_{\pmb{\theta}}\mathcal{L}(\pmb {\theta}))\approx \mathcal{L}(\pmb {\theta}) + \epsilon \nabla_{\pmb{\theta}}\mathcal{L}(\pmb {\theta})^{\top}\nabla_{\pmb{\theta}}\mathcal{L}(\pmb {\theta}),$
$$
\Delta \mathcal {L} (\boldsymbol {\theta}) \approx \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}), \tag {12}
$$
expanded by chain rule  $\nabla_{\pmb{\theta}}\mathcal{L}(\pmb {\theta}) = \nabla_{\pmb{\theta}}f(\pmb {\theta})\nabla_{\mathcal{Z}}\mathcal{L}(\pmb {\theta})$  into
$$
\Delta \mathcal {L} (\boldsymbol {\theta}) = \left(\nabla_ {\boldsymbol {\theta}} f (\boldsymbol {\theta}) \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta})\right) ^ {\top} \nabla_ {\boldsymbol {\theta}} f (\boldsymbol {\theta}) \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) \tag {13}
$$
$$
\begin{array}{l} \left(\nabla_ {\boldsymbol {0}} f (\boldsymbol {0}) \cdot \mathcal {L} \mathcal {L} (\boldsymbol {0})\right) \cdot \nabla_ {\boldsymbol {0}} f (\boldsymbol {0}) \cdot \mathcal {L} \mathcal {L} (\boldsymbol {0}) \\ = \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} \left(\nabla_ {\boldsymbol {\theta}} f (\boldsymbol {\theta}) ^ {\top} \nabla_ {\boldsymbol {\theta}} f (\boldsymbol {\theta})\right) \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}), \end{array} \tag {13}
$$
then following Eq. 1 and matrix decomposition of mNTK,
$$
\begin{array}{l} \Delta \mathcal {L} (\boldsymbol {\theta}) = \nabla_ {\boldsymbol {\mathcal {Z}}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} (\boldsymbol {\Theta}) \nabla_ {\boldsymbol {\mathcal {Z}}} \mathcal {L} (\boldsymbol {\theta}) \\ = \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} \left(\sum_ {l = 1} ^ {L} \Theta^ {l}\right) \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) \\ = \mathcal {Y} ^ {\top} \left(\sum_ {l = 1} ^ {L} \sum_ {i = 1} ^ {n k} \lambda_ {i} ^ {l} \mathbf {u} _ {i} ^ {l} \mathbf {u} _ {i} ^ {l} ^ {\top}\right) \mathcal {Y} \tag {14} \\ = \sum_ {l = 1} ^ {L} \sum_ {i = 1} ^ {n k} \lambda_ {i} ^ {l} \left(\mathbf {u} _ {i} ^ {l} ^ {\top} \mathcal {Y}\right) ^ {\top} \left(\mathbf {u} _ {i} ^ {l} ^ {\top} \mathcal {Y}\right) \\ = \sum_ {l = 1} ^ {L} \sum_ {i = 1} ^ {n k} \lambda_ {i} ^ {l} \left(\mathbf {u} _ {i} ^ {l} ^ {\top} \mathcal {Y}\right) ^ {2}. \\ \end{array}
$$
The directional derivative of the loss function is closely related to the eigenspectrum of mNTKs. As for the latter projection item  $\left(\mathbf{u}_i^{l^\top}\mathcal{Y}\right)^2$ , Arora et al. (2019) have studied the relationship between the projection norm and labels, and they demonstrate that true labels generate better alignment with top eigenvectors, especially for the maximum one. Therefore, we can assume that  $\left(\mathbf{u}_i^{l^\top}\mathcal{Y}\right)^2\propto \lambda_i^l$  in a real dataset.
A.2 Generalization
We denote the set of  $n$  training instances as  $\mathcal{X} = (\mathbf{x},\dots ,\mathbf{x}_n)$  and the target set as  $\mathcal{Y} = (\mathbf{y}_1,\dots ,\mathbf{y}_n)^\top$ . A two-layer neural network with the ReLU activation function can be written as
$$
f _ {\boldsymbol {\theta}, \boldsymbol {a}} (\mathcal {X}) = \frac {1}{\sqrt {m}} \sum_ {r = 1} ^ {m} a _ {r} \sigma \left(\boldsymbol {\theta} _ {r} ^ {\top} \mathcal {X}\right) \tag {15}
$$