SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 3: (Main result). Comparing our method TempBalance (TB) to CAL and SNR. Our method TempBalance outperforms CAL and SNR in 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 D for the details in all hyperparameters.
Datasets. We consider CIFAR100, CIFAR10, SVHN and Tiny ImageNet (TIN) [72-75]. CIFAR100 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) [76-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 SGD training with a cosine annealing learning rate schedule (CAL), which follows the formula:  $\eta_t = \frac{\eta_0}{2}\left(1 + \cos\left(\frac{t\cdot\pi}{T}\right)\right)$ , where  $t$  is the current epoch,  $T$  represents the total training epochs, and  $\eta_0$  is the initial learning rate. We grid search the optimal initial (base) learning rate  $\eta_0$  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  $5\times 10^{-4}$ , respectively, which are both standard choices.
Another baseline is spectral norm regularization (SNR). Prior work [37] uses the SNR objective:
$$
\min  _ {\Theta} \frac {1}{n} \sum_ {i = 1} ^ {n} l \left(f _ {\Theta} \left(\boldsymbol {x} _ {i}\right), \boldsymbol {y} _ {i}\right) + \frac {\lambda_ {s r}}{2} \sum_ {l = 1} ^ {L} \sigma \left(W _ {l}\right) ^ {2}, \tag {4}
$$
where  $\lambda_{sr}$  is the SNR coefficient,  $\sigma(W_l)$  is the largest eigenvalue, i.e., spectral norm of weight matrix  $\mathbf{W}_l$ , and  $L$  is the number of layers. We use the power iteration method to calculate  $\sigma(W_l)$  in our experiments. For SNR, we grid search the optimal regularization coefficient  $\lambda_{sr}$ , and we again adopt the CAL schedule for SNR, similar to the CAL baseline.
To make our results fully reproducible, we report in Appendix D all 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 a CAL learning rate schedule, and SNR means SGD trained
(a) ResNet18, CIFAR100
(b) ResNet34, CIFAR100
Figure 4: (More baseline optimizers). Comparing our method TempBalance (TB) to cosine annealing (CAL) baseline and other state-of-the-art optimizers and learning rate schedulers for ResNet18 and ResNet34 trained on CIFAR100. 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 TB combined with SGDP.
with spectral norm regularization. TB means our method TempBalance, and TB + SNR means TempBalance combined with SNR. All error bars are obtained from five random seeds. From Figure 3, 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 leads to better test accuracy.
Second, we compare our method to a number of optimizers and learning rate schedulers that are not necessarily related to ESD of weights. These include SGDR [10], SGDP [41], 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 SGD 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 in Figure 4 show that TempBalance outperforms these baselines, and TempBalance combined with SGD 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 [11], AdamP [41] and LAMB [79]. However, we 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 E.
4.3 Corroborating results and ablation studies.
In addition to the main results (Figures 3 and 4), we provide corroborating results and ablation studies.
Experiment one: tuning initial learning rate  $\eta_0$ . 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 ResNet18, ResNet34, VGG16 and VGG19 as our architectures and show results on CIFAR100. Results in Figure 5 show that TempBalance achieves a higher test accuracy than CAL for both ResNet and VGG.
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, VGG16 and VGG19, and trained on CIFAR100, and we grid search for the optimal learning rate for each width to get the best accuracy. Results in Figure 6 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_Alpha_Hi11 in (2) with other HT-SR metrics including SpectralNorm and AlphaWeighted [22]. Results in Figure 7 show that PL_Alpha_Hi11 achieves the optimal test accuracy.
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 outperforms CAL for both ResNet and VGG trained on CIFAR100. All results are obtained by running five random seeds.
Figure 6: (Different widths). Comparing TempBalance and the CAL baseline for different network widths. Our method TempBalance consistently outperforms the CAL baseline across various network widths for both ResNet and VGG trained on CIFAR100. All results are obtained by running five random seeds.
Figure 7: (Different HT-SR metrics). Comparing PL_Alpha_Hill with multiple HT-SR metrics. PL_Alpha_Hill achieves the best test accuracy among these metrics. All results are obtained by running five random seeds.
Experiment four: varying PL fitting methods. The HT-SR metric PL_Alpha_Hill is derived through PL fitting, which is influenced by the choice of hyperparameter  $\lambda_{\mathrm{min}}$ . More specifically, this involves determining the adjustable parameter  $k$  as per Equation 3. Past research has employed various methods to select  $\lambda_{\mathrm{min}}$  based on the task, such as performance prediction. For instance, Martin et al. [22], Clauset et al. [39] choose  $\lambda_{\mathrm{min}}$  that aligns with the best fit according to the Kolmogorov-Smirnov statistic [40], a method termed Goodness-of-fit. Meanwhile, Yang et al. [32] adopted the Fix-finger approach, which identifies  $\lambda_{\mathrm{min}}$  at the peak of the ESD. In our study, we designate  $\lambda_{\mathrm{min}}$  as the median of all eigenvalues present in the ESD for TempBalance. As depicted in Figure 8, our 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 of learning rate schedulers that demand low computation overhead.
Empirical analysis results. We conduct an empirical analysis of TempBalance to discuss why it provides improvement. Our first analysis involves visualization to demonstrate how TempBalance
Figure 8: (Varying PL fitting method to determine the  $\lambda_{\mathrm{min}}$ ). 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 (Median) 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 is averaged over ten times.
(b) ResNet34, CIFAR100
(c) VGG16, CIFAR100
(d) VGG19, CIFAR100
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 excursions [80] (see Appendix F).
Corroborating results on other tasks. We extend our evaluation of TempBalance to 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.
Our extensive empirical evaluations demonstrate that TempBalance offers a straightforward yet effective layer-wise learning rate schedule. Our approach for balancing layer-wise temperature confirms the following: (i) HT-SR-motivated metric PL_A1pha_Hi11 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 HT-SR 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 of layer-wise tuning approaches and load-temperature balancing in deep NN training, as 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 to 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 to achieve a more adaptive learning rate scheduler? Currently, we calculate layer-wise PL_Alpha_Hill once per epoch, resulting in a minimal increase in computational cost. Consider the example of training ResNet18 for 200 epochs on CIFAR100. Calculating layer-wise PL_Alpha_Hill takes 1.14 seconds for each epoch, leading to 3.8 minutes in total. Training CIFAR100 on 1 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 on computation overhead is detailed in Appendix H.
Our research centers around developing a 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 a 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 a theory that provides vacuous upper bounds and then relies on extremely expensive hyperparameter tuning to obtain good results.
Acknowledgements. WeightWatcher is a publicly available tool distributed under Apache License 2.0 with a copyright held by Calculation Consulting. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
[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.02083, 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.02218, 2020.
[92] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[93] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting 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-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dóllá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 a large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2):313-330, 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.
A Heavy-tail phenomena in different DNN matrices are closely related
Recently, several papers have separately studied HT structures in different types of matrices, including the Hessian, the Fisher Information Matrix (FIM), and input/output covariance matrices [35, 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_Hi11), or HT-shaped decaying eigenvalues [33-35]. 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 to 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.1. We find that if one of these matrices shows the PL phenomenon, the other two matrices have a high chance to exhibit a similar PL phenomenon. Then, in subsection A.2, we derive the connection between our metric PL_Alpha_Hi11 and the PL exponent on decaying eigenvalues, showing a simple reciprocal relationship between these two.
A.1 Connections between different matrices
Consider a NN  $f_{\theta}:\mathbb{R}^{d}\to \mathbb{R}^{C}$ , where  $\theta \in \mathbb{R}^P$  is the vectorized weights,  $d$  is the input dimension, and  $C$  is the output dimension. When the NN is used for a classifying task,  $C$  is also the number of classes. We denote the input data as  $\{(x_i,y_i)\}_{i = 1}^n$ , where  $x_{i}\in \mathbb{R}^{d}$ , and the number of samples is  $n$ . We denote the loss function as  $L(\theta) = \frac{1}{n}\sum_{i = 1}^{n}l(y_{i},f_{\theta}(x_{i}))$ .
Covariance matrices. We denote the output covariance matrix as  $\mathbb{E}[f_{\theta}(x)f_{\theta}^{\top}(x)]$ , where the expectation is taken over the input distribution. We tend to consider the following empirical covariance matrix:
$$
C (\theta) := \frac {1}{n} \sum_ {i = 1} ^ {n} f _ {\theta} \left(x _ {i}\right) f _ {\theta} ^ {\top} \left(x _ {i}\right) \in \mathbb {R} ^ {C \times C}. \tag {5}
$$
Fisher Information Matrices. We denote the (output) FIM as
$$
\mathbb {E} \left[ \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {\top} \right] = \sum_ {k = 1} ^ {C} \mathbb {E} \left[ \nabla_ {\theta} f _ {\theta} ^ {(k)} (x) \nabla_ {\theta} f _ {\theta} ^ {(k)} (x) ^ {\top} \right], \tag {6}
$$
where  $f_{\theta}^{(k)}(x)$  is the  $k$ -th entry of the vector function  $f(x)$ . We also consider the empirical version of the FIM:
$$
F (\theta) := \sum_ {k = 1} ^ {C} \frac {1}{n} \sum_ {i = 1} ^ {n} \nabla_ {\theta} f _ {\theta} ^ {(k)} \left(x _ {i}\right) \nabla_ {\theta} f _ {\theta} ^ {(k)} \left(x _ {i}\right) ^ {\top} \in \mathbb {R} ^ {P \times P}. \tag {7}
$$
Note that (7) can be equally written as
$$
F (\theta) := \frac {1}{n} \nabla_ {\theta} \tilde {f} _ {\theta} (x) \nabla_ {\theta} \tilde {f} _ {\theta} (x) ^ {\top}, \tag {8}
$$
where  $\nabla_{\theta}\tilde{f}_{\theta}(x)$  has the following form:
$$
\left[ \begin{array}{c c c c} \frac {\partial f _ {\theta} ^ {(1)} (x _ {1})}{\partial \theta_ {1}} \dots \frac {\partial f _ {\theta} ^ {(1)} (x _ {n})}{\partial \theta_ {1}} & \dots & \frac {\partial f _ {\theta} ^ {(C)} (x _ {1})}{\partial \theta_ {1}} \dots \frac {\partial f _ {\theta} ^ {(C)} (x _ {n})}{\partial \theta_ {1}} \\ \vdots & \ddots & \vdots \\ \frac {\partial f _ {\theta} ^ {(1)} (x _ {1})}{\partial \theta_ {P}} \dots \frac {\partial f _ {\theta} ^ {(1)} (x _ {n})}{\partial \theta_ {P}} & \dots & \frac {\partial f _ {\theta} ^ {(C)} (x _ {1})}{\partial \theta_ {P}} \dots \frac {\partial f _ {\theta} ^ {(C)} (x _ {n})}{\partial \theta_ {P}} \end{array} \right] \in \mathbb {R} ^ {P \times C n}.
$$
Hessian Matrices. We denote the Hessian as  $\mathbb{E}\left[\frac{\partial^2l(y,f_\theta(x))}{\partial\theta^2}\right]$ , and we tend to consider the empirical Hessian Matrices:
$$
H (\theta) := \frac {\partial^ {2} L (\theta)}{\partial \theta^ {2}} \in R ^ {P \times P}, \tag {9}
$$
where  $L(\theta)$  is the empirical loss function  $L(\theta) = \frac{1}{n}\sum_{i = 1}^{n}l(y_i,f_\theta (x_i))$
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:
$$
s F I M := \mathbb {E} \left[ \nabla_ {\theta} \log P (y | x; \theta) \nabla_ {\theta} \log P (y | x; \theta) ^ {T} \right], \tag {10}
$$
where  $P(y|x;\theta)$  represents the likelihood. After simple derivations, one can show that sFIM also has the following form [83, 84]:
$$
s F I M = - \mathbb {E} \left[ \frac {\partial^ {2} \log P (y | x ; \theta)}{\partial \theta^ {2}} \right]. \tag {11}
$$
Therefore, when the loss function is defined as the negative log-likelihood, the sFIM in (11) is equivalent to Hessian defined in (9).
Why is the FIM defined in (6) equivalent to (10). Back to deep learning, the FIM is often defined as (6). It is thus meaningful to derive the equivalence between these two forms. Suppose  $P(y|x;\theta)$  here means the conditional probability distribution of output  $y$  given input data  $x$ . If  $P(y|x;\theta)$  is assumed to take the following form:
$$
P (y | x; \theta) = \frac {1}{\sqrt {2 \pi}} \exp \left(- \frac {1}{2} \| y - f _ {\theta} (x) \| ^ {2}\right), \tag {12}
$$
then the MSE estimator  $\min_{\theta} \frac{1}{2} \| y - f_{\theta}(x) \|^2$  is equivalent to the maximum likelihood estimation of  $P(y|x; \theta)$ . Then, plugging (12) into (10), we have:
$$
s F I M _ {m s e} = \mathbb {E} [ \| y - f _ {\theta} (x) \| ^ {2} \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {T} ]. \tag {13}
$$
We now expand  $sFIM_{mse}$  by the definition of expectation, and we have the following [81]:
$$
\begin{array}{l} s F I M _ {m s e} = \int_ {\mathbb {R}} \int_ {\mathbb {R}} \| y - f _ {\theta} (x) \| ^ {2} \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {T} p (x, y; \theta) d y d x (14) \\ = \int_ {\mathbb {R}} \int_ {\mathbb {R}} \| y - f _ {\theta} (x) \| ^ {2} \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {T} P (y | x; \theta) q (x) d y d x (15) \\ = \int_ {\mathbb {R}} \left[ \int_ {\mathbb {R}} \frac {1}{\sqrt {2 \pi}} \| y - f _ {\theta} (x) \| ^ {2} \exp \left(- \frac {1}{2} \| y - f _ {\theta} (x) \| ^ {2}\right) d y \right] \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {T} q (x) d x (16) \\ = \int_ {\mathbb {R}} \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {T} q (x) d x (17) \\ = \mathbb {E} \left[ \nabla_ {\theta} f _ {\theta} (x) \nabla_ {\theta} f _ {\theta} (x) ^ {T} \right], (18) \\ \end{array}
$$
where (14) follows from the definition of expectation,  $q(x)$  is input distribution, and (17) holds because the integral of  $y$  in the brackets [] equals 1 due to the property of Gamma function  $\Gamma(\cdot)$ .
Therefore, from (18), we find that  $sFIM_{mse}$  is just equal to  $FIM$ , defined in (6). Also, plugging (12) into  $\mathbb{E}\left[\frac{\partial^2\log P(y|x;\theta)}{\partial\theta^2}\right]$  and taking the loss function  $L(\theta)$  as the mean-square loss, we will again find that  $\mathbb{E}\left[\frac{\partial^2\log P(y|x;\theta)}{\partial\theta^2}\right]$  is equal to  $H(\theta)$ . Therefore, jointly considering (11), we can see that FIM is equal to the Hessian  $H(\theta)$ .
PL in the covariance matrix and PL in Hessian are tightly correlated. Next, we consider the relationship between the covariance matrix and the Hessian. Suppose the NN function  $f_{\theta}$  is a Lipschitz 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.
Figure 9: We train a MLP for 50 epochs and fit PL exponent  $s$  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 a strong correlation.