Temperature Balancing, Layer-wise Weight Analysis, and Neural Network Training

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.

Strengths

1. The article proposes a simple yet effective layer-wise learning rate schedule TempBalance based on HT-SR theory. 2. 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). 3. 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. 4. 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. 2. 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. 3. The paper only conducts experiments on classification tasks and does not conduct experimental validation on downstream tasks or other areas.

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? 2. 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?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

This article provides necessary discussions on the limitations and future directions that can be explored.

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.

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). The 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. A 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. In 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. There 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. Overall, 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).

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. Conceptually: Would it be possible that the whole effect observed boils down to implicitly addressing exploding or vanishing gradients? EDIT: 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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).

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.

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. 2. 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? 2. The authors do not provide a convergence analysis. 3. 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. 4. Missing baseline: LAMB [1]. No hyper-parameter study on $s_1$ and $s_2$. 5. 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. [1] You, Yang, et al. "Large batch optimization for deep learning: Training bert in 76 minutes." arXiv preprint arXiv:1904.00962 (2019).

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? 2. See weakness.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have address the limitations.

Summary

The paper proposes a way of modulating the learning rate, independently for each layer, when training deep networks via gradient descent. This 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. That 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. A 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).

Strengths

Originality -------------- Layer-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. Quality ---------- The 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. Experiments have been replicated with 5 different seeds, and error bars reported. Experimental settings are clearly reported in the appendix. Overall really solid experimental validation. Clarity --------- The paper is clearly organized, and explained really clearly the theoretical bases it builds on, existing algorithms, as well as the proposed new method. Significance ----------------- A 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. 1. It would have been interesting to see one series of experiments on larger-scale data (full ImageNet?), or maybe non-image data. 2. No explicit comparison with parameter-wise scaling schemes (e.g., a variation of Adam) 3. No mention of the algorithmic complexity, or overhead of computing these scales, before the penultimate paragraph.

Questions

Clarification questions 1. 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. 2. Should $(s_1, s_2)$ be mentioned in the Hyperparameters section (l.235)? Further information / curiosity 3. 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. 4. 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`? 5. 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. 6. 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? *Update after rebuttal* I believe all my questions were addressed. The additional experiments are thorough and lead me to increase my score.

Rating

9: Very Strong Accept: Technically flawless paper with groundbreaking impact on at least one area of AI/ML and excellent impact on multiple areas of AI/ML, with flawless evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Properly addressed within the body of the article.

A gentle reminder about the discussion phase

Dear reviewers, thanks for your constructive feedback. Please tell us if we should include anything further in the revised draft. We are more than happy to clarify if anything is unclear.

Thank you!

Thanks for the update and additional thorough experiments. I believe all my questions are addressed.

Thanks for your response

The new experiment results and analysis have addressed my concerns. My score has been updated.