Train Faster, Perform Better: Modular Adaptive Training in Over-Parameterized Models
Authors: Yubin Shi
Yixuan Chen
Mingzhi Dong
Xiaochen Yang
Dongsheng Li
Yujiang Wang
Robert P. Dick
Qin Lv
Yingying Zhao
Fan Yang
Tun Lu
Ning Gu
Li Shang
Venue: NeurIPS 2023 poster (2023)
Keywords: Modular Adaptive Training
Efficient Training
Over-parameterized Model
Neural Tangent Kernel.
TL;DR: This work analyses modular learning behaviors and proposes modular adaptive method for efficient training.

Abstract: Despite their prevalence in deep-learning communities, over-parameterized models convey high demands of computational costs for proper training. This work studies the fine-grained, modular-level learning dynamics of over-parameterized models to attain a more efficient and fruitful training strategy. Empirical evidence reveals that when scaling down into network modules, such as heads in self-attention models, we can observe varying learning patterns implicitly associated with each module's trainability. To describe such modular-level learning capabilities, we introduce a novel concept dubbed modular neural tangent kernel (mNTK), and we demonstrate that the quality of a module's learning is tightly associated with its mNTK's principal eigenvalue $\lambda_{\max}$. A large $\lambda_{\max}$  indicates that the module learns features with better convergence, while those miniature ones may impact generalization negatively. Inspired by the discovery, we propose a novel training strategy termed Modular Adaptive Training (MAT) to update those modules with their $\lambda_{\max}$ exceeding a dynamic threshold selectively, concentrating the model on learning common features and ignoring those inconsistent ones. Unlike most existing training schemes with a complete BP cycle across all network modules, MAT can significantly save computations by its partially-updating strategy and can further improve performance. Experiments show that MAT nearly halves the computational cost of model training and outperforms the accuracy of baselines.

Train Faster, Perform Better: Modular Adaptive Training in Over-Parameterized Models
Yubin Shi
Yixuan Chen'
Mingzhi Dong
Xiaochen Yang?**
Dongsheng Li
Yujiang Wang'
Robert Dick Qin Lv
Yingying Zhao?
Fan Yang"
Tun Lu Ning Gu Li Shangh
China and Shanghai Key Laboratory of Data Science, School of Computer Science, Fudan University
School of Mathematics Statistics, University of Glasgow
Microsoft Research Asia, Shanghai, China
'Department of Engineering Science, University of Oxford
Department of Electrical Engineering and Computer Science, University of Michigan " Department of Computer Science, University of Colorado Boulder School of Microelectronics, Fudan University
Abstract
Despite their prevalence in deep-learning communities, wee-parameterized models convey high demands computational costs for proper training. This work studies the fine-grained, modular-level learning dynamics of over-parameterized models to attain a more efficient and fruitful training strategy. Empirical evidence reveals that when scaling down into network modules, such heads in self-attention models, we can observe varying learning patterns implicitly associated with each module's 's trainability. To describe such modular-level learning capabilities, we introduce a novel concept dubbed modular neural tangent kernel (mNTK), and we demonstrate that the quality of module's 's learning tightly associated with its mNTK's principal eigenvalue Amax- A large Amax indicates that the module learns features with better convergence, while those miniature ones may impact generalization negatively. Inspired by the discovery, we propose a novel training strategy termed Modular Adaptive Training (MAT) to update those modules with their Amax exceeding dynamic threshold selectively, concentrating the model on learning common features and ignoring those superfluous ones. Unlike most existing training schemes with complete BP cycle across all network modules, MAT can significantly save computations by its partially-updating strategy and can further improve performance. Experiments show that MAT nearly halves the computational cost model training and outperforms the accuracy of baselines.
Introduction
The proliferation of large-scale pre-trained models (Devlin al. 2018 Brown al. 2020 Doso- vitskiy et al. 2020) demonstrates the advantages over-parameterized models. Those models adequately fit the training data with their sizeable parameters far exceeding the data scale (Du ct al. 2018). and such over-parameterization. are known to be essential in model optimization and generalization and are also crucial their success in deep learning (Liu et al. (022). However, challenges remain. Over-parameterized models are expensive to train; for example, training large language models (LLMs) such GPT-3 (Brown al. 2020 Zaheer et 2020) may take weeks to months, even on a large collection of powerful GPUs.
"Corresponinng authors.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Modular Adaptive Training (MAT)
Effletive Kank
Baad wi BP Head
BERT
BERT with
Thrasbold
Willlaten Lano 291.
Epeadi
Epech 40
Epoch
Avea ofNTXX
g
Heat
Howe Nowe
dasTK 1o
Heef
Iiframn Tarte
00
0
20
69
300
25 35 20 25
Tpock
Upoch
1poch
Comporotion (FFLORS))
(a)
Figure 1: Characteristics of training BERT on WikiText-2. Figure (a) demonstrates the joint variation of effective rank and Amax across the attention heads. Head; refers to the 1th attention head in the 1h layer. Figure (b-left) illustrates the idea of MAT, which governs the heads training by a dynamic threshold. Using MAT speeds up convergence and achieves lower validation loss (b-right).
Generally, most paraaaatttterze models consist of various functioning modules connected layer- wisely, e.g., heads in self-attention models like Transformer (Vaswani et al. 2017), experts in mixture-of-experts (MoE) blocks (Shazeer et al. 2017b), or filters in convolutional neural networks (CNN). This module- resource been spent efficiently and effectively across modules during model optimization? To this end, we dive into the modular-level training behaviors of eeee-parameterized models and pursue more efficient and fruitful training strategy. We propose modular Neural Tangent Kernel (mNTK), which is derived from the vanilla NTK (Jacot et al. 2018 Fort et al. 2020), powerful tool for understanding the fine-grained learning dynamics of each network module. We compute mNTK as the first-order approximation of network module's ' evolution during training, and it reflects the gradient correlation among the training data module-wisely. The eigenspectrum of mNTK describes the learning dynamics of a module, and its principal eigenvalue Amax indicates the degree of consistency in the gradient direction for learning the most common features in data.
architecture inevitably leads question: Has the computational
To illustrate the module-Ievel training dynamics intuitively, we employ the standard BERT (Devlin eet al. 2018) as an example and consider each attention head as specific module. In addition to Amax- adopt the effective rank (Roy & Vetterli 2007) measure the effective dimensionality of a feature matrix, like an attention matrix, and those metrics are computed at the end of each training epoch. Figure [ a) presents the trend of training dynamics of BERT, and we can observe asynchronous learning behaviors across modules and training iterations. We explain those observations using Amax as an indicator of trainability parameters. As discussed in Bowman & Montular (2022), the network is biased to learn the top eigenfunctions of the NTK over the entire imput space, while the parameters with very low eigenvalues barely update. In Figure a), the attention head in blue marker has large Amax. meaning that it has large consistency over data samples on feature learning and is thus apt to converge with better trainability. Indeed, these heads contain richer information as reflected through effective rank. In contrast, Amax of the red attention head oscillates at a low level, likely due to learning diverse features or even noise that may be hard to converge, and thus updating it brings little benefit to training loss.
measure how effectively gradient descent can optimize network
Our analysis (Section 2) on trainability and generalization demonstrates that aver-parameterized models exhibit modularly and temporally varying asynchronous learning behaviors. Modules with large mNTK principal eigenvalues learn features that are more rapidly learned, and those miniature ones with limited impact on training loss negatively influence generalization. Such asynchronous modularly and temporally training dynamics result in inefficient resource utilization by existing training schemes. Motivated by these analyses, we propose simple yet effective training method termed Modular Adaptive Training (MAT), aiming at dynamically training partial modules to improve training efficiency. As shown in Fīg, we design dynamic threshold (the dash-dotted line), according to modular difference and temporal variation of mNTK Amax- MAT only back pook-propagates the parameters of those modules whose Xmax is larger than the threshold. This forces the network to

---

concentrate on learning the common features and ignore the inconsistent ones, preventing it from fitting superfluous features or noises. This work makes the following contributions.
We empirically and theoretically reveal the associations between modular-level training dynamics and the proposed modular Neural Tangent Kernel (mNTK).
We propose Modular Adaptive Training (MAT) with a dynamic threshold to selectively update modules during prok-propagation to improve learning efficiency.
Experimental results verify that MAT can significantly reduce the training computation, improve performance, and generalize well to different over seeer-prmmeterizee models.
2 Analysis of Modular Neural Tangent Kernel in Over-Parameterized Models
In this section, we first introduce the definition of Neural Tangent Kernel (NTK) and modular Neural Tangent Kernel (mNTK). Next, we present empirical analysis revealing the eigenspectrum characteristics the mNTK, including the imbalanced distribution of mNTK eigenvalues and the modular and temporal diversity of mNTK eigenvalues during training. Finally, we provide theoretical analysis show that trainability of wer-paramettriied structured models is closely related to the largest mNTK eigenvalues; on the other hand, learning features associated with the miniature eigenvalues have negative impact on generalization.
2.1 Modular Neural Tangent Kernel (mNTK)
Let X' denote the set of n training instance, which are i.i.d. drawn from an unknown distribution D, and let y E gun denotes the targets. We study the network function f parameterized by O aiming to map the input vector X' to output vector 2, termed as 2 [X'; 0) where O E IR and Z Rn* Following the original definition of Neural Tangent Kernel (NTK) (Jacot et al. 2018 Wang et al.) 2020). we introduce modular NTK for fine- fine-graine grained analysis of structured deep networks as below.
Definition 1 (Modular Neural Tangent Kernel (mNTK))- Suppose model f contains L disjoint mod- ules 0 = "'....,.'"' ,''} where O' denote the parameters of th module. We define mNTK as matrix by e ('X,') = Jo Jo:) where Jo V V f(X: 6') denotes the Jacobian of the function f at the points X with respect the 1 module's 's parameters 0¹.
O' is a positive semi-definite real symmetric matrix and can be digo-deeo as O¹ U'A'U' En Nufu!!" with nk non-negative eigenvalues X()') (X1 N2..... X nk Note that a network can be partitioned into modules in flexible ways, for example, by dividing layers within a deep network components of network with inherent structures (e.g., attention heads in Transformer or convolutional filters a CNN). Since each module has its own mNTK, the integral NTK of a model can be computed as the sum of mNTKs of each module (Yang & Littwin 2021):
(X) Σ Jo, (X) (63)
Σ Σ
Joi(
(168) Σ
(2, X),
(1)
OpEo'
where (i) decomposes the matrix multiplication into the sum of vector multiplication; (tt) gathers addends by each module; (iii) follows the definition of mNTK.
2.2 Empirical Analysis
In this subsection, apply BERT to WikiText-2 for language modeling and regard each layer module to analyze the properties of cigenvalue distribution of all mNTKs, as well as the modular and temporal variation of eigenvalues during the training process.
The principal eigenvalue of mNTK dominates the eigenspectrum of mNTK. Figure (a) shows the first 16 eigenvalues of all layer-wise mNTKs at the 10th epoch. Apparently, the eigenvalue distribution of each layer-wise mNTK is inbalanced where there exists single large eigenvalue X (A, will be denoted by Amax hereafter) and a large number of small eigenvalues. In particular, the largest eigenvalue often has to 4 orders of magnitude larger than that of the rest eigenvalues. These findings indicate that the principal eigenvalue of mNTK dominates the spectrum of the mNTK.
3

---

Magaitude of coch eigeavaloe
Arax of layer-wiic MNIK
15000
08
10
12300
06-
100
Layer
5000
2300
20
600
103
10
35
Eigenvalue index
(4)
Epoch
(k)
Epoch
Figure 2: Training dynamics characterized by layer-wise mNTK of BERT trained on WikiText-2. (a): The first 16 eigenvalues distribution of layer-wise mNTKs the 10th epoch. (b): Variation of Xmax of layer-wise mNTKs. (c): Variation of x of layer-wise mNTKs during the first 20 epochs.
The phenomenon is also indicated by Xiao al. (2020) in terms of NTK. Furthermore, according to the connection between NTK and mNTK (Eq. D. the spectrum NTK is dominated by the principal eigenvalues of all mNTKs. Below, utilize Amax for fine- fine-grained optimization analysis of over-parameterized model.
The variation of principal eigenvalues of mNTKs is highly asynchronous across modules. Figure shows the variation Xmax of all layer-wise mNTKs during the training process. is worth noting that the principal eigenvalues of mNTKs exhibit different trends during the training process. Intuitively speaking, NTK (also mNTK) captures the training dynamics throughout the training process, and its principal eigenvector corresponds to the direction of the steepest descent in loss. Specifically, for mNTKs of shallow layers, the magnitude of Amax is large after model initialization, and it decreases and converges fast in the early training stages. A potential explanation that shallow layers learn simple and common features of data samples (Yosinski et al.) 2014 Zhang model convergence in the early training stages. In
al. 2019), which contribute significantly contrast, for mNTK of the deep layer, the magnitude of Amax is small after initialization and gradually increases during training. The reason is that deep layers learn complex and unique features, relying on the stability of simple features in shallow layers, which occurs in the later training stages.
To analyze the training dynamics, condition number is also commonly used (Xu et al. 2021 Xiao et al. 2020), defined as к = Amax/ Amiu- Figure ( shows the variation of for different modules during training. We observe that the layer-wise condition number exhibits asynchronization among different modules only in the early training stages (especially the first 10 epochs) and к of all modules quickly converges to the same value. This reveals that the condition number is hard to provide sustained information for fine-grained optimization analysis, unlike the principal eigenvalue which is capable of capturing the effective information throughout the training process.
The temporal variation of the principal eigenvalue of mNTK indicates the generalization ability. In addition to the above analysis depicting the asynchronization between different modules, examine the temporal variation of the principal eigenvalue given a specific module and establish its connection with the generalization ability. In detail, train the same model using masked modeling on half sequence length, which makes the model prone to overfitting. As shown in Figure a), the validation loss starts increase from the t5th epoch while the training loss continually decreases. Figure 3(b-c) present the temporal variation of the principal eigenvalue for the 3rd, 6th 9th and 12th layers of the model during the training process, respectively. During the first 45th epochs, i.e., the non-overfitting stage, the principal eigenvalues of these layers continually decrease. However, after the thth epoch when the model enters the overfitting stage, the principal eigenvalues converge to excessively low values. A possible reason for overfitting is that the optimization direction of the model starts to align with the direction of small eigenvalues, which means that the model is learning superfluous features that are too specific or even noisy.
2.3 Analysis of Trainability and Generalization
Using the tool of mNTK, show that the optimization properties of over-parameterized structured models are closely related to eigenvalues of mNTKs.
Suppose for any module 0¹, Amin (e') X 0. denote the loss function VoC(() the gradient w.r.t parameter 0. For step of gradient descent, loss reduction can be characterized by the

---

1d15
lel3
Log Perplecity
Layer
Layer
Layer
Layer 12
Validation
Amax
Amex
Amex
Amox
2.0
2.5
0.5
L4
1.-
10-
100
300
110
Epoch
Epoch
Epoch
Epoch
Epoch
(e)
Figure 3: Training dynamics in the overfitting Case of 4-layer BERT trained by 64 token MLM task.
directional derivative of the loss function
Wang et al.
2020):
L(e + <Vo((())) - (()) (1i)
AC
lim
Vol(O)\voC(C)
(128)
VzC(0) (VoS(0) (VV (Vo Vol(())
(2)
VzL(O) Σ e')=2((()) (1) 4 nk
(30)
1=1
11
where (i) follows the definition of directional derivative; (ii) follows the first-order Taylor expansion: (izi) follows chain rule of derivative; (iv) follows E and (v) follows the cigen-dccomposition of
mNTK under the assumption of squared error loss (Arora et al. 2019). Following the work of Arora
et al. (2019), we assume that true labels align well with top eigenvectors, i.e., (u: )) is large for large XT Thus, the directional derivative of the loss function can be regarded as closely related to the eigenspectrum mNTKs
(1) The relation between trainability and principal eigenvalues of mNTKs.
Based on the first observation of empirical analysis, i.e., the spectrum of each mNTK is dominated by principal eigenvalues, we utilize Amax of mNTKs as the nearly equivalent proxy of the spectrum of msTKs. Therefore, Eq. 2 is simplified as:
AC
Σ
Xmax u ))²
The above equation suggests that the loss decreases faster along the eigenspaces that correspond to larger eigenvalues. Given the fact that the principal eigenvalues of mNTKs are highly asynchronous across modules, we suggest selectively training the modules that are active with larger X to achieve efficient learning with limited computational resources.
(2) The relation between generalization ability and eigenvalues of mNTKs.
Figure4shore shows the eigen-spectrum distribution of 4ayyr BERT model the 10th epoch, by regarding each head as a module and cal- culating 256 eigenvalues for each head. Among
Number Eigenvalues
threshold
Nuisance
Info
all 4x x2x25xx 256 eigenvalues, the distribution of eigen-spectrum exhibits distinct regions: the left region with a large number of small eigen- values and the right region with small number of large cigenvalues. The two regions are widely separated. According Eq. 3 the eigenspaces corresponding to miniature eigenvalues in the left region have negligible impact on training loss minimization, while those in the right re- gion make dominant contributions. Following the definition of Oymak et al. (2019), we refer to the right region as the information space where the training error decreases rapidly, and the left region as the nuisance space where training is slow and may cause overfitting.
Space
Space
102
101
20.5
20.5
22-4
20-1
10:7
10:3
10a
Normalized Eigenvalue
Figure 4: The normalized eigen-spectrun distribu- tion exhibits two distinct regions, termed informa- tion space and nuisance space.
5

---

Model training with more parameter updates tends to have longer distance to the initializations 0n (Hoffer et al., 2017). Based on Lemma 5.4 in Arora et al. (2019). longer distance results in relatively larger Rademacher complexity, as R. < ||0 0ol[r. and larger R. results in relatively worse generalization. As a consequence, we suggest selectively stopping the training in some modules so as to stop the increase of the distance to the initializations in these directions. Specifically, we propose to stop training the modules with principal cigenvalues falling into the nuisance space so as to enjoy better generalization and efficient learning simultaneously.
Let's consider the inter-module training dynamics on model generalization. During the training process, given modules with variant principal eigenvalues, it often occurs that even though a small number of modules still lie in the information space, the most of others already fall into the nuisance space. The modules located in the information space, i.e., those with principal eigenvalues exceeding the threshold, are trained to learn informative features that facilitate loss minimization. Conversely, the modules falling into the nuisance space, i.e., those with principal eigenvalues below the threshold, are prone to learn superfluous or even noise features, which significantly increase the Rademacher complexity, and deteriorate their generalization ability. This analysis indicates that the existing training scheme overlooks the asynchronization among modules, resulting in inefficient utilization of the computational resource for loss minimization during model training.
3 Modular Adaptive Training
Based on the analysis in the previous section, advocate dynamically optimizing partial module parameters in back propagation, termed as Modular Adaptive Training (MAT). The forward propa- gation is computed as usual; during back propagation, MAT quantifies module parameters in terms of modular policy and temporal policy such that only subsets of modules are updated. Below, we present the proposed modular policy and temporal policy. Algorithm | summarizes the proposed MAT algorithm.
Modular Policy. Following the analysis that trainability and generalization of an over- weerparameterized model are related to the principal eigenvalues of module parameters, we propose to only update partial module parameters with Amax greater than certain threshold during the model training; this policy is termed as modular policy.
Specifically, suppose the model parameters consist of L module parameters and each mNTK has nk positive eigenvalues, and let X denote 7th eigenvalue of the mNTK computed over the parameters
0'. We define Ä = {, nk.I. as the set of eigenvalues calculated over all mNTKs, and Amax. Amin
as the maximum and minimum value of the set A, respectively. To select the appropriate modules, we introduce hyperparameter ratio (0, 1), thus the corresponding eigenvalue threshold X, is computed as:
d Amiu + Amaxx Amin) a.
(4)
According to the eigenvalue threshold A., all modules are split into information modules and nuisance modules as:
Iz (0") Amma (0)) 2 Aa), Amao (0)) < AaJ.
(5)
As result, the computation cost of back propagation is reduced.
We empirically observe that a E (0, 1) is relatively stable over the training process for a certain model and dataset. Therefore, 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 0', its principal eigenvalue after warming up is recorded as the initial value Amax (e)). During each training episode t, the temporal variation of principal eigenvalue is defied as A! I>mmax (0)) Amax ( J. 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 pre-defined threshold 8:
A 1/-1 <, A
(6)
/ is empirically bounded by eigenvalue variation within the nuisance region.
Related Work
Optimization Analysis using NTK. Neural Tangent Ker- nel (NTK) (Jacot et al.) 2018). which calculates the Gram matrix of Jacobian, is known as powerful tool to analyze convergence and generalization properties (Arora et al. (2019). Modern neural networks are gen- erally over-parante peeeparrmtteee with positive definite NTKs, meaning that theoretically the gradient flow always converges to the equilibrium where the training loss zero (Chizat & Bach
Algorithm 1 Modular Adaptive Training
Require: training data Yx ) sample number S, modules set M, modular threshold a, temporal threshold B.
1: I,N M, Ø
2: while not converged and [Z] > 0 do
3: Compute Vof (x: 0) for samples (x.,
4: Ä
5:
for 0 in M do
6:
O' = Vo f(x 0J ozx
Compute eigen-decomposition O¹ U U'A
7:
8: A Amaax ( Amax (0))
9:
AUA
10:
end for
A = min(A) + (max(() mi)(A)) a
2018, Bombari al., (002)). A line of works study the training dynamics of over-pore parameterized models (Li et al., 20200 Liu et al., 2022) based NTK, mainly under the assumption of two-layer infinite-width net- works. Many papers (Xiao et al.] 2020) study the spectrum of the NTK, and find in particular the largest cigenvalue dominates the training regime (Jacot et al.
12: N = (0)) Mmmaa (0)) < >a (Modular policy Eq. 4)
13: N = NU A4-A!_11
< B} (Temporal policy Eq. 6)
A
I M -
15:
for each batch x, in yx do
16:
for 0' in I do
17:
0 end for
v D( ().)))
18:
19: end for
20:
end while
2018 Bowman & Montular (2022). Empirical NTK, developed by many practical tools Novak
et al. 2022;
et al. (022). succeeds in deep learning applications such as neural architecture
search (Xu et al. 2021| Chen all
202T) and network pruning (Chen al. 2023 Wang et al. 2023).
Our work takes a deeper look into the modules of an parameterizeed model, demonstrating module-level training asynchronization of the components by the ppatio-temporn distribution of the modular NTK. To the best of our knowledge, this is the first work that proposes 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, which parameters in the slow part are updated less frequently. Our work analyzes the fine-grained training dynamics of mncoreticalinnspired method to dynamically update parameters. Besides, there is a wide range of methods that can save computational resources, such as network pruning (Lee) al. (2118a) Rachwan al., 2022) and dynamic sparse training (Liu al.) [2220] Diang et al., (022). 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.
modules and proposes
5 Experiments
This section presents experimental results on various model architectures. The models under con- sidcration are both verrpaaameeeerized and highly structured, including BERT (Devlin et al. 2018)
7

---

with multi-head self-attention (MHSA), Switch- Transformer (Fedus et al. 2022) with both MHSA and cexture-of-experte (MoE) (Shazeer 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 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 over- werr-aarameterized model, due to the computational and memory cost when calculating the Jacobian matrix J E Rnkxm 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 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 measure it using the DeepSpeed Profile All experiments are conducted on 8 x NVIDIA GeForce RTX 3090 GPUs. For further experimental details, please refer to the Appendix.
5.1 Main Result
Results 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 (Merity et al. 2016). We compare the proposed MAT method with directly training 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 0.1, 3 10
Table Results of BERT on WikiText-2. FLOPs is measured per GPU without embedding. Compu- tation refers to the FLOPs model used until achieving best test loss. Best results are in boldface.
Validation Loss
Test Loss (Log PPL.) (PFL.Ps)
Computation
Method
(Log PPL.
@ 10 PFLOPs
@ 15 PFLOPs
@ 20 PFLOPs
@ Convergence
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
4.82 + 0.06
21.53
Multirate
4.93 + 0.10
4.57 + 0.03
4.52
4.48 + 0.03
19.46
MAT
4.46 + 0.04
4.41 0.02
4.35=0002
4.27 - 0.03
16.50
Table shows the performance of all methods used in training BERT on 2ikiText-2. MAT out- performs all the basclines training efficiency. In particular, MAT achieves the almost identical validation loss of 4.46 using 10 PFLOPs (1 PFLOPs 10 FLOPs) of computation whereas the vanilla BERT achieves 4.48 using 20 PFLOPs. Our method saves half of the computation by per- forming 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 Amax 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 reduction of about 15% in test perplexity, and also saves 41 % of computational resources.
Results of Transsfrrrrnsoommer. To further evaluate the scalability of our method, we apply MAT to the feed-forve forward network (FFN) of transformer layer. Switch-Trie Transformer (Fedus et al. 2021 is a representative implementation of Mixture-of-Expee -xpeert (MoE) architecture (Shazeer et al. 2017a).
Phups://githubbbo Phtnps//gituuh.co//mmllliaa htt.://hhhccccmmm////m///ccnnkkk hhngs//githbbboom/mnncsss///ssss
8

---

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 Sritich-TTansffrmee using vanilla, Multirate and Switch-Rand training methods on WikiText-109 (Merity et 2016). Specifically, MAT regards both experts and attention heads as the modules with Orp = 0.1, 0.2, Sh. 3e 10-3.
Table 2: Results of Switch-Transforme Transformer on WikiTextt1033 FLOPSs is measured per GPU without embedding. Best results are in boldface.
Validation Loss
Test Loss (Log PPL)
Computation (PFLOPs)
Method
(Log PPL)
@ 200 PFLOPs
@ 400 PFLOPs
@ 600 PFLOPs
@ Convergence
Switch
3.16 0.12
2.31 0.05
1.93 0.03
1.74 1 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 1 0.08
1.83 1 0.05
1.77 + 0.04
740.4
MAT
2.34 1 0.06
1.92 1 0.03
1.69 1 0.02
1.68 1 0.02
614.8
trained on a large dataset with a large number of parameters. demonstrates that MAT significantly reduces the computation required for training. Switch-Rand and Multirate also improve training efficiency to some extent. Compared to them, method precisely identifies the appropriate heads and experts trained with features that easy to learn and generalize, reducing validation perplexity by 52.3%, 25.2%, and 18.1 1 % (25. 95,, 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 MAT larger models and datasets.
Table 2 shows that the performance
Results of VGG. further validate the versatility MAT, we deploy it on computer vision tasks. We take classic convolutional network VGG16 example, which is wver-parameterized for the CIFAR-10 dataset. MAT splits convolutional filters into two parts of modules by a 0.2, 3 10 10-6 Experimental results listed in with better test accuracy compared with the vanilla model.
that method helps VGG16 converge faster (47.3%)
Table 3: Results of VGGI6 on CIFAR-10. Best results are in boldface.
Computation until Train Acc n (PFLOPs)
Test Accuracy
Computation (PFLOPs)
Method
(Top-1, %)
@ 95% @ 99%
@ Convergence
VGGT6
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 more than half of the heads were trained with backpropaga- tion 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 back- propagation induced by MAT enhances training ef- ficiency. Furthermore, approximately 20% of the heads were never updated, indicating the sparse ac- tivation over-parament prerpearaeterrieed structured models. This finding suggests the possibility of further re- ducing the computation of stable modules during the forward pass.
17.5%
15.0%
Ratio Heads
12.5%
7.5%
5.0%
2.5%
25
100
125
150
175
Number of Eppccs
Figure Histogram of epochs where the heads are trained using back-propagation.

---

Conclusion
We have analyzed the modular-level and temporal training characteristic of structured parameterized 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.
Acknowledgements
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.
References
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-parameetrrizatoon. In Advances in Neural Information Processing Systems, 2022.
Benjamin Bowman and Guido 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 few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search on imagenet in
four gpu hours: A theoretically inspired perspective. In International Conference 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- ojasiewicz condition. The Eleventh International Conference on Learning Representations, 2023.
Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable program- (2018. arXiv preprint arXiv:1812.07990 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.04800 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. Torch- ntk: A library for calculation of neural tangent kernels of pytorch models. arXiv preprint 123.: 2022.
10

---

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): 5232-5270, 2022.
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.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: a033vv1803.03635, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. (026-1034, 2015.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap large batch training of neural networks. Advances neural information processing systems, 30, 2017.
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.
Peng Jiang, Lihan Hu, and Shihui Song. Exposing and exploiting fine-grained block structures for fast and accurate sparse training. Advances in Neural Information Processing Systems, 35: 38345-3883 38357, 2022.
Daniel Justus, John Brennan, Stephen Bonner, and Andrew Stephen McGough. Predicting the computational cost of deep learning models. In 2018 IEEE international conference on big data (Big Data), pp. 3873-3882, IEEE, 2018
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Torr. Snip: Single-shot network pruning based on connection sensitivity. In International Conference on Learning Representations, 2018a. 8a.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning based connection sensitivity, arXiv preprint arX1v:1810.02340. 2018b.
Mingehen Li, Mahdi Soltanolkotabi, and Samet Oymak. Gradient descent with early stopping is provably robust to label noise for overparameterized neural networks. In International conference on artificial intelligence and statistics, pp. 4313-4324. PMLR, 2020.
Chaoyue I Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameterried non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59: 85-116, 2022.
Junjie Liu, Zhe Xu, Runbin Shi, Ray CC Cheung, and Hayden KH So. Dynamic sparse train- ing: Find efficient sparse network from scratch with trainable masked layers. arXiv preprint 0:Xi1:2005.06870, 2020.
Liyang Liu, Shilong Zhang, Zhanghui Kuang, Aojun Zhou, Jing-Hao Xue, Xinjiang Wang, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. Group fisher pruning for practical network compression. In International Conference on Machine Learning, pp. 7021-7032. PMLR, 2021.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692, 2019.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXiv: 2016.

---

Mohamad Amin Mohamadi and Danica Sutherland. A fast, well-founded approximation to the empirical neural tangent kernel, arXiv preprint arXiv:2206. 125:43, 2022.
Roman Novak, Jascha Sohl-Dickstein, and Samuel S Schoenholz. Fast finite width neural tangent kernel. In International Conference on Machine Learning, pp. 17018-17044. PMLR, 2022.
Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi. Generalization guaran- tees for neural networks via harnessing the low-rank structure of the jacobian. arXiv preprint arXiv:1906.05392, 2019.
John Rachwan, Daniel Zügner, Bertrand Charpentier, Simon Geisler, Morgane Ayle, and Stephan Günnemann. Winning the lottery ahead time: Efficient early network pruning. International Conference on Machine Leurning, pp. 18293-18308 PMLR, 2022.
Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimiza- tions enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery Data Mining, pp- 3505-3506, 2020.
Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007 15th European signal processing conference, pp. 606-610. IEEE, 2007.
Victor Sanh, Thomas Wolf, and Alexander Rush. Movement pruning: Adaptive sparsity by fine-tuning. Advances in Neural Information Processing Systems, 33:20378-20333 2020.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gan gated mixture-of-exppee layer. arXiv preprint arXiv: 06538, 2017a.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017b.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1556, 2014.
Julia Ture, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv: 1rXiv:1908.08922 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Tiffany Vlaar and Benedict Leimkuhler. Multirate training of neural networks. In International Conference Machine Learning, pp. 22342-223222 22360. PMLR, 2022.
Chaoqi Wang, Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. In International Conference on Learning Representations, 2020.
Yite Wang, Dawei Li, and Ruoyu Sun. Ntk-sap: Improving neural network pruning by aligning 2023.
training dynamics. arXiv preprint ar
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generaliza- tion deep neural networks. In International Conference on Machine Learning, pp. 10462-10472. PMLR, 2020.
Jingjing Xu, Liang Zhao, Junyang Lin, Rundong Gao, Xu Sun, and Hongxia Yang. Knas: green neural architecture search. In International Conference Machine Learning, pp. 11613-11625. PMLR, 2021.
Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. In International Conference on Machine Learning, pp. 11762-11772. PMLR, 2021.
12

---

Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances neural information processing systems, 27, 2014.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, 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.01196, 2019.
13

---

Appendix
A Trainability and Generalization
A.1 Trainability
Following the previous work (Jacot et al. 2018) training neural networks function space instead of parameter space, analyze the trainability of wver-parameterized model by investigating the evolution of its predictions. Specifically, after a step of gradient descent
0++1 = 0. - V,
the updated predictions can be approximated using the first-order Taylor expansion as
((x:0++1 (.00--Vo..)
f(x;0() (V..f(X;0.) TV.. Vo.6
(8)
f(X;0) 'VV.f((X 0,) (o
f(x;0) ().
Motivated by Arora et al. (2099) rewrite the equation into the form of norn and eigenpairs:
1f(X; f(X; ~ (1)
(X, (O,(X,
vL (X, (X,
mk =
nn )
nnn V2
nk
(6) n
V
X uu u
V 2 T
uu V2C
nk
(u; VzL)² T u: T V
n nk
(u; (u: VzL}²
where (i) follows u; u, 0, * j and (ii) follows u Tu, | || 1. Then we can derive Eq. into:
nk indu ZZZ2
(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):
L(e + cV L())
AC lim
(1I)
c++0
14

---

using Taylor's first-order approximation = (VVc
Voc(()) expanded by chain rule VoC(0)= (ef(0)v=[(0) into
(12)
AC =
(13)
V=((())(Vof(0) (V vof(0))V2((())
then following Eq. and matrix decomposition of mNTK,
AC = (
V=C(O) T Σ 0')V
1=1
y
Muluf y
(14)
L nk
Ml y) Tuf y)
\u v)
The directional derivative of the loss function is closely related to the eigenspectrum of mNTKs. As for the latter projection item (uf })) Arora et (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 (uf) T >) X in real dataset.
A.2 Generalization
We denote the set of n training instances as X (x,..., A two-layer neural network with the ReLU activation function can be written as
X,) .,,,,..,,,
and the target set as y = (T)..
(x) = Σ a, 0, x)
(15)
where 02....... "e.....,,. 6, are the weight vectors in the first layer, a1, Q2 are weights in the second layer. For simplicity, we denote 0 (OI...., Om) and a (a), am Following (Arora et al. 2019), we assume the network is over-parameterized and is trained by gradient descent on the quadratic loss over dataset (X.)), We freeze the second layer and optimize the first layer through gradient descent. Let (0)) and (tt) denote the weights initialized from scratch and the weights after iterations, respectively.
Under these settings, according to the Lemma 5.3 of (Arora et al. 2019), the embedding function fo.a learned from GD is in a restricted class of neural networks with weights close to initialization (0). Therefore, the Rademacher complexity of the two-layer function can be bounded as follows:
Lemma 2 Given R O, with probability at least 1 8 over the random initialization (0(0), a). simultaneously for every B > 0, the following function class
F 0 {So.a 110, (-)0))l s R (Vr [m]),
R.B
(16)
||0 ((0)]lF S B}
15

---

has empirical Rademacher complexity bounded as:
Rss
R.
a0(0),a))
-
1 [t
sup
Σ
oif (xi)
71
SETR R.B
(0))a
(17)
B V2n
2 log 2
1/4
2R
1
+ Ry
log
8
m
K
Lemma indicates that the Rademacher complexity is proportional the weight distance from its initialization. Specifically, a greater weight distance implies a more significant amount of Rademacher complexity and is thus associated with weaker generalization ability. Following those intuitions, extend the connection between weight distance and Rademacher complexity, initially portrayed in two-layer network, into more generalized deep neural networks.
For deep models, as mentioned in (Hoffer et al. 2017). the weight distance from its initialization grows logarithmically with the number of iterations, which be described
[[O(t) ( logt.
(18)
discover that training iterations increase, the model's Rademacher complexity also grows with its weights more deviated from initializations, which subsequently impairs the model's generalization ability. However, solutions remain. As trainability varies across various modules, we can prevent the significant growth of Rademacher complexity and thus retain satisfying model generalization by ignoring the updates of those modules with undesirable mNTK eigenvalues.
Combining Lemma and Eq. 18
B Experiments
B.1 Experimental Settings
BERT and Switch-Trane Transformer. We generally follow the settings of Liu et al. (2019) to train BERT and Transforrrn from scratch using self-supervised Masked Language Modeling (MLM). In detail, we uniformly sample approximately 15% input tokens to replace them with a mask, and the learning objective is predict those masked tokens accurately with cross-entropy loss.
All experiments are conducted on NVIDIA GeForce RTX 3090 GPUs, and we list those key hyperparameters in Table
Table Hyperparameters configuration in BERT and Switch-Transforme
Hyperparameter
BERT
Number of layers
12
12
Attention heads
12
12
Hidden dimension size
768
768
Dropout
0.1
0.1
Attention dropout
0.1
0.1
Sequence length
512
512
Batch size
8
8
Warmup steps
0
12k
Max steps
12k
300k
Weight decay
0
0.01
Peak learning
2e-4
le-4
Learning rate decay
Linear
Cosine
Adam [e, _ 82]
[1e-6, 0]
[1e-6. 0.9, 0.99]
Number of experts
4
Capacity factor
1.5
The baseline methods consist of: 1) the vanilla model, 2) the Multirate (Vlaar & Leimkuhler 2022
algorithm, and 3) the adaptive training with random selection of the updated modules (abbreviated as Rand).
16

---

For BERT, Multirate randomly partitions the heads into fast and slow groups into 50% 40%/50% 50% split, with fast heads updated every step of stepsize 1) and slow ones updated every l 5 step of stepsize kon.
BERT-Rand randomly selects and updates 50% heads in every epoch. MAT first randomly samples 64 data samples to approximate 64 eigenvalues for each head-based mNTK. The modular policy and temporal policy are set with a 0.1, 3 10 The adaptive training algorithm starts from the 5th epoch and performs every 8 epochs.
For Switch-Transformn heads/experts into fast and slow sets at 50%/50% ratios. The fast heads and experts are updated every step with current stepsize 17. and the slow ones are updated every k 5 steps with stepsize kn. Switch-Rand randomly selects 50% heads or experts for updating at the beginning of each epoch. MAT first randomly samples 64 data samples and separately approximate 64 eigenvalues for each head-based mNTK and expert-basee mNTK. The modular policy and temporal policy are set with ag 0.1, 3r. 10-3 and ae 0.2, 3e 10-3 3 for head and expert, respectively. The adaptive training algorithm performs every 2 epochs after warming up.
Transformer, Multirate considers each head/expert single module and partitions
To avoid the case that no heads or experts updated in any layer, we utilize protecting strategy that maintains the gradient of least one head or expert in each layer for all baselines and our model.
VGG. All baselines of VGG are initialized with Kaiming initialization (He et al. 2015) and are trained with SGD for 200 epochs with an initial learning rate of 0.1 and batch size of 128. We decay the learning rate by 0.1 at 1/2 and 3/4 of the total number of epochs. Specifically, Multirate randomly partitions the filters into fast and slow parts by 50%/50% randomly selects 50% filters to update filters in each epoch. MAT approximates eigenvalues for each filter-based mNTK using random 64 data samples. The modular policy and temporal policy are set with = 0.2, 8 10-6, All experiments are repeated by three runs and the final computation costs are calculated on average.
each layer. VGG-Rand
B.2 MAT and Network Pruning
Network pruning (Frankle & Carbin 2018; Sanh et al. 2020 Liu et al. 2021 applies various criteria determine the importance of different components and prunes those that most redundant to compress model size, which often results in slight performance drops. The proposed MAT distinguishes from network pruning in that MAT only sparsifies modules during backward propagation. while pruning methods eliminate the forward and backward propagations of pruned modules. Besides, MAT is the first work to employ the principal eigenvalue of mNTK as the module selection criterion.
Our empirical study reveals some modules that are never or rarely selected by the proposed adaptive training method (MAT), showing potential for being pruned achieve further computation savings. However, due to the complete removal of modules, existing network pruning methods may negatively impact the model's performance and generalizations. Based those observations, we have explored whether we can apply MAT to network pruning methods to accelerate the training process or to improve performance.
We employ a BERT model for this experiment, and we prune 50% attention heads according to the ranking of Amax of head-based mNTKs the 15th epoch. Across the training session, use MAT with a = 0.2 to introduce sparsity in the backward pass, resulting in approximately 25% sparsity of weight gradients.
Table Results BERT on WikiText-2 by pruning methods.
Sparsity (Pruning Ratio)
Test Loss (Log PPL)
Computation (PHL.PS))
Method
Forward Pass
Backward Pass
@ Convergence
Vanilla
0%
0%
4.41
27.80
SNIP (50%)
50%
50%
4.39
19.02
SNIP (75%)
75%
75%
4.70
15.22
MAT
50%
755
4.32
12.03
Table 5 compares the extended MAT, the vanilla BERT model, and SNIP (Leee al. 2018b) in terms of forward and backward sparsity.
17

---

(nkmm² SNIP is widely used pruning method that operates based on connection sensitivity, enabling sparsity in over-parameterized models. In our implementation, we apply SNIP modular manner by calculating the connection sensitivity of each module. As shown in the Table, SNIP achieves 31.6% reduction in computation when pruning 50% of the attention heads without any performance degradation. However, when the pruning ratio (sparsity) is increased to 75%, SNIP fails to achieve comparable performance with the vamilla model. This suggests that a large ratio sparsity can have negative impact model performance.
In contrast, using the criteria of MAT, we prune 50% of the attention heads while training the remaining ones by MAT. This approach leads to further acceleration of computations by 56.7% while slightly improving the overall performance. This experiment serves an example highlighting the potential of MAT in network pruning, showcasing the trade-off between computation savings and performance maintenance.
B.3 Computational Complexity and Overhead of MAT
The proposed MAT can introduce additional computational and memory overheads as it involves the calculation of mNTKs and their principal eigenvalues. However, demonstrate in this subsection that MAT only yields negligible proportion of extra computations, and we also report the numeric overhead results from experiments to support this claim.
To clarify, we have employed two strategies to accelerate the computation of mNTKs significantly: (1) sampling subset of size S( < n) for the NTK approximation instead of using the entire set, and (2) computing the gradient using the sum-of-logits approach instead of considering all output units. With those strategies, we approximate the Jacobian matrix J R with the approximated Jacobian matrix J RSxm and we only need to perform S additional gradient propagations and concatenate the gradients together. mNTKs are then computed with the approximated Jacobian matrix, and we perform the eigen- oegeppiocooooiiiioon to mNTKKs to obtain the principal eigenvalue.
Apart from those techniques mentioned above, the module division strategy also accelerates the matrix multiplication of the Jacobian. Unlike the integral NTK, MAT applies a lightweight NTK estimation by modular NTK that significantly reduces the computation time required, and this estimation can be scalable to deeper structured networks. The complexity of computing integral NTK is O(Sm²); O(Sm² L) in MAT, assuming
however, the overall time complexity reduces to O (LS(mm/))) we are computing L mNTKs. As for the singular value decomposition (SVD), since S & n < m, its complexity, O(LS*), can be far lower than others. Table(6illuustaate the comparison of computational complexities. showcasing MAT's significant computational advantages for NTK approximation. In short, the overhead produced by MAT is negligible.
Table 6: Complexity comparison, where n denotes the number training data, m the number of parameters, & the output dimension, L the number of components, and S the sample number.
Complexity
Full
Our Approximation
NTK computaion
O(Sm² O
SVD computaion
O(
We also measure the actual computation overhead by MAT. Following the experimental setting of Ture et al. (2019), we apply the proposed MAT to BERT models with different network scales, namely BERT-Mini (L.=4, H=256), BERT-Small (L=4, H=512), BERT-Medium (L==8, H=512), BERT-Base (L=112, H=768), and BERT-Large (L.224, H=1024). Table Table?7caa demonstrates the computational costs across varying multi-scale BERT models, and we can see that MAT can save 26.4%. 33.4%, 39.4%, 40.6%, and 50.9% computations for BERT-Mini, BERT-Small, BERT-Medium, BERT-Base and BERT-Large, respectively. These observations indicate that applying MAT to larger models can better improve training efficiency. Notably, when compared to the overall computational costs, the overheads introduced by MAT are extremely small and can be arguably ignored.
18

---

Table 7: The computation (PFLOPs) required for training to convergence of models of different scale on WikiText-2.
Model
BERT-Mini
BERT-Small
BERT-Meduim
BERT-Bass
BERT-Large
Vanilla computation
1.44
4.28
9.52
27.80
61.25
MAT computation
1.06
2.85
5.77
16.50
30.09
MAT overhead
0.01
0.04
0.08
0.24
0.45
19


Reviews:
[{"review_id": "Evv0SSQ3GV", "reviewer": "Reviewer_TyiH", "summary": "Inspired by the Neural Tangent Kernel literature, this paper analyzes the training dynamics of over-parameterized neural networks from the perspective of modularity. In particular, the paper proposes to decompose the Jacobians of the NTK into those of modules. Here, the concept of \"module\" is a bit ill-defined, referring to some coherent computation block. The paper proposes to analyze the NTK within each module and in particular look at the eigenspectrum. It shows that empirically the first eigenvalues dominate within each module and that those eigenvalues are very different across modules. This provides an avenue for focusing mostly on those modules with a high MNTK eigenvalue. Results on different architectures and datasets show noticeable speedups. Finally, a connection between optimizing only high eigenvalue parts of the NTK and generalization is cited and corresponding empirical results are provided in this direction as well.", "questions": "- I'm a bit surprised at the lack of role of interactions between Jacobians or similar terms across modules in the entire derivation. Given how little constraints there are in what constitutes a module, I would've expected to see some effect. Could you explain this in more detail?\n- As mentioned under weaknesses, I would really appreciate a deeper comment on the entire computation time when accounting for the algorithm itself.", "limitations": "- The added compute should've been mentioned more clearly in my opinion.\n\n- No ethics review necessary.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper reads well\n- It's an interesting topic and potentially impactful in improving our understanding of training dynamics.", "weaknesses": "- The biggest limitation for me is that the method of computing the NTK and then computing its eigenvalues seems quite expensive. Speedup experiments are performed using flops, which bypasses the actual running time (which would have to include the algorithm itself), defeating the purpose if it takes longer than the speed gain.\n\n- As presented, the algorithm mostly claims speedups by not updating some modules, but if we need to update one of the earliest layers we still need to perform most backpropagation computations. Furthermore, at best, this algorithms provides a 2x (or 1.5x depending on how you count) improvement since it basically only affects the backward step.\n\n- I found the connection of eigenvalues of mNTKs and generalization a bit week both conceptually/theoretically and in terms of empirical results. It's enough to show promise, but the connection with the distance to initialization is only tangential and improvements in generalization are not necessarily directly linked to the NTK eigenspectrum in the experiments.\n\nMinor suggestion: I would put % of heads instead of number of heads in figure 5 to make it more immediately parseable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> `Q1` The biggest limitation ... the speed gain. \n> I would really appreciate a deeper comment on the entire computation time when accounting for the algorithm itself.\n\n`A1` Thanks for your comment. The proposed method MAT will introduce additional computational overheads as it involves calculation of empirical NTK and eigen-decomposition. However, numerical overhead results from *Table 7* demonstrate that MAT only requires a negligible proportion (1%~1.5%) of extra computations. The complete computational complexity analysis and numerical results can be found in *Appendix B.3*.\n\nTo accelerate the computation, we employed two strategies to approximate the Jacobian matrix $J\\in\\mathbb{R}^{nk\\times m}$ with $\\tilde{J}\\in\\mathbb{R}^{S\\times m}$, where we sample $S$ data and use the sum-of-logits approach [A] instead of considering all output units. Furthermore, MAT applies a lightwight NTK estimation by modular NTK instead of integral NTK, which reduces the complexity from $O(Sm^2)$ to $O(Sm^2/L)$, assuming we are computing $L$ mNTKs.\n\nFLOPs comparison gives the upper bound of the time savings, and the actual wall-clock time is limited by the hardware performance and optimization. Under our experimental conditions, the wall-clock time is reduced by 31.5% compared by vanilla model. Note that computional cost of the algorithm itself is also included in the statistics.\n\n| Method | Test PPL @ Final | Computation (PFLOPs) | Wall-Clock Time (s) |\n|---|---:|---:|---:|\n| vanilla | 4.41 | 28.70 | 8021 |\n| MAT | 4.27 | 16.50 | 5494 |\n\n> `Q2` As presented, the algorithm ... the backward step.\n\n`A2` Thanks for your comments. When using MAT, most gradient computation of fixed modules can be left out even when earliest layers need to be updated.\n\nTake a linear layer $y=Wx$ in an arbitrary model as an example, it goes through two computation steps during the standard backward process: the gradient of parameter weight W $\\frac{\\partial L}{\\partial W}=\\frac{\\partial L}{\\partial y}x^{\\top}$ and the gradient of the input vector $\\frac{\\partial L}{\\partial x}=W^{\\top}\\frac{\\partial L}{\\partial y}$, where $\\frac{\\partial L}{\\partial x}$ will be used to compute $\\frac{\\partial L}{\\partial W}$ of the front layer. If we split $W$ into equal-sized modules $W=cat(W_1,W_2,...,W_n)$, which satisfies $y=cat(y_1, y_2,...,y_n)=cat(W_1x,W_2x,...,W_nx)$, we can sparsify the two steps into: $\\frac{\\partial L}{\\partial W}=cat(0,...,\\frac{\\partial L}{\\partial y_i}x^\\top,...,0)$ and $\\frac{\\partial L}{\\partial x}=\\sum_{i}W_i^\\top\\frac{\\partial L}{\\partial y_i}$, where $W_i$ is the informative module. In other words, if $k$ modules in a layer are in nuisance space, we reduce the backward computation to $(n-k)/n$ of the origin. \n\nWith multiple layers, as long as the latter layer produces non-zero $\\frac{\\partial L}{\\partial x}$, i.e., by keeping at least one module in each layer active, the gradient can continue to propagate to the earlier layers. In our empirical study, we have not observed instances where later layers end their learning before the earlier ones, and thus the backpropagation is performed uninterruptedly. In addition, once the shallow layers are all in nuisance space, we can stop the gradient computation in those layers altogether.\n\nThis example is congruent with the multi-head attention structure, and the associated gradient sparsification is proved to guarantee the convergence [B].\n\nFurthermore, our approach is not limited to sparsification on gradient backpropagation computation. We can further save the forward computation if we consider mNTK $\\lambda_{\\max}$ as network pruning criteria. *Table 5* demonstrates the comparison of applying MAT as pruning method. The complete experimental results can be found in *Appendix B.2*.\n\n> `Q3` I found the connection ... eigenspectrum in the experiments.\n\n`A3` Thanks for your comments. Please refer to the **General Response**.\n\n> `Q4` Minor suggestion: I would put % of heads instead of number of heads in figure 5 to make it more immediately parseable.\n\n`A4` Thanks for your suggestion. We will improve this figure in the revised version.\n\n> `Q5` I'm a bit surprised at ... expected to see some effect.\n\n`A5` Thanks for your question. While it might seem counterintuitive, mNTKs are calculated independently across modules, and its sum of mNTKs equals to the integral NTK; the derivation is shown below. $\\boldsymbol{\\Theta}(\\mathcal X, \\mathcal X)=\\sum_{l=1}^L \\sum\\_{\\boldsymbol{\\theta}_p \\in \\boldsymbol{\\theta}^l}J\\_{\\boldsymbol{\\theta}^l}(\\mathcal X) J\\_{\\boldsymbol{\\theta}^l}(\\mathcal X)^{\\top}=\\sum\\_{l=1}^L \\boldsymbol{\\Theta}^l(\\mathcal X, \\mathcal X)$. Different from MLP, modern neural networks are inherently structured. Nevertheless, modules are correlated for model optimization. As presented in the paper, we observe significant inter-module variation during training, specifically, some modules already converge and fall into the nuisance space, while others are still in the information space. Our work indicates that mNTK is a good metric to measure the inter-module training dynamics, and presents an mNTK-based adaptive training method to optimize network training process.\n\nIn addition, the granularity of the modular division will have an impact on our adaptive strategy, the following table evaluates the impact of division granularities. As we can see, attention head is a suitable granularity.\n\n| Method | Valid PPL @ 10 PFLOPs | Test PPL @ Final | Computation (PFLOPs) |\n|---|---:|---:|---:|\n| MAT (layer) | 4.64 | 4.37 | 19.40 |\n| MAT (head) | 4.46 | 4.27 | 16.50 |\n| MAT (half-head) | 4.53 | 4.29 | 17.26 |\n\n---\n**References:**\n\n[A] Mohamadi et al. \"A fast, well-founded approximation to the empirical neural tangent kernel.\" International Conference on Machine Learning. PMLR, 2023.\n\n[B] Alistarh, Dan, et al. \"The convergence of sparsified gradient methods.\" Advances in Neural Information Processing Systems 31 (2018)."}, {"review_id": "iAjYzvLcrE", "reviewer": "Reviewer_pEnu", "summary": "In this work the authors propose a method for analyzing modules of a neural network, in particular their utility in test-time generalization, via a modular neural tangent kernel (mNTK). The mNTK is eigenspectrum the NTK of the network modules, e.g. different attention heads. In this work it was found that the larger the largest eigenvalue for a module's mNTK, (I will call this lam-max-mNTK for convenience) the more that module contributes features which are useful for generalizability. For example when, during training, a neural network's lam-max-mNTKs plateauing on a low value is indicative that the network has overfit. The authors also demonstrate that the modules' lam-mat-mNTKs evolve asynchronously during training, indicating that optimization of certain modules is more important during different epochs during training. Using this intuition the authors propose modular adaptive training (MAT), which weights modules' learning rates during training in proportion to the lam-max-mNTK of that module. The authors find that this improves training, achieving a minimum loss faster (around 1-5% faster). These experiments, including the analysis of generalizability, were done on BERT, although MAT was also performed on VGG, where the computational savings were even greater.", "questions": "Do we have some more evidence that your interpretation of mNTK is correct? Is there some reason we are more interested in languange models? Do the VGG lam-max-mNTK results align with those you got for BERT?", "limitations": "none", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is clearly written and pleasant to read.\n\nThe proposed analysis is interesting and would be of general interest to the deep learning community.\n\nThe computational improvements are nice.", "weaknesses": "The mNTK isn't super novel, as it is very similar to spectral analysis work which as been performed before, although I do think it is sufficiently novel enough for to ML venues.\n\nThe performance improvements, while nice, aren't large enough (at least for BERT) that I would expect MAT to see a large deployment in the future considering implementation seems somewhat involved. Perhaps more work is needed to make the improvement more substantial.\n\nThe experimental results are somewhat limited. In particular its hard to conclude that the observations in Section 2.2 generalize to neural networks in general considering this was just applied to one network. Also I think concluding that two regions in Figure 4 correspond to \"info space\" and \"nuisance space\" seems like a bit of a large jump and again, its just for one network. The split is interesting, but can we really conclude that the small eigenvalues are nuisances? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> `Q1` The mNTK isn't super novel, as it is very similar to spectral analysis work which has been performed before, although I do think it is sufficiently novel enough for to ML venues. \n\n`A1` Thanks for your comments. Indeed, the spectrum of NTK has been applied to two-layer networks to analyze the optimization of the overall network [A,B]. However, one key contribution of our work is that we observe significant inter-module variance during training, specifically, some modules already converge and fall into the nuisance space, while others are still in the information space. To the best of our knowledge, our work, for the first time, indicates that mNTK is a good metric to measure the inter-module training dynamics, and presents an mNTK-based adaptive training method to optimize the network training process.\n\n> `Q2` The performance improvements, while nice, aren't large enough (at least for BERT) that I would expect MAT to see a large deployment in the future considering implementation seems somewhat involved. Perhaps more work is needed to make the improvement more substantial.\n\n`A2` Thanks for your comments. To further verify the potential of MAT, we consider the following improvements:\n\n- Scalability: In Appendix Table 7, we present the computational costs of varying multi-scale BERT models. MAT can save 26.4%, 33.4%, 39.4%, 40.6%, and 50.9% computations for BERT-Mini, BERT-Small, BERT-Medium, BERT-Base, and BERT-Large (L=24, H=1024), respectively. These findings indicate that applying MAT to larger models can better improve training efficiency. In addition, larger datasets are more likely to contain more superfluous or noisy features, making more modules unnecessary to update.\n- Network Pruning: Our approach is not limited to sparsification on gradient backpropagation computation. Our empirical study reveals some modules that are never or rarely selected by the proposed adaptive training method (MAT), showing potential for being pruned to achieve further computation savings. In other words, $\\lambda_{\\max}$ of mNTK can be criteria of structured pruning. The experimental results can be found in Appendix B.2.\n- Training Stability: Large model pre-training is often unstable and requires careful tuning of the learning rate to maintain stable convergence. A common practice is to use a learning rate scheduler to gradually reduce the learning rate. MAT prompts the model to update in a consistent direction, which to some extent avoids the instability caused by noise, allowing the model to keep a large learning rate and achieve higher training efficiency. Experimental results are shown in following table.\n\n| Method | Valid PPL @ 10 PFLOPs | Test PPL @ Final | Computation (PFLOPs) |\n|---|---:|---:|---:|\n| MAT (w/ lr scheduler) | 4.46 | 4.27 | 16.50 |\n| MAT (w/o lr scheduler) | 4.37 | 4.30 | 14.75 |\n\n> `Q3` Also I think concluding that two regions in Figure 4 correspond to \"info space\" and \"nuisance space\" seems like a bit of a large jump and again, its just for one network. The split is interesting, but can we really conclude that the small eigenvalues are nuisances? Do we have some more evidence that your interpretation of mNTK is correct?\n\n`A3` Thanks for your questions. These two terms, information space and nuisance space, are originated from Oymak et al. [C], and the idea of splitting is followed by other works. For example, Li et al. proposes clean and corrupted information and proves that the clean residual is aligned with the top singular direction of the Jacobian matrix whereas label noise is aligned with the small singular directions [D]. Please refer to the **General Response** for detailed explanation.\n\nTo further verify the discrepancy between the two spaces, we conduct an ablation study which performs adaptive training on modules lying in the nuisance space. The comparison is listed below, where further verify the modules entered nuisance space learn superfluous or noisy features.\n\n| Method | Valid PPL @ 10 PFLOPs | Valid PPL @ 15 PFLOPs | Test PPL @ Final | Computation (PFLOPs) |\n|---|---:|---:|---:|---:|\n| vanilla | 5.39 | 4.75 | 4.41 | 28.70 |\n| MAT (nuisance) | 5.78 | 5.08 | 4.63 | 23.79 |\n| MAT (information) | 4.46 | 4.41 | 4.27 | 16.50 |\n\n\n> `Q4` The experimental results are somewhat limited. In particular its hard to conclude that the observations in Section 2.2 generalize to neural networks in general considering this was just applied to one network. Do the VGG lam-max-mNTK results align with those you got for BERT?\n\n`A4` Thanks for your questions. Prior works have found that parameter updates in over-parameterized models are sparse and low-rank, which opens up the possibility of neural network pruning and sparse training. Following your suggestion, we conduct the experiments in VGG. The result of eigen-spectrum distribution is similarly splitted, as shown in *Figure C* of PDF in General Response.\n\n---\n**References:**\n\n[A] Arora, Sanjeev, et al. \"Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.\" International Conference on Machine Learning. PMLR, 2019.\n\n[B] Xiao, Lechao, Jeffrey Pennington, and Samuel Schoenholz. \"Disentangling trainability and generalization in deep neural networks.\" International Conference on Machine Learning. PMLR, 2020.\n\n[C] Oymak, Samet, et al. \"Generalization guarantees for neural networks via harnessing the low-rank structure of the jacobian.\" arXiv preprint arXiv:1906.05392 (2019).\n\n[D] Li, Mingchen, Mahdi Soltanolkotabi, and Samet Oymak. \"Gradient descent with early stopping is provably robust to label noise for overparameterized neural networks.\" International conference on artificial intelligence and statistics. PMLR, 2020."}, {"review_id": "vkJEqePDyi", "reviewer": "Reviewer_afqQ", "summary": "This paper introduces modular adaptive training based on the largest eigenvalue of the NTK. In particular, based on the assumption that  the principal eigendirection matters the most for generalization, the authors propose to dynamically turn off gradient updating on certain modules if they have small principal eigenvalues. Experiments are conducted on NLP and CV tasks to show the superiority of MAT.", "questions": "1. In Figure 2(a), can we also see a picture showing the distribution of all eigenvalues? Now even the 16th eigenvalue seems pretty large, is there very small eigenvalue?\n2. In definition 1, can you write down the size of each matrix? Like $J\\in\\mathbb{R}^{nk\\times p}$.\n3. In figure 3, I wonder if there is a phenomenon of benign overfitting in this BERT task. If you keep training beyond 100 epochs to make preplexity even smaller, will you see a double descent on validation?\n4. In figure 2(a), $\\lambda_1$ of the 12th layer is still large close to 100 epoch. If you keep training will that eigenvalue eventually decays to $0$? What is special about this layer?\n5. Around line 220, you mentioned the computation cost of backprop is reduced. Say at some point my first layer is in the information modules but the last layer is in the nuisance modules. In this case, the computation cost is as large as all modules are active, is this right? Because you need to backprop through the later layers.\n6. Can you give more intuition on eq(6)? If my $\\lambda_1(\\Theta_1)$ and $\\lambda_1(\\Theta_0)$ are small, but $\\lambda_1(\\Theta_t)$ and $\\lambda_1(\\Theta_{t-1})$ are both large, then this equation can still be satisfied. Do you want to stop gradient in this scenario?\n7. In algorithm 1, how large did you use for sample number $S$? How accurate is this estimation on the eigenvalues?\n8. When you compute FLOPs, do you include the computations used for eigendecomposition?\n9. Can we see the downstream performance of the MAT-trained BERT on a downstream task?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "It is very important to reduce computation resources for large model training. This paper focuses on an important topic, and the approach is principled. The presentation is very organized and the delivery is clear.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> `Q1` In Figure 2(a), can we also see a picture showing the distribution of all eigenvalues? Now even the 16th eigenvalue seems pretty large, is there very small eigenvalue?\n\n`A1` Thanks for your question. The eigenvalues can be refered to *Figure A* of PDF in General Response. The computation of NTK considers each output unit, so the magnitude of the eigenvalues is related to the specific task. The number of output units in BERT's MLM task is equal to the vocubulary size, which makes the magnitude of the smallest eigenvalue in BERT MLM task can be as large as $10^9$ ~ $10^{10}$. Therefore, we observe their relative or normalized value instead, as *Figure 4* for all attention heads.\n\n> `Q2` In definition 1, can you write down the size of each matrix?\n\n`A2` Sure. Suppose the parameter vector is partitioned into $L$ disjoint modules $\\boldsymbol{\\theta}=\\{\\boldsymbol{\\theta}^1,\\boldsymbol{\\theta}^2,...,\\boldsymbol{\\theta}^L\\}$, each of which contains $m^1,m^2,...,m^L$ parameters respectively. The size of each matrix is annotated below:\n\n$J_{\\boldsymbol{\\theta}^l}(\\mathcal{X})\\in\\mathbb{R}^{nk\\times m^l}$, $J_{\\boldsymbol{\\theta}_p}(\\mathcal{X})\\in\\mathbb{R}^{nk\\times 1}$, $\\boldsymbol{\\Theta}^l(\\mathcal{X},\\mathcal{X})\\in\\mathbb{R}^{nk\\times nk}$,$\\boldsymbol{\\Theta}(\\mathcal{X},\\mathcal{X})\\in\\mathbb{R}^{nk\\times nk}$\n\n> `Q3` In Figure 3, I wonder ... double descent on validation?\n\n`A3` Thanks for your question. In our experiment, the validation perplexity keeps increasing when training for an additional 1000 epochs without descent on validation. Thanks for pointing it out, and we will further study on it.\n\n> `Q4` In Figure 2(a), ... about this layer?\n\n`A4` Thanks for your question. The 12th layer is the last Transformer layer, closest to the output and classifier. In our experiment, $\\lambda_{\\max}$ of the last layer is still large even it's trained more(~500) epochs. A possible explanation is that, in the later training process, due to the noisy dataset, the training error is not necessarily reduced to zero, so the residual makes its $\\lambda_{\\max}$ keep stable. In comparison, we conduct an experiment on MLP with MNIST, whose training loss and $\\lambda_{\\max}$ of the last layer are both converged to small value. Thanks for pointing out this interesting phenomenon, we will do more experiments to further explain why it occurs.\n\n> `Q5` Around line 220, ... backprop through the later layers.\n\n`A5` Thanks for your comments. When using MAT, most gradient computation of fixed modules can be left out even when earliest layers need to be updated.\n\nTake a linear layer $y=Wx$ in an arbitrary model as an example, it goes through two computation steps during the standard backward process: $\\frac{\\partial L}{\\partial W}=\\frac{\\partial L}{\\partial y}x^{\\top}$ and $\\frac{\\partial L}{\\partial x}=W^{\\top}\\frac{\\partial L}{\\partial y}$. If we split $W$ into equal-sized modules $W=cat(W_1,W_2,...,W_n)$, we can sparsify the two steps into: $\\frac{\\partial L}{\\partial W}=cat(0,...,\\frac{\\partial L}{\\partial y_i}x^\\top,...,0)$ and $\\frac{\\partial L}{\\partial x}=\\sum_{i}W_i^\\top\\frac{\\partial L}{\\partial y_i}$, where $W_i$ is the informative module. In other words, if $k$ modules in a layer are in nuisance space, we reduce the backward computation to $(n-k)/n$ of the origin. \n\nWith multiple layers, as long as the latter layer produces non-zero $\\frac{\\partial L}{\\partial x}$, i.e., by keeping at least one module in each layer active, the gradient can continue to propagate to the earlier layers. In our empirical study, we have not observed instances where later layers end their learning before the earlier ones, and thus the backpropagation is performed uninterruptedly.\n\nThis example is congruent with the multi-head attention structure, and the associated gradient sparsification is proved to guarantee the convergence [A].\n\n> `Q6` Can you give more intuition on eq(6)? ... in this scenario?\n\n`A6` Thanks for your question. Our intuition is to stop the module whose $\\lambda_1(\\boldsymbol{\\Theta})$ has small variation. Referring to Figure 3, near the time point of overfitting, certain modules begin to satisfy this condition. As for the situation you mentioned, where $\\lambda\\_1(\\boldsymbol{\\Theta}\\_{n})$ is larger than the initial one $\\lambda\\_1(\\boldsymbol{\\Theta}\\_{0})$, is avoided by recording the initial value $\\lambda\\_1(\\boldsymbol{\\Theta}\\_{0})$ after warming up or at the time $\\lambda\\_1(\\boldsymbol{\\Theta})$ reaches its largest value.\n\n>  `Q7` In algorithm 1, how large did you use for sample number $S$ ? How accurate is this estimation on the eigenvalues?\n\n`A7` Thanks for your question. The sample number $S$ is set according to the learning task and GPU memory. In our experiment, $S=64$ for BERT/Switch-Transformer and $S=128$ for VGG-16. We demonstrate the mNTK $\\lambda_{\\max}$ distribution of approximate NTK and true empirical NTK in *Figure B* of PDF in General Response, which are nearly identical.\n\n> `Q8` When you compute FLOPs, do you include the computations used for eigendecomposition?\n\n`A8` Thanks for your question. We demonstrate that MAT only yields a small proportion (<1.5%) of extra computations. *Table 7* shows the computational costs of varying multi-scale BERT models, and we can see that the overheads introduced by MAT are negligible. The complete computational complexity analysis and numerical results can be found in *Appendix B.3*.\n\n> `Q9` Can we see the downstream performance of the MAT-trained BERT on a downstream task?\n\n`A9` Thanks for your question. We compare performance of vanilla and MAT pre-trained model on the SST-2 (sentiment recognition) from GLUE Benchmark in the following table. \n\n| Method | MLM (perplexity) $\\downarrow$ | SST-2 (accuracy) $\\uparrow$ |\n|---|---:|---:|\n| vanilla | 4.41 | 84.36 |\n| MAT | 4.27 | 85.47 |\n\n---\n**References:**\n\n[A] Alistarh, Dan, et al. \"The convergence of sparsified gradient methods.\" Advances in Neural Information Processing Systems 31 (2018)."}, {"review_id": "ApeVw13wGg", "reviewer": "Reviewer_1e7h", "summary": "The authors introduce a novel modification of the neural tangent kernel (NTK) called modular NTK (mNTK).\nIt is essentially the decomposition of the NTK as a sum of tangent kernels for the disjoint modules\nthat make up the network. This allows for module-level analysis of training dynamics. They find that\nprinciple eigenvalue of mNTKs tends to be orders of magnitude larger than the others and that the variation\nof these principle eigenvalues is a synchronous across modules. Since the directions associated with\nthese principle eigenvalues dominates gradient updates during training, they can be utilized to selectively\nupdate subsets of the modules to better allocate computational resources during training. The authors\nfurther characterize modules with smaller max eigenvalues as being more likely to learn \"nuisance\" features.\nThey introduce a novel optimization schema called modular adaptive training (MAT) that stops back propagation\nto modules whose principle eigenvalue falls below a threshold or if variation of the principle eigenvalue\nfalls below a threshold.\nThey demonstrate that MAT produces models with equivalent performance to vanilla trained model while requiring\nfewer FLOPs during training. Furthermore, they show that MAT trained models tend to generalize better.\n", "questions": "\n- Line 221 mentions how the hyperparameter $\\alpha$ is relatively stable throughout training. This doesn't really make sense to me because how could a constant not be stable throughout training. Is this a typo? I think you may have meant to say something like $\\lambda_\\alpha$ is stable throughout training.\n- Have you explored using MAT for the process of fine-tuning a pretrained model? Typically, fine-tuning takes far less time than pretraining, but it can still take a while to converge if the target task has many examples.", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Paper is well-written.\n- Method is novel and theoretically well motivated.\n- Supported by experiments.\n- The method is very general and has the potential for high impact.\n", "weaknesses": "\nFor the language models, the paper only looks at perplexity scores of the pretrained model for text. I'm not the most\nup to date on this, but I think typically practitioners are more concerned about the performance of the model\nwhen fine-tuned on down stream tasks. I think comparing performance on (vanilla) fine-tuned downstream tasks of MAT vs vanilla pretrained models would strengthen the results.\n\nTying a bit into the previous point, papers such as RoBERTa indicate that continuing to train pretrained models after\nthe train perplexity has mostly stopped decreasing can lead to significant boosts in performance when fine-tuned on\ndownstream tasks. If the low $\\lambda_\\textrm{max}$ heads mostly learn diverse features, then this might be negatively impacted by MAT pretraining. If they mostly learn noise, then there shouldn't be much negative impact on this. Success\nof MAT in this setting (as indicated by increased performance on downstream tasks at a given computational cost) would\ngreater strengthen the method.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> `Q1` For the language models, the paper only looks at perplexity scores of the pretrained model for text. I'm not the most up to date on this, but I think typically practitioners are more concerned about the performance of the model when fine-tuned on down stream tasks. I think comparing performance on (vanilla) fine-tuned downstream tasks of MAT vs vanilla pretrained models would strengthen the results.\n\n`A1` Thanks for your suggestion. We compare performance of vanilla and MAT pre-trained model on the SST-2 (sentiment recognition) from GLUE Benchmark in the following table. \n\n| Method | MLM (perplexity) $\\downarrow$ | SST-2 (accuracy) $\\uparrow$ |\n|---|---:|---:|\n| vanilla | 4.41 | 84.36 |\n| MAT | 4.27 | 85.47 |\n\n\n> `Q2` Tying a bit into the previous point, papers such as RoBERTa indicate that continuing to train pretrained models after the train perplexity has mostly stopped decreasing can lead to significant boosts in performance when fine-tuned on downstream tasks. If the low $\\lambda_{\\max}$ heads mostly learn diverse features, then this might be negatively impacted by MAT pretraining. If they mostly learn noise, then there shouldn't be much negative impact on this. Success of MAT in this setting (as indicated by increased performance on downstream tasks at a given computational cost) would greater strengthen the method.\n\n`A2` Thanks for your comments. It's in line with what we are thinking. Many works demonstrate similar conclusions that continuing pre-training the model whose perplexity has converged boosts the downstream task performance. However, a major challenge in this process is how to distinguish learning diverse, useful features and learning noise. The proposed MAT shows potential for indicating which modules are learning the informative features that are common and likely to generalize.\n\nWe've conducted related experiments. We pretrain BERT from scratch by MLM task on AGNews tilL the validation perplexity converges, and save the snapshot called $A$. After that, we obtain $B$ and $C$ by continuing to pretrain $A$ for 10 epochs with vanilla method and MAT seperately. We fine-tuned the three models on AGNews by classification task and the experimental results are shown in Table.\n\n| Accuracy | A (Full-Trained) | B (Continue-Trained) | C (Adaptive-Trained) |\n|---|---:|---:|---:|\n| Full Data | 74.73 | 75.02 | **75.54** |\n| Long-tail Data | 73.69 | 74.23 | **75.00** |\n\n> `Q3` Line 221 mentions how the hyperparameter $\\alpha$ is relatively stable throughout training. This doesn't really make sense to me because how could a constant not be stable throughout training. Is this a typo? I think you may have meant to say something like $\\lambda_\\alpha$ is stable throughout training.\n\n`A3` Thanks for your question. The modular policy is motivated by the analysis of the eigen-spectrum distribution in Section 2.3. We find it exhibits two distinct regions in Figure 4 with a significant threshold, and we try to locate the threshold, which corresponds to the $\\lambda_\\alpha$. In fact, due to the magnitude variation of eigenvalues during the training process, $\\lambda_\\alpha$ should also be dynamic. However, we empirically observe the threshold is relatively stable in the normalized distribution as Equation 4, leading us to assign the normalized threshold $\\alpha$ as a constant. Thus, the hyperparameter $\\alpha$ keeps constant and the threshold $\\lambda_\\alpha$ is located dynamically according to the magnitude of overall eigenvalues.\n\n> `Q4` Have you explored using MAT for the process of fine-tuning a pretrained model? Typically, fine-tuning takes far less time than pretraining, but it can still take a while to converge if the target task has many examples.\n\n`A4` Thanks for your question. We do consider MAT as a potential method for fine-tuning or transfer learning, because parameters updating tends to be sparse or low-rank in these tasks. Following your suggestion, we conduct the experiments that fine-tune a pre-trained BERT(obtained from Huggingface) on SST-2 data. Experimental results show MAT achieves better performance with less computation.\n\n| Method | Accuracy | Computation (PFLOPs) |\n|---|---:|---:|\n| vanilla | 92.96 | 5.71 |\n| MAT | 93.41 | 3.35 |"}]

Author rebuttal:
# General Response

Thank you very much for reviewing our manuscript and providing detailed and constructive comments, which have been very helpful for us to improve the quality of our work. Please see our answers to address the comments of individual reviewers. Additionally, enclosed pdf with figures is used for `Q4` of reviewer `pEnu`, `Q1` and `Q7` of reviwer `afqQ`.

Here, we provide a general response to the common concerns of reviewers (especially `Q3` of reviewer `TyiH`, `Q1` and `Q6` of reviewer `pEnu`), in terms of the key contributions of our work, mNTK and its implications, the split of information and nuisance space, and their relationship to generalization.

Firstly, we would like to point out that analyzing the eigenvalues of NTK is equivalent to analyzing the singular values of the Jacobian matrix, specifically $\lambda_i (JJ^\top) = \sigma_i^2(J)$, where $\lambda$ denotes the eigenvalue and $\sigma$ denotes the singular value. Next, several prior works have analyzed the relationship between generalization and the singular values of the Jacobian matrix. Oymak et al. notice the low-rank structure of the Jacobian matrix, and the features that fall on the lower part of the Jacobian singular value spectrum are hard to generalize [A]. Li et al. consider clean and noisy labels, splits the residual accordingly as presented below, and proves that the clean residual is aligned with the top singular vectors whereas label noise is aligned with the small singular vectors [B].

$$
\underbrace{\boldsymbol y-f(\boldsymbol W_t)}\_{\text {corrupted residual }}=\underbrace{\tilde{\boldsymbol y}-f(\boldsymbol W_t)}\_{\text {clean residual }}+\underbrace{\boldsymbol y-\tilde{\boldsymbol y}}\_{\text {label corruption }}
$$

Intuitively speaking, mNTK measures the correlation of the gradient that different data samples produce on a certain module. The eigenspectrum of mNTK measures how frequent  data features exisit in the dataset. Features occur frequently are related to the large eigenvalues, while data-specific features, generally considered as noise, related to the small eigenvalues. As the dataset becomes larger, the small eigenvalues, corresponding to the data-specific noise, become even smaller. Therefore, the gap between large and small eigenvalues becomes more significant.

One key contribution of our work is that we observe significant inter-module variance during training, specifically, some modules already converge and fall into the nuisance space, while others are still in the information space. Therefore, we use mNTK $\lambda_{\max}$ as an indicator of modular training. 

We further introduce Radmacher complexity as evaluation of generalization ability from Lemma 2 in Appendix A.2:

**Lemma 2** Given $R>0$, with probability at least $1-\delta$ over the random initialization $(\boldsymbol{\theta}(0), \boldsymbol{a})$, simultaneously for every $B>0$, the following function class

$$
\mathcal F_{R, B}^{\boldsymbol{\theta}(0), \boldsymbol{a}}=\\{f_{\boldsymbol{\theta}, \boldsymbol{a}}:\|\theta_r-\theta_r(0)\|_2 \leq R~(\forall r \in[m]), \\
\|\boldsymbol{\theta}-\boldsymbol{\theta}(0)\|_F \leq B\\}
$$

has empirical Rademacher complexity bounded as:

$$
\mathcal R_S(
\mathcal F_{R, B}^{\boldsymbol{\theta}(0), \boldsymbol{a}})=
\frac{1}{n} \mathbb E_{\boldsymbol{\varepsilon} \in\{ \pm 1\}^n}[\sup_{f \in 
\mathcal F_{R, B}^{\boldsymbol{\theta}(0), \boldsymbol{a}}
} \sum_{i=1}^n \varepsilon_i f(\mathbf{x}_i)] \\
\leq \frac{B}{\sqrt{2 n}}(1+(\frac{2 \log \frac{2}{\delta}}{m})^{1 / 4})+\frac{2 R^2 \sqrt{m}}{\kappa}+R \sqrt{2 \log \frac{2}{\delta}}.
$$

Lemma 2 indicates that Rademacher complexity is proportional to the weight distance from its initialization, where $\|\boldsymbol{\theta}-\boldsymbol{\theta}(0)\|_F \leq B$. Parameters updating of those modules falling into the nuisance space slightly contribute to the loss reduction but increase the Rademacher complexity.

In summary, during model training, modules already fallen into the nuisance space (with low mNTK $\lambda_{\max}$) are prone to learn superfluous or even noise features, which increase the Rademacher complexity with weak trainability, and deteriorate the model's generalization ability.

To the best of our knowledge, this work indicates for the first time that mNTK is a good indicator of inter-module training dynamics, and presents an mNTK-based adaptive training method to optimize the network training process.

---
**References:**

[A] Oymak, Samet, et al. "Generalization guarantees for neural networks via harnessing the low-rank structure of the jacobian." arXiv preprint arXiv:1906.05392 (2019).

[B] Li, Mingchen, Mahdi Soltanolkotabi, and Samet Oymak. "Gradient descent with early stopping is provably robust to label noise for overparameterized neural networks." International conference on artificial intelligence and statistics. PMLR, 2020.


Meta-review: The paper proposes to study the modular-level training dynamic of transformer models from a neural tangent kernel view. It shows an interesting connection between the mNTK's principal eigenvalue and the generalization ability of modular. Reviewers generally appreciate the contribution of this paper.
Decision: Accept (poster)