SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yubin Shi
Yixuan Chen
Mingzhi Dong $^{1,*}$
Xiaochen Yang $^{2,*}$
Dongsheng Li<sup>3</sup> Yujiang Wang<sup>4</sup>
Qin Lv
Yingying Zhao<sup>1</sup>
Fan Yang
Tun Lu<sup>1</sup>
Ning Gu
Li Shang $^{1,*}$
<sup>1</sup>China and Shanghai Key Laboratory of Data Science, School of Computer Science, Fudan University
$^{2}$ School of Mathematics Statistics, University of Glasgow
<sup>3</sup>Microsoft Research Asia, Shanghai, China
$^{4}$ Department of Engineering Science, University of Oxford
$^{5}$ Department of Electrical Engineering and Computer Science, University of Michigan
$^{6}$ Department of Computer Science, University of Colorado Boulder
$^{7}$ School of Microelectronics, Fudan University
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_{\mathrm{max}}$ . A large  $\lambda_{\mathrm{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_{\mathrm{max}}$  exceeding a dynamic threshold selectively, concentrating the model on learning common features and ignoring those superfluous 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.
The proliferation of large-scale pre-trained models (Devlin et al., 2018; Brown et al., 2020; Dosovitskiy et al., 2020) demonstrates the advantages of over-parameterized models. Those models can adequately fit the training data with their sizeable parameters far exceeding the data scale (Du et al., 2018), and such over-parameterizations are known to be essential in model optimization and generalization and are also crucial to their success in deep learning (Liu et al., 2022). However, challenges remain. Over-parameterized models are expensive to train; for example, training large language models (LLMs) such as GPT-3 (Brown et al., 2020; Zaheer et al., 2020) may take weeks to months, even on a large collection of powerful GPUs.
Figure 1: Characteristics of training BERT on WikiText-2. Figure (a) demonstrates the joint variation of effective rank and  $\lambda_{\mathrm{max}}$  across the attention heads.  $\mathrm{Head}_i^l$  refers to the  $i^{\mathrm{th}}$  attention head in the  $l^{\mathrm{th}}$  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).
Modular Adaptive Training (MAT)
Generally, most over-parameterized models consist of various functioning modules connected layerwisely, 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-based architecture inevitably leads to a question: Has the computational resource been spent efficiently and effectively across modules during model optimization? To this end, we dive into the modular-level training behaviors of over-parameterized models and pursue a 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), as a powerful tool for understanding the fine-grained learning dynamics of each network module. We compute mNTK as the first-order approximation of a 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  $\lambda_{\mathrm{max}}$  indicates the degree of consistency in the gradient direction for learning the most common features in data.
To illustrate the module-level training dynamics intuitively, we employ the standard BERT (Devlin et al., 2018) as an example and consider each attention head as a specific module. In addition to  $\lambda_{\mathrm{max}}$  we adopt the effective rank (Roy & Vetterli, 2007) to 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 1(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  $\lambda_{\mathrm{max}}$  as an indicator of trainability to measure how effectively gradient descent can optimize network parameters. As discussed in Bowman & Montufar (2022), the network is biased to learn the top eigenfunctions of the NTK over the entire input space, while the parameters with very low eigenvalues barely update. In Figure 1(a), the attention head in blue marker has large  $\lambda_{\mathrm{max}}$ , 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,  $\lambda_{\mathrm{max}}$  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.
Our analysis (Section 2) on trainability and generalization demonstrates that over-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 a simple yet effective training method termed Modular Adaptive Training (MAT), aiming at dynamically training partial modules to improve training efficiency. As shown in Figure 1(b), we design a dynamic threshold (the dash-dotted line), according to modular difference and temporal variation of mNTK  $\lambda_{\mathrm{max}}$ . MAT only back-propagates the parameters of those modules whose  $\lambda_{\mathrm{max}}$  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 back-propagation to improve learning efficiency.
- Experimental results verify that MAT can significantly reduce the training computation, improve performance, and generalize well to different over-parameterized 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 of 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 to show that trainability of over-parameterized 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  $\mathcal{X}$  denote the set of  $n$  training instance, which are i.i.d. drawn from an unknown distribution  $\mathcal{D}$ , and let  $\mathcal{V} \in \mathbb{R}^{nk}$  denotes the targets. We study the network function  $f$  parameterized by  $\pmb{\theta}$  aiming to map the input vector  $\mathcal{X}$  to output vector  $\mathcal{Z}$ , termed as  $\mathcal{Z} = f(\mathcal{X};\pmb{\theta})$  where  $\pmb{\theta} \in \mathbb{R}^m$  and  $\mathcal{Z} \in \mathbb{R}^{nk}$ . Following the original definition of Neural Tangent Kernel (NTK) (Jacot et al., 2018; Wang et al., 2020), we introduce modular NTK for fine-grained analysis of structured deep networks as below.
Definition 1 (Modular Neural Tangent Kernel (mNTK)). Suppose model  $f$  contains  $L$  disjoint modules  $\pmb{\theta} = \{\pmb{\theta}^1, \pmb{\theta}^2, \dots, \pmb{\theta}^L\}$  where  $\pmb{\theta}^l$  denote the parameters of  $l^{th}$  module. We define mNTK as a matrix by  $\Theta^l(\mathcal{X}, \mathcal{X}) = J_{\pmb{\theta}^l}(\mathcal{X}) J_{\pmb{\theta}^l}(\mathcal{X})^\top$ , where  $J_{\pmb{\theta}^l} = \nabla_{\pmb{\theta}^l} f(\mathcal{X}; \pmb{\theta}^l)$  denotes the Jacobian of the function  $f$  at the points  $\mathcal{X}$  with respect to the  $l^{th}$  module's parameters  $\pmb{\theta}^l$ .
$\Theta^l$  is a positive semi-definite real symmetric matrix and can be eigen-decomposed as  $\Theta^l = \mathbf{U}^l\pmb{\Lambda}^l\mathbf{U}^{l^\top} = \sum_{i=1}^{nk}\lambda_i^l\mathbf{u}_i^l\mathbf{u}_i^{l^\top}$  with  $nk$  non-negative eigenvalues  $\lambda(\Theta^l) = \{\lambda_1^l, \lambda_2^l, \dots, \lambda_{nk}^l\}$ . Note that a network can be partitioned into modules in flexible ways, for example, by dividing layers within a deep network or components of a network with inherent structures (e.g., attention heads in Transformer or convolutional filters in 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):
$$
\Theta (\mathcal {X}, \mathcal {X}) \stackrel {(i)} {=} \sum_ {p = 1} ^ {m} J _ {\boldsymbol {\theta} _ {p}} (\mathcal {X}) J _ {\boldsymbol {\theta} _ {p}} (\mathcal {X}) ^ {\top} \stackrel {(i i)} {=} \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} \stackrel {(i i i)} {=} \sum_ {l = 1} ^ {L} \Theta^ {l} (\mathcal {X}, \mathcal {X}), \tag {1}
$$
where  $(i)$  decomposes the matrix multiplication into the sum of vector multiplication;  $(ii)$  gathers addends by each module;  $(iii)$  follows the definition of mNTK.
2.2 Empirical Analysis
In this subsection, we apply BERT to WikiText-2 for language modeling and regard each layer as a module to analyze the properties of eigenvalue 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 2(a) shows the first 16 eigenvalues of all layer-wise mNTKs at the  $10^{\mathrm{th}}$  epoch. Apparently, the eigenvalue distribution of each layer-wise mNTK is imbalanced where there exists a single large eigenvalue  $\lambda_1^l$  ( $\lambda_1$  will be denoted by  $\lambda_{\mathrm{max}}$  hereafter) and a large number of small eigenvalues. In particular, the largest eigenvalue often has 2 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.
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 at the  $10^{\mathrm{th}}$  epoch. (b): Variation of  $\lambda_{\mathrm{max}}$  of layer-wise mNTKs. (c): Variation of  $\kappa$  of layer-wise mNTKs during the first 20 epochs.
The phenomenon is also indicated by Xiao et al. (2020) in terms of NTK. Furthermore, according to the connection between NTK and mNTK (Eq. 1), the spectrum of NTK is dominated by the principal eigenvalues of all mNTKs. Below, we utilize  $\lambda_{\mathrm{max}}$  for fine-grained optimization analysis of over-parameterized model.
The variation of principal eigenvalues of mNTKs is highly asynchronous across modules. Figure 2(b) shows the variation of  $\lambda_{\mathrm{max}}$  of all layer-wise mNTKs during the training process. It 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  $\lambda_{\mathrm{max}}$  is large after model initialization, and it decreases and converges fast in the early training stages. A potential explanation is that shallow layers learn simple and common features of data samples (Yosinski et al., 2014; Zhang et al., 2019), which contribute significantly to model convergence in the early training stages. In contrast, for mNTK of the deep layer, the magnitude of  $\lambda_{\mathrm{max}}$  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  $\kappa = \lambda_{\mathrm{max}} / \lambda_{\mathrm{min}}$ . Figure 2(c) shows the variation of  $\kappa$  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  $\kappa$  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, we examine the temporal variation of the principal eigenvalue given a specific module and establish its connection with the generalization ability. In detail, we train the same model using masked modeling on half sequence length, which makes the model prone to overfitting. As shown in Figure 3(a), the validation loss starts to increase from the  $45^{\text{th}}$  epoch while the training loss continually decreases. Figure 3(b-e) present the temporal variation of the principal eigenvalue for the  $3^{\text{rd}}$ ,  $6^{\text{th}}$ ,  $9^{\text{th}}$ , and  $12^{\text{th}}$  layers of the model during the training process, respectively. During the first  $45^{\text{th}}$  epochs, i.e., the non-overfitting stage, the principal eigenvalues of these layers continually decrease. However, after the  $45^{\text{th}}$  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, we show that the optimization properties of over-parameterized structured models are closely related to eigenvalues of mNTKs.
Suppose for any module  $\pmb{\theta}^l$ ,  $\lambda_{\mathrm{min}}(\Theta^l) = \lambda_{n,k}^l > 0$ . Let  $\mathcal{L}$  denote the loss function and  $\nabla_{\pmb{\theta}}\mathcal{L}(\pmb{\theta})$  the gradient w.r.t parameter  $\pmb{\theta}$ . For a step of gradient descent, loss reduction can be characterized by the
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):
$$
\begin{array}{l} \Delta \mathcal {L} \stackrel {(i)} {=} \lim  _ {\epsilon \rightarrow 0} \frac {\mathcal {L} (\boldsymbol {\theta} + \epsilon \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta})) - \mathcal {L} (\boldsymbol {\theta})}{\epsilon} \stackrel {(i i)} {\approx} \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}) \\ \stackrel {(i i i)} {=} \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} \left(\nabla_ {\boldsymbol {\theta}} f (\boldsymbol {\theta}) ^ {\top} \nabla_ {\boldsymbol {\theta}} f (\boldsymbol {\theta})\right) \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) \tag {2} \\ \stackrel {(i v)} {=} \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) ^ {\top} \left(\sum_ {l = 1} ^ {L} \boldsymbol {\Theta} ^ {l}\right) \nabla_ {\mathcal {Z}} \mathcal {L} (\boldsymbol {\theta}) \stackrel {(v)} {=} \sum_ {l = 1} ^ {L} \sum_ {i = 1} ^ {n k} \lambda_ {i} ^ {l} \left(\mathbf {u} _ {i} ^ {l} ^ {\top} \mathcal {Y}\right) ^ {2}, \\ \end{array}
$$
where  $(i)$  follows the definition of directional derivative;  $(ii)$  follows the first-order Taylor expansion;  $(iii)$  follows chain rule of derivative;  $(iv)$  follows Eq. 1; and  $(v)$  follows the eigen-decomposition 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.,  $(\mathbf{u}_i^l^\top \mathcal{V})^2$  is large for large  $\lambda_i^l$ . Thus, the directional derivative of the loss function can be regarded as closely related to the eigenspectrum of 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 the principal eigenvalues, we utilize  $\lambda_{\mathrm{max}}$  of mNTKs as the nearly equivalent proxy of the spectrum of mNTKs. Therefore, Eq. 2 is simplified as:
$$
\Delta \mathcal {L} \approx \sum_ {l = 1} ^ {L} \lambda_ {\max } ^ {l} \left(\mathbf {u} _ {1} ^ {l} ^ {\top} \mathcal {Y}\right) ^ {2}. \tag {3}
$$
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  $\lambda_{\mathrm{max}}^l$  to achieve efficient learning with limited computational resources.
(2) The relation between generalization ability and eigenvalues of mNTKs.
Figure 4 shows the eigen-spectrum distribution of a 4-layer BERT model at the  $10^{\mathrm{th}}$  epoch, by regarding each head as a module and calculating 256 eigenvalues for each head. Among all  $4 \times 12 \times 256$  eigenvalues, the distribution of eigen-spectrum exhibits two distinct regions: the left region with a large number of small eigenvalues and the right region with a small number of large eigenvalues. The two regions are widely separated. According to Eq. 3, the eigenspaces corresponding to miniature eigenvalues in the left region have negligible impact on training loss minimization, while those in the right region make dominant contributions. Following the definition of Oymak et al. (2019), we refer
Figure 4: The normalized eigen-spectrum distribution exhibits two distinct regions, termed information space and nuisance space.
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.
Model training with more parameter updates tends to have longer distance to the initializations  $\theta_0$  (Hoffer et al., 2017). Based on Lemma 5.4 in Arora et al. (2019), longer distance results in relatively larger Rademacher complexity, as  $\mathcal{R} \propto \| \pmb{\theta} - \pmb{\theta}_0\|_F$ , and larger  $\mathcal{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 eigenvalues 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, we advocate dynamically optimizing partial module parameters in back propagation, termed as Modular Adaptive Training (MAT). The forward propagation 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 1 summarizes the proposed MAT algorithm.
Modular Policy. Following the analysis that trainability and generalization of an over-parameterized model are related to the principal eigenvalues of module parameters, we propose to only update partial module parameters with  $\lambda_{\mathrm{max}}$  greater than a 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  $\lambda_i^l$  denote  $i^{\mathrm{th}}$  eigenvalue of the mNTK computed over the parameters  $\pmb{\theta}^{l}$ . We define  $\widetilde{\Lambda} = \{\lambda_i^l\}_{i,l=1}^{nk,L}$  as the set of eigenvalues calculated over all mNTKs, and  $\widetilde{\lambda}_{\mathrm{max}}, \widetilde{\lambda}_{\mathrm{min}}$  as the maximum and minimum value of the set  $\widetilde{\Lambda}$ , respectively. To select the appropriate modules, we introduce a hyperparameter ratio  $\alpha \in (0,1)$ , thus the corresponding eigenvalue threshold  $\lambda_{\alpha}$  is computed as:
$$
\lambda_ {\alpha} = \widetilde {\lambda} _ {\min } + \left(\widetilde {\lambda} _ {\max } - \widetilde {\lambda} _ {\min }\right) \times \alpha . \tag {4}
$$
According to the eigenvalue threshold  $\lambda_{\alpha}$ , all modules are split into information modules and nuisance modules as: