Temperature Balancing, Layer-wise Weight Analysis, and Neural Network Training
Authors: Yefan Zhou
Tianyu Pang
Keqin Liu
charles h martin
Michael W. Mahoney
Yaoqing Yang
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Heavy-tail self-regularization
learning rate schedule

Abstract: 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.

Temperature Balancing, Layer-wise Weight Analysis, and Neural Network Training
Yefan Zhou Tianyu Pang*, Keqin Liu?, Charles H. Martin', Michael W. Mahoney', and Yaoqing Yang'
Department of Computer Science, Dartmouth College
National Center for Applied Mathematics and Department of Mathematics, Nanjing University CCalculation Consulting
*ICSI, LBNL, and Department of Statistics, University of California at Berkeley
{yefan.zhou.ge gr, yaoqing -aaianngggaaee ydan))darrmooutt edu tianyupang6280gaail com, kqliuβnju edu cn
clarlle@callolaattooooooattttnni com, rmahoney@stat, berkeley.ee edu
Abstract
Regularization modern machine learning crucial, and 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 temperature-like parameter within the statistical mechanics of learning, plays a crucial role 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 learning rate that varies for each parameter. This
rate (for the entire model) paper proposes lempBalance, a straightforward yet effective layer-wise learning rate method. TempBalanee ance is based Heavy-Tailee Self- Regfaaeeulaiizttionn (HT-SR) Theory, an approach which characterizes the implicit sell-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 ance on CIFARIO CIFAR 100, SVHN, and ImaylmageNet datasets using ResNets, VGGs and WideResNets with various depths and widths. Our results show that TempBalance significantly outperforms ordinary SGD and carefully- -uuned spectral norm regularization. We also show that TempBalance ance outperforms number of state-of-the-ar optimizers and learning rate schedulers.
Introduction
Having learning rate schedule that gradually decreases 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 progression of parameter updates, realized by different learning rate schedules [[4]] Common schedules assign a global learning rate per epoch, where the same learning rate is used for all layers the model. This includes the family of cyclical learning rates (3), and parameter-wise learning rate schedules like Adam 2 and its variants [5] [). However, such global learning rate schedule does 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
First two authors contributed equally.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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-TTurin NL. requires 10 terabytes of aggregate memory, and the Adam optimizer's 's first and second-order moments
2 consume 40% of the memory. Nonetheless, improving parameter-wist learning rate schedules is an active field of study 50 11]
largely under-explored idea to reconcile the two extremes of setting 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 [2] 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-specifte 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 [1S.. AutoLR also automatically tunes its layer-wise learning rates according to the "role" of each layer (6). 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 [7 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.
layer-wise
In this paper, we propose TempBalance, a simple yet effective layer- learning rate assignment (and regularization) method, grounded in Heavy-Tail Self Regularization (HT-SR) Theory 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-lin like parameter to control each layer's '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 a temperature- 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- 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- temperature.
HT-SR Theory. HT-SR Theory (4] 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 I. layers and the corresponding weight matrices W1, W.,... W. with shape n x (where n S m). For the i-th layer, we calculate the eigenvalues of its correlation matrix X; = W W., 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 power law (PL) distribution to the HT part of the ESD, and extract its exponent as, namely, following formula:
The fitted PL will have the
p(A) x x= Amin <> X Amax-
The PL_Alpha metric measures the PL exponent of the weight matrices' ESD. Its underlying motiva- tion 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_Alpha metric has been shown to predict the trends in the test accuracy of state-of of-the--rt 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_Alpha's '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_ PI__Apph makes it a fine-grained
metric that can be used to assess the quality of individual layers of the network. Thus, in this paper,
extend and apply HT-SR Theory (originally designed as a predictive diagnostic for analyzing
2

---

PL Alpha Hill 1.669
PL Alpha Hill 2.115
Amin
Amin
10
10²
ESD
ESD
10
10
10"
10
10
10
10
10
10
10"
Eigenvalues of correlation matrix
Eigenvalues of correlation matrix
Figure 1: Examples of PL fitting. Blue histograms depict the ESDs. Vertical black lines indicate the lower threshold Ain used to truncate the full ESDs and extract the tail portion. Solid red curves represent the tail part of the ESDs truncated by Amin- while dashed red curves represent the fitted HT distributions. The left shows more HT ESD, which requires 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 find the "optimal" PI. exponent. (Thus, we are less interested in obtaining a precise estimate than in obtaining a robust estimate.) Instead, we use the PI. exponent to rank ESDs to find layers that need higher/lower learning rates. These two ESDs correspond two layers of 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-wist learning rates for better test accuracy.
We note that, while 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 33335 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_A1 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 Al_hlpha tend 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 more overtrained layers and a larger learning rate to more undertrained layers, using PL_Alpha (see Figure l). In implementing this learning rate balancing approach, we use scale-free method to map the PL_Alpha value of each layer predefined learning rate range. This range is established in relation to 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_ 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 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-scarching the base global learning rate to avoid having to decide the "optimal" PL exponent, as this can be tricky due to different ways of mcasuring HT properties. Indeed, there are different ways to measure PL Alpha [9), and we use the Hill estimator (36). While not necessarily the best estimate (see [19 2 it shows stable performance in our experiments. We refer to our version of the PL_Alpha metric as the PL Altpaalpaaa__a 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 different form of regularization, compared to HT-SR, because regulates the largest eigenvalue instead of the ESD slope (i.c., the PL_ PL_Alpha_Hii Hill metric). It has been demonstrated that the spectral norm
3

---

and PL_AIpha_HHe 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 (33 [2], the spectral norm and PL_Alpha__te 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_ PL_Alpha_Hail Hill performs the best among them.
Our main contributions. The following summarizes our main contributions
We propose a simple yet effective layer-wise learning rate schedule, TempBal ance, 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, _Hlll 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 Amin- a standard practice in PL
fitting 19 39 401, leads unstable training. To improve stability, we instead fix Amin 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 CIFARIO, CIFARI00, SVHN TinylmageNet, and (3) ablation studies, such as varying widths, depths, initial learning rates, HT-SR layer-wise metrics, and PI. 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 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 TempBal ance to a range of state-of-the-art optimizers and learning rate schedulers, including SGDR [10), SGDP [T], Lookahead 42) and LARS [17 18] on ResNetl8 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 PL_Alpha Hill provides the best test accuracy among several layer-wise metrics considered by HT-SR 222 32]. We also show that TempBalance maintains stable performance over SCD baselines when the model size changes. Furthermore, we show visualization results in Appendix B verifying that TempBalance controls ESDs during training.
2
Related Work
Here, we give 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 temperature- -lie 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-availabe
pre-trained NNs, Martin and Mahoney [19] introduce HT-SS Theory, showing that the weight matrices of deep NNs exhibit HT ESDs, and they show that a decay coefficient of ESD, PL_ PL_Alpha, effectively gauges model quality. Subsequently, [31 50-444 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 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, BJity-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
"Our code is open-sourced. http://gggaaloocmo///ee/z//uu/on

---

More heavy-tailed
Less heavy-tailed
PL fitting
Btting
PL fitring
ESD
ESD
ESD
(a) Analysis and PL fitting
Less heavy-tailed, larger LR
heavy-tailed, smaller LR
ESD PL exponent
Learning Rate (TempBalunce)
Learnnng Rute (Uniform)
(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 PL_Alpha_ _ilpa__iil, and (b) Using the layer-wise PL_ Aalpa_ppaa_H_____ _illl to assign learning rates to each layer using assignment function.
"three-regime model" in network pruning, unifying multiple practical hyperparameter tuning methods 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 [1-66] 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 666 However, most generalization metrics, such as those based on the PAC-Bayes bounds 167/70) do not straightforwardly transfer layer-wise quality metries, 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, has been mentioned the literature
[99] 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- norm-baser based metries. cannot be expected to correlate with test accuracy causally [60], making the link between these generalization metrics and practical training methods nuanced. will be clear in the next section that we do not regularize ESD metrics directly. Instead, we change learning rates modify ESDs.
The TempBalance Algorithm
In this section, we introduce our simple yet effective method TempBalance, based on the PI_Alpha_Hiil __ill metric from HT-SR Theory. For NN, different layers tend to have different values for PL_ Al_Alpha Hill, A1lha_Hill, [9 24) a layer with larger PL_Alpha_Hill AapphaaHHlllllllp indicates that layer is relatively undertrained, while a layer with smaller PL_Alpha_Hiil _iill indicates that layer is relatively overtrained. natural idea for layer whose PL_Alpha_Hal _ailllllll- is too large, could assign a larger learning rate to accelerate its learning, and vice versa. The intuition of our method transferring one layer's learning rate to another and hence, TempBalance. The pipeline is in Figure 2
adjust the degree learning among different layers to get a balance:
5

---

Algorithm 1: TempBalance
Input: M: Deep NN, T: Total training epoch, t: Current epoch, oj in layer's
PL_Alpha, Hill at epoch T/: Baseline global learning rate at epoch 82: Minimum and maximum scaling ratio,
ft: Learning rate schedule function
Initialize model M; for i 4- 0 T' do
Compute of for all layers using the Hill estimator;
Leverage all of and adopt ft in (2) to assign per-layer learning rate fi(i) between f)e and
527) for the next epoch;
Update the optimizer for the next epoch;
end
We provide the details of TempBalance in Algorithm 1 Based on PL_Alpha A1l_A_phaaHill111111111111111111 in different layers, use the learning rate schedule function fo map the i-th layer to a particular learning rate fi(i) in epoch t. We adopt fi as a linear map between the layer-wise PL_Alpha_ Hill and the final layer-wise learning rate, which has the following formula:
aj Ot
min
(e(1)
1)
(82 81)
max
where means the base global learning rate in epoch f, (51, $2) are the minimum and maximum learning rate scaling ratio relative to Tk. aj represents the layer i's @o_Alpha_Hill at the beginning of epoch f, and (a. a denote the minimum and maximum PL_ Ml__Apph__Hie
across all the layers in epoch Using 2 we ensure that the new learning rate fe(i) scaled version of the original base learning rate Ne and is always inside the interval (t SyT]). Note that ($1, 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 PI. distribution p(A) defined in ①.. we use the Hill estimator [36] (1]. (It is not the best estimator for fine-scale diagnostics based on HT-SR Theory 119 222. but it is robust, and it suffices for our purposes.) For the i-th layer, suppose the weight matrix is W, and the correlation matrix W. T W, has ascending eigenvalues in Then, the Hill estimator calculates PL_Alpha_Hie Hill using the following:
k
PL AL_hapha_HHe
Hilpha_Hillle Hill
+ E i=1 In An-it1
where k is the adjustable parameter. We adopt k 2 in our experiments. Note that changing essentially changes the lower eigenvalue threshold Amin for (truncated) PL estimation, as shown by the vertical black line in Figure Choosing & means using the largest half of the eigenvalues to estimate the slope. We empirically find that fixing i for all layers leads to more stable performance than searching & for different layers (e.g.. optimizing & using the Kolmogorov-Smme Smiiovv test [40], as is needed for other applications of HT-SR Theory 19 (2)).
One advantage of mapping PL PL_Alpha_Har _Hill to learning rates using 2 that the scale of PI__lppaa_Han _Hill is unimportant, i.e., linearly scaling PL_ PI_Alpha_Haal _illlllll 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
4 Empirical results
In this section, we give full details of the experimental setup (Section 4.1 and compare our method TempBalance to a few baselines (Section p.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.

---

CAL
SNR
TB
TB + SNR
Tea Avcumccy
70.5
Teat Accuncy
16.5
Teat Accunaty
760
Text Accurcy
15.0
5000
15.5
19.5
750
18.5
74.5
(a) ResNet18, CIFARIO0 (b) ResNer34, CIFARI00
(c) VGGI6, CIFARI00
(d) VGG19, CIFARIO0
Tok Accurage
Tex Aesuure
Tos Axcurayy
Toe Accuroey
68.5
(c) R8 TIN
(I) ResNet34, TIN
WRN16-8, TIN
(h) WRN28-6, TIN
Test Acganey
Tes Accunayy
Tox Axcuracy
Tex Accurey
(i) ResNet18, CIFARIO
(i) VGGI6, CIFARIO
(k) ResNet18, SVHN
(I) VGGI6. SVHN
Figure 3: (Main result). Comparing method TempBalance (TB) to CAL and SNR. Our method TempBalance outperforms CAL and SNR almost all the settings except for VGG19 and ResNet 34 on CIFAR 100. For all experiments, combining TempBalance and SNR (TB+SNR) yields the best performance. All baselines are carefully tuned. All results are obtained by running five random seeds. See Appendix for the details in all hyperparameters.
4.1 Experimental setup
Datasets. We consider CIFAR100, CIFARIO. SVHN and Tiny ImageNet (TIN) 75 CIFARIO0 consists of 50K pictures for training and 10K pictures for testing with 100 categories. CIFAR10 consists of 50K pictures for training and 10K pictures for testing with 10 categories. SVHN consists of around 73K pictures for training and around 26K pictures for testing with 10 categories. Tiny ImageNet consists of 10K pictures for training and 10K images for testing with 200 classes.
Models. We mainly consider three types of NNs: VGG, ResNet, and WideResNet (WRN) 776 78 For each network, we consider two different size options. For VGG, we consider VGG16 and VGG19. For ResNet, we consider ResNet18 and ResNet34. For WideResNet, we consider WRN16-8 and WRN28-6. Also, for ResNet and VGG, we consider three different widths for ablation studies.
Hyperparameters. One baseline is ordinary SCD training with a cosine annealing learning rate schedule (CAL), which follows the formula: π)e 2 (1 + cos ((4)), where the current epoch, T represents the total training epochs, and t/o is the initial learning rate. We grid search the optimal initial (base) learning rate To for the CAL baseline, using the grid {0.05, 0.1, 0.15) for ResNet and (0.025, 0.05, 0.1} } for VGG. The momentum and weight decay are 0.9 and x 10 respectively, which are both standard choices.
Another baseline is spectral norm regularization (SNR). Prior work 37 uses the SNR objective:
Aar 2
min 0
Σ ((fe (x;), y.) + 1m
o (W))?,
T
1=1
where Agr is the SNR coefficient, (W;) is the largest eigenvalue, i.e., spectral norm of weight matrix Wr. and L is the number of layers. We the power iteration method calculate (Wi) in our experiments. For SNR, we grid search the optimal regularization coefficient Ase, and we again adopt the CAL schedule for SNR, similar the CAL baseline.
To make our results fully reproducible, we report in Appendix Diall hyperparameters, random seeds, and all numerical values of experimental results shown in the figures.
4.2 Comparing TempBalance and multiple baseline methods.
First, we compare TempBalance to two baseline training methods. See results in Figure 3 In the figure, CAL means SGD training with CAL learning rate schedule, and SNR means SGD trained
7

---

79.5
80.0
Test Accuracy
Test Accuracy
79.0
79.5
78.5
79.0
88 033
78.0
78.5
77.5
78.0
77.0
CAL
SGDR
LARS
Lookahead
SGDP
TB
TB SGDP
77.5
CAL
SGDR
LARS
Lookahead
SGDP
TB
TB + SGDP
(a) R8,Net18, CIFARIO0
(b) ResNe134, CIFARI00
Figure 4: (More baseline optimizers), Comparing our method TempBalance (TB) cosine anneal- ing (CAL) baseline and other state-of-the-art optimizers and learning rate schedulers for ResNet18 and ResNet34 trained on CIFARIO0. Crosses for the same method represent different hyperparameter settings. Each cross represents the mean test accuracy of five random seeds. The best performing model thus far is combined with SGDP.
with spectral norm regularization. means our method TempBalance, and TB SNR means TempBalance combined with SNR. All eтror bars are obtained from five random seeds. From Figure we see that TempBalance outperforms the CAL baseline in all settings. In almost all cases, it performs better than SNR baseline. When TempBalance does not outperform SNR, combining SNR with TempBalance ance leads to better test accuracy.
Second, we compare our method to number of optimizers and learning rate schedulers that are not necessarily related to ESD of weights. These include SGDR [10], SGDP [4]] Lookahead [42] and LARS [17 18], and we compare these baselines with TempBalance for ResNet18 and ResNet34 trained on CIFAR100. SGDR is stochastic gradient descent with warm restarts. SGDP modifies the ordinary SGD to compensate for the effect of increasing weight norm. Lookahead 42 modifies SCD by letting each gradient update approximate the future trajectory of multiple updates. LARS assigns layer-wise learning rates based on the so-called "trust-ratio" and is the closest to our method. Results F show that TempBalance ance outperforms these baselines, and TempBal ance combined with SGDP is the best-performing method. The crosses on each column represent training runs with different hyperparameters. Note that there are several other methods based on modifying the Adam optimizer [2], such as AdamW [II), AdamP [4]] and LAMB (79). However, do not find them to provide better results than the SGD baseline with cosine annealing (CAL in Figure 4). The results are detailed in Appendix
4.3 Corroborating results and ablation studies.
In addition to the main results (Figures) and we provide corroborating results and ablation studies.
Experiment one: tuning initial learning rate rjo- We train models from scratch using TempBalance versus CAL with various initial learning rates, comparing TempBalance and the CAL baseline when both methods are allowed to search for the optimal hyperparameters. We again use ResNet1 8, ResNet34, VGG16 and VGG19 Figure that TempBalane ance achieves a higher test accuracy than CAL for both ResNet and VGG.
architectures and show results CIFAR 100. Results in
Experiment two: varying channel width. We view the fraction of model width in Experiment one as "100%," and we experiment with models with varied widths in [50%, 100%, 150%]. We again used ResNet18, ResNet34, VGGI6 and VGG19, and trained on CIFAR 100, and we grid search for the optimal learning rate for each width to get the best accuracy. Results in Figure show we find that TempBalance outperforms the baseline for all widths.
Experiment three: varying HT-SR metric. We use different HT-SR metrics to assign layer-wise learning rates. That is, we replace the layer-wise PL_ Aa_App_p_HH11 Hillllll in (2) with other HT-SR metrics including SpectralNorm and AlphaWeighted
22 Results Figure show that PL__llpha_Hal Hill
achieves the optimal test accuracy.
8

---

AEARES
CAL
TB
Test Test Accuracy
Test Accuracy
Test Accuracy
Test Accuracy
79
75
74-
78-
79-
73
74
77
78
72
0.05
0.1
0.15
0.05
0.1
0.15
0.025
0.05
0.1
0.025
0.05
0.1
Initial Learning Rate
Initial Learning Rate
Initial Learning Rate
Initial Learning Rate
(a) ResNer18 CIFARIO0 (b) ResNer34, CIFARI00
(c) VGGI6, CIFAR100
VGG19, CIFARIO0
Figure 5: (Tuning initial learning rate). Comparing the test accuracy of TempBalance (red) and CAL baseline (blue) for varying initial learning rate. Our method TempBalance ance outperforms CAL for both ResNet and VGG trained on CIFAR100, All results are obtained by running five random seeds.
CAL
TB
81
77
Test Accuracy
Test Accuracy
Test Accuracy
Test Accuracy
79
75
74
79-
72
77
73
77-
75
71
256
512
768
256
512
768
Model Width
Model Width
Model Width
Model Width
(a) ResNet18, CIFARI00 (b) ResNet34, CIFAR100
(c) VGGI6, CIFARI00
(d) VGG19, CIFARI00
Figure 6: (Different widths). Comparing TempBalance and the CAL baseline for different network widths. Our method TempBalanee consistently outperforms the CAL baseline various network widths for both ResNet and VGG trained on CIFAR100. All results are obtained by running five random seeds.
SpectralNorm
AlphaWeighter
Alphapha Hill
Test Accumey
79.0
Test Accuracy
80
Test Lccurayy
75
74
78.5
79
74
78.0
78
0.05
0.1
0.15
0.05
0.1
0.15
0.025
0.05
0.1
0.025
0.05
0.1
Leurning Rate
Initial
Learning Rate
Initial
Learning Rate
Initial Learning Rate
(a) ResNet18, CIFARI00 (b) ResNet34, CIFARI00
(c) VGGI6, CIFARI00
(d) VGG19, CIFARI00
Figure 7: (Different HT-SR metries). Comparing PL_ PL_AlphaaHHil _Hill with multiple HT-SR metrics. Pl__Alpha_s
Milha_Hillll___ Hill achieves the best test accuracy among these metrics. All results are obtained by running five random seeds.
varying PL fitting methods. The HT-SR metric PL_Alpha_HHi _iill is derived through PL fitting, which is influenced by the choice hyperparameter Amin- More specifically, this involves determining the adjustable parameter k as per Equation 3 Past research has employed various methods to select Amin based on the task, such as performance prediction. For instance, Martin et al. 22), Clauset et al. [39] choose Amio that aligns with the best fit according the Kolmogorov- Smirnov statistic [40], a method termed Goodness-of of -itt. Mcanwhile, Yang et al. [32] adopted the Fix-finger approach, which identifies Amin at the peak of the ESD. In our study, designate Amin as the median all eigenvalues present the ESD for TempBalance. As depicted in Figure 8 fitting method, termed Median, not only ensures optimal test accuracy but also notably decreases computation time. This shows that this PL fitting method is suited for the design learning rate schedulers that demand low computation overhead.
Experiment four:
Empirical analysis results. We conduct an empirical analysis of TempBalance ance to discuss why provides improvement. Our first analysis involves visualization to demonstrate how TempBalance

---

mnn (oe))
Test Accuncy
7,5
Time of fring TD onte
Text Accunacy
80-
Test Accuracy
Time (ftll TB ofeic
Test Accurncy
74.5
Usioa elcc
10
75.0-
74.0-
5.00
74.5
73.5
Grodrana-
Fix- Fic-dnger
Modian (Uuns)
Goodress-
Fia-dinges
Media (Oum)
Goodrew-
Flaefnger
Medien (Dun)
Goodsase-
Fooffnger
Medion (Durg)
of-fi
of-Et
ofift
of/01
(a) ResNet18, CIFARI00 ResNet34, CIFARTO0
(c) VGGI6, CIFAR100
(d) VGG19, CIFARIO0
Figure 8: (Varying PL fitting method to determine the Amin). Results of using different PL. fitting methods. The blue bar plot and the left y-axis label denote the test accuracy (higher the better), and the red line plots and the right y-axis label denote the time in seconds of using TempBalance once (lower the better). Our design (M) used in the proposed method achieves higher test accuracy and takes lower computation times compared to Goodness-of -fit and Fix-finger The test accuracy is averaged over five random seeds and computation time averaged over ten times.
effectively regularizes ESDs by scheduling the learning rate (see Appendix B). The second analysis strengthens the connections between TempBalance and HT structure, illustrating that the observed improvements are not due to indirectly addressing other training issues, such as gradient excur- sions [0 (see Appendix Fb.
Corroborating results on other tasks. We extend evaluation of TempBalance two additional tasks: object detection and language modeling, the details of which can be found in Appendix G Across these tasks, TempBalance consistently outperforms the baseline CAL in terms of generalization.
Conclusion
Our extensive empirical evaluations demonstrate that TempBalance offers straightforward yet effec- tive layer-wiee wiie learning rate schedule. Our approach for balancing layer-wise temperature confirms the following: (i) -TTSS--motivater metric PL A___pppp__H_____ a_tallllllll11 helps layers achieve temperature balance during training, exhibits strong correlations with model quality, and yields improved performance during testing; (ii) temperature balancing is a novel and essential aspect of NN training, and HTT-S Theory provides a strong theoretical support for balancing temperatures; and (iii) layer-wise learning rate schedules are cheap and effective to apply, and it is useful to study these layer-wise learning rate schedules further. Our method provides insights into the study layer-wise tuning approaches and tea-reemeerature balancing in deep NN training, it serves both as a layer-wise learning rate schedule and an effective regularization technique based on HT-SR Theory.
Future directions, limitations, and societal impacts. Our paper leaves many future directions to explore, of which we mention just a few.
Can HT-SR metrics be extended to parameter-wise learning rate schedules, global learning rate schedules, or other hyperparameters? It would be of interest observe how HT-SR can assist in acquiring a comprehensive set of hyperparameter tuning tools.
Is it possible to accelerate the computation of ESDs and PL_Alpha_ Hill achieve more adaptive learning rate scheduler? Currently, we calculate layer-wise PL_Alpha_Hill _iill once per epoch, resulting
minimal increase computational cost. Consider the example of training ResNet18 for 200 epochs on CIFAR1ON 100. Calculating layer-wise PI_AApha_Hill takes 1.14 seconds for each epoch, leading to 3.8 minutes total. Training CIFARIO0 on Quadro RTX 6000 takes 59 minutes, and thus using TB increases 6% of training time. However, if we can significantly decrease the expense of computing ESDs, it might enable an optimizer that adjusts the learning rate every few gradient updates. A study computation overhead is detailed in Appendix H
Our research centers around developing generic algorithm for optimizing NNs. Although TempBalance could be applied to learning models with adverse applications. we do not see any immediate negative societal impacts stemming from the algorithm itself. Indeed, we see lot of societal value in using a practical, predictive, and quantitative theory, such as HT-SR Theory, as opposed to developing a method that relies on theory that provides vacuous upper bounds and then relies extremely expensive hyperparameter tuning to obtain good results.
10

---

Acknowledgements. WeightWatcher is publicly available tool distributed under Apache License 2.0 with copyright held by Calculation Consulting. Our conclusions do necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
References
[1] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61)2121-2159, 2011.
[2] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2014.
[3] Leslie Smith. Cyclical learning rates for training neural networks. In 2017 IEEE Winter Conference on Applications of Computer Vision, pages 464-472, 2017.
[4] Zhewei Yao, Amir Gholami. Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
[5] Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikkonaa, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelicf optimizer: Adapting stepsizes by the belief in observed gradients. Advances in Neural Information Processing Systems, 2020.
[6] Liyuan Liu, Haoming Jiang. Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations, 2020.
[7] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, 2017.
[8] Bharat Singh, Soham De, Yangmuzi Zhang, Thomas Goldstein, and Gavin Taylor. Layer- specific adaptive learning rates for deep networks. In IEEE 14th International Conference on Machine Learning and Applications, 2015.
[9] Shaden Smith, Mostofa Patwary. Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper. Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminahadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, and Bryan Catanzaro. Using deepspeed and megatron to train tegatron-turing nig 530b, a large-scale generative language model. arXiv preprint arXiv:2201. 2tt.:2201..//999 2022.
[10] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In International Conference Learning Representations, 2017.
[11] Uyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations, 2019.
[12] Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classi- fication. In Proceedings of the 56th Anmual Meeting of the Association for Computational Linguistics.
[13] Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In International Conference on Machine Learning, 2015.
[14] Yixiong Chen, Jingxian Li, Hua Jiang, Li Liu, and Chris Ding. Metalr: Layer-wise learning rate based meta-learnng for adaptively fine-tuning medical pre-trained models. In Medical Image Computing and Computer Assisted Intervention, 2023.
[15] Luca Franceschi, Paolo Frasconi, Saverio Salzo, Riccardo Grazzi, and Massimiliano Pontil. Bilevel programming for hyperparameter optimization and meta-learning. International Conference Machine Learning, 2018.

---

[16] Youngmin Ro and Jin Young Choi. Autolr: Layer-wise pruning and auto-tuning of learning rates in fine-tuning of deep networks. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
[17] Yang You, Igor Gitman, and Boris Ginsburg. Scaling SGD batch size 32k for ImageNet training. arXiv preprini arXiv: / 708.03888, 6(12);6, 2017.
[18] Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th Interaational Conference on Parallel Processing, 2018.
[19] Charles H Martin and Michael W Mahoney. Implicit self- reeflfegularizaiionn in deep neural networks: Evidence from random matrix theory and implications for learning. Journal of Machine Leurning Research, 22(165): 2021.
[20] Charles H Martin and Michael W Mahoney. Traditional and heavy tailed self regularization in neural network models. In International Conference on Machine Learning. 2019.
[21] Charles H Martin and Michael W Mahoney. Heavy-tailed universality predicts trends in test accuracies for very large pre-trained deep neural networks. In SIAM International Conference on Data Mining, 2020.
[22] Charles H Martin, Tongsu Serena Peng, and Michael W Mahoney. Predicting trends in the quality of state-of-the-arr neural networks without access to training or testing data. Nature Communications, 12(1):1-13, 2021.
[23] Charles H Martin and Michael W Mahoney. Post-mortem on deep learning contest: Simpson's paradox and the complementary roles of scale metrics versus shape metrics. Technical Report Preprint: 2XXi:21066000734 2021.
[24] Charles H Martin and Michael W Mahoney. Rethinking generalization requires revisiting old ideas: statistical mechanics approaches and complex learning behavior. Technical Report Preprint: M.Kiv:1110.09555 2017.
[25] Hyunjume Sebastian Seung, Haim Sompolinsky, and Naftali Tishby. Statistical mechanics of learning from examples. Physical Review A, (5(8))60566608 1992.
[26] Timothy L. H. Watkin, Albrecht Rau, and Michael Biehl. The statistical mechanics of learning rule. Rev. Mod. Phys., 65:499-556, Apr 1993.
[27] David Haussler, H. Sebastian Seung, Michael Kearns, and Naftali Tishby. Rigorous learning curve bounds from statistical mechanics. In Proceedings of the Seventh Anmual Conference on Computational Learning Theory, page 76-87, 1994.
[28] Andreas Engel and Christian Van den Broeck. Statistical mechanics of learning. Cambridge University Press, 2001.
[29] Samuel L Smith and Quoc Le. A bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
[30] Samuel L. Smith, Pieter-Jan Kindermans, and Quoc V. Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018.
[31] Mert Gurbuzbalaban, Umut Simsekli, and Lingjiong Zhu. The heavy-tail phenomenon in sgd. In International Conference on Machine Learning, 2021.
[32] Yaoqing Yang, Ryan Theisen, Liam Hodgkinson, Joseph E Gonzalez, Kannan Ramchandran, Charles H Martin, and Michael W Mahoney. Test accuracy vs. generalization gap: Model selection in ulp without accessing training testing data. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 3011-3021, 2023.
[33] Kumar K Agrawal, Arnab Kumar Mondal, Arna Ghosh, and Blake Richards. Stalphaš-req Assessing representation quality self-supervised learning by measuring eigenspectrum decay. In Advances in Neural Information Processing Systems, 2022.
12

---

[34] Josue Nassar, Piotr Sokol, Sue Yeon Chung, Kenneth D Harris, and I1 Memming Park. On 1/n neural representation and robustness. In Advances Neural Information Processing Systems, 2020.
[35] Zeke Xie, Qian- Yuan Tang. Yunfeng Cai, Mingming Sun, and Ping Li. On the power-law hessian spectrums in deep learning. arXiv preprint arXiv:2201.13011 2022.
[36] Bruce M Hill. A simple general approach to inference about the tail of a distribution. The Annals of Statistics, pages 1163-1174, 1975.
[37] Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generaliz- ability of deep learning. arXiv preprint arXiv: 10941, 2017.
[38] Takeru Miyato, Toshiii Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018.
[39] Aaron Clauset, Cosma Rohilla Shalizi, and Mark EJ Newman. Power-law distributions in empirical data, SIAM review, 51(4):661-703, 2009.
[40] Jeff Alstott, Ed Bullmore, and Dietmar Plenz. Powerlaw: python package for analysis of heavy-tailed distributions. PloS one, (1))x85777, 2014.
[41] Byeongho Heo, Sanghyuk Chun, Seong Joon Oh, Dongyoon Han, Sangdoo Yun, Gyuwan Kim, Youngjung Uh, and lung-Woo Ha. Adamp: Slowing down the slowdown for momentum optimizers on ecale-invariam weights. In International Conference on Learning Representations, 2021.
[42] Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey Hinton. Lookahead optimizer: steps forward, 1 step back. In Advances in Neural Information Processing Systems, 2019.
[43] John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the National Academy of Sciences, 1982.
[44] Haim Sompolinsky. Statistical mechanics of neural networks. Physics Today, 41(21):70-80, 1988.
[45] LM Rasdi Rere, Mohamad Ivan Fanany, and Aniati Murni Arymurthy. Simulated annealing algorithm for deep learning. Procedia Computer Science, 72: 137-1444 2015.
[46] Adriano Barra and Francesco Guerra. About the ergodic regime in the analogical hopfield neural networks: moments of the partition function. Journal of Mathematical Physics, 49(12): 125217, 2008
[47] Adriano Barra, Alberto Bernacchia, Enrica Santucci, and Pierluigi Contucci. On the equivalence of hopfield networks and boltzmann machines. Neural Networks, 34:1-9, 2012.
[48] Stephen G. Brush. History of the lenz-ising model. Reviews of Modern Physics, 39:883-89 1967.
[49] Yaoqing Yang, Liam Hodgkinson, Ryan Theisen, Joe Zou, Joseph E Gonzalez, Kannan Ram- chandran, and Michael W Mahoney. Taxonomizing local versus global structure in neural network loss landscapes. In Advances in Neural Information Processing Systems, 2021.
[50] Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827-5837, 2019.
[51] Umut Simsekli, Ozan Sener, George Deligiannidis, and Murat A Erdogdu. Hausdorff dimension. heavy tails, and generalization in neural networks. Advances in Neural Information Processing Systems, 33:5138-51511 2020.
[52] Liam Hodgkinson and Michael W Mahoney, Multiplicative noise and heavy tails in stochastic optimization. In International Conference on Machine Learning, pages 4262-4274, 2021.
13

---

[53] Liam Hodgkinson, Umut Simsekli, Rajiv Khanna, and Michael Mahoney. Generalization bounds using lower tail exponents in stochastic optimizers. In International Conference on Machine Learning, pages 8774-8795, 2022.
[54] Anant Raj. Lingjiong Zhu, Mert Gurbuzbalaban, and Umut Simsekli. Algorithmic stability of heavy-tailed sgd with general loss functions. In International Conference on Machine Learning, pages 28578-28597, 2023.
[55] Marco Baity-Jesi, Levent Sagun, Mario Geiger, Stefano Spigler, Gérard Ben Arous, Chiara Cammarota, Yann LeCun, Matthieu Wyart, and Giulio Biroli. Comparing dynamics: Deep neural networks versus glassy systems. In International Conference on Machine Learning, pages 314-323, 2018.
[56] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 116(12)5551--4400, 2019.
[57] Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In Advances in Neural Information Processing Systems, 2022.
(58) Yefan Zhou, Yaoqing Yang, Arin Chang, and Mahoney W Michael. A three-regime model of network pruning. In International Conference on Machine Learning, 2023.
[59] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishman, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2019.
[60] Gintare Karolina Dziugaite, Alexandre Drouin, Brady Neal, Nitarshan Rajkumar, Ethan Ca- ballero, Linbo Wang. loannis Mitliagkas, and Daniel M Roy. In search of robust measures of generalization. In Advances in Neural Information Processing Systems, 2020.
[61] Peter Bartlett, Dylan Foster, and Matus Telgarsky. pnoctrallynnormalizer margin bounds for neural networks. In Advances in Neural Information Processing Systems, 2017.
[62] Huanrui Yang, Xiaoxuan Yang, Neil Zhenqiang Gong, and Yiran Chen. Hero: Ressian-enhaneed robust optimization for unifying and improving generalization and quantization performance. In Proceedings of the 59th ACM/IEEE Design Automation Conference, 2022.
[63] Zhen Dong, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. HAWQ: Hes- sian aware quantization of neural networks with sexxed-rrecsion. In /EEE/CVF International Conference on Computer Vision, 2019.
[64] Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Q-BERT: Hessian based ultra low precision quantization of bert. In AAAI Conference on Artificial Intelligence, 2020.
[65] Amartya Sanyal, Philip H. Torr, and Puneet K. Dokania. Stable rank normalization for im- proved generalization in neural networks and gans. In International Conference on Learning Representations, 2020.
[66] Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search on imagenet in four gpu hours: A theoretically inspired perspective. In International Conference on Learning Representations, 2021.
[67] David A McAAlester. Some pac-bayesian theorems. In Proceedings of the Eleventh Annual Conference on Computational Learning Theory, 1998.
[68] John Langford and John Shawe-Taye Taylor.. Pac-bayes & margins. Advances in Neural Informa- tion Processing Systems, 2002.
[69] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Annual Conference Uncertainty in Artificial Intelligence, 2017.
14

---

[70] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-Bayesian approach to spectrally- normalized margin bounds for neural networks. In International Conference on Learning Representations, 2018.
[71] Xuanzhe Xiao, Zeng Li, Chuanlong Xie, and Fengwei Zhou. Heavy-tailed regularization of weight matrices in deep neural networks. arXiv preprint arXiv:2304.02911, 2023.
[72] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 and cifar- 100 datasets. 2009.
[73] Pierre Sermanet, Koray Kavukcuoglu, and Yann LeCun. Traffic signs and pedestrians vision with multi-scale convolutional networks. In Snowbird Machine Learning Workshop, 2011.
[74] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition, 2009.
[75] Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
[76] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409 1556, 2014.
[77] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEÉE Conference on Computer Vision and Pattern Recognition, 2016.
[78] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference, 2016.
[79] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020.
[80] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International Conference on Machine Learning, pages 1310-1318, 2013.
[81] Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Pathological spectra of the fisher infor- mation metric and its variants in deep neural networks. Neural Computation, 33:2274-2300 2019.
[82] Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Universal statistics of fisher information in deep neural networks: Mean field approach. In the 22nd International Conference on Artificial Intelligence and Statisties, 2019.
[83] Ka- Veng Yuen. Bayesian methods for structural dynamics and civil engineering. John Wiley Sons, 2010.
[84] Yudi Pawitan. In all likelihood: statistical modelling and inference using likelihood. Oxford University Press, 2001.
[85] Bohang Zhang, Du Jiang, Di He, and Liwei Wang. Rethinking lipschitz neural networks and certified robustness: A boolean function perspective. In Advances in Neural Information Processing Systems, 2022.
[86] Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 2916):141-142, 2012.
[87] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. PyHessian: Neural networks through the lens of the hessian. In IEEE International Conference Big Data, pages 581-590, 2020.
[88] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, 2018.
15

