Super Consistency of Neural Network Landscapes and Learning Rate Transfer

We thank the Reviewer for the strong overall score and for marking “excellent” our paper across all the three evaluation axes.

On the weakness:

1. Larger scale experiment: We managed to scale the Hessian computation up to 300 million parameters in a Transformer model trained on Wikitext with Adam (width scaling), surpassing the scale reached in the reference literature on the EoS (Cohen et al, 2021). Please consult Figure 1 in the one-page pdf attached. We find that Super Consistency still holds at this scale. See also the global response on this point.

Questions:

1. Hessian vs NTK dynamics. Under $\mu$P, both the Hessian and NTK evolve from initialization, a distinctive feature of feature learning scaling limits. These two quantities are indeed related by the Gauss-Newton decomposition of the Hessian, which for MSE loss reads $H = K^T K + R$. Indeed, for MSE loss the dynamics of the NTK and Hessian are connected, in that the change of the Hessian is largely driven by the change in the NTK (Figure 5). However, notice that the matrix $R$ still plays a role and cannot be neglected when the network is far from convergence (more subtly, the interaction between K and R gives the self-stabilization property that preserves edge of stability for $\lambda_{\max}$). Secondly, under cross-entropy loss, the first term of the decomposition is not $K^T K$ but $K^T D K$, where $D$ contains the second derivative of the loss for different datapoints. Thus, there is an additional term that potentially makes the dynamics of the Hessian different from the NTK. Indeed, this is what we observe in Figure 2, where the NTK eigenvalues accumulate significant finite-size effects, while the Hessian top eigenvalues are Super Consistent.

We hope that the new experiments and clarification will make the Reviewer more confident about their evaluation.

We thank the reviewer for acknowledging the importance and novelty of our results. Here we address the concerns:

1. A more diverse set of experiments: we would like to gently push back on this point. We have already performed experiments on CIFAR-10, Imagenet (vision), and wikitext (language) with two different architectures: ResNets and Transformers (including Vision Transformers). We have also performed experiments on several parametrizations, including NTK, $\mu$P and SP in width, and Depth-$\mu$P and residual networks without $1\sqrt{L}$ scaling for depth. We have trained models with up to 96 layers in depth (e.g. Figure 19), and 8172 units in width. We have other interesting ablations, including batch size, data augmentation, and different optimizers (SGD, Adam, now including AdamW as well, see the one page pdf included in the global response).

2. New Evidence on Transformers at scale and convolutional networks. We have performed several new experiments on width scaling of Transformers (with Adam) and convolutional networks (with Adam and AdamW). In contrast to depth-scaling, where Super Consistency was not great due to the nature of the Transformer layer, in this case we observe Super Consistency to hold from very small to very large width Transformers’ (scaling up to about 300 million parameters for the largest model) and the learning rate to transfer significantly better. Please see Figures 1 and 2 of the one-page pdf. In particular, our results are based on Post-LN (like GPT-2) models trained on wikitext (with Adam) and Covnets on CIFAR-10. We hope that this extra experimental evidence addresses the Reviewer's concerns about the presence of Super Consistency across more models and datasets.

3. Absence of Super Consistency in Transformer experiment (with $\mu$P-Depth): as the Reviewer mentions, we do hypothesize that the observed breach of super consistency is due to the nature of the Transformer block design, having multiple layers per residual block. This breach is thus expected, and we have also reproduced it in the ResNet case with 2 layers per residual block, as in Figures 4, 19 (a,b), and 20 (a,b,c). Furthermore, the absence of Super Consistency correlates with worse learning rate transfer, and it is thus compatible with our claims.

Questions:

1. Hessian Computation. We use PyHessian to compute the eigenvalues, which adopts the Power iteration method to get the eigenvalue spectrum. We use a large fixed batch of 1024 datapoints for Hessian estimation. Models with more than 124 million parameters would certainly fit a single A100 GPU, but the estimation becomes very slow and thus cannot be performed at the same frequency as a smaller model. However, as shown in Figure 1 of the one-page pdf, we have run a Post-LN Transformer with up to about 300 million parameters and observed that super consistency still holds at this larger scale.

