SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yefan Zhou\* 1, Tianyu Pang\*2, Keqin Liu2, Charles H. Martin3, Michael W. Mahoney4, and Yaoqing Yang1
$^{1}$ Department of Computer Science, Dartmouth College
$^{2}$ National Center for Applied Mathematics and Department of Mathematics, Nanjing University  $^{3}$ Calculation Consulting
$^{4}$ ICSI, LBNL, and Department of Statistics, University of California at Berkeley
{yefan.zhou.gr, yaoqing.yang}@dartmouth.edu
tianyupang628@gmail.com, kqliu@nju.edu.cn
charles@CalculationConsulting.com, mmahoney@stat.berkeley.edu
Regularization in modern machine learning is crucial, and it can take various forms in algorithmic design: training set, model family, error function, regularization terms, and optimizations. In particular, the learning rate, which can be interpreted as a temperature-like parameter within the statistical mechanics of learning, plays a crucial role in neural network training. Indeed, many widely adopted training strategies basically just define the decay of the learning rate over time. This process can be interpreted as decreasing a temperature, using either a global learning rate (for the entire model) or a learning rate that varies for each parameter. This paper proposes TempBalance, a straightforward yet effective layer-wise learning rate method. TempBalance is based on Heavy-Tailed Self-Regularization (HT-SR) Theory, an approach which characterizes the implicit self-regularization of different layers in trained models. We demonstrate the efficacy of using HT-SR-motivated metrics to guide the scheduling and balancing of temperature across all network layers during model training, resulting in improved performance during testing. We implement TempBalance on CIFAR10, CIFAR100, SVHN, and TinyImageNet datasets using ResNets, VGGs and WideResNets with various depths and widths. Our results show that TempBalance significantly outperforms ordinary SGD and carefully-tuned spectral norm regularization. We also show that TempBalance outperforms a number of state-of-the-art optimizers and learning rate schedulers.
Having a learning rate schedule that gradually decreases over time is crucial for the convergence and performance of state-of-the-art machine learning algorithms. Indeed, many optimization algorithms essentially boil down to designing a progression of parameter updates, as realized by different learning rate schedules [1-4]. Common schedules assign a global learning rate per epoch, where the same learning rate is used for all layers in the model. This includes the family of cyclical learning rates [3], and parameter-wise learning rate schedules like Adam [2] and its variants [5, 6]. However, such a global learning rate schedule does not take into account the structural characteristics of neural networks (NNs). At the same time, parameter-wise learning rate schedules are sometimes used, but they have long been conjectured to have worse generalization performance than carefully tuned
stochastic gradient descent (SGD) optimizers [7], and storing both first and second-order moments for each parameter can lead to substantially increased memory consumption [8]. As mentioned in [9], storing the whole Megatron-Turing NLG requires 10 terabytes of aggregate memory, and the Adam optimizer's first and second-order moments [2] consume  $40\%$  of the memory. Nonetheless, improving parameter-wise learning rate schedules is an active field of study [4, 5, 10, 11].
A largely under-explored idea to reconcile the two extremes of setting a single global learning rate or assigning fine-grained parameter-level learning rates is to assign layer-wise learning rates. Such a learning rate assignment method does not require much storage cost, and it can assign very different training speeds to different layers. However, existing layer-wise schemes are often introduced as an additional part of hyperparameter sweeping, thus substantially increasing computational cost; and most lack a strong (or any) theoretical foundation. For instance, layer-wise learning rates can increase test accuracy in transfer learning [12] and domain adaptation [13], but these learning rates are often empirically tuned. More recently, motivated by the intuition that lower-level layers should be domain-specific and higher-level layers should be task-specific, [14] automates the search for an optimal set of learning rates. However, the authors find the nested, bi-level optimization scheme to be too computationally expensive in practice [15]. AutoLR also automatically tunes its layer-wise learning rates according to the "role" of each layer [16]. The method is validated almost entirely by empirical results, further explained by layer-wise weight variations. While the authors attempt to assign a different initial learning rate to each layer, the learning rate for each layer continues to stay largely constant throughout training. LARS [17, 18] is another method to assign layer-wise learning rate. It is based on the "trust ratio," defined as the ratio of weight norm to gradient update norm of each layer, and it is specifically used in large batch training to avoid gradient divergence.
In this paper, we propose TempBalance, a simple yet effective layer-wise learning rate assignment (and regularization) method, grounded in Heavy-Tail Self Regularization (HT-SR) Theory [19-24]. Our approach leverages HT-SR Theory to assess the quality of each network layer. This is achieved through an analysis of the heavy-tail (HT) structure present in the Empirical Spectral Density (ESD) of NN weight matrices. Given this information, TempBalance meticulously adjusts the temperature-like parameter to control each layer's quality, with the objective of ensuring consistently high quality across all layers of the network. From the statistical physics viewpoint on learning and optimization [24-28], a temperature-like parameter refers to some quantity related to the empirical noise/stochasticity of the learning process. This is the noise scale described by [29, 30], and it can be written as a function of learning rate, batch size, and momentum. Prior research [19, 31] has shown that temperature-like parameters significantly influence HT structure in the ESD. Our approach, TempBalance, focuses on the strategic adjustment of the learning rate as the temperature-like parameter, thereby facilitating accurate control of the quality across each network layer, as characterized by its HT ESD structure. The following paragraph will delve deeper into the importance of HT-SR, highlighting its connection to the concept of layer-wise temperature.
HT-SR Theory. HT-SR Theory [19-24] relies on the empirical fact that very well-trained models tend to exhibit strong correlations, resulting in HT structure in the ESD of each layer. To obtain this ESD, we take a NN with  $L$  layers and the corresponding weight matrices  $\mathbf{W}_1, \mathbf{W}_2, \dots, \mathbf{W}_L$  with shape  $n \times m$  (where  $n \leq m$ ). For the  $i$ -th layer, we calculate the eigenvalues of its correlation matrix  $\mathbf{X}_i = \mathbf{W}_i^T \mathbf{W}_i$ , and then we plot the ESD for that layer. Upon training, the ESD will typically gradually change to have an HT structure [19, 22]. We can then fit a power law (PL) distribution to the HT part of the ESD, and extract its exponent as, namely, PL_A1pha. The fitted PL will have the following formula:
$$
p (\lambda) \propto \lambda^ {- \alpha}, \quad \lambda_ {\min } <   \lambda <   \lambda_ {\max }. \tag {1}
$$
The PL_Alpha metric measures the PL exponent of the weight matrices' ESD. Its underlying motivation stems from random matrix theory and statistical physics, as well as the empirical observation that HT ESDs are ubiquitous in well trained NN models [20, 22].
The PL_A1pha metric has been shown to predict the trends in the test accuracy of state-of-the-art models in computer vision (CV) and natural language processing (NLP), without even the need for access to training or testing data [22, 32]. According to [22], one can aggregate PL_A1pha's for different layers either by simple averaging or weighted averaging, and each can predict test accuracy in different cases [22, 32]. Furthermore, the layer-wise nature of PL_A1pha makes it a fine-grained metric that can be used to assess the quality of individual layers of the network. Thus, in this paper, we extend and apply HT-SR Theory (originally designed as a predictive diagnostic for analyzing
Figure 1: Examples of PL fitting. Blue histograms depict the ESDs. Vertical black lines indicate the lower threshold  $\lambda_{\mathrm{min}}$  used to truncate the full ESDs and extract the tail portion. Solid red curves represent the tail part of the ESDs truncated by  $\lambda_{\mathrm{min}}$ , while dashed red curves represent the fitted HT distributions. The left shows a more HT ESD, which requires a relatively lower learning rate. The right one shows a less HT ESD, which requires a relatively higher learning rate. Unlike prior work, we do not aim to find the "optimal" PL exponent. (Thus, we are less interested in obtaining a precise estimate than in obtaining a robust estimate.) Instead, we use the PL exponent to rank ESDs to find layers that need higher/ lower learning rates. These two ESDs correspond to two layers of a ResNet18 model trained on TinyImageNet.
pre-trained NN models) to NN training, and we exploit the layer-wise information provided by PL_Alpha to determine the layer-wise learning rates for better test accuracy.
We note that, while it provides perhaps the most principled approach, the PL_Alpha metric is not the only way to try to measure the HT structure in NN models. Several recent papers [33-35] use different HT metrics to try to measure the spectral properties of other matrices (such as input/output covariance matrices, Fisher Information Matrices, and the Hessian). We show in Appendix A that these HT phenomena, measured in different ways on different matrices, are closely related to each other. On the other hand, this also means that (for the problems considered in this paper) the absolute numerical value of PL_Alpha is less important, as optimal PL exponents estimated by different algorithms can be different [22, 33]. What matters the most, as we show in this paper, is the layer-wise quality ranked by the PL exponent: layers with a smaller PL_Alpha tend to be relatively more "overtrained," and layers with a larger PL_Alpha tend to be relatively more "undertrained." (We emphasize that this is true for the training problem we consider in this paper—for prior HT-SR work, the actual numerical value of PL_Alpha mattered a lot.)
This observation leads to a simple and efficient way to balance layer-wise learning rates: assign a lower learning rate to more overtrained layers and a larger learning rate to more undertrained layers, using PL_Alpha (see Figure 1). In implementing this learning rate balancing approach, we use a scale-free method to map the PL_Alpha value of each layer to a predefined learning rate range. This range is established in relation to a global learning rate. Rather than depending on the absolute numerical values of PL_Alpha for each layer, this method emphasizes the importance of their relative differences and quality ranking. As a result, the learning rates assigned to individual layers remain stable and unaffected by arbitrary linear scaling of PL_Alpha estimates, whether they arise from the choice of the estimator or the presence of noisy measurements. On top of this, we can perform a grid search on the global learning rate. This is standard practice, and it is more efficient than grid-searching the layer-wise learning rates. We use this combination of assigning layer-wise learning rates using PL_Alpha and grid-searching the base global learning rate to avoid having to decide the "optimal" PL exponent, as this can be tricky due to different ways of measuring HT properties. Indeed, there are different ways to measure PL_Alpha [19], and we use the Hill estimator [36]. While not necessarily the best estimate (see [19, 22]), it shows stable performance in our experiments. We refer to our version of the PL_Alpha metric as the PL_Alpha_Hill metric, and we use it for the remainder of the paper.
Another popular way to change the ESD of weights is to constrain the spectral norm (i.e., the largest eigenvalue) using spectral norm regularization (SNR) [37, 38]. SNR provides a different form of regularization, compared to HT-SR, because it regulates the largest eigenvalue instead of the ESD slope (i.e., the PL_Alpha_Hi11 metric). It has been demonstrated that the spectral norm
and PL_Alpha_Hill serve distinct roles in evaluating model quality, and their combined form yields optimal predictions for test accuracy trends [19, 22, 23, 32]. To complement this, our results demonstrate that TempBalance outperforms SNR in training deep NNs in most cases. Moreover, when these two regularization methods are combined during training, they result in optimal test accuracy, thereby confirming their complementary roles. As described in [23, 32], the spectral norm and PL_Alpha_Hill measure the scale and the shape of a ESD, respectively; and regulating both the scale and shape is crucial for achieving better ESD regularization. We provide ablation studies on several layer-wise metrics for assigning layer-wise learning rates, including spectral norm, and we show that PL_Alpha_Hill performs the best among them.
Our main contributions. The following summarizes our main contributions.2
- We propose a simple yet effective layer-wise learning rate schedule, TempBalance, which is motivated by HT-SR Theory. Based on our empirical results, we obtain two main high-level insights. First, the mapping from PL_Alpha_Hi11 to learning rates should be scale-free, meaning that arbitrary linear scaling on the estimated PL exponent should not change the learning rate assignment. Second, searching for the minimum eigenvalue  $\lambda_{\mathrm{min}}$ , a standard practice in PL fitting [19, 39, 40], leads to unstable training. To improve stability, we instead fix  $\lambda_{\mathrm{min}}$  as the median of the ESD.
- We compare TempBalance to ordinary SGD and SNR on various training tasks. This includes (1) different network architectures, such as ResNet, VGG, WideResNet, (2) different datasets, such as CIFAR10, CIFAR100, SVHN, TinyImageNet, and (3) ablation studies, such as varying widths, depths, initial learning rates, HT-SR layer-wise metrics, and PL fitting methods. Compared to ordinary SGD, TempBalance achieves higher test accuracy by setting layer-wise learning rates. Compared to SNR, TempBalance performs better by providing a more fine-grained regularization on shape/slope instead of norm. We also show that combining TempBalance and SNR leads to further improved accuracy, verifying their complementary roles in informing deep learning training.
- We compare TempBalance to a range of state-of-the-art optimizers and learning rate schedulers, including SGDR [10], SGDP [41], Lookahead [42] and LARS [17, 18] on ResNet18 and ResNet34 trained on CIFAR100. We show that TempBalance achieves the highest test accuracy. We do careful hyperparameter tuning for all baselines. All results are obtained from five random seeds.
- We use ablation studies to show that PL_Alpha_Hill provides the best test accuracy among several layer-wise metrics considered by HT-SR [22, 32]. We also show that TempBalance maintains stable performance over SGD baselines when the model size changes. Furthermore, we show visualization results in Appendix B, verifying that TempBalance controls ESDs during training.
Here, we give an overview of the statistical mechanics of learning and recent progress in theoretical and empirical studies on generalization metrics and their applications.
Statistical mechanics of learning and HT-SR. Our paper is motivated by statistical mechanics of learning [43-45], and especially by works that connect load-like [43, 46, 47] and temperature-like parameters [25, 48] to NNs. According to prior works in this area [24, 49], a temperature-like parameter represents the amount of noise/variance in an iteration of SGD, such as learning rate, weight decay parameters, and batch size. A load-like parameter represents the quantity and/or quality of data relative to the size of the learning model. To measure the quality of publicly-available pre-trained NNs, Martin and Mahoney [19] introduce HT-SR Theory, showing that the weight matrices of deep NNs exhibit HT ESDs, and they show that a decay coefficient of ESD, PL_A1pha, effectively gauges model quality. Subsequently, [31, 50-54] provide rigorous bounds for HT phenomenon and generalization, further adding support to HT-SR Theory. HT-SR has also been applied to predicting trends in test accuracy of large-scale NNs, in both CV and NLP [22, 23, 32], but it has yet to be systematically incorporated to novel training algorithms. Recently, more papers realize the important connections between deep NNs and statistical mechanics of learning [24]. To name a few, Yang et al. [49] use load and temperature parameters to study a wide range of loss landscapes, providing a taxonomy from the perspective of global structure of a loss landscape. On the theory side, Baity-Jesi et al. [55] investigates the glassy behavior of NNs, and Barbier et al. [56] derives the optimal generalization error of generalized linear systems. More recently, Sorscher et al. [57] studies easy versus hard samples used in training and design a "data-pruning" method; and Zhou et al. [58] establishes a
$^{2}$ Our code is open-sourced: https://github.com/YefanZhou/TempBalance.
(a) ESD Analysis and PL fitting
(b) Layerwise Learning Rate Assignment
Figure 2: The pipeline diagram of TempBalance. In each epoch, TempBalance undergoes two steps: (a) Performing ESD analysis on all layers and employing PL fitting to derive the layer-wise PL_Alpha_Hill, and (b) Using the layer-wise PL_Alpha_Hill to assign learning rates to each layer using an assignment function.
"three-regime model" in network pruning, unifying multiple practical hyperparameter tuning methods in a principled way.
Generalization measures. The search for effective and robust generalization metrics (which, importantly, can be very different than model quality metrics [32]) has been the focus of several recent theoretical and empirical works [22, 32, 49, 59-61]. Several recent papers apply metric-informed training and architecture search, such as those based on the Hessian [4, 62-64], spectral norm [37, 38], stable rank [65], and the spectrum of the neural tangent kernel [66]. However, most generalization metrics, such as those based on the PAC-Bayes bounds [67-70], do not straightforwardly transfer to layer-wise quality metrics, because such generalization metrics often study the whole NN as an architecture-free function, and they lack the fine granularity to unveil the quality of each layer. Also, it has been mentioned in the literature [59] that (1) directly regularizing generalization metrics can lead to difficulty in training, (2) evaluating these regularization methods may be hard due to the existence of implicit regularization in SGD, and (3) these metrics, especially norm-based metrics, cannot be expected to correlate with test accuracy causally [60], making the link between these generalization metrics and practical training methods nuanced. It will be clear in the next section that we do not regularize ESD metrics directly. Instead, we change learning rates to modify ESDs.
3 The TempBalance Algorithm
In this section, we introduce our simple yet effective method TempBalance, based on the PL_Alpha_Hill metric from HT-SR Theory. For a NN, different layers tend to have different values for PL_Alpha_Hill, [19, 24]: a layer with a larger PL_Alpha_Hill indicates that layer is relatively undertrained, while a layer with a smaller PL_Alpha_Hill indicates that layer is relatively overtrained. A natural idea is to adjust the degree of learning among different layers to get a balance: for a layer whose PL_Alpha_Hill is too large, we could assign a larger learning rate to accelerate its learning, and vice versa. The intuition of our method is transferring one layer's learning rate to another and hence, TempBalance. The pipeline is in Figure 2.
Algorithm 1: TempBalance
Input:  $M$ : Deep NN,  $T$ : Total training epoch,  $t$ : Current epoch,  $\alpha_{t}^{i}$ :  $i_{\mathrm{th}}$  layer's PL_Alpha_Hill at epoch  $t$ ,  $\eta_{t}$ : Baseline global learning rate at epoch  $t$ ,  $s_{1}, s_{2}$ : Minimum and maximum scaling ratio,  $f_{t}$ : Learning rate schedule function
1 Initialize model  $M$
2 for  $t\gets 0$  to  $T$  do
3 Compute  $\alpha_{t}^{i}$  for all layers using the Hill estimator;
4 | Leverage all  $\alpha_{t}^{i}$  and adopt  $f_{t}$  in (2) to assign per-layer learning rate  $f_{t}(i)$  between  $s_1\eta_t$  and  $s_2\eta_t$  for the next epoch;
5 Update the optimizer for the next epoch;
We provide the details of TempBalance in Algorithm 1. Based on PL_Alpha_Hill in different layers, we use the learning rate schedule function  $f_{t}$  to map the  $i$ -th layer to a particular learning rate  $f_{t}(i)$  in epoch  $t$ . We adopt  $f_{t}$  as a linear map between the layer-wise PL_Alpha_Hill and the final layer-wise learning rate, which has the following formula:
$$
f _ {t} (i) = \eta_ {t} \cdot \left[ \frac {\alpha_ {t} ^ {i} - \alpha_ {t} ^ {\operatorname* {m i n}}}{\alpha_ {t} ^ {\operatorname* {m a x}} - \alpha_ {t} ^ {\operatorname* {m i n}}} \left(s _ {2} - s _ {1}\right) + s _ {1} \right], \tag {2}
$$
where  $\eta_{t}$  means the base global learning rate in epoch  $t$ ,  $(s_1, s_2)$  are the minimum and maximum learning rate scaling ratio relative to  $\eta_{t}$ ,  $\alpha_{t}^{i}$  represents the layer  $i$ 's PL_Alpha_Hill across all the layers in epoch  $t$ . Using (2), we ensure that the new learning rate  $f_{t}(i)$  is a scaled version of the original base learning rate  $\eta_{t}$  and is always inside the interval  $[s_1 \eta_t, s_2 \eta_t]$ . Note that  $(s_1, s_2)$  serves as tunable hyperparameters in our method. We conducted ablation studies on it, which are detailed in Appendix C. The hyperparameter values used across all experiments can be found in Appendix D. Our studies reveal that the optimal results are usually achieved around (0.5, 1.5).
To fit the PL distribution  $p(\lambda)$  defined in (1), we use the Hill estimator [36, 71]. (It is not the best estimator for fine-scale diagnostics based on HT-SR Theory [19, 22], but it is robust, and it suffices for our purposes.) For the  $i$ -th layer, suppose the weight matrix is  $\mathbf{W}_i$  and the correlation matrix  $\mathbf{W}_i^\top \mathbf{W}_i$  has ascending eigenvalues  $\{\lambda_i\}_{i=1}^n$ . Then, the Hill estimator calculates PL_Alpha_Hill using the following:
$$
\text {P L} _ {\text {A l p h a}} \text {H i l l} = 1 + \frac {k}{\left(\sum_ {i = 1} ^ {k} \ln \frac {\lambda_ {n - i + 1}}{\lambda_ {n - k}}\right)}, \tag {3}
$$
where  $k$  is the adjustable parameter. We adopt  $k = \frac{n}{2}$  in our experiments. Note that changing  $k$  essentially changes the lower eigenvalue threshold  $\lambda_{\mathrm{min}}$  for (truncated) PL estimation, as shown by the vertical black line in Figure 1. Choosing  $k = \frac{n}{2}$  means using the largest half of the eigenvalues to estimate the slope. We empirically find that fixing  $k$  for all layers leads to more stable performance than searching  $k$  for different layers (e.g., optimizing  $k$  using the Kolmogorov-Smirnov test [40], as is needed for other applications of HT-SR Theory [19, 22]).
One advantage of mapping PL_Alpha_Hill to learning rates using (2) is that the scale of PL_Alpha_Hill is unimportant, i.e., linearly scaling PL_Alpha_Hill arbitrarily does not change the learning rate assignment because the linear scaling cancels each other in (2). This can maximally reduce the artifact of estimating the ESD PL exponent/slope due to estimation noise, which has been found to be a tricky issue in practice [19, 23].
In this section, we give full details of the experimental setup (Section 4.1) and compare our method TempBalance to a few baselines (Section 4.2), and then (Section 4.3) we perform ablation studies on varied initial learning rates, model widths, HT-SR layer-wise metrics, and PL fitting methods.