---

[89] Lingyu Gu, Yongqi Du, Yuan Zhang, Di Xie, Shiliang Pu, Robert Qiu, and Zhenyu Liao. "lossless" compression of deep neural networks: A high-dimensional neural tangent kernel approach. In Advances in Neural Information Processing Systems, 2022.
[90] Zeke Xie, Qian- Yuan Tang, Zheng He, Mingming Sun, and Ping Li. Rethinking the structure of stochastic gradients: Empirical and statistical evidence. arXiv preprint arXiv:2212.0200
2022.
[91] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022188 2020.
[92] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[93] Peter I. Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitring in linear regression. In Proceedings of the National Academy of Sciences, 2020.
[94] Romain Couillet and Zhenyu Liao. Random Matrix Methods for Machine Learning. Cambridge University Press, 2022.
[95] Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The pascal visual object classes (voc) challenge. International Journal of Computer Vision, 88: 303-338, 2010.
[96] Joseph Redmon, Santosh Divvala, Ross Girshick. and Ali Farhadi. You only look once: Unified, real-time object detection. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 779-788, 2016.
[97] Tsung-Yr Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6- 12, 2014, Proceedings, Part V 13, pages 740-755, 2014.
[98] Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2):313-33,, 1993.
[99] Xindian Ma, Peng Zhang, Shuai Zhang, Nan Duan, Yuexian Hou, Ming Zhou, and Dawei Song. A tensorized transformer for language modeling, Advances in neural information processing systems, 32, 2019.
16

---

Appendix
A Heavy-tail phenomena in different DNN matrices are closely related
Recently, several papers have separately studied HT structures in different types matrices, including the Hessian, the Fisher Information Matrix (FIM), and input/output covariance matrices [5 81 82] The results confirm that when NNs are well-trained, various matrices have HT properties. Among these works, there are two major ways to characterize the HT spectrum, namely the HT-shaped ESDs (such as PL_Alpha_Hiiie __ill,, or HT-shaped decaying cigenvalues [33 Our paper mainly uses the first way of characterizing the HT structure. On the other hand, the second way is to sort eigenvalues from largest to smallest and study the PL phenomena between the ordered eigenvalues and their index. Our experiments show fruitful connections between the PL phenomena manifested in different DNN matrices; if one matrix shows a PL spectrum, the other matrices often show something similar [35] Thus, it is meaningful ask why and how the PL phenomena in different prior works correlate.
This section first establishes the connections between input/output covariance matrices, the FIM and the Hessian in subsection A. We find that if one of these matrices shows the PI. phenomenon, the other two matrices have a high chance to exhibit a similar PL phenomenon. Then, in subsection [A.2 derive the connection between metric PL Alpha_Hilll Hill and the PL exponent on decaying cigenvalues, showing a simple reciprocal relationship between these two.
A.1 Connections between different matrices
Consider a NÑ fe : Rd RC, where 0 E P is the vectorized weights, d is the input dimension, and C is the output dimension. When the NN is used for classifying task, C is also the number of classes. We denote the input data as { (I;, 3/) J where x; E IR', and the number of samples is n. We denote the loss function as L(0)= = N/2 Σ Uys. fx.)))
Covariance matrices. We denote the output covariance matrix as (lfo(=) (x)), where the expec- tation is taken over the input distribution. We tend consider the following empirical covariance matrix:
C(0) fo(x))fo (x;) E BxC
(5)
Fisher Information Matrices. We denote the (output) FIM as
E [[oef(())
V k
(x) (1
(6)
where ((k) (z) is the k-thh entry of the vector function f(x). We also consider the empirical version of the FIM:
C
F(0)= =
Vale
(1.) (1) (x))
T PxP.
(7)
Note that 7 can be equally written
F(0) - Vofo(2]Voof(111
(8)
where (ofo(x) has the following form:
(x))
os (zn)
05 (x1)
(1))) (.))
89,
3e,
09,
80,
E R PXCB
05 (1) (m1)
a5 (3m)
85. (m1) 00
a5 (m.))
50
00
17

---

Hessian Matrices. We denote the Hessian IE
('ll..f(xx)) 002
and we tend to consider the empirical
Hessian Matrices:
H(0) = 82((0)
E RXx..
(9)
00
where (()) is the empirical loss function I(0) = E Mys fo(z:)).
Hessian and FIM are equivalent under certain conditions. FIM can be defined in alternative ways different from (6). For instance, from classic statistical knowledge, we have the standard FIM (sFIM) in the following form:
FFM V log P( V) log P(y|=;;))' 0)"}
(10)
where P(yla;; 0) represents the likelihood. After simple derivations, can show that sFIM also has the following form 83 84]
82 log P(ylar; 0)
sFIM = -E
(11)
082
Therefore, when the loss function is defined as the negative log-likelihood, the sFIM in equivalent to Hessian defined in (9).
Why is the FIM defined in (6) equivalent to [TO.. Back to deep learning, the FIM often defined
as (6). It is thus meaningful to derive the equivalence between these two forms. Suppose P(ylæ; 0)
here means the conditional probability distribution of output y given input data x. If P(yla; ; 8) assumed to take the following form:
exp (5lly-So(o)]
Ply|=;0)
(12)
2m
then the MSE estimator ming P(y|=;0) 0). Then, plugging 122 into 10
fo (x)]l is equivalent the maximum likelihood estimation of
we have:
sFMM Eilly (o.o
(13)
We now expand sFIMmss by the definition of expectation, and we have the following [1]
/ [lv (z)) (ooo thhhu dyddd
s sFIM
=
(14)
2t R
- (1o= ()
(15)
R. 2m 11 f0(x))1 exp (1 (oo))o()))))) dy Volo(e)Volo(x)" q())de
(16)
= Vefa(z) g(
(17)
= EoVefa
(18)
where 14 follows from the definition of expectation, g(x)) is input distribution, and 17 holds because the integral of y in the brackets [] equals 1 due the property of Gamma function (.).
Therefore, from (8)
we find that & FT Mmae just equal to FTM, defined (6. Also, plugging 12)
into E
8log Pay (m:0))
and taking the loss function T-()) the mean-square loss, will again find
002
that E
Pogg P(yla99
is equal (0)) Therefore, considering
we can see that FIM is
802
equal to the Hessian (()).
PL in the covariance matrix and PL Hessian are tightly correlated. Next, we consider the relationship between the covariance matrix and the Hessian. Suppose the NN function fo is Lipchitz function [85). Then, it can be seen that the covariance matrix (5) may be controlled and estimated by FIM defined in (6). which is equivalent to being controlled by Hessian.
18

---

Output covariance PL exponent
linear fitting
1.80
1.75
1.70
1.65
1.60
0.9
1.0
1.1
1.2
1.3
Hessian PL sxponent
We train MLP for 50 epochs and fit PL exponent for both the output covariance and the Hessian. For models trained with epochs (1, 10, 20, 30, 40, 50), we see their PL exponents s show strong correlation.
Figure 9:
Although deriving an exact equivalent between these two can be hard, we numerically show that the PL in matrix informs the PL in the other. To visualize their relationship in the presence of PL, we train a simple MLP on MNIST [6 with one hidden layer and 2000 neurons for 50 epochs. We leverage the spectral regularization from Nassar et al. 34] to make the output covariance matrix exhibit PL spectrum. Meanwhile, we calculate the top eigenvalues of the covariance and the Hessian (87), fit the PL. exponent for each matrix, and compare the PL exponents against each other. More specifically, we take trained NNs from epochs [1, 10, 20, 30, 40, 50) and plot the Hessian PL exponent versus the output covariance PL exponent s. From the results shown in Figure | we can see that their PL exponent s shows a strong correlation, which supports our claim that the PL phenomena in one matrix can inform the other.
Connections to the NTK matrix. Interestingly, we ignore the constant X and switch the two (ofo().) This matrix equal to Neural Tangent Kerne((NTK) 18), which kernel used approximate the deep NN when NN's width is infinite. We thus conjecture that NTK should show PI. when the NN well trained 89] Indeed, Karakida al. 81) and Karakida al.
matrices multiplied together, we obtain
[2 study the eigenvalues of NTK, showing PL trend. Some other work on stochastic gradient [90] claim that the so-called "stochastic gradient matrix" (which is similar the NTK matrix) shows a PL spectrum well, which matches our expectations. Also, Lewkowycz et al. 191), Dyer and Gur-Ari [92] show that the eigenvalues of NTK are similar to those in the Hessian, which again meets our expectation because the Hessian tends be PL when NNs well-trained [35].
In summary, this section investigates different "important matrices" and shows that they are tightly correlated to each other in terms of the PI. trends: one matrix shows a PI. spectrum, there a high chance that the other ones show something similar.
A.2 Connections between PL ESD and PL in decaying eigenvalues
Next, we derive the connection between our PL_Alpha_Hi Hill metric and the exponent of PL distribu- tion on decaying cigenvalues. Take the covariance matrix (5) as an instance. According to Nassar al. 34 the HT phenomenon in the output covariance matrix is similar the layer-wise covariance matrices. Thus, without the loss of generality, we can consider the case when there is only one layer in the NN. We assume the weight matrix L is in R NxQ, According to prior works, when is well-trained, the ESD follows a PL distribution:
p(A)
_ x- H
Amin <
À
<
Amax-
(19)
Here, H is normalizing constant, and is the PL exponent.
Another way to characterize the PL. phenomennn to consider eigenvalues directly following a PI. series. For example, Xie et al. 35 show that the decaying eigenvalues follows the following PL
19

---