2. SGD along Hessian directions and learning rate transfer Indeed, there is compelling evidence (Gur-Ari et al, 2018) that SGD happens in the subspace spanned by the top $k$ Hessian eigenvectors, where k is the number of classes. In Figure 2 of the paper, we report how a subset of the top 10 eigenvalues behaves across time and model scale. In most of the other experiments of the paper, we restrict to top eigenvalue $\lambda_{\max}$ due to computational reasons and due to its importance in learning rate choice. Please see Appendix A for a thorough discussion on the importance of the sharpness in the choice of the learning rate, before and after the EoS era. Regarding the specific wording, by that sentence, we meant that “lambda_{\max} is super consistent across the training trajectory”. We made this clear in the revised version of the paper.

   • New Evidence on Gradient-Hessian alignemnt. We performed an experiment where we computed $g^THg/||g||^2$, where g are the weight gradients and H is the loss Hessian, thus taking into account the curvature along the gradient direction. The results are in Figure 3. We observe that Super Consistency largely holds in this case as well. Finally, please notice that all our experiments show a clear correlation between sharpness super consistency (in general, top k eigenspectrum) and learning rate transfer. Due to the importance of the sharpness in determining the trainability of the model at large learning rates, our results provide strong empirical evidence for a potential explanation of the phenomenon of learning rate transfer. However, more work has to be done in neural network optimization to establish (i.e. prove) how the optimal learning rate depends on the Hessian spectrum (together with other quantities).

We sincerely thank the reviewer for the detailed feedback and for the further discussion that we anticipate. Also, we thank the reviewer for acknowledging certain strengths of the paper, such as the importance and validity of our findings in the context of understanding neural networks’ loss landscape at different scales.

The reviewer’s main concern is regarding how Super Consistency is defined (Definition 3.1). The main purpose of putting forward the definition was to have an actionable measure of Super Consistency that would reflect the precise quantitative results of Figure 3 and it is not supposed to be used in a mathematical theory. This seems to be in line with the Reviewer’s suggestions to “cut back on formalising things like “super consistency”. However, we do agree with the Reviewer that certain aspects of the presentation may seem to introduce mathematical formalism, which is not intended due to the experimental nature of the presented results. Thus, we decided to describe superconsistency outside of a formal definition environment, taking into account the Reviewer's valuable feedback. We hope that the Reviewer agrees that this is a more sensible choice in the context of our paper.

To clarify the reviewer’s confusion on the concept of Super Consistency: we started from the concept of Consistency in Vyas et al (2023) which refers to how certain aspects of the network (such as the logits and the loss function) are the same at different scales early in training. We extend this concept to Super Consistency by requiring that Consistency is maintained for a longer period of training time, as we state in the abstract (lines 8-9) and introduction (lines 34-36). In the revised version of the paper, we try to achieve this objective once again, hoping to resolve the confusion. For your reference, here is the new phrasing of Section 3, replacing Definition 3.1:

Super Consistency refers to when certain aspects of the loss landscape and of the predictor $S_N(t)$ (in this paper $S_N(t)$ refers to the NTK's and loss Hessian's eigenvalues or the loss itself) exhibit the following two properties:

1. At large $N$, $S_N(t)$ does not deviate significantly from its limit $S_{\infty}(t) := \lim_{N\to\infty}S_N(t)$. This is what is referred to as consistency in Vyas et al, 2023.

2. $S_N(t)$ does not accumulate significant finite size effects over time, i.e. the curves of $S_N(t)$ and $S_\infty(t)$ remain close over a sustained period of training.

With respect to the experiment illustrated in Fig. 1, notice that the curves of the loss (center) at different widths show progressive and significant deviations, thus violating Super Consistency. On the other hand, the sharpness dynamics for $\mu$P qualitatively exhibit little-to-no deviations. Also, notice that we assume the existence of the limit $\lim_{N\to\infty}S_N(t)$. For those parametrizations (e.g. standard parametrization) that do not have a well-defined limit, $S_N(t)$ diverges at large $N$ and Super Consistency is trivially violated.

And a few lines later:

To give a quantitative measure to the finite-size accumulation property, we measure deviations over time by estimating the following quantity: $g(t) := |S_{N}(t) - S_{\infty}(t)|.$ When $g(t)$ is an increasing function (up to fluctuations), Super Consistency is violated.

The current phrasing simplifies the notation (addressing the first 5 points in the weakness Section), and clarifies the intended meaning of Super Consistency in the abstract and intro (addressing the 6th point). We hope that the clarifications on the goal of the definition, together with the new simplified version of the corresponding Section solve the Reviewer’s concerns about the formalization of the Definition and its underlying meaning.

Others:

1. “I find the phrase “optimization happens along trajectories of super consistent top Hessian eigenvalues” quite confusing”. This wording was intended and comes from a bias on related literature, such as “Optimization happens at the edge of stability” or “Gradient Descent Happens in a Tiny Subspace”. However, we agree that it might create slight confusion and have changed it according to the Reviewer’s recommendation.

2. The typo in line 161 has been fixed.

Questions:

1. Why are there two sets of lines for each width in Figure 1, one solid and one dashed? Referring to the left column, different line styles indicate different learning rates, we plot 3 learning rates for $\mu$P and 2 learning rates for NTK to not make the plots too cluttered.

2. Appendix G. We thank the reviewer for pointing this out. We have now completed the section.

3. How come the authors decided to look at a fixed N-th largest eigenvalue? The reason for this is that earlier work established that SGD dynamics largely happens in the subspace spanned by the top $k$ Hessian’s eigenvectors, where $k$ is the number of classes. In our setting, $k$ is a constant independent of the width. Also, due to EoS-type of results, the largest eigenvalue would be converging to $2/\eta$ regardless of the width. However, in the direction proposed by the Reviewer, we have performed a new experiment where we measure the curvature along the gradient direction, which is a more global measure in the sense that it takes into account the whole Hessian. Super Consistency still holds (See Figure 3 of the one-page pdf)

Other remarks:

1. More figures like Figure 3. The Reviewer appreciated the experiments on the distance between finite and infinite models over time as in Figure 3, suggesting that more plots like that should be included for other experiments. Please notice that Figures 20 and 22 also show this kind of analysis for other experiments.

2. On the time horizon. We operate at large training time, but not to complete convergence, as we show in Appendix E.1), thus we avoided talking about time asymptotics of the sharpness dynamics.

Finally, we thank the reviewer for their valuable and precise feedback, and we would like to gently push back on the score of 1 given to the paper presentation. The general consensus amongst the other reviewers is that the main body of work is well-written and well-presented and we would kindly ask the reviewer to reconsider their score.

We deeply appreciate that the reviewer is now in favor of acceptance.

To wrap it up, we will include the additional experiments on the hessian's eigenspectrum, and clarify the meaning of Super Consistency at all points in the paper according to the meaning and phrasing agreed here. We will also run a coord check experiment to verify that features updates are $\mathcal{\Theta}(1)$.

We thank the Reviewer for the careful assessment of our paper, and for highlighting the extensive suite experiments that we run, including different scaling regimes ($\mu$P, Depth $\mu$P, SP, NTK, etc..) and architectures. On the weaknesses:

1. Theoretical limitation: please notice that our paper is the first that analyzes a two-layer network in this generality. The closest work to our paper is Kalra et al, 2023. In contrast to this work we allow multiple datapoints instead of a single data point study, learning targets different than 0. Additionally, we use the standard definition of sharpness (largest eigenvalue of the loss Hessian) instead of the Hessian trace as a proxy. Extending our framework to nonlinear (and deeper) networks is certainly a fascinating avenue for future work. However, it will likely incur significant difficulties in computing the sharpness of deep nonlinear networks during training. Finally, please notice that our paper is mainly empirical, and the scope of the theory is mainly to provide intuition and justification of the empirical results. Under this perspective, the two-layer linear network is an interesting tractable model that expresses the phenomenology observed in practice for more complex networks.

2. AdamW: We have performed new width-scaling experiments with Adam and AdamW on 3 layers Convolutional Networks, extending Super Consistency to these settings. The weight decay scaling in AdamW follows the settings of Wang, X. and Aitchison, L., 2024 and Yang et al., 2022. Note that in the added experiments on Convolutional Networks, the learning transfers across width, and we have a Super Consistent sharpness evolution during training. We also believe that our results could be empirically extended to other optimization algorithms as well. See Figure 2 of the one-page pdf for our training runs on CIFAR-10. We note that a very interesting future work avenue would involve understanding the role of adaptivity in Super Consistency from a theoretical point of view.