PL_Alpha, Hill VS
Alphha Hill VS 5 M==2
PIL. Alphapppaaaaa Hill VSS
PiL Alphalppaaaaa Hill VS 5
M- 16
M-64
120
2
Alsha Hll
Alpha
Alsha HIl
Alpha.
PL_ Alpha Hill VS.
PL_ Alphhh Hill V
PL_ Alpha Hill VSS
PL. Pl_hip Hill VS
(4 M-256
512
M-7688
M-1024
PL_Agpos_II HI
PL_AiphaLH
PLAAipho_Hi
PLAMiphajHI HI
Figure 10: We show the connection between PL Alpha_Hilllll and the PL exponent of the decaying eigenvalues (denoted as s) satisfy Al. 8111 Results are shown for different matrix size Q In partieular, we see that PL_ Ph_Alpha Hill 2
19 is equivalent to 8 = 3 the linear case.
series:
A A,k= 100==1, 1,22, 22,
(20)
where d1 is the same Amax used in the main paper.
Now, we will analytically and empirically show that these two ways of characterizing PL are strongly related. Furthermore, the two PL. coefficients satisfy s 1
0-1-
An analytical way to show that - The derivation is actually quite simple. Consider the case that A = λ ~s (i.e.. (20) holds), and suppose A is a random variable distributed according to the empirical distribution from these eigenvalues λ = Ayk * Now, from (20), we can see that the distribution function takes the following form:
k Q
P(A > A,k
(21)
By changing variables d,k = X, we get the cumulative distribution function of A:
F(A > A) ~ =
(22)
After that, we take the derivative with respect to X, and we get the ESD:
p(A) (1++1)
(23)
In other words, we have A -G+++1 = a which means s a-1'
An empirical way to show that === T- Q in {16, 32, 64, 128, 256, 512, 768, 1024), and we assign the parameters such that the decaying eigenvalues obey the formula X k= for s in (0.2, 0.3, 0.4..... 3.2 J. Then, we fit the ESD and get our estimate PL TL_AAppa__Hij Hill. We plot the relationship between PL_Alpp Alpha_Hill and s in T0 From Figure 10 we find that the connection between PL Ao__ApppaHHiip Hill and s shows a good fit with With increasing matrix size Q, the fitting becomes increasingly accurate.
1_ We consider matrices of size Q x Q, where we choose
the formula $
When T -1 corresponds to = 2. Some prior works Nassar et al. 34 Xie et al. Bartlett et al. 993 measure the HT phenomena from the perspective of decaying eigenvalues with PI. exponent 8, and they show either theoretically or empirically that 8 = is the optimal exponent. Now that we have s -- in the linear case, and from the theory of NTK [TK(88) the infinite wide NN is approximated as linear model, we tend to believe that a = 2 satisfies a similar property. Indeed, one of the main contributions of Martin and Mahoney [19] is to establish different HT families of ESDs, and a = 2 is believed to be the boundary between "moderately HT" and "very HT." corresponding to the best models. Martin and Mahoney [19 further argue that the optimal exponent for PL_Alpha is in the range [2,4]. Combining the perspective from Nassar et al. (34], Xie et al. (35), Bartlett et al. (93] and those from Martin and Mahoney [19, it reasonable to believe that the optimal exponent for PL_ Pl_Alpha is around 2. When PL Al_ha is much higher lower than 2,
35
20

---

the NN probably has some issue in training, Although we argued in the main paper that the absolute numerical value of PL_Alpha is unimportant in implementing our TempBalance algorithm. it is, however, helpful to have an "optimal" PL Alpha value to test if our algorithm actually works in controlling the ESDs. We will show visualization results in Appendix that TempBalance leads to a better distribution of our estimated PL_Alpha_HHil
In summary, this section explores two distinct methods for determining PL fit. We demonstrate that, although these two methods yield numerically distinct PL exponents, they essentially capture the same underlying phenomenon. Moreover, it is noteworthy that the "optimal" values of the PL exponents reported in various papers are consistent with one another 19
35 93
B Visualization results: how does TempBalance control ESDs
We demonstrate that the proposed method, TempBalance, effectively controls the shape of ESDs, resulting in a more favorable distribution of PL_ PL_Alpha_ _Hlll among the layers of NNs compared to the baseline method CAL. This observation elucidates the superior performance of TempBalance over CAL in our main experiment, as presented in Section 4.2
We evaluate the models reported in the main paper. For each individual NN, we compute and aggregate PL__Alpha_Hill __illllllllll values across all layers, excluding the first and last layers that have an extremely small number of eigenvalues and thus cause inaccurate PL PL_Alpha__He _Hill estimation. We aggregate the PL_Alpha__ __illlllllllllllllllllllllll values from five models trained using different random seeds for each method. Figure 11 shows the distribution of PL_Alpha_Hiil baseline CAL. Comparing TempBalance with CAL, we see that TempBalance consistently yields more concentrated distribution. Furthermore, TempBalance causes the median and mean of the distribution to approach 2 (shown in each subplot respectively as the middle vertical line and the red star). The value 2 represents the theoretically optimal PL_Alpha__ie Hill value, as we have justified in Appendix A
of TempBalance and the
Next, Figure 12, we group the models into different subgroups based on their architectures and/or datasets, aggregating the PL__Appha_Hi Hill values and comparing the distributions of the two methods TempBalance and CAL. Once again, we observe that TempBal ancce ance results in a more concentrated distribution, with a larger number of samples (layers) having PL_Alpha_HHi _Hill values closer to 2.
We provide visualization to demonstrate how the learning rates are distributed over layers during the training. In Figure 13] we report the learning rate and PL_Alpha_ Alpha_Hihllll every epoch throughout the 200-epoch training duration. The key observation includes the following.
1. How does the learning rate vary across layers? We observed correlation between the layer-wise learning rate and the layer-wise PL_Alpha_Hill distribution: layers with larger PL PL_Alpha__Ht Hill are allocated larger learning rates, whereas those with smaller PL_Alpha_Hiil Hill receive smaller learning rates.
2. How does the layer-wise learning rate evolve during training? The variations in layer- wise learning rates closely reflect shifts in the layer-wise PL_Alpha_Hill Hill distribution. Initially, the PL_ PL_Alpha__ Hill distributes uniformly across layers but eventually converge to a layer-wise pattern where earlier layers have smaller PL PL_Alpha Hill and later layers have larger ones.
We present visualizations of how PL_ PL_Alpha_Hipe _Hill and learning rate evolve through training. In Figure [4] we show PL_ Alpha_ppaa___c Alpha_Hill Hill and learning rate of two layers within the same ResNet18 during the training process. The two layers are layerl. conv2 (index=1) and layer4. conv2 (index=15). From Figure [4b] and [4d] we can see that with the baseline CAL scheduler (blue curves), the earlier layer (index=1) achieves a smaller PL_Alpha_Hie Hill value compared to the larger PL_Alpha_ Hill value of the later layer (index=15). In contrast, TempBalance (orange curves) narrows this gap, indicating our approach balances the underrainin//overrraiingg levels (as signified by PL PL_Alpha_Hii Hilll of different layers. This balancing effect is further corroborated by Figures I1 and 12). where our method consistently refines the layer-wise PL_AAlp Aalpha_Hil) distribution. Regarding the learning rate plots in Figure 14a and 14c TempBalance allocates lower learning rate for earlier layers and a higher one for later layers than the baseline does. This leads to a more balanced PL Alpha Hill distribution between layers mentioned above. Additionally, noted
21

---

CAL
TB
CAL
CAL
CAL
CAL
TB
TB
TB
TB
1.0
2.0
25
2.0
2.0
25
200
Pl_Alpha_Iie Hill
PL PL__alpha_iae Hill
PL_ Alpha Hill
(a) ResNet18, CIFARIO0 (b) ResNe(34, CIFARIO0
(c) VGGI6, CIFAR100
(d) VGG19, CIFARIO0
CAL
CAL
CAL
CAL-
T1
TB
TB
TB
1.0
25
2.0
25
1.5
2.0
PL Alpha Hill
PL Alpha
PL Alpha Hill
Alpha
(e) ResNet18, TIN
(I) ResNet34, TIN
WRN16-8, TIN
(h) WRN28-6, TIN
CAL
CAL
CAL
CAL
TB
TB
Plphlphaa Hill
Algha
Alpha
Alpho
(i) ResNet18, CIFARIO VGGI6, CIFARIO
ResNet18, SVHN
(I) VGGI6, SVHN
SVHN
Figure 11: Comparing the distribution of PL PL_Alpha_Hil Hill of NNs trained by our method TempBalance (TB) and CAL. The mean each distribution indicated by a star marker. Each distribution aggregates the PL_Alpha_ palpaaalllll____l_llllllllllllll values from models trained using five different random seeds. Across all experiments, our method TempBalance consistently yields more concentrated distribution, resulting in the mean and median approaching the theoretically optimal PL Al_hap [ipha_Hill] value 2, as supported in Appendix
CAL
TB
Density
Density
Density
Dencity
9.0
1.5
2.0
2.5
3,0
0
1.5
20
2.5
1.0
1.3
2.0
2,5
0.0
3,0
15
PL. Awphap_paa___0 Hill
PL. Alpha Hill
PL. Alpha Hill
Alpha Hill
(a) Total
(h) ResNet
(c) VGG
(d) WRN
Density
Density
Density
Dossity
1.0
9.0
1.5
2.0
2.5
2.5
0
1.5
2.0
PL, T1__Alpha_Ha Hill
PL Alpha Hill
PL. Alpha Hill
PL Alpha Hill
(e) TIN
(I) CIFARIO0
(g) CIFARIO
(h) SVHN
Figure 12: Comparing our method TempBalance (TB) to CAL in terms of the distribution of PL PL_Alpha_Hip Hill of aggregating NNs into different architectures and datasets. Each distribution aggregates the PL PL_Alpha Hill of models trained with five random seeds. Across all subgroups. our method TempBalan ance consistently exhibits a more concentrated distribution, accompanied by higher number of layers approaching PL_AP1 PT__Alpha_Hi Hill value close This value of corresponds to the theoretically optimal PL. Alpha Hill value, as justified in Appendix A
22

---

Training Epochs
50
100
130
200
2.5
Leiniag Ratt
2,5
0.10
0.10-
2.0
$2.0
0.00
0.00
10
15
0
10
15
10
20
30
10
20
Layer BAdes
Layer Index
Layns IaDox
Layer index
(a) LR, ResNet18
(b) Alpha, ResNet18
(c) LR, ResNet34
(d) Alpha, ResNet34
Figure 13: (Visualization of layer-wise learning rate (LR) and Al_happa_Hill (Alpha) over training). (a-b) The layer-wise LR and PL_Alpha Aa__lpaa__aap Hill ResNet18 over training. (c-d) The
layer-wise LR and PL_Alpha_Hai Hill of ResNet34 over training.
2.00
Layer
Layer indox 15
2.00
Leaming Rate
0.10
Loaning Rite
0.10
Alpha
1.75
Alpha
1.75
0.05
0.05
1.50
1.50
0.00
0.00
50
100
150
200
0
50
100
150
200
50
100
150
200
0
50
100
150
200
Epoch
Epoch
Epoch
Opoch
(a) LR, layer index=1
(b) Alpha, layer index=1
LR, layer index=15
(d) Alpha, layer index=15
Figure 14: (Visualization of learning rate (LR) and PL PL_Alpha__t Hill (Alpha) of two layers during training) (a-b) LR and PL_Alpt [ipha_Hill of one layer with index = in ResNet18. (c-d) LR and Hill of one layer with index 15 ResNet18. The ResNet18 is trained on CIFARI00.
PL_Alpha _ill
instability in the learning rate curves during early training phases, while smoother transitions emerge in later phases.
C Ablation studies
We provide additional ablation studies on the choices of learning rate assignment function, assignment hyperparameters.
Varying LR assignment function. For TempBalance, we selected the linear interpolation (Equation 2 for learning rate assignment function fi, based on its superior performance in our ablation study.
We evaluated three alternative learning rate assignment functions: Square root (Sqrt), Log2, and Step:
Sqri: fe(1) -
Log2: fi(i) 1)t I. E log(a;)
Iog(a?)
Step: For layer with k-th minimum PL Alpha_ Biccha_HHllll_ among all the layers,
82 - f(1i (s(s1 + 1-11
Here, 7)t denotes the base global learning rate epoch t, (s), 92) represents the minimum and maximum learning rate scaling ratios relative to nh, of is the PL P1__Apha_HHi Hill estimate of the layer i epoch t, and L is the total number of model layers. All these notations consistently used in the main paper.
TempBalance (TB), with the current assignment function, surpasses the other designs when tested on VGG and ResNet architectures on CIFAR100. All hyperparameters are consistent with the main paper. Each experiment conducted with five random seeds.
As depicted in Figure 15
23

---

Sqrt
Log2
Step
79.0
75.0
79.5-
74.5
(a) R8 CIFARI00 (b) ResNe134, CIFARIO0 (c) VGG16, CIFARI00
(d) VGGI9, CIFARI00
Figure 15: (Different designs for learning rate assignment function.) Results of using different learning rate assignment functions on different architectures and CIFAR-100 100. Our design the main paper TempBalance (TB) outperforms others. Reporting mean/std over five random seeds.
Varying LR assignment function hyperparameters. We provide additional results of a hyperparameter study on ($1, in which we consider five different settings for (s1, 82): [(0.5, 1.5), (0.6, 1.4), (0.7, 1.).. (0.8, 1.2), (0.9, 1.1)]. We run tasks on CIFAR 100 with four VGG and ResNet architectures, each with five random seeds. Our results in Figure 16 show that larger learning rate scaling range (0.5, 1.5) performs best. This hyperparameter setting is the default setting used in paper. All hyperparameters consistent with those described in the main paper.
RB5
700
$20
TED
TED
19.5
13.5
19-1
(a) ResNet18, CIFARIO0 (b) ResNet34, CIFARI00 (c) VGG16, CIFAR100 (d) VGGI9, CIFAR100
Figure 16: (Hyperparameter study on ($1, $2)). Search for hyperparameters ($1 82) with different architectures on CIFARI00. The current hyperparameter choice (0.5. 1.5) used in the paper performs best among the cases. Reporting mean/std over five random seeds.
D Hyperparameter settings for reproducing our results
We report all hyperparameters, random seeds and numerical values of experimental results shown in the main paper (in Section 4).
First, we report the common hyperparameters shared by all the experiments: the default optimizer is SGD, trained with batch size 128. number of training epochs 200, weight decay 5e-4, and momentum 0.9. The default HT-SR metric used in TempBalance is PL_ PL_Alpha__s Hill. For each experimental setting, we use five random seeds, which are always 43, 37, 13, 51, 71, and we report the mean and standard deviation of the test aceuracy across these seeds.
First, Table reports the details of experiments shown in Figure 3 We carefully tune the initial learning rate ?/ and Asr for the two baseline methods CAL and SNR. Then, Table / reports the detailed hyperparameter settings of the experiments shown in Figure 4 We again carefully tune the hyperparameters of various baseline optimizers and schedulers, specified in their papers. Finally, Table 3 Table 4 Table and Tabel 6 respectively report the details of the experiments shown in Figure Figure
Figure and Figure
24

---

Table 1: Parameter settings of the experiment reported in Section 4.2 Figure 3
The hyperparameter
in bold the best hyperparameter selection reported in the main paper. The five random seeds for each setting are {43, 37, 13, 51, 71,, and the means and standard deviations of the test accuracy among the five seeds are reported.
Initial learning rale 7j
Test Acc (best hyperparam.
scaling ratio
Index
Dataset
Model
Method
(81, s2)
ResNet18
CAL
0.05, 0.1. 0.15
78.31 = 0.05
ResNet18
0.005, 0.01,
78.65
ResNet18
TB
0.1
78.97 + 0.29
(0.5, 1.5) (0.6, 1.4)
ReaNet18
TB
SUR
0.001
79.06
ResNet34
CAL
0.05, 0.I, 0.15
78.98 + 0.14
ReaNer34
SNR
0.001, 0.005, 0.01, 0.015
79.97
ReaNer34
TB
0.1
79.89 = 0.15
(0.5. 1.5)
CIFARIOD
ReaNet34
TB
SMR
0.005
80.09 1: 0.35
(0.6, 1.4)
VGG16
CAL
0.025, 0.05, 0.1
74.59 +: 0.23
VGG16
SNR
0.05
0.001, 0.005, 0.01, 0.015
74,80 + 0,28
10
VGG16
0.05
74.96 + 0.15
(0.5, 1.5)
11
VGG16
TB SMR
0.05
0.005
75.52 + 0.46
(0.6, 1.4)
VGG19
0.025, 0.05,
73.26
13
VGG19
SMR
0.05
0.001, 0.005, 0.01, 0.015
74,37 1 0.16
VGG19
73.77
1.5)
15
VGG19
TB SMR
0.05
0.01
74.74 1 0.10
(0.5, 1.5)
ResNet18
0.1,
17
ResNet18
SNR
0.1
0.001, 0.005, 0.01, 0.015
66.20 + 0.22
18
ReaNet18
TB
66.77 =
(0.6, 1.4)
19
ReaNet18
TB
SHR
0.1
0.001
66.86 =: 0.22
(0.6, 1.4)
20
ReaNo134
CAL
0.05,
0.1. 0.15
68.19 =
21
ReaNer34
SHR
0.1
0,001, 0.005, 0.01. 0.015
68.69 + 0.13
22
ResNet34
TB
69.12 0.16
(0.6. 1.4)
23
ResNer34
TB SNR
0.001
69.27 0.21
(0.6, 1.4)
24
TinylmageNe
WRN16-8
CAL
0.05, 0.1. 0.15
63.67 = 0.09
25
WRN16-8
SNR
0.00005, 0.0001, 0.001
63.98 + 0.23
26
WRN16-8
TB
0.1
64.09 = 0.17
(0.6, 1.4)
27
WRN16-8
1.4)
28
WRN28 6
CAL
0.05,
0.1, 0.15
65.88 1. 0,20
29
WRN28-6
SNR
0.00005, 0.0001, 0.001
30
WRN28-6
TB
0.1
66.58 = 0.23
(0.6, 1.4)
31
WRN28-6
SNR
0.0001
(0.6, 1.4)
32
ResNet18
CAL
0.05, 0.1, 0.15
95.53 = 0.12
33
ReaNet18
SUR
0.001, 0.005, 0.01. 0.015
95.57 = 0.06
34
ResNei18
TB
0.1
95.63 1
(0.5, 1.5)
35
CIFAR
ResNet18
TB
SNR
0.001
95.66 + 0.09
(0.6, 1.4)
36
VGG16
CAL
0.025, 0.05, 0.1
93.98 1 0.12
37
VGG16
SNR
0.05
0.001, 0.005, 0.01, 0.015
94.04 + 0.07
38
VGG16
0.05
94,
(0.5, 1.5)
39
VGG16
TB SUR
0.05
0.005
94.26 = 0.10
(0.6, 1.4)
ResNei18
0.1.
41
ResNet18
SNR
0.1
0.001, 0.005, 0.015, 0.01
96.63 = 0.12
ResNei18
1.5)
43
ResNet18
TB SMR
0.1
0.01
96.67 1: 0.09
(0.6, 1.4)
44
VGG16
CAL
0.025,
0.05,
96.28
45
VGG16
SNR
0.05
0.001, 0.005, 0.015. 0.01
96.32 = 0.07
VGG16
1.5)
47
VGG16
TB SUR
0.05
0.005
96.40 0.08
(0.6, 1.4)
Table 2: Parameter settings of the experiment reported in Section 4.2 Figure 4 The hyperparameter in bold the best hyperparameter selection reported in the main paper. The five random seeds for each setting are {43, 37, 13, 51, 71), and the means and standard deviations of the test accuracy among the five seeds are reported.
Inirial
SCDRR
Lonkshead
Looksherd
Test Aze (best hyperparams.
scaling mato
Lndes
Dosaset
Model
Method
learning rale
(Te. Zoue)
(o 99)
ResNeil3
G&l.
6.1,
0.00
ResNei18
scp/t
0.1,
(100,0) (10. 20/1.
0.20
ResNei18
LASS
28, 30,
78.44
0.12
ReaNei18
Lanksbead
0.05, 0.1, 0.15
10.5
0.8 0.5
75.46
0.11
4
ResNet18
S0D?
0.1,, 0.05, 0.1. 0..01,,,0000
78.74
0.11
5
ReaNeil8
TB
0.05, 0.1. 0.15
75.97
0.20
(0.5, 1.5)
CIFARIO0
ResNet18
ResNo134
TB
CAL.
800F
RUIS, 0.1. 0.15
0.0S, 0.1. 0.15
78.08 0.14
79.13 2: 0.15
(0.1, 1.)
8
ResNo:34
9228
8JI5, 0.1. 0.15
(TOBLE), (10, (1.
78.61 0.20
ReaNe134
LARS
30, 32,34
73.04 0.19
ResNe134
Lookahoad
0.1.
79.19 0.11
ResNe134
9222
0.01.
0.155
79.34 0.21
ResNo14
7B
70.89 0.15
79.94
15)
25

---

Table 3: Parameter settings of the experiment reported in Section 4.3 for each setting are (43, 37, 13, 51, 71 and the means and standard deviations of the test accuracy among the five seeds are reported.
Figure The five random seeds
Initial learning rute Yp
scaling rutio
Index
Dataset
Model
Method
Test Acc
($1.
0
ResNet18
CAL
0.05, 0.1, 0.13
78.08 + 0.19.
78.31 1 F 0.05,
77.72 = 0.44
ResNet18
TB
0.05, 0.1, 0.15
78.48
0.27,
78.97
+ 0.29, 78.69 +: 0.11
(0.5, 1.5)
ResNer34
0.05, 0.15
78.98
78.89
78.51
COFAR100
ResNet34
TB
0.05, 0.1. 0.15
79.36
+.18 79.89
0.15, 79.09
0.64
(0.5, 1.5)
VGG16
CAL
0.025, 0.05,
73.96 +: 0.27, 74.59
0.23,
74.46
0.12
VGG16
TB
0.025, 0.05,
74,40
0.31, 74.96
1.
0.16
(0.5, 1.5)
VGG19
0.025. 0.05,
72.57
0.45,
73.26
0.37,
72.98
0.16
VGG19
TB
0.025, 0.05,
73.47
1 0.16,
73.77
0.4, 73.40
0.38
(0.5,
Table 4: Parameter settings of the experiment reported in Section 4.3 Figure The five random seeds for each setting are [43, 37, 13, 51, 71), and the means and standard deviations of the test accuracy among the five seeds reported.
Initial
scaling
Index
Dataset
Model
Method
learning rate
Width
Test Ace
ResNer18
256, 512, 768
75.05
0,26 78.31
0,05,
79.44
0.26
R
тB
256, 512, 768
75.63
0.12, 78.97
0.29,
80.47
0.18
(0.5, 1.5)
ResNet34
CAL.
256, 512. 768
76.79
0.34, 78.89
0,24,
79.94
1.333 0.31
CIFARIOD
ResNet34
TB
256, 512. 768
77.25 0.14, 79.89
0.15
80.23
0.553
(0.5, 1.5)
VGG16
CAL
0.05
256, 512, 768
71.04
0.14, 74.59
75,53
0.32
VGG16
768
71.26
74.96
0.15
(0.5,
VGG19
CAL
0.05
256, 512, 768
69.58
0.39, 73.26
0.37,
74.39
0.33
VOG19
0.35
(0.5,
Table 5: Parameter settings of the experiment reported in Section 4.3 for each setting are {43, 37, 13, 51, 71), and the means and standard deviations of the test accuracy among the five seeds are reported.
Figure The five random seeds
Initial learning rate
ccaling
Index
Dataset
Model
Method
HT-SR Metric
Test Acc
[*), ax)
ResNet
1B
Spectra Ihorn
0.05. O.I. 0.15
77.83
0.21,
78.30
0.32
78.27
1 0.25
(0.5, I.5)
ResNett
Al.hhaVaighted
0.05, 0.15
78.18
0.27
78.67
78
ResNei18
Al._Alpha_Htt H1ll
0.05, 0.1, 0.15
78.48
0.27,
78.97
0.29.
78.69
0.11
(0.5, 1.5)
ResNer34
SpectralKort
0.05, 0.1. 0.15
78.25
+ 0.16
78.71
0.15,
78.92
0.28
(0.5, 1.5)
giphaMeighted
0.05,
0.34
78.83
(0.5,
CIFARI00
ResNer34
VGG16
TB
Pt._AApha_Bii Hi11
SpectralNorm
0.025. 0.05, 0.1
0.05, 0.1, 0.15
73.58
79.36
0.18
0.19,
74.29
79.89
0,16,
0.15
74.17
79:09
0.228
0.64
(0.5, 1.5)
VGGI6
AlphaVeighted
0.025, 0.05, 0.1
73.97
1 0.22
74.19
0.11
0.311
(0.5,
VGG16
7B
AL__a_pa__111
0.025, 0.05, 0.1
14.40 + 0.31
74.96
0.15 74.94
0.16
(0.5, 1.5)
VGG19
Spectralloom
73.04
VGG19
7B
AlphaVeighted
0.025, 0.05, 0.1
72.85
4 0.16,
73.41
4
0.17,
73.33
0.21
(0.5. 1.5)
VGG19
FL_Alpha_B111
0.025, 0.05,
73.47
0.16,
73.77
0.43,
73.40
0.338
(0.5,
Table 6: Parameter settings of the experiment reported in Section 4.3 Figure The five random seeds for each setting are {43, 37, 13, 51, 71), and the means and standard deviations of the test accuracy among the five seeds, the means and standard deviations of the computation time of using TB among the 10 times are reported.
Initial
Scaling ratlo
Index
Darasct
Model
Method
litting method
Jearning
Test Acc
Computation Time (scc)
ResNet13
Goodcees-od
U.I
75.39 0.11
8.20
U.33
(0.3, 1.5)
ResNet18
TB
Fix-finger
79.06 0.22
0.74
(0.5
ResNet18
TB
Median
0.1
78.97 1: 0.29
1.14 = 0.04
(0.5, 1.5)
ResNer34
Coodnesa-of
79,13 0.21
16455
0.48
(0.5.
ResNe134
TB
Fix-fingor
79.64 0.222
15.13 1 1.05
(0.5, 1.5)
CIFAR100
Re<Net34
TB
Madian
0.1
79,89 0.15
2.27 = 0.06
(0.5. 1.5)
VGG16
Guodnees-
7.226
VGG16
TB
Fix-ringer
0.05
74.48 0.20
8.45-. 0.59
(0.5 1.5)
VGG16
TB
Median
0.05
74.96
1.37=005
(0.5
VGG19
Goodnesa-ol -o -f
0.16
(0.5
VGG19
TB
Piz- Siz-finge finger
0.05
73.52
11.15=007
(0.5
VGG19
Hedian
0.05
73.77 0.43
1.85 1.05
00.5.
26

---

E Comparison with more baselines
In Figure T7 provide additional results by comparing TempBalance with LAMB and Adam. We found that our method outperforms both baseline methods. Furthermore, we also found that the Adam-based methods do not provide better results than the SGD baseline with cosine annealing (CAL) in our experiment setting, which was mentioned in Section 4.2 For Adam, we searched the initial learning rate over {0.00005, 0.0001, 0.001, 0.01, 0.13, and used
10 8, For LAMB,
searched the initial learning rate over {0.005, 0.01, 0.02), and we used e = 10-6. Both methods used weight decay 5.0 10 31 0.9, 32 0.999, learning rate decay with cosine annealing. Each experiment was conducted with five random seeds.
We also discuss the difference between TempBalance ance and these two types of learning rate scheduling.
Compared layer-wise learning rate scheduling (e.g., LARS): TempBalance uses a more precise model quality metric, PL_Alpha_HHHal _ill from HT-SR Theory, to enhance the performance of deep models during training. This "shape-based" metric estimates the shape of the eigenspectrum of weight matrices. In contrast, LARS uses "norm-based" metric, such as the layer-wise gradient norm. A recent study HT-SR [22 has shown that the shape-based metrics surpasses norm-based ones in assessing model quality and
performance. Figure 3 confirms that our method outperforms the layer-wise scheduler LARS in test accuracy.
Compared to parameter-wise learning rate scheduling (e.g., Adam): Similarly, our method employs the "shape-based" metric PL_Alpha_Hiil Hill improve the generalization, an approach not incorporated in traditional parameter-wise methods.
79.5
80.5
Test Accuracy
1896
Test Accuracy
78.0
79.0
76.5
77.5
75.0
73.5
88
76.0
72.0
Adamn
LAMB
CAL
SODR
LARS
Lodosbad
SCDP
7B
TB soDp
74.5
Adaan
LAMI
CAL
SoDR
LARS
Lodahoad
SCDP
78ll
TB scDe
(a) R8 CIFAR100
(b) ResNet34, CIFAR
Figure 17: (Comparison with additional baselines). Comparing method, TempBalance (TB), with other baselines such as parameter-wise learning rate schedulers Adam and LAMB, using ResNet18/34 trained on CIFARI00. Each cross represents the mean test accuracy of five random seeds.
Maximum Layorwise Gradicot Norm
Minimum of Layerwise Gradient Norn
Mean Layerwise Gradient Norm
Godliaa Nomt (ogalo
TB
CAL
Chviters (eemlinggg)
Oradiens Namn (logi))
10'
CAL
10"
100
10"
-00
9.10
4
6
Epozha
Epochs
Epodh:
(a) Maximum
(b) Minimum
(c) Mean
Figure 18: (Layerwise gradient norm during training). From left right: maximum, minimum,
and mean of the layerwise gradient norm at every 30 iterations for the first 10 epochs. ResNet18 on CIFAR-100.
F Does addressing other training issues lead to TempBalance's improvement?
We discuss whether the improvement from the proposed method, TempBalance, due to indirectly addressing another fundamental training issue that could distort the ESD, specifically the gradient
27

---

CAL
TB
Iniiallzation
Epochtl Bon 10
Epochel Iren 60
Epochtt
Epochel Bert 120
Spoch: 150
Epoch Bear 180
0it2 aidee Exts
2/6
8.0
25
1.
Spochel 210
Epochs 1 her 240
Epoch:1 hen 270
Upoch:1 Iber:
Epochel Joer
Spoch:1 Ier 350
Epoch: Iter: 350
Dogg
Gudded
Figure 19: (Histogram of gradient distribution during first epoch). ResNet18 on CIFAR- 100.
Epachi 1. Ber: 30
Epoch: 1. Ber 150
10'
107
10
10
Epoch: 2
Epoch: 5
10
Figure 20: (Impact of large rank-1 1 updates on ESD). Large rank- updates result in the spikes of ESD, observed exclusively during the first epoch. From the second epoch onward, the ESD exhibits heavy-tailed distribution. R on CIFAR-100.
CAL
TB
79.0
78.5
118
11
(sact opece
Figure 21: (Varying the starting epoch of applying TempBalance (TB)). Postponing the usage of TempBalance to Epochs 2, 5, and 10 doesn't affect the performance of TempBalance (originally starting from Epoch l).
magnitude excursions [80] (explosion/vanishingg) This discussion further strengthens the connection between our method and the HT structure, as discussed in the Sections [
A and B
We first summarize the questions and the corresponding primary findings, with subsequent detailing of our experiment and supporting results.
Does gradient excursion exist? We discovered that gradient explosion does exist, but it is confined to the first epoch out of a total of 200 training epochs, leading us to believe it does not significantly impact the test accuracy. We observed no gradient vanishing. Does the observed gradient explosion impact the estimation PI__Appha_Ha Hill? We discovered that the large rank-1 updates resulting from the gradient explosion do indeed
28

---

affect the ESD as well as the A__Appaa_Hill estimation. However, this effect is again
restricted to the first epoch.
Does TempBalance ance boil down to addressing gradient explosion? We found that postpon- ing the use of TempBalance ance until the epoch when neither the gradient explosion nor the PL_A1pha_HHa Hill estimation is affected does not compromise the test accuracy.
To support the above answers, we conducted three experiments. We discuss the setup of these experiments first and then analyze the results.
(Figures 18 19) We aim detect gradient excursion by tracking the gradient norm across layers during training. We examine the model every 30 iterations over the first 10 epochs, calculating the L2 norm of each gradient update across layers using the training batches of size 128. This produces an empirical gradient norm distribution with a total sample size of update numbers layer numbers. Figure [8] presents the maximum/minimun/mmea of the distribution, while Figure 19 visualizes these distributions for several iterations of Epoch 1. (Figure20) We aim to assess the impact of gradient explosion on PL_Alphs Aha__ppa___i Hilll estimation by monitoring the ESD. Figure 20 examines the change of the ESD of single weight matrix over several iterations, tracked in the experiment depicted in Figures [8 and [9
(Figure[IIp We aim to see if TempBalance enhances generalization by implicitly addressing gradient explosion. Since the gradient explosion and its effect on PL_Alpha_HHin estima- tion only transpire in the first epoch, we postpone the starting epoch of TempBalance to Epochs 2, 5, and 10 and see it affects the test accuracy.
Our answers to the above questions are supported by the results obtained from the three experiments:
We observed that the notable exploding gradients only occur in the initial 200 iterations of the first epoch. In Figure[[8] we pinpoint a singular peak of maximum gradient norm within the first epoch. This aligns with the abnormal distribution with a large gradient norm in the subfigure of Figure 19 titled "Epoch 1, iteration 30."
First question (Figure and 19:
Second question (Figure 20): Note that large rank-one updates have been studied in random matrix theory, which manifests
"bulk+spike" pattern. This has been analyzed in, e.g.. Theorem 2.13 of [94). Figure 20 shows this "bulk+spike" pattern, but only in the first epoch. The ESD exhibits a heavy-tail distribution in subsequent epochs, suggesting the influence of rank-one updates is limited.
Third question (Figure 21): Delaying the application of TempBalance until after the first epoch does not adversely affect the test accuracy. Figure 21 illlssrates that applying TempBal ance from Epochs 2, 5, and 10 results in test performance comparable to when TempBal ance is applied from Epoch 1. Since the gradient explosion only occurs in the first epoch and its effect PL_ ____lppp__HHa __ill estimation diminishes after this, the effectiveness of TempBal ance does not rely on addressing gradient explosion or biased PL_ PI_Alpha__ Hill estimation from large rank-one updates.
Third question: We compare TempBalance with the baseline method LARS, which uses gra- dient norms to determine layer-wise learning rates in combating gradient vanishing/cxplosion issues. As illustrated in Figure TempBal Bancee outperforms LARS in terms of generalization performance.
G Corroborating results on other tasks
We provide corroborating results of applying TempBalance to two different tasks: object detection (OD) and language modeling (LM). In both tasks, TempBal ance consistently improves generaliza- tion, outperforming the baseline scheduler cosine annealing (CAL) when both are combined with Adam/AdamW optimizers.
For OD, we studied the PASCAL VOC2007 95 dataset with YOLO series 96 pre-trained model. We compared TempBalance ance with the baseline scheduler CAL with both applied to Adam/Adami optimizer. For both scheduler methods, we trained for 200 epochs with batch size 64, and we set the same hyperparameter for the optimizers: 33 = 0.9, 32 = 0.999, = 1 weight decay= 5.0 x 10 We searched the initial learning rate for all methods among (7.5 x 10-6, x 10-5 2.5 10-5). For metrics we use the COCO (97 version mean Average Precision (mAP, higher better), which is calculated for 10 IOUs varying in range of 0.5 to 0.95 with steps of 0.05. We
29

---

report the mean of mAP over five random seeds on the test set. We set the scaling factors (s1, $2) of TempBalance to be (0.6, 1.4).
Here are the experimental settings for L.M. We studied the Penn Treehank (PTB) dataset [ using three-layer "tensorized transformer core-1" [99]. We compared TempBalance with the baseline scheduler CAL with both applied to Adam optimizer. For both scheduler methods, we trained the models for 40K iterations with a batch size of 120, and dropout rate of 0.3. We searched the initial learning rate for baseline methods among 11.25 x 10 2.5 10 x 10 1 x 10 1.25 10-,, 10- for baseline CAL. The hyperparameters for Adam are 31 0.9, B3 0.999, = 10 The mean of perplexity (PPL, lower is better) across five random seeds on the test set reported. We observed improved performance of TempBalance in this task when extending our hyperparameter search to include the scaling factors ($1, 82) E {00.5, 1.5). (1.0, 2.0) }, the power-law fitting hyperparameter Amin index 50) iterations.
2.5 10
( 1.25 ). and the TempBalance update interval over (10, 25,
Table 7: (a) Object Detection (OD): mean Average Precision (mAP) on PASCAL VOC 2007 using model Yolov8n. (b) Language Modeling (LM): test perplexity (PPL) on Penn TreeBank (PTB) using
the tensorized transformer. TempBalance (TB) consistently outperforms the CAL in different tasks.
CAL Adam TB Adam
CAL + Adamä TB + Adami
CAL + Adam
TB + Adam
59.59
60.03 (+0.44)
59.68
59.96 (40.28)
49.94
47.30 (-2.4)
(a) OD, VOC2007, mAP (f)
(b) LM, PTB, PPL (4)
We present additional results in Figure 22, showing the application of our method TempBalane ance to ResNet 101 on CIFAR- 100, and we compare it with the baseline (CAL). We searched the initial learning rate among (0.05, 0.1, 0.15) for both the baseline and our method. The results report the mean and standard deviation across five seeds. We found that FempBalance ance offers improvements for the larger ResNet101 model comparable to those observed for ResNet18/34 18344, demonstrating its potential for larger models.
CAL
TB
Test Accuraey
80
79
78
ResNet18
ResNe:34
ReaNer101
Figure 22: (Applying the TempBalance (TB) to different sizes of ResNets). TempBalance consistently outperforms the baseline CAL method in the larger model ResNet101. The dataset is CIFAR100. Reporting mean/std over five random seeds.
H Analysis of computation overhead
We conducted a study on the computational overhead of TempBalance, demonstrating that our method is both applicable and scalable for large models. To do so, we conducted a scaling experiment to demonstrate that the computational cost remains low for different sizes of models. We recorded the duration of a single training epoch and the time taken apply our method once. From this, we calculated the percentage increase in time when using TempBalance Bance once per epoch, using this as an indicator of computational overhead. The experiment setup is based on ResNet-series on CIFAR100. We studied models of depth in [18, 34, 50, 101) and ResNet1 8 models of width in 5512, 768, 1024, 2048). We report the mean and the standard deviation of the results over runs. The test platform was one Quadro RTX 6000 GPU with Intel Xeon Gold 6248 CPU. The results are presented in Figure 23 Our findings reveal that the computational overhead remains low (less than 9%) even when applied exceptionally wide or deep models (ResNet18 with width 2048 or ResNet101). The
30

---

Time incremtr Na
Tine locresete (5)
150
75
and
15
34
412
700
1004
2046
512
Dopth
Scaling depth
Scaling depth
(c) Scaling width
(d) Scaling width
Figure 23: (Computation overhead of TempBalance (TB) in scaling the model depth/width). (a)(c) Time duration (second) of one training epoch (blue) and using TempBalance once (red).
(b)(d) Time increment of using TempBalance once per epoch. The dataset is CIFARI00, reporting mean/std over 10 epochs. The computational overhead using TempBalance remains low (less than 9%) even when applied to exceptionally wide or deep models.
CAL
Test Accuracy
79.2
79.0
78.8
78.6
78.4
390
300
200
100
50
TB update interval (Iters)
Figure 24: (Varying the TempBalance (TB) update interval). Reducing the update interval from 390 iters (used in the paper) brings mild improvement in test accuracy. Both use ResNet18 on CIFAR 100. Reporting mean/std over five random seeds.
computation overhead is not large because: 1) we select the efficient PL fitting method to obtain PL_Alpha_ Hi 11, which is demonstrated in Figure and the most computation-inneeniive part of our method is SVD decomposition, which we have optimized using GPU implementation and batch processing.
We conducted an experiment on reducing the update interval the learning rate schedule to see if it affects the test accuracy of TempBalance. Figure 24 shows the experiments conducted with ResNet18 on CIFAR-100. We reduce the update interval from 390 iterations used in our paper (equivalent to epoch) to 300, 200, 100, and 50. We observed that there indeed exists trade-off between the computation time and test accuracy, but reducing the update interval only brings mild improvement.
31


Reviews:
[{"review_id": "HZszuvxQyd", "reviewer": "Reviewer_zFJE", "summary": "The paper proposes a way of modulating the learning rate, independently for each layer, when training deep networks via gradient descent.\nThis modulation keeps the average learning rate (over all layers) on a predefined path (e.g., cosine decay), but \"balances\" it according to the relative training stage of each layer: layers comparatively \"overtrained\" get a smaller scaling factor, layers comparatively \"undertrained\" get a larger one.\nThat comparison is done by leveraging the theory of \"heavy-tail self-regularization\" (HT-SR), estimating the heavy-tail characteristics of the empirical spectrum density (ESD), the spectrum of eigenvalues of weight correlation matrix, specifically the alpha coefficient of a power law fitting the heavy tail.\n\nA wide range of experiments on small-scale image datasets, across architectures (ResNet, VGG) and variants (width and depth) show this method leads to better generalization error, compared to single-learning-rate optimizers, and can be further improved when combined with spectral norm regularization (SNR).", "questions": "Clarification questions\n1. l. 129, $\\lambda_\\mathrm{min}$ as the medium (median?) of the ESD. Is that what is shown on Figure 1? It was unclear to me how $\\lambda_\\mathrm{min}$ was selected based on the figure.\n2. Should $(s_1, s_2)$ be mentioned in the Hyperparameters section (l.235)?\n\nFurther information / curiosity\n\n3. In structured architectures like ResNet, is there a pattern of under / overtraining, either within each block, or between them? For instance, are layers within a block usually at the same \"stage\", or do you see correlations between the first layers of each stage? This might suggest ways in which to share a scaling factor within a block, for instance.\n4. Is there a usual way in which the $\\alpha_t^i$ usually evolve during regular SGD training, or a pattern? If there is, is it disrupted or modified when using `TempBalance`?\n5. Similarly, how do the $f_t(i)$ evolve through time? For instance, smoothness may indicate that increasing the frequency of update may not be beneficial, but instabilities may suggest the opposite.\n6. Is there evidence that convergence with `TempBalance` could be faster or slower than SGD? Could it compensate for the overhead of computing `PL_Alpha_Hill`, or worsen it?\n\n\n*Update after rebuttal*\n\nI believe all my questions were addressed. The additional experiments are thorough and lead me to increase my score.", "limitations": "Properly addressed within the body of the article.", "rating": 9, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Originality\n--------------\nLayer-wise scaling factors for learning rates is an under-explored area. Exploring it by applying the results of theoretical models of deep network training is novel, and welcome.\n\nQuality\n----------\nThe paper presents extensive experiments supporting the proposed methods, and the papers conclusion. They feature a breadth of variants for the VGG and ResNet architectures, and hyper-parameters that have been carefully tuned.\n\nExperiments have been replicated with 5 different seeds, and error bars reported. Experimental settings are clearly reported in the appendix. Overall really solid experimental validation.\n\nClarity\n---------\nThe paper is clearly organized, and explained really clearly the theoretical bases it builds on, existing algorithms, as well as the proposed new method.\n\nSignificance\n-----------------\nA method that improves the generalization of existing architectures is extremely interesting, especially when the computation overhead is reasonable. This method could be quite impactful, either directly or through further improvements or further research in a similar direction.", "weaknesses": "No major weaknesses, just a few limitations.\n\n1. It would have been interesting to see one series of experiments on larger-scale data (full ImageNet?), or maybe non-image data.\n2. No explicit comparison with parameter-wise scaling schemes (e.g., a variation of Adam)\n3. No mention of the algorithmic complexity, or overhead of computing these scales, before the penultimate paragraph.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Experiments with non-image data \nIn Table 6 (b) of our rebuttal PDF, we present a new experiment using a language dataset. Our TempBalance (TB) method performs better than the baseline cosine annealing learning rate schedule (CAL) when both used the Adam optimizer for language modeling.\n\nHere are experimental settings for language modeling. We studied the Penn Treebank (PTB) dataset [1] using a three-layer \"tensorized transformer core-1\" [2]. We compared our method TB with the baseline scheduler CAL with both applied to the Adam optimizer. For both methods, we trained the models for 40000 iterations with a batch size of 120, and a dropout rate of 0.3. We searched the initial learning rate for all methods among {0.000125, 0.00025, 0.0005, 0.001, 0.00125, 0.0025, 0.005}. The hyperparameters for Adam are $\\beta_{1}$ = 0.9, $\\beta_{2}$ = 0.999, $\\epsilon$ = $10^{−8}$. The mean and standard deviation of perplexity (PPL, lower is better) across five random seeds on the test set are reported.\n\n## Explicit comparison with parameter-wise scaling schemes\nIn Figure 18 of the rebuttal PDF, we compared our method with parameter-wise learning rate schedulers including Adam and LAMB. We show that our method outperforms them with ResNet18/34 on CIFAR100. \n\nHere is the experimental setup. For Adam, we searched the initial learning rate over {0.00005, 0.0001, 0.001, 0.01, 0.1}, and used $\\epsilon$ = $10^{−8}$. For LAMB, we searched the initial learning rate over {0.005, 0.01, 0.02}, and used $\\epsilon$ = $10^{−6}$. Both methods used weight decay = $5.0×10^{−4}$, $\\beta_{1}$ = 0.9, $\\beta_{2}$ = 0.999, learning rate decay with cosine annealing. All results were obtained by running five random seeds.\n\n## $\\lambda_{min}$ as the median of the ESD (l.129)?\nIn line 129, we state that the $\\lambda_{min}$ is fixed as the median of all eigenvalues in the Empirical Spectral Density (ESD), represented by the black vertical line in Figure 1. The histogram plot's log scales on both axes might make this less intuitive. We thank the reviewer for pointing out the typo, and we will fix and clarify it in the revised version of the paper.\n\n## Should $(s_1,s_2)$ be mentioned in the hyperparameters section (l.235)?\nWe have listed the settings for $(s_1,s_2)$ for each experiment of the main paper in Appendix C's Table 1-5 (Hyperparameter settings) under the last column. We will mention it in line 235 in the revised draft.\n\n## No mention of computation overhead\nWe provide an additional study on the computation overhead of our TB method, with results presented in rebuttal PDF Figure 17. We conducted a scaling experiment to demonstrate that the computational cost remains low for different sizes of ResNet models. Our findings reveal that the computational overhead remains low (less than 9%) even when applied to exceptionally wide or deep models (ResNet18 with width 2048 or ResNet101).\n\nThe experiment setup is based on ResNet-series on CIFAR100. We studied models of depth in {18, 34, 50, 101} and ResNet18 model of width in {512, 768, 1024, 2048}. For each model size, we recorded the duration of a single training epoch and the time taken to apply our method once per epoch. From this data, we calculated the percentage increase in time when using TB once per epoch, using this as an indicator of computational overhead. We report the mean and standard deviation of the results over 10 runs. The test platform used was a Quadro RTX 6000 GPU with an Intel Xeon Gold 6248 CPU. We will be sure to include these discussions in the updated draft of the paper.\n\n## Patterns of under/overtraining between ResNet blocks\nIn the rebuttal PDF, Figure 21 (b, d) illustrate the layer-wise alpha of ResNet 18 and 34 during training. Note that the ResNet 18/34 architecture is organized into four stages, with each stage comprising multiple residual blocks. Our primary focus is on the blue curves, which represent the alpha value at the end of training. One can recognize the patterns between stages: earlier layers (closer to the input) tend to have a greater number of layers with lower alpha values (indicative of overtraining) compared to later layers. Within one stage of the network, the second convolutional layer typically exhibits a lower alpha compared to the first convolutional layer. We will be sure to include more discussions about over/undertraining patterns in network architectures, as well as the future work on saving the computations by implementing a shared scaling factor for certain layers in the updated draft of the paper.\n\n\n## How $\\alpha_{t}^{i}$ and $f_{t}(i)$ evolve through training\nIn Figure 22 of the rebuttal PDF, we present visualizations of the $\\alpha_t^i$ (alpha) and $f_t(i)$ (learning rate) of two layers within the same ResNet18 during the training process.\n\nFrom Figure 22 (b, d), we can see that with the baseline CAL scheduler (blue curves), the earlier layer (index=1) achieves a smaller alpha value compared to the larger alpha value of the later layer (index=15). In contrast, our TB method (orange curves) narrows this gap, indicating our approach balances the undertraining/overtraining levels (as signified by alpha) of different layers. This is further corroborated by Figures 10 and 11 in the submitted paper, where our method consistently refines the layer-wise alpha distribution.\n\nRegarding the learning rate plots in Figure 22 (a, c), our TB method allocates a lower learning rate for earlier layers and a higher one for later layers than the baseline does. This leads to a more balanced alpha distribution between layers as mentioned above. Additionally, we noted instability in the learning rate curves during early training phases, while smoother transitions emerge in later phases. \n\n## Comparison of convergence rate \nWe observed that our method converges at the same rate as the SGD baseline, but it achieves higher test accuracy upon convergence.\n\n## Reference\n[1] Mikolov et al, 2011.\n\n[2] Ma et al, 2019. "}, {"review_id": "olmyyNw9vB", "reviewer": "Reviewer_MZpk", "summary": "This paper propose TempBalance, an adaptive lr schedule that assigns lr to each layer based on its heavy-tail characterization. The authors estimate PL_Alpha, the exponent of the power law distribution that fits the heavy tail part of the empirical spectral density, for the weight at each layer. They propose to assign a higher lr to the layer with larger PL_Alpha and lower lr to the layer with smaller PL_Alpha as a larger PL_Alpha often indicates a layer is under-trained while a smaller PL_Alpha often indicates a layer is over-trained. \n\n", "questions": "1. Could authors provide a visualization of the layerwise lr with some insights on how lr varies across layers and how layerwise lr changes through training?\n2. See weakness. ", "limitations": "The authors have address the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The proposed TempBalance is novel in that it adjusts lr based on metrics from HT-SR theory, providing a new perspective from statistical mechanics of learning to neural networks optimization.\n2. The experiments results demonstrate that TempBalance exhibits regularization effect and improves the generalization performance upon some existing lr schedulers, optimizers and spectral norm regularization methods. ", "weaknesses": "1. Eq. 2: The authors design the relation between the lr and the value of PL_Alpha_Hill to be linear, which seems a bit arbitrary to me. Have the authors tried other designs?\n2. The authors do not provide a convergence analysis. \n3. Eq. 2: TempBalance requires the computation of eignevalues of the weight at each layer, inducing some computational overheads. Therefore, it can be difficult to scale TempBalance to models with larger widths and depths. The authors reduce the lr update frequency to alleviate the cost, but this might compromise the model performance. \n4. Missing baseline: LAMB [1]. No hyper-parameter study on $s_1$ and $s_2$. \n5. The authors only evaluate on small models. It would be good if the authors can provide an experiment on a ResNet-101 to demonstrate its potential for larger models.\n\n[1] You, Yang, et al. \"Large batch optimization for deep learning: Training bert in 76 minutes.\" arXiv preprint arXiv:1904.00962 (2019).\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Other design of learning rate assignment?\nOur selection of the linear interpolation design for learning rate assignment in our TempBalance (TB) method was based on its superior performance in our ablation study, as provided in rebuttal PDF Figure 16.\n\nWe evaluated three alternative learning rate assignment functions: Square root (Sqrt), Log2, and Step:\n* Sqrt : $f_t(i)=\\eta_t\\frac{\\sqrt{\\alpha_t^i}}{\\frac{1}{L} \\sum_{j=1}^{L}\\sqrt{\\alpha_t^j}}$,\n* Log2: $f_t(i)=\\eta_t\\frac{log(\\alpha_t^i)}{\\frac{1}{L} \\sum_{j=1}^{L}log(\\alpha_t^j)}$,\n* Step: For layer $i$ with $k$-th minimum alpha among all the layers, \n   $f_t(i)=  \\eta_t (s_1 + (k-1)\\frac{s_2 - s_1}{L-1}) $ .\n\nHere, $\\eta_t$ denotes the base global learning rate at epoch $t$, $(s_1, s_2)$ represents the minimum and maximum learning rate scaling ratios relative to $\\eta_t$, $\\alpha_t^i$ is the PL_Alpha_Hill estimate of the layer $i$ at epoch $t$, and $L$ is the total number of model layers. All these notations are consistently used in the main paper.\n\nAs depicted in Figure 16, our method, TB, surpasses the other designs when tested on VGG and ResNet architectures on CIFAR100. All hyperparameters are consistent with the main paper. Each experiment was conducted with five random seeds.\n\n## Computation problem\n1. **Is our method difficult to scale to large models?**\n    * In rebuttal PDF Figure 17, we conducted a scaling experiment to show that our method is applicable to large models. The experiment setup is based on ResNet-series on CIFAR100. We studied models of depth in {18, 34, 50, 101} and ResNet18 models of width in {512, 768, 1024, 2048}. For each model size, we recorded the duration of a single training epoch and the time taken to apply our method once. From this, we calculated the percentage increase in time when using TB once per epoch, using this as an indicator of computational overhead. Our findings reveal that the computational overhead remains low (less than 9%) even when applied to exceptionally wide or deep models (ResNet18 with width 2048 or ResNet101). We report the mean and the standard deviation of the results over 10 runs. The test platform was one Quadro RTX 6000 GPU with Intel Xeon Gold 6248 CPU.\n    * The computation overhead is not large because the most computation-intensive part of our method is SVD decomposition, which we have optimized using GPU implementation and batch processing. \n\n2. **Does reducing the SVD computation compromise the test accuracy?**\nWe conducted an experiment on reducing the update interval of learning rate schedule to see if it affects the test accuracy of TB. Figure 15 (b) shows the experiments conducted with ResNet18 on CIFAR-100. We reduce the update interval from 390 iters used in our paper (equivalent to one epoch) to 300, 200, 100, and 50. We observed that there indeed exists a trade-off between the computation time and test accuracy, but reducing the update interval only brings mild improvement.\n\n## Missing baseline and hyperparameter study\n1. **Missing baseline.** In rebuttal PDF Figure 18, we provide additional results by comparing our method with Adam and LAMB. We found that our method outperforms both baseline methods. Furthermore, we also found that the Adam-based methods do not provide better results than the SGD baseline with cosine annealing (CAL) in our experiment setting, which was mentioned in line $267-268$ in the paper. For Adam, we searched the initial learning rate over {0.00005, 0.0001, 0.001, 0.01, 0.1}, and we used $\\epsilon$ = $10^{−8}$. For LAMB, we searched the initial learning rate over {0.005, 0.01, 0.02}, and we used $\\epsilon$ = $10^{−6}$. Both methods used weight decay = $5.0×10^{−4}$, $\\beta_{1}$ = 0.9, $\\beta_{2}$ = 0.999, learning rate decay with cosine annealing. Each experiment was conducted with five random seeds.\n\n2. **Missing hyperparameter study.** In Figure 19, we provide additional results of a hyperparameter study on $(s_1, s_2)$, in which we consider five different settings for $(s_1,s_2)$: $[(0.5,1.5), (0.6,1.4),(0.7,1.3), (0.8,1.2), (0.9,1.1)]$. We run tasks on CIFAR100 with four VGG and ResNet architectures, each with five random seeds. Our results show that a larger learning rate scaling range $(0.5,1.5)$ performs best. This hyperparameter setting is the default setting used in our paper. All hyperparameters are consistent with those described in the paper.\n\n## Experiments on ResNet 101\nWe present additional results in rebuttal PDF Figure 20, showing the application of our method (TB) to ResNet 101 on CIFAR-100, and we compare it with the baseline (CAL). We searched the initial learning rate among {0.05, 0.1, 0.15} for both the baseline and our method. The results report the mean and standard deviation across five seeds. We found that TB offers improvements for the larger ResNet101 model comparable to those observed for ResNet18/34, demonstrating its potential for larger models. We used the same hyperparameters as those for ResNet18/34 in Appendix C.\n\n## Visualization of layer-wise learning rate\nIn rebuttal PDF Figure 21, we visualize the layer-wise learning rates for ResNet 18/34 trained on CIFAR-100. We report the learning rate (or alpha) every epoch throughout the 200-epoch training duration.\n1. **How does the learning rate vary across layers?** We observed a correlation between the layer-wise learning rate and the layer-wise alpha distribution: layers with larger alphas are allocated larger learning rates, whereas those with smaller alphas receive smaller learning rates.\n2. **How does the layer-wise learning rate evolve during training?** The variations in layer-wise learning rates closely reflect shifts in the layer-wise alpha distribution. Initially, the alpha distributes uniformly across layers but eventually converge to a layer-wise pattern where earlier layers have smaller alphas and later layers have larger ones.\n"}, {"review_id": "v23zevX08O", "reviewer": "Reviewer_QLkn", "summary": "The paper proposes a layer-wise learning rate scheduler that adapts the learning rate to the eigenvalue distribution of the Weight-covariance matrix. As explained in the appendix, output and weight covariance matrices as well as the Hessian and Fischer-Information matrix are closely related, and they often show a power-law decay in their eigenvalues. The hypothesis is that layers with a smaller exponent in the power-law-decay are more \"overtrained\" than those with larger values, and thus should receive smaller updates. This behavior is implemented by interpolating learning rates linearly by (fitted) decay-exponents. Experiments show an improved generalization performance over base-line methods for a variety of standard computer-vision benchmarks and (CNN) architectures.", "questions": "Are the new results state-of-the-art in terms of generalization performance (compared with other LR-schedulers)? That does not need to be the case to have an interesting paper, but it would be good to know which characteristics of other approaches reach similar goals.\n\nConceptually: Would it be possible that the whole effect observed boils down to implicitly addressing exploding or vanishing gradients?\n\nEDIT: Post rebuttal, I have raised my score to accept (and soundness and contribution accordingly), as the response of the authors shows very clearly that the new approach has effects beyond gradient magnitude excursions, which was my main concern in terms of empirical evidence for the effect described. I would encourage the authors to clarify this in a revision of the paper. ", "limitations": "My main concern in terms of limitations is the possibility that other primary effects case the spectral criterion to trigger as a secondary feature. This could be excluded by additional experiments (or even some theory). One could also discuss limitations of the empirical study a bit more in detail, but I would think that the limitations are obvious to an attentive reader, so there would be no serious issues in this regard (I would not understand this as a \"new technique\" paper but rather a paper that explores a novel methodological approach, not yet arriving at a deployable method).", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper follows an interesting and less explored approach to understanding the training dynamics of networks. The observation of an emergence of power laws as such is fascinating, and studying the connection to training dynamics adds an interesting and novel perspective. Further, the empirical results are encouraging: It seems that the proposed criterion for adaptive training does indeed improve generalization performance, which is an unexpected, non-trivial result that on its own raises some eyebrows (i.e., should be discussed and explored).\n\nThe paper is well-written and in particular the discussion of the background in the appendix provides an interesting and insightful read.", "weaknesses": "My main concern with the submission in its current form is that it is light on the analytical side. In short, the paper does not try to explain why there is a connection between training success and the proposed criterion/schedule and there is a risk of overlooking unmodelled or indirect effects: In my experience, the training process of multi-layer networks is affected by a mix of numerical and fundamental issues that are usually difficult to separate. \n\nA very significant problem are for example gradient magnitude excursions (explosion/vanishing). For networks that contain batch normalization layers (or most variants thereof), the standard He initialization leads to exploding gradients at initialization. This causes early layers to \"learn\" at an exponentially larger rate than later layers, but the effect vanishes over time. It seems likely, that large rank-1 updates from initially exploding gradients distort the estimation of the decay coefficients (\"alpha\"), and likely in the \"correct\" direction of dampening exploding layers. Residual architectures still suffer from this effect within each stack in each residual block (but to a lesser degree). In contrast, non-normalized networks such as the traditional VGG suffer from vanishing gradients for other reasons. Then, there is the problem of concentration of the overall singular value spectrum for stacks of random matrices (higher-order powers of Wigner spectra), which again affects the effective rank of the updates.\n\nIn order to understand better if and how the proposed method could (or does) work, it would be useful to take such effects into account, for example by measuring gradient magnitudes over layers, or by a theoretical model. This could inform the reader better with respect to the  technical justification the proposed scheme and might help removing (or justifying) ad-hoc choices, such as linear interpolation of learning rates. It would also be helpful to carefully consider architectural aspects such as normalization and residual connections, and (at least) monitor the layer-wise statistics of gradient magnitudes and decay coefficients in order to understand better what is going on.\n\nThere are some smaller issues that could be improved, such as experimenting on a more diverse set of architectures and data sets and comparing against a larger set of base-line schedulers. In the former case, I would consider the current effort sufficient (as computational costs are skyrocketing easily, and the effect does seem quite pronounced already), in the latter case a more comprehensive study (maybe only for one or a few representative examples) could help in solidifying the result further.\n\nOverall, I found this paper very interesting and the results are unexpectedly strong. The only reason for my slightly negative overall assessment is that (for a paper at NeurIPS) more effort could be put into understanding the results better (empirically and/or analytically).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Does TempBalance (TB) boil down to addressing gradient excursions\nWe first summarize the reviewer's questions and our primary responses, with subsequent detailing of our experiment and supporting results.\n\n1. **Does gradient excursion exist?**\nWe discovered that gradient explosion does exist, but it is confined to the first epoch out of a total of 200 training epochs, leading us to believe it does not significantly impact the test accuracy. We observed no gradient vanishing.\n2. **Does the observed gradient explosion impact the estimation of alpha?**\n  We discovered that the large rank-1 updates resulting from the gradient explosion indeed affect the Empirical Spectral Density (ESD) as well as the alpha estimation. However, this effect is again restricted to the first epoch.\n3. **Does TB boil down to addressing gradient explosion?**\n   We found that postponing the use of TB until the epoch when neither the gradient explosion nor the alpha estimation is affected does not compromise the test accuracy.\n\nTo support the above answers, we conducted three experiments (see figures in the rebuttal material). We discuss the setup of these experiments first and then analyze the results.\n\n* (Figures 12,13) We aim to detect gradient excursion by tracking the gradient norm across layers during training. We examine the model every 30 iterations over the first 10 epochs, calculating the $L_2$ norm of each gradient update across layers using the training batches of size 128. This produces an empirical gradient norm distribution with a total sample size of update numbers $\\times$ layer numbers. Figure 12 presents the maximum/minimum/mean of the distribution, while Figure 13 visualizes these distribution for several iterations of Epoch 1.\n* (Figure 14) We aim to assess the impact of gradient explosion on alpha estimation by monitoring the ESD. Figure 14 examines the change of the ESD of a single weight matrix over several iterations, tracked in the experiment depicted in Figures 12 and 13.\n* (Figure 15 (a)) We aim to see if TB enhances generalization by implicitly addressing gradient explosion. Since the gradient explosion and its effect on alpha estimation only transpire in the first epoch, we postpone the starting epoch of TB to Epochs 2, 5, and 10 and see if it affects the test accuracy.\n\nOur responses to the questions are supported by the results obtained from the three experiments:\n1. **First question (Figure 12 and 13):** We observed that the notable exploding gradients only occur in the initial 200 iterations of the first epoch. In Figure 12, we pinpoint a singular peak of maximum gradient norm within the first epoch. This aligns with the abnormal distribution with a large gradient norm in the subfigure of Figure 13 titled \"Epoch 1, iteration 30\".\n2. **Second question (Figure 14):** Note that large rank-one updates have been studied in random matrix theory, which manifests as a \"bulk+spike\" pattern. This has been analyzed in, e.g., Theorem 2.13 of [1]. Figure 14 shows this \"bulk+spike\" pattern, but only in the first epoch. The ESD exhibits a heavy-tail distribution in subsequent epochs, suggesting the influence of rank-one updates is limited. \n3. **Third question (Figure 15 (a)):** Delaying the application of TB until after the first epoch does not adversely affect the test accuracy. Figure 15 (a) illustrates that applying TB from Epochs 2, 5, and 10 results in test performance comparable to when TB is applied from Epoch 1. Since the gradient explosion only occurs in the first epoch and its effect on alpha estimation diminishes after this, the effectiveness of TB does not rely on addressing gradient explosion or biased alpha estimation from large rank-one updates.\n4. **Third question**: We compare TB with the baseline method LARS, which uses gradient norms to determine layer-wise learning rates in combating gradient vanishing/explosion issues. As illustrated in Figure 4 of the submitted paper, TB outperforms LARS in terms of generalization performance.\n\n## Light on the analytical side \n1. Our method is founded on the HT-SR theory, detailed in both the introduction and Appendix A. This explains: **(1) our use of alpha for better generalization:** The paper [2] showed that modern neural networks' Empirical Spectral Densities (ESDs) typically exhibit a heavy-tail distribution. Alpha, the decay coefficient of ESD, effectively gauges generalization. [3] provided a rigorous bound for this, and [4] pinpointed an optimal alpha value close to 2. **(2) our approach to learning rate schedule based on alpha:** Based on insights from [2] and [5], we set layer-wise learning rates according to alpha values, as a higher learning rate decreases alpha. Thus, layers with larger alpha get higher learning rates. Our linear interpolation design for learning rate assignment was based on its better performance over other designs (e.g., square root, log2, step). See Figure 16 in the rebuttal PDF.\n\n2. The success of TB and its relationship with HT-SR are elucidated in Appendix B (Figures 10, 11). It reveals that TB effectively controls the shape of the ESDs. Compared to CAL, TB consistently attains a more concentrated distribution, with both the mean and median approaching the theoretically optimal PL_Alpha_Hill value of 2 [4]. This is consistently observed across various settings.\n\n## Smaller issues: comparison to other schedulers\n\nFigure 4 in our paper demonstrates our method's advantage over multiple baselines. Our rebuttal's Figure 18 further confirms its consistent improvement over parameter-wise schedulers such as Adam and LAMB. Also, our rebuttal PDF Table 6 reveals our method's performance gains in object detection and language modeling.\n\n## Reference\n[1] Couillet and Liao, 2022\n\n[2] Martin and Mahoney, 2021\n\n[3] Simsekli et al, 2020\n\n[4] Bartlett et al, 2020\n\n[5] Gurbuzbalaban et al, 2020\n\n"}, {"review_id": "UE2mwCZVML", "reviewer": "Reviewer_Gitq", "summary": "This paper proposes TempBalance for temperature balancing based on the theory of heavy-tail self-regularization (HT-SR), which is a simple yet effective layer-wise policy applicable to general global temperature allocations in deep learning regularization. This paper proposes learning rate balancing across layers, which has received less attention compared to global or parameter-wise learning rate assignment. The HT-SR motivated capacity control metrics characterize the layers to achieve maximum temperature balance during model training, resulting in improved performance during testing. Extensive experiments show that TempBalance significantly outperforms ordinary SGD and carefully tuned spectral norm regularization, as well as a number of state-of-the-art optimizers and learning rate schedulers.", "questions": "1. What are the advantages of layer-wise learning rate schedule strategy based on HT-SR theory over other layer-wise learning rate schedule strategies? What are the advantages over other global and parameter-wise learning rate schedule strategies?\n2. Optimizers are closely related to the model architecture, e.g. Transformers typically use the second-order optimizers, while CNNs typically use the first-order optimizers. In this paper, we do not analyze and experimentally verify this common phenomenon, and only use CNNs in combination with first- and second-order optimizers for classification tasks, which may lead to \"We do not find them to provide better results than the SGD baseline with cosine annealing\". Can more extensive experimentation on this issue be done when conditions permit?", "limitations": "This article provides necessary discussions on the limitations and future directions that can be explored.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The article proposes a simple yet effective layer-wise learning rate schedule TempBalance based on HT-SR theory.\n2. The article compares TempBalance with SGD and SNR on various training tasks, including different network architectures (such as ResNet, VGG, WideResNet), different datasets (such as CIFAR10, CIFAR100, SVHN, TinyImageNet), and extensive ablation studies (such as varying widths, depths, initial learning rates and HT-SR layer-wise metrics).\n3. TempBalance outperforms a range of state-of-the-art optimizers and learning rate schedulers, and maintains stable performance over SGD baselines when the model size changes.\n4. The designed algorithm is simple and easy to understand, and the paper is well-written.", "weaknesses": "1. The algorithm implementation is too simple and lacks theoretical innovation.\n2. The motivation is not clear enough, and it does not explain why it is necessary to design layer-wise learning rate schedule strategy based on HT-SR theory.\n3. The paper only conducts experiments on classification tasks and does not conduct experimental validation on downstream tasks or other areas.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Experiments on other areas\nIn Table 6 of the rebuttal PDF, we provide experiments of applying our method TempBalance (TB) to two different tasks: object detection and language modeling. In both tasks, TB consistently improves generalization, outperforming the baseline scheduler cosine annealing (CAL) when both are combined with Adam/AdamW optimizers.\n\nHere are experimental settings for object detection. We utilized the PASCAL VOC2007 [1] dataset and the YOLO-v8n [2] model with the pre-trained weights from the official website. The baseline methods are CAL combined with Adam/AdamW, while our method is TB with Adam/AdamW. For both methods, we trained for 200 epochs with batch size 64 and set the same hyperparameter for the optimizers: weight decay = $5.0×10^{−4}$, $\\beta_{1}$ = 0.9, $\\beta_{2}$ = 0.999, $\\epsilon$ = $10^{−8}$. We searched the initial learning rate for all methods among {$7.5×10^{-6}$, $1×10^{-5}$, $2.5×10^{-5}$}. We report the mean and standard deviation of mean Average Precision (mAP, higher is better) over five random seeds on the test set. \n\nHere are experimental settings for language modeling. We studied the Penn Treebank (PTB) dataset [3] using a three-layer \"tensorized transformer core-1\"[4]. We compared our method TB with the baseline scheduler CAL with both applied to the Adam optimizer. For both methods, we trained the models for 40000 iterations with a batch size of 120 and a dropout rate of 0.3. We searched the initial learning rate for all methods among {0.000125, 0.00025, 0.0005, 0.001, 0.00125, 0.0025, 0.005}. The hyperparameters for Adam are $\\beta_{1}$ = 0.9, $\\beta_{2}$ = 0.999, $\\epsilon$ = $10^{−8}$. The mean and standard deviation of perplexity (PPL, lower is better) across five random seeds on the test set are reported.\n\n## Experiments with Transformer and Adam\nWe agree that optimizers are closely related to the architecture and using the Adam with Transformer instead of CNN is a more suitable experimental setup. In our rebuttal PDF Table 6 (b), we trained a tensorized transformer using Adam. Further details are in \"Experiments on other areas.\" Our results show that TB outperforms the Adam baseline with the CAL scheduler, demonstrating its compatibility with Transformers and Adam.\n\n## Advantages over different learning rate scheduling\n1. **Compared to layer-wise learning rate scheduling (e.g., LARS):** Our method utilizes a more precise generalization metric, the alpha metric from heavy-tailed self regularization (HT-SR) theory, to enhance the generalization performance of deep models during training. This \"shape-based\" metric estimates the shape of the eigenspectrum of weight matrices. In contrast, LARS uses a \"norm-based\" metric, such as the layer-wise gradient norm. A recent study in HT-SR [5] has shown that the shape-based metric alpha surpasses norm-based ones in assessing model generalization performance. Figure 4 of the submitted paper confirms that our method outperforms the layer-wise scheduler LARS in test accuracy.\n\n2. **Compared to parameter-wise learning rate scheduling (e.g., Adam):** Similarly, our method employs the \"shape-based\" metric alpha to improve the generalization, an approach not incorporated in traditional parameter-wise methods. Our updated results in Figure 18 of the rebuttal PDF confirm that our method outperforms the parameter-wise scheduler Adam and LAMB in test accuracy. Moreover, our updated experiments in Table 6 show that combining our method with Adam/AdamW further improves the generalization.\n\nHere are the experimental setups for Figure 18. For Adam, we searched the initial learning rate over {0.00005, 0.0001, 0.001, 0.01, 0.1}, we used $\\epsilon$ = $10^{−8}$. For LAMB, we searched the initial learning rate over {0.005, 0.01, 0.02}, and used $\\epsilon$ = $10^{−6}$. Both used a weight decay of $5.0×10^{−4}$, $\\beta_{1}$ of 0.9, $\\beta_{2}$ of 0.999, learning rate decay with cosine annealing. Each experiment was conducted with five random seeds. All other hyperparameters were consistent with those described in the paper. The experimental details of Table 6 can be found in our response titled \"Experiments on other areas\".\n\n## Lack of theoretical innovation and simple algorithm design\nWe wish to emphasize that this is the first study to design learning rate schedulers based on the heavy-tailed self regularization (HT-SR) theory. We draw the theoretical insights from HT-SR, noting that the weight matrices of layers in well-trained models typically exhibit a heavy-tailed eigenspectrum. We propose to use the Hill estimator to quantify the heavy-tailed pattern of each neural network layer, utilizing this value for more steady layer-wise learning rate scheduling. This theory-guided method is novel, efficient in improving test accuracy, and straightforward to implement, as affirmed by the reviewer in the \"Strengths\" section.\n\n## Unclear motivation: why design layer-wise learning rate schedule strategy based on HT-SR theory\nThe heavy-tailed self regularization (HT-SR) theory suggests that the empirical spectral densities of weight matrices in a well-trained neural network typically display a heavy-tailed pattern. This heavy-tailed characteristic can indicate the quality of each layer, signifying whether it is overtrained or undertrained. Such theoretical insights motivate us to balance the overtrained and undertrained levels of different layers by designing a layer-wise learning rate schedule based on the heavy-tailed pattern measurement. Specifically, we assign higher learning rates to undertrained layers (as indicated by less pronounced heavy-tailed patterns) and lower rates to overtrained layers (indicated by more pronounced heavy-tailed patterns). We dynamically monitor these measurements and adjust the scheduling throughout the training process.\n\n## Reference\n[1] Everingham et al, 2010.\n\n[2] Redmon et al, 2016.\n\n[3] Mikolov et al, 2011.\n\n[4] Ma et al, 2019. \n\n[5] Martin and Michael, 2021."}]

Author rebuttal:
We want to thank all the reviewers for the constructive feedback, which helps us improve our paper. Please refer to the attached PDF for our new experiments and see below for our responses to each comment.

Meta-review: The paper proposes a theoretically-motivated procedure for automatically balancing learning rates across different layers of a neural network.  Experimental results demonstrate the empirical effectiveness of the layer-wise learning rates derived from the method.  After the rebuttal and discussion phase, all reviewers recommend accept, with some notable enthusiasm over the results achieved by the proposed optimizer.  The AC agrees with the reviewer consensus.
Decision: Accept (spotlight)