Super Consistency of Neural Network Landscapes and Learning Rate Transfer
Authors: Lorenzo Noci
Alexandru Meterez
Thomas Hofmann
Antonio Orvieto
Venue: NeurIPS 2024 poster (2024)
Keywords: mup
deep learning theory
optimization theory
edge of stability
NTK
Primary area: optimization_for_deep_networks
TL;DR: We show the some properties of the loss landscape are preserved across different model sizes under the right scaling of the architecture.

Abstract: Recently, there has been growing evidence that if the width and depth of a neural network are scaled toward the so-called rich feature learning limit ($\mu$P and its depth extension), then some hyperparameters --- such as the learning rate --- exhibit transfer from small to very large models. From an optimization perspective, this phenomenon is puzzling, as it implies that the loss landscape is consistently similar across very different model sizes. In this work, we study the landscape through the lens of the Hessian, with a focus on its largest eigenvalue (i.e. the sharpness), and find that certain spectral properties under $\mu$P are largely independent of the width and depth of the network along the training trajectory. We name this property *super consistency* of the landscape. On the other hand, we show that in the Neural Tangent Kernel (NTK) and other scaling regimes, the sharpness exhibits very different dynamics at different scales. But what causes these differences in the sharpness dynamics? Through a connection between the Hessian's and the NTK's spectrum, we argue that the cause lies in the presence (for $\mu$P) or progressive absence (for the NTK scaling) of feature learning.
We corroborate our claims with a substantial suite of experiments, covering a wide range of datasets and architectures: from ResNets and Vision Transformers trained on benchmark vision datasets to Transformers-based language models trained on WikiText.

Super Consistency of Neural Network Landscapes and Learning Rate Transfer
Lorenzo Noci"
Alexandru M 344
Thomas Hofmann Antonio Orvieto 234
Abstract
Recently, there has been growing evidence that if the width and depth of neu- ral network are scaled toward the so-called rich feature learning limit (μP and its depth extension), then some hyperparameters hibit transfer from small to very large models. From an optimization perspective, this phenomenon is puzzling, similar across very different model sizes. In this work, we study the landscape through the lens of the loss Hessian, with focus on its largest eigenvalue (i.e. the sharpness), and find that certain spectral properties under μPP are largely inde- pendent of the size of the network, and remain consistent training progresses. We name this property Super Consistency of the landscape. On the other hand, we show that in the Neural Tangent Kernel (NTK) and other scaling regimes, the sharpness exhibits very different dynamics at different scales. But what causes these differences in the sharpness dynamies? Through connection between the Hessian's is and the NTK's spectrum, we argue that the cause lies in the presence (for μP) or progressive absence (for the NTK scaling) of feature learning. We corroborate our claims with substantial suite of experiments, covering wide range of datasets and architectures: from ResNets and Vision Transformers trained on benchmark vision datasets to Transformers-based language models trained on WikiText.
such the learning rate
it implies that the loss landscape consistently
Introduction
Recent trends in deep learning research have unmistakably shifted towards an increase in model sizes, with networks comprising of billions of parameters emerging as the standard [1]. However, models enlarge, does the cost incurred in hyperparameter tuning which has led researchers look for ways to scale up the architecture both in terms of width and depth - while preserving the optimal hyperparameters (such as the learning rate).
While there exist several ways (a.k.a parametrizations) to scale up the width and depth of the net- work, not all of them facilitate learning rate transfer. For standard deep learning practices, such as networks parametrized with LeCun/K oeCun/Kaiming initializations [2, 3], a significant shift the optimal learning rate is usually observed as the width and the depth of the model are increased. Similarly, under the Neural Tangent Kernel (NTK) parametrization [4]. which provides theoretical insights into the behavior of very wide neural networks during training, the optimal learning also varies as the width and depth of the network change. Alternatively, Yang and Hu [5] and Yang et al. [6] propose the μuP framework, designed maximize the gradient update the representations of the intermediate layers (i.e. feature learning) as the width increases. Under μP scaling, and its depth ex- tension for residual networks Depth-uP [7, 8]. it has been empirically demonstrated that the learning
Equal contribution. Correspondence to: lorenzo. noci@inf. www.fettz.ccmm ch, ameterezOfas harvard. edu "ETH Zürich, "ELLIS Tübingen, MMI Intelligent Systems, "Tübingen AI Center, "Harvard University
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Width
128
256
512
1024
2048
4096
JP
J/P
UP
x 10"
Sharpness X
Train Loss
109
Leamiing Rate
0.19
0.53
10a
9 10-1
10
20 Epoch
30
40
50
10
20
30
40
50
10-1
10a Learning Rate
Epoch
NTK
2x 10
Sharpness A
-0-1
Train Loss
Learning Rate
10
6.81
52.75
10a
90-1
0
20
40
20
40
10a
10a
10a
Epoch
Epoch
Learning Rate
Figure 1: Top row. Under JP, (left) the sharpness dynamics are largely identical for the whole training dynamics across different widths, phenomenon that we call Super Consistency. The dashed horizontal lines are the Edge of Stability thresholds. Center: The loss dynamics similar early in training, but accumulate finite-size effects over time, thus violating Super Consistency. Right: The learning rate transfers from small to large model, suggesting that the loss landscape is Super Consis- tent across different model sizes. Bottom row. Under NTK parameterization (NTP), the sharpness dynamics show large discrepancies. Also, the learning rate does not transfer. The architecture is a two-layer convolutional network trained on CIFAR- with data augmentation, where the width corresponds to the number of filters in the convolution. (See App. J). Other parameters: B 128, epochs 50.
rate transfers across both width and depth. Vyas al. [9] is observed that in feature learning parametrizations (e.g. μP) the model's dynamics are consistent across model sizes, but for harder tasks or longer training times there are progressive and significant deviations across different model sizes. We give example in Figure 1 (top center), where the training losses exhibits an increasing gap. The fact that the learning rate exactly preserved, however, suggests that some properties of the landscape do not exhibit these finite-size deviations, and must be precisely preserved across different model sizes for the whole training trajectory.
Motivated by this, in the present work we identify the notion of Super Consistency to describe the function
properties of the neural network's loss landscape that are preserved across training of the model width and depth, thus not accumulating finite-size effects. In particular, we analyze the landscape through the lens of the loss Hessian. It provides insights into the landscape's 's local curvature, and its structure for neural networks has been studied in several works [10-14]. Of great interest in optimization theory is sharpness, i.e. the top Hessian eigenvalue, which for neural networks exhibit rapid increase (progressive shurpening) towards threshold called Edge of Stability (EoS) [15, 16]. However, although a few works have provided carly insights [10, 16, 9]. the scaling properties of the sharpness and Hessian's 's dynamics under different scaling limits remain unexplored. In this work we first present evidence of Super Consistency in the Hessian's 's largest eigenvalues, which have been shown to control the curvature along the optimization subspace [17]. We then focus on the sharpness dynamics, and find that the presence (resp. absence) of Super Consistency correlates well with presence (resp. absence) of learning rate transfer under jaP, NTP and other scaling limits. These results suggest that learning rate transfer happens in super consistent landscapes, the geometry of the landscape does significantly change with the network's size.
Then, we investigate the role of feature learning in the progressive sharpening phase, argue that while in jP feature learning causes progressive sharpening to reach a inttheindeeendent sharpness,

---

in the NTK regime the progressive lack of feature learning when the width is increased prevents the Hessian from adapting, and its largest eigenvalue from reaching the convergence threshold.
More concretely:
We define Super Consistency, and show that under μP and Depth- it holds for the largest eigen- values of the loss Hessian (Fig. 2), which converge to largely indeehindeeendent threshold and remains there for the rest of training. On the other hand. we show that other quantities, such as the training loss and the NTK cigenvalues accumulate significant finite-size effects. We quantify the rate of divergence of these quantities with power law fits (Fig. 3).
We analyze the relationship between Super Consistency of the sharpness and learning rate transfer across JaP, Depth-ju,, NTP and other parametrizations (Fig. 1, Fig. 4 and Sec. B). For juP and Depth-uuP, which do transfer, the sharpness stays super consistent, stabilizing to threshold (Fig. 1, top left), which in some cases corresponds Edge of Stability [16], and oscillates around it for a sustained period of training time. On the other hand, under NTP, Standard Parametrization (SP), or Depth-uP with multiple layers per residual block, the sharpness dynamics significantly separate during training for different widths, albeit in different ways. Also, here do not observe transfer.
We reproduce some of these results at realistic scale, including ResNets and Vision Transformers (ViTs) trained on Imagenet and GPT-2 on text data. Also, we analyze the effect of batch size, learning rate warm-up. and long training times.
In Sec. 4.1 we show that the progressive sharpening phase is mainly driven by the NTK's largest eigenvalue, which is asymptotically fixed to its initial value for NTP, while it evolves at any width under μpP In Sec. 5 we provide intuition with a theoretical analysis on a two-layer linear network.
Finally, in Sec. we discuss to what extent Super Consistency of these properties explains learning rate transfer, and the relevance of our results in the existing literature on optimization and scaling limits. Due to page limitations, we defer the discussion on related work the appendix (App. A).
2 Background and Definitions
We consider a neural network with residual connections, defined by the following recursive equa- tions over the layer indexes & E [L]:
h = (x) + VNIa
w w' ol (x)),
(I)
where N and L are the width and depth of the network, W' E NNN for t= f I...., L- 1, and r is factor that enables (r = 1) or disables (7 0) the skip branch. We denote the output with f(z) W (zh((() where W E gIxN and y scales the network output. Similarly, a has the role of interpolating between different depth limit regimes. At the first layer, we define 1 (x)) = J/ Wo,, | where 8 E R-NxD_ All the weights 0 = ( are initialized independently from ( 1) and we denote with P the total number of parameters. We stress that the fully connected layer can be replaced with any type of layer (our experiments include convolutional and attention layers). Given a dataset D p yu) yp)) [P] of datapoints xu IR." and labels y1+ E R, we train the network with stochastic gradient descent (SGD) with batch size B and learning rate 17 E R,
B
0c+1 = 02 - Σ (oL((((E)))
(2)
where the loss L is a twice differentiable loss function. Defining Jr == (fi(w))) E R da| to be the vector of network's 's outputs at time t, if one considers continuous time, the correspond- ing gradient descent dynamics in function space df// dt take the following form [18]: de = -((f))(ff), where A(fe); == (c(f((x)) f)).))(((( i E [[]]] is the vector of residuals, and ((f:)i := (Vof((z.), (of((z)) for i,j E [[D]] is the Neural Tangent Kernel (NTK).
de
Infinite Width. The parameters 7, ar, 7) E R determine the nature of the scaling limit. y = 70. 7) Vo are ((1) constants with respect to N. L (neural tangent parameterization, or NTP). then the network enters the NTK regime [4]. Here, in the limit of infinite width, the NTK remains constant to its value at initialization throughout training, i.c O(f) O(fo) for all t 2 0. Thus,
3

---

the network's 's dynamics become equivalent to a linear model trained on the first order term of the Taylor expansion of the model initialization [19]. The fact that the NTK is fixed to its value at initialization is associated with the lack of feature learning of the model in the large width limit. If "o N, and ? no? (μP, or mean-field parameterization), the features evolve in the limit (i.e. the NTK (ff) evolves), and the richer model's 's dynamics can described using either Tensor Programs [5] dynamical mean field theory [20]. Under μP, Yang et al. [6] show that the learning rate 'jo as well as other hyperparameters transfer across width, in contrast to kernel limits, which we reproduce for our residual network in Fig.
Infinite Depth. If on top of the μP framework, the residual branches are scaled with a = 1/2 (Depth-uPP), then Bordelon al. [7] and Yang al. [8] show that the infinite width dynamics also admit feature-learning infinite depth limit. Under Depth-uu, the learning rate yo transfers with both width and depth. this paper, we compare NTP and μP regimes as the width is increased, and show that our results extend to depth-scaling using the Depth-uP model. We summarize the feature learning parametrizations and report Depth uP for Adam in Appendix K.
10
AOO)
Ámax
Aa
10a
Width
4096
10
2
10
10
Epoch
Epoch
(a) Largest Hessian eigenvalues JPP
(b) Largest NTK eigenvalues JuP
Figure 2: (a) The top Hessian eigenvalues exhibit progressive increase to a threshold, with larger cigenvalues showing precise Super Consistency, while lower eigenvalues show finite-size accumu- lation at small width in the initial phase of training. (b) Top eigenvalues of the NTK matrix O. As opposed to the top eigenvalues of the Hessian, these exhibit evident finite-size accumulation during training. Model: 3-layer ConvNet, 0, nu = 0.7 (optimal). Details in Sec. J.
3 Super Consistency of the Optimization Landscape
In this work, we analyze the landscape through the lens of the preconditioned Hessian 2², where H, V2.((,) == En VzL(f((zz)) E RPxP, as Or evolves with gradient descent. The Hes- sian is key object in optimization theory [21]. information geometry [22, 23J, and deep learning theory [17, 24, 16, 13] and its relation to optimal step sizes is often used to design second-order optimizers [23, 25-27). In Figure I, we can observe that learning rate transfer correlates with strong alignment across model sizes of the Hessian top eigenvalue dynamics, a property which we name Su- per Consistency. The choice of the preconditioning factor 7² ensures the right scaling with respect to the width N, as the theory will justify. We also provide an intuition and an extension to Adam in Appendix. J.1. Unless stated otherwise, every experiment is conducted with the this preconditioning factor 72 set according to the corresponding parametrization.
More concretely, Super Consistency refers to when certain aspects of the loss landscape and of the predictor SN (in this paper (1) refers to the NTK's and loss Hessian's eigenvalues or the loss itself) exhibit the following two properties:
At realistically large N, SN(() does not deviate significantly from its limit Sx) =
1im (1). This is what is referred consistency in Vyas al. [9].
(w) does accumulate significant finite-size effects over time, i.e. the curves of Sx) and Sze (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 ,P qualitatively exhibit little-to-no no deviations. Also, notice that we assume existence of the limit lim (t(t). For those parametrizations (e.g. stan- dard parumetrization(( (6] that do not have well-defined limit, (t) diverges at large N and Super Consistency is trivially violated.
We now turn to the analysis of the Hessian spectrum and observe the following:
Observation: in uP (and Depth-uP), Hessian eigenvalues are super consistent along the optimisa- tion trajectory. Smaller eigenvalues have progressively different dynamics.
In Fig. 2 (a), we train residual network on CIFAR-10 10 (a 10 classes image classification task) us- ing cross-entropy loss, and show super consistent dynamics of three of the top 10 eigenvalues. The choice of k = 10 is guided by Gur-Ari et al. [17], where it is observed that stochastic gradient descent happens in small subspace where the gradient lies in the space spanned by top Hessian eigenvectors (where k is the number of classes). Thus, our results show that the curvature along the training trajectory preserved super consistently at different scales, thus suggesting that the geom- etry of the landscape across the trajectory is preserved across model size. Lower order eigenvalues tend to accumulate finite-size effects in the first phase of training, and stabilize at lower thresholds for smaller width models. We discuss the effect of lower order eigenvalues through the Hessian trace in Appendix G. Finally, to make sure that Super Consistency holds along the training trajec- tory regardless of the suny-sube assumption, in Appendix I we track the directional sharpness, that measures the curvature along the gradient direction.
To give quantitative measure to the finite-size accumulation property, we measure deviations over time by estimating the following quantity:
gt) == Stx (t)l-
When g(t) increases over time (up to fluctuations), Super Consistency violated. We illustrate this in Fig. 3 (b, c), where we compute the left hand side of Eq. 3 for the loss L(fe) and the NTK's largest eigenvalue (9). To estimate the infinite width limit, we use very large-wide width model as proxy. Notice how under yPP the the loss dynamics progressively diverge from the infinite width model, indicating a finite-size accumulation over time. The same holds for Amaz(O) To study the rate of divergence gt),, we fit a power law of the form y at to the observations. A larger 33 indicates a higher divergence rate. Notice how B > 0.6 for the loss, and B =e 2 for Amaze (m)z(()). indicating quadratic divergence.
comparison, in Fig. (left), we show Super Consistency of the sharpness, in that finite-size effects are not accumulated over time (10 epochs). Finally, both in Fig. 2 and 3 notice how the sharpness is not just converging in width, but fact width- hiitth-independen.. This might be due to the fact that the threshold is a stable attractor of the dynamics [28].
With 32
201
4x30
Denie. A-w. Jevil
30-3-
2-5
2.33
20rl
Esoch
Bpoch
Epoch
(a) Sharpness vs Time PP
(b) Loss vs Time pP
(c) NTK Amax vs Time JPP
Figure 3: (a) Convergence rate of the sharpness at finite width N to the infinite limit proxy. Note that the distance approaches as the training time increases. (b) Convergence rate of the loss at finite width N to the infinite limit proxy. Note that the loss accumulates finite-size effects over time and the distance to the proxy increases. (c) Convergence rate of the top NTK eigenvalues over time to the infinite limit proxy. Similar to the loss, this also accumulates finite-size effects over time. Details: infinite limit proxy width 4096, model ConvNet, O, rjo = 0.7.
4 Super Consistency and Learning Rate Transfer
Sharpness and Edge of Stability. We now focus on the sharpness X := X defined as the largest eigenvalue of the Hessian. In the theory of smooth convex [21], nonconvex [29], and stochastic [30] optimization, the sharpness plays crucial role in in the guarantees of convergence
5

---

of gradient methods and selection of the optimal step size. For instance, for quadratic objective,
A, constant and gradient descent would diverge if the learning rate satisfies ? > 2/A.
and training speed is maximized for % 1// (LeCun et al. [2], page 28). Beyond this classical
example, which assumes constant Hessian, the descent lemma [21] states that [(0,+1) S ((),)
S where 3 == sup, V 2. and [|V is the sharpness at 0. When it comes to deep neural networks, X is generally observed increase during training (progressive sharpening): in the early phase of training it increases [31, 32] and then decreases close to convergence [10]. Under full batch gradient descent training, the sharpness consistently rises above the EoS threshold of 2/1o [16].
Conditions for hyperparameter transfer. The empirical success of hyperparameter transfer cru- cially relies on the following two observations, constituing "theoretical puzzle" [8].
1. The optimal learning rate is preserved across widths/depths, indicating very fast conver- gence with respect to the scaling quantity.
2. The models show consistent improvement in training speed with respect the scaling quan- tity (i.e. there is a clear "wider/deeper is better" effect), indicating that the loss dynamics have not yet converged to the limiting behaviour predicted by the theory.
In this section, we study the role of Super Consistency in learning rate transfer. We focus on the dynamics of sharpness Amax across training, due to its well-established connection to optimization theory and step size selection, as well as better computational tractability than the full Hessian spectrum. We provide extensive studies of other relevant spectral quantities (i.e. Hessian and NTK eigenvalues) in Appendix G.
Observation: in pP (and Depth-uP), the sharpness X is super consistent along the training tra- jectory. while for NTP the sharpness decreases in width. This correlates with presence/absence of hyperparameter transfer.
In Fig. we train a two-layer convolutional network under the jP and NTP scalings with cross entropy loss, while keeping track of the sharpness at fixed gradient step intervals. The top row shows the dynamics of X. Notice how the sharpness' behaviour is qualitatively different in the two parameterizations: in JaP it reaches a width-indeependet value which is close to the EoS threshold of 2/mo- On the other hand, in NTP we observe progressive diminishing of the sharpness with width, as previously observed for Mean-Squaree-rroo loss by Cohen et al. [16].
We then study the effect of depth under the Depth- jPP model of Eq. 1. In Fig. 4 (left), show that the sharpness' dynamics are also super consistent across depth, although progressively diminishing from the EoS threshold. This suggests that EoS is not necessary for the learning rate to transfer, but the consistency of the sharpness dynamics is.
Other feature learning parameterizations. Finally, we study the effect of other feature parameter- izations that do not exhibit learning rate transfer. In particular, we study the Depth-uP scaling the residual branches in residual networks with multiple layers per residual branch denoted by k (i.e. each branch has multiple weight matrices and non linearities). A typical example is the Transformer architecture, which has multiple layers per block in both the attention and fully connected blocks. This parameterization, although it learns features in the infinite depth limit, it is lazy within each residual branch [8, 7]. The results are in Fig. 4. Notice how the sharpness dynamics consistent, in that they accumulate finite-size effects over time. We study other parametrizations, including those without stable limit in Appendix B, showing compatible results with those pre- sented here. The observation that sharpness dynamics exhibit greater consistency compared to loss dynamics suggests that under jP scaling, although models with larger capacity fit the data faster, the paths taken by various models through the optimization landscape show a surprisingly uniform curvature.
not super
4.1 Feature Learning and Progressive Sharpening
We now study the effect of feature learning in the sharpness dynamics. Following the Gauss-Newton decomposition [25, 12], the Hessian can be decomposed as sum of two matrices H G RR where

---

6x
Depth
Depth
10'
To Uos
Thei Loos
24
12
xxx 10
Ex
10
96
60
4x10-
10-1 3-30-1 4x
10a1
10 Learning Rate
10
10
Learning
Learring Rate
Depth
Depth
Depth
10
12
24
30
0.16
20
40
60
60
1.57
10a
10
15
20
Epocn
Epoch
Epach
(a) Depth- yuP ConvNet it
(b) Depth-pul puPth--PPoa ConvNet
(c) Depth-juP VIT
Figure 4: Depth-uP extensions with top row showing transfer plots and bottom row the sharpness evolution. (a) ConvNets with 1 layer per block exhibit both hyperparameter transfer and sharpness Super Consistency. (b) ConvNets with 2 layers per block. The model has a lazy behavior within each block, and no transfer. The sharpness starts accumulating finite-sizz size effeets during training, violating Super Consistency. (c) ViTs also have lo > blocks per layer by design, and thus have a similar behaviour. Details: (a), (b) are trained with SGD, with widths 128 and 32 respectively; (c) is trained with Adam, with the learning rate scaled by 1 1/VL [8]. See Fig. 22 for convergence rates.
G is the Gauss-Newton (GN) matrix and R depends on the Hessian of the model. For MSE loss,
[D]
g Σ Vof(::)Vof(a)) = K K and R Σ Vāf(x:)(1 f(x))),
i=1
where K E RIPiXP is matrix where each row is (x; (i.e. the Jacobian of f(x;), and yi E R is the label. One can readily see that the NTK matrix can be written as ((fe) = KKT thus the NTK and share the same nonzero eigenvalues. In Figure 5, we show that under jPP the sharpness evolution is dominated by the G matrix consistently across different widths, while for NTP the sharpness evolution slows down when increasing the width. Since in the limit the NTK matrix fixed for NTP, while it evolves with time for j.P these results provide further supporting evidence for the role of feature learning in the evolution of the hessian. While this argument strictly holds for MSE loss, it can be generalized to any twice differentiable loss function, albeit with some caveats. In Appendix D, we generalize the setting, analyze the cross-entropy loss and perform validating experiments, confirming the conclusions drawn here. Finally, in Appendix H, we show that our results remains valid a random feature model, where the NTK matrix is fixed at initialization at any finite width. In Section 5 we revisit the above claims more preciscly in simplified setting, providing further intuition on the sharpness dynamics and learning rate transfer.
uP
NTP
Width
Width
10
50
50
(R)
Amg
200
Amax
200
20
500
500
1000
10
1000
10-9
0.2
0.4
0.6
0.8
1.0
0.30
0.35
0.40 Anaxld)
0.45
Figure 5: Evolution of the top cigenvalues of the Hessian components G and R for two-layer linear network trained on random data under MSE loss. The vector field characterizes the evolution during training for fixed learning rate. Top: jP. Note how G drives the initial change super consistently. Bottom: NTP. For wider networks the sharpening phase reduces, since the network is approaching the limit where the NTK is fixed to its value at initialization.
Large seale experiments. In App. F, we perform more experiments to validate the connection between the consistency of the sharpness' dynamics and learning rate transfers across datasets (Tiny-
7

---

ImageNet, Wikitext), architectures (VIT, GPT-2 [33]), and optimizers (Adam [34] and Adam [35]). We find these results to be consistent with those in the main text.
End of training dynamics. In App. E,1 (Fig. 12), we study the width dependence of the sharpness the late phase of training. It is well-known that for cross-entropy loss, a phase transition happens where the sharpness starts to decrease [16]. We found that even for jPP this transition point is width- dependent, with a consequent slight shift in optimal learning rates during this late phase. Again, these results are in line with our results that super consistent sharpness facilitates transfer.
Batch size ablation. We repeat the experiment in Fig. I with increasing batch size, observing that the threshold is reached across all the tested batch sizes, thus not affecting learning rate transfers. For larger batches, close-to-EB EoS threshold is reached across more learning rates. Results are summarized in Fig. 13 and 14 in App. E.
5 Case study: Two-Layer Linear Network
We now revisit and validate our intuition and empirical findings in Sec. 4 through the lens of a two-layer neural network with linear activations and L2 loss. Our purpose is to characterize the dynamics of uP and NTP at the edge of stability through the lens of a simple example that shares similar phenomenology with the more complex scenarios observed in the last section (sce Fig. 10, App. C). In particular, the theory justifies the preconditioned Hessian 7 H ox NH as the right object of study when it comes the sharpness computations (see Prop. 5.3). Also, it provides an intuition to the -ndhhiinddeendent evolution of the sharpness. Our setting is similar to the one leading to the insightful analysis of EoS in [28, 36). Compared to these works, we do not limit the analysis to single datapoint or to vanishing targets
Notation and assumptions. Consider dataset of [D] datapoints in D dimensions X R (P) 2 D). and labels generated through a latent ground-truth vector w. E RD, that is Y = =x... The neural network we use here is parametrized by weights 0 E R DXN, W E R x where N is the width. To simplify the notation in our setting, we name E == W and V == W XEV, C(E, V) = 11l -1 We initialize each entry of E, V i.i.d. Gaussian
f(X) = 7VND
11.
with mean zero and variance 1. Recall that "Nu 1, TrP vN. We train with gradient de- scent (GD) with a learning rate = ? Empirically, we observed (Fig. 10, App. C) that picking [D] D and data X = In (ID is the D D identity matrix) is sufficient to track most of the crucial features of μP/ NTP explored in this paper, except the "wider is better" effect which here is less apparent due the simple hypothesis class. The loss function reduces to:
1
C(E.V)
|| w.ll", with
EV. NVND
(4)
Finally, we reparametrize the model by defining:
1 V V VEED Rz
= ND EE E DDD D= ND
(5)
Note that using learning rate ? when optimizing L is equivalent to using a learning rate foo when optimizing C. Next, we characterize how e, evolve through time, and give conclusions for IuP.
5.1 Dynamics and Edge of Stability in Latent Space
We now show that at any value of the width N, under GD on the original network parameters (E, V). the dynamics of w, e, v, can be described completely through a self-contained dynamical system in (1 D + D2) dimensions. This property is surprising because the original dynamical system described by GD on the variables E, V lives in N(D + dimensions. Concretely, this means
study the Hessian dynamics at different network widths in the same space.
Theorem 5.1 (Evolution Laws), Let (E.V) evolve with GD at stepsize т TT= on the loss of Eq. 4. The evolution of (w, e, v) is completely described by the following self-contained equation:
If our dataset has cardinality 1, then the NTK is a scalar. If targets vanish, for 2-layer linear networks with 1.2 loss, NTP and P induce the same loss on the parameters (y cancels).
8

---

let the + denote updated quantitiess
2 ND
wt w- - -no(v ID + +e)( 110-w,4 w.)
(ww) w.w )w W.).
et e + ?o0?
+ n
-2w +wa T + ww. _
ND [oww] T vww www." T. + UwW, w
ND
1
+ 1172 ND
v+ +
-2w TouTur w + 2w, w|
w 2w, + w ew.)
ND
While the system above describes the evolution laws (wk, Ck. U&) + (wk+1, Ck+1, ux+1), the dy- namics influenced also by initialization. In Prop. C.1 in the Appendix, show that the only dependence in width in the evolutions laws are in the initial conditions.
Last, by analyzing the stability of the dynamical system in Theorem 5.1, we can characterize the edge of stability using tools from dynamical systems [37]. First of all, we need the following Lemma, which implies that at the minimizer (w w.)) the Hessian has the same non-zero eigenvalues the NTK O, which only depends on e and v.
Lemma 5.2 (GN bound). Let 22 L 9 + R be Gauss-Newton decompasition (see Sec. 4.1) of the Hessian for the loss in Eq. 4, with G = KT K, where K E Dx
Let us denote the NTK matrix O K KKT E DDD Then
OEE, = = v.v
and
Dmmax h V b D(E Amax (O(E,
22 || w. 12-
ND
We stress that this result implies that evolution of the NTK feature learning) goes hand in hand with the evolution of the sharpness, as we empirically show in Sec. 4.1. We are now ready to state the result on the sharpness at convergence.
noy on the loss of Eq. 4 towards a minimizer (E., V.). Assume the corresponding solution in latent space (w,, C., U.) is marginally stable 3. Then, Amax 2 V.)] E
Proposition 5.3 (EoS). Let (E, V) evolve with GD with stepsize
+
No no
ND
Implications for NTP. Consider ? in Thm. 5.1. The dynamics of (w, e, u) width- dependent. Let us take N - oc in the equation above to amplify this effect: the system becomes linear
w+ = w - no ID + e)(w - w.), = v vー=== =.
While evolves from wo expected from standard NTK theory [4], v stay clamped at initial- ization. Applying Lemma 5.2 with 7 = O(1), the Hessian and the NTK have the same largest eigenvalue at large width (at rate O(vN)). This agrees with our intuition, as under NTP the predic- tor converges to linear model in the large N limit, and thus R vanishes. Also, this proves that the sharpness has no dependency on the learning rate in the width limit (we observe this, e.g., in Fig. and throughout all our experiments). This derivation is also in line with our discussion in See. 4.1: we only have sharpening under feature learning, and for the same reason we cannot observe NTP at the edge of stability N + de (see stepsize dependency in Prop. 5.3), also noted empirically by [16].
Implications for μP. The following result immediately follows by inspection of the equations in Thm 5.1, combined with Prop. C.I.
Corollary 5.4. Consider uP (7 = /N) and let (E, V) evolve with GD with stepsize n 7o on the loss of Eq. 4. Then, the equations governing the evolution of (w, e, v) (defined in Thm. 5.1) in latent space have no width dependency this holds at any finite width and not just at the limit, Initialization of (w,e, e, v) is instead dependen yet the error from N + 30 case scales in expectation like 1/VN.
Recall:: [D] = D in our simplified setting.
In dynamical system sense: eigenvalues of the Jacobian have unit others have L.

---

The corollary shows that μP trajectories at different widths align in the latent space (w, e, e), albeit with a vanishing perturbation in the initial condition (see Prop. C.I). While NTP's dynamics for and become slower as the width increases, for μP their evolution laws are width independent. This implies that the dynamics converge towards a minimizer for e and v, this will at the sharpness value predicted by Prop. 5.3. Under jP,, where y² x N, this value will be width--ineepnnent. as Su- per Consistency would suggest. We stress that Prop. 5.3 characterizes the sharpness at convergence (i.e. at infinite time). At finite time, there is still a discrepancy between the Amax(O) and the sharp- ness of the order of the residual term 1/VD|[ -wwl (Lemma 5.2). Finally, we stress that Prop. 5.3 prescribes the right scaling for the Hessian by including the preconditioning factor of 7²2 Thus, we do not prove that at any finite time, the whole sharpness trajectory is width-independent nor we estimating converge rates in N at finite time. Indeed, there will be a finite-size dependence coming from the initial conditions. We leave a precise characterization of the whole sharpness dynamics across the training trajectory for future work.
6 Discussion & Conclusions
On Feature Learning Parametrizations. In this paper, we have shown how certain properties of the loss Hessian evolve almost identically across training for different model sizes, and named this property Super Consistency. We have also compared the sharpness dynamics under different scaling limits and parameterizations, and related Super Consistency of the landscape to learning rate trans- fer. Beyond being able to distinguish feature learning (rich) and kernel (lazy) parametrization, we have also shown how other suboptimal feature learning parametrizations have sharpness dynamics violating Super Consistency through finite-size accumulations. This seems to suggest that Super Consistency of the landscape is an important discriminant when it comes to hyperparameter transfer beyond the rich/lazy regimes. We foresee that our paper could spark further research interest the intersection between the scaling limits of neural networks and optimization theory.
On the NTK and Hessian dynamics. In Section 4.1 we have drawn the connection between pro- gressive sharpening and NTK evolution in the early phase of training. However, Figure (b), shows how the NTK cigenvalues at different widths accumulate finite-size effects over time and diverge from each other, while the Hessian eigenvalues are Super Consistent. This suggests that other forces play after progressive sharpening, such Self--tahilll Stellittaalizztiinn [38]. In fact, progressive sharpen- ing on one hand, and self-Stabilization on the other, make the stability threshold a stable attractor of the sharpness dynamics. Gaining theoretical understanding for these complex interactions in the context of scaling limits is an exciting area of future research.
Design of Step-size Tuners. In most of our experiments, we rely on constant step size 7/o How- ever, an alternative is to use a step-size tuner, i.c. to automatically choose 77o based on some criteria of the local landscape [39]. Our results open directions into some possible investigations and design choices for new step size tuners. For instance, do step size tuners transfer with the width and depth of the architecture? Given results the role of warmup schedule to improve transfer, it seems plausible to design step size tuners that use EoS results to achieve optimal learning rate transfer under different parameterizations.
Limitations. One of the underlying assumptions of the argument presented here is that the sharpness is an important property of the landscape when it comes to step size selection. Indeed, the results in Cohen et al. [16] establish a more intricate relationship between sharpness and learning rate. We discuss this in Sec. A.2. Overall, our theory on Super Consistency does not exclude the existence of other factors that might influence the optimal learning rate. Hyperparameter transfer requires Super Consistency of the landscape, thus expect other potential factors to have this property. Finally, we note that due to the high computational cost of Hessian estimation, we do not perform experiments at larger scale than presented here. It would be interesting see if Super Consistency still holds at an even larger scale.
Acknowledgements
The authors would like to thank Bobby He, Imanol Schlag, Dayal Kaira, Tiago Pimentel and Gregor Bachmann for providing insightful feedback on early versions of this manuscript. L.N would also like to thank Blake Bordelon, Boris Hanin and Mufan Li for the stimulating discussions on the topic
10

---

of scaling limits that helped inspiring this work. AO acknowledges the financial support of the Hector Foundation. LN acknowledges the support Google PhD fellowship. AM acknowledges the support of a Kempner Fellowship.
References
[1] W. X. Zhao, K. Zhou, J. Li, Tang, X. Wang, Y. Hou, Y. Min, B. Zhang, Zhang, Z. Dong,
Y. Du, C. Yang, Y. Chen, Z. Chen, J. Jiang. R. Ren, Y. Li, X. Tang, Z. Liu, P. Liu, J.Y. Nie,
and J.-R. Wen. "A Survey of Large Language Models". In: arXiv preprint 1xXi::3303.1822
(2023). URL: atttp://arxis org/ abs hbtp//bb//333.././//// 18223.
[2] Y. LeCun, L. Bottou, G. B. Orr, and K.-R. Müller. "Efficient backprop". In: Neural networks:
Tricks of the trade. Springer, 2002. pp. 9-50.
[3] K. He, X. Zhang, S. Ren, and J. Sun. "Delving deep into rectifiers: Surpassing human-level performance on imagenet classificationn. In: Proceedings of the IEEE international confer- ence computer vision. 2015, pp. 1026-1034.
[4] A. Jacot, F. Gabriel, and C. Hongler. "Neural tangent kernel: Convergence and generalization
in neural networks". In: Advances in neural information processing systems 31 (2018).
[5] G. Yang and E. J. Hu. "Tensor programs iv: Feature learning in infinite-width neural net- works". International Conference on Machine Learning. PMLR, 2021, pp. 11727-11737.
[6] G. Yang. E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, Pachocki, W. Chen, and J. Gao. "Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer". In: arXiv preprint wxXiv22233.0340 (2022).
[7] B. Bordelon, L. Noci, M. B. Li, B. Hanin, and C. Pehlevan. "Depthwise hyperparameter trans- fer in residual networks: Dynamics and scaling limit". In: arXiv preprint arXiv:2309.1662) (2023).
[8] G. Yang, D. Yu, C. Zhu, and S. Hayou. "Tensor Programs VI: Feature Learning in Infinite- Depth Neural Networks". In: arXiv preprint 1rX0v:2310002244 (2023).
[9] N. Vyas, A. Atanasov, B. Bordelon, D. Morwani, S. Sainathan, and C. Pehlevan. "Feature- Learning Networks Are Consistent Across Widths At Realistic Scales". In: arXiv preprint arXiv:2305.1844 (2023).
[10) I. Sagun, I. Bottou, and Y. LeCun. "Figenvalues of the hessian in deep learning: Singularity and beyond".. In: arXiv preprint 161 (2016).
[11] L. Sagun, U. Evei, V. U. Guney, Y. Dauphin, and L. Bottou. "Empirical analysis of the hessian of over wver-parametrized neural networks". In: arXiv preprint arXiv:/706 1706.04454 (2017).
[12] J. Martens. "New insights and perspectives on the natural gradient method". In: The Journal of Machine Learning Research 21.1 (2020), pp. 7776-5851.
[13] S. P. Singh, G. Bachmann, and T. Hofmann. "Analytic insights into structure and rank of neural network hessian maps". In: Advances in Neural Information Processing Systems 34 (2021), pp. 23914-23927.
[14] A. Orvieto, Kohler, Pavllo, T. Hofmann, and A. Lucchi. "Vanishing curvature and the power of adaptive methods in randomly initialized deep networks".. In: arXiv preprint trXiv:2106.083 03766 (2021).
[15] Lewkowyez, Y. Bahri, E. Dyer, J. Sohl-Dickstein, and G. Gur-Ari. "The large learning rate phase of deep learning: the catapult mechanism". In: arXiv preprint arXiv:2003.02218 (2020).
[16] J. M. Cohen, S. Kaur, Y. Li, J. 7. Kolter, and A. Talwalkar. "Gradient descent on neural net- works typically occurs at the edge of stability". In: arXiv preprint arXiv:2113. 103.00065 (2021).
[17] Gur-Ari, D. A. Roberts, and E. Dyer. "Gradient descent happens in a tiny subspace". In: arXiv preprint Pr.Xiv1812.000 047554 (2018).
[18] S. Arora, S. S. Du, W. Hu. Z. Li, R. R. Salakhutdinov, and R. Wang. "On exact computation with an infinitely wide neural net". In: Advances in neural information processing systems 32 (2019).
[19] J. Lee, L.. Xiao, Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. "Wide neural networks of any depth evolve as linear models under gradient descent", In: Advances in neural information processing systems 32 (2019).

---

[20] B. Bordelon and C. Pehlevan. "Selffconsistent dynamical field theory of kernel evolution in wide neural networks". In: Advances in Neural Information Processing Systems 35 (2022), pp. 32240-32256.
[21] Y. Nesterov. Introductory lectures convex optimization: A basic course. Vol. 87. Springer Science & Business Media, 2013.
[22] S. Amari. "Information geometry". In: Contemporary Mathematics 203 (1997), pp. 81-96.
[23] S.-I. Amari. "Natural gradient works efficiently in learning". In: Neural computation 10.2 (1998), pp. 251-276.
[24] B. Ghorbani, S. Krishnan, and Y. Xiao. "An investigation into neural net optimization via hes- sian eigenvalue density". In: International Conference Machine Learning. PMLR. 2019,
pp. 2232-2241.
[25] J. Martens. Second-order optimization for neural networks. University of Toronto (Canada),
2016.
[26] J. Martens and R. Grosse. "Optimizing neural networks with kronecker-factored approximate curvature". In: International conference machine learning. PMLR. 2015, pp. 2408 -417. [27] A. Botev, H. Ritter, and D. Barber. "Practical Kauss-Newton optimisation for deep learning".
In: International Conference on Machine Learning. PMLR. 2017, pp. 557-565,
[28] D. S. Kaira, T. He, and M. Barkeshli. "Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos". In: arXiv preprint arXiv:2311 1.320100076 (2023).
[29] G. Garrigos and R. M. Gower. "Handbook of convergence theorems for (stochastic) gradient methods". In: arXiv preprint arXiv:2301.1123: (2023).
[30] R. M. Gower, N. Loizou, X. Qian, A. Sailanbayev. E. Shulgin, and P. Richtárik. "SGD: Gen- eral analysis and improved rates". In: International conference on machine learning. PMLR.
2019, pp. 5200-5209.
[31] S. Jastrzkebski, Z. Kenton, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey. "On the relation between the sharpest directions of DNN loss and the SGD step length". In: arXiv preprint 0rXiv:1807.05033 (2018).
[32] S. Jastrzebski, M. Szymczak, S. Fort, D. Arpit, J. Tabor, K. Cho, and K. Geras. "The break-even point on optimization trajectories of deep neural networks". In: arXiv preprint arXiv:2002.0955 09570 (2020).
[33] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. "Language models are unsupervised multitask learners", In: OpenAl blog 1.8 (2019), p. 9.
[34] D. P. Kingma and Ba. "Adam: A method for stochastic optimization". In: arXiv preprint urXiv:1412.666 (2014).
[35] I. Loshchilov. "Decoupled weight decay regularization". In: arXiv preprint arXiv: (2017).
[36] M. Song and C. Yun. "Trajectory alignment: understanding the edge of stability phenomenon via bifurcation theory". In: arXiv preprint arXivv2307.04224 (2023).
[37] E. Ott. Chaos in dynamical systems. Cambridge university press, 2002.
[38]
Damian, E. Nichani, and D. Lee. "Seelfssabbilization: The implicit bias of gradient descent at the edge of stability"". In: arXiv preprint arXiv:2209. 15594 (2022).
[39] V. Roulet, A. Agarwala, and F. Pedregosa. "On the Interplay Between Stepsize Tuning and Progressive Sharpening". In: OPT 2023: Optimization for Machine Learning. 2023.
[40] M. Claesen and B. D. Moor. Hyperparameter Search in Machine Learning. 2015. arXiv:
1502. 02127 [cs LG].
[41] J. Bergstra and Y. Bengio. "Random search for hyper-parameter optimization.." " In: Journal of machine learning research 13.2 (2012).
[42] K. Jamieson and A. Talwalkar. "Non-stochastic best arm identification and hyperparameter optimization". In: Artificial intelligence and statistics. PMLR. 2016, pp. 240-248.
[43] I. Li, K. Jamieson, G. DeSalvo, A. Rostamizadeh, and A. Talwalkar. "Hyperband: A novel bandit-based approach to hyperparameter optimization". In: Journal of Machine Learning Research 18.185 (2018), pp- 1-52.
[44] J. Snoek, H. Larochelle, and R. P. Adams. "Practical bayesian optimization of machine leamn ing algorithms". In: Advances in neural information processing systems 25 (2012).
12

---

[45] J. Snock. O. Rippel, K. Swersky, R. Kiros, N. Satish, N. Sundaram, M. Patwary, M. Prabhat, and R. Adams. "Scalable bayesian optimization using deep neural networks". In: Interna- tional conference on machine learning. PMLR. 2015, pp. 2171-2180.
[46] A. Shaban, C.-A. Cheng, N. Hatch, and B. Boots. "Truncated back-propagation for bilevel optimization". In: The 22nd International Conference on Artificial Intelligence and Statistics. PMLR. 2019, pp. 1723-1732.
[47] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil. "Forward and reverse gradient-based hyperparameter optimization". In: International Conference on Machine Learning. PMLR. 2017, pp. 1165-1173.
[48] D. Maclaurin, D. Duvenaud, and R. Adams. "Gradient-based hyperparameter optimization through reversible learning". In: International conference on machine learning. PMLR. 2015, pp. 21 2113-2122.
[49] Yogatama and G. Mann. "Efficient transfer learning method for automatic hyperparameter tuning". In: Artificial intelligence and statistics. PMLR. 2014, pp. 1077-1085.
[50] D. Stoll, J. K. Franke, D. Wagner, S. Selg, and F. Hutter. "Hyperparameter transfer across developer adjustments". In: arXiv preprint arXiv:2010.13111 (2020).
[51] V. Perrone, R. Jenatton, M. W. Seeger, and C. Archambeau. "Scalable hyperparameter transfer learning". In: Advances in neural information processing systems 31 (2018).
[52] S. Horváth, A. Klein, P. Richtárik, and C. Archambeau. "Hyperparameter transfer learning with adaptive complexity". . In: International Conference on Artificial Intelligence and Statis- tics. PMLR. 2021, pp. 1378-1386.
[53] G. Yang and E. Littwin. "Tensor programs ivb: Adaptive optimization in the infinite-widhe width limit". In: arXiv preprint arXiv:2308.018/ (2023).
[54] B. Bordelon and C. Pehlevan. "Dynamics of Finite Width Kernel and Prediction Fluctuations
in Mean Field Neural Networks". In: arXiv preprint arXiv:2304.
(2023).
[55] S. Jelassi, B. Hanin, Z. Ji, S. J. Reddi, S. Bhojanapalli, and S. Kumar. "Depth Dependence of mup Learning Rates in ReLU MLPs". In: arXiv preprint Stfi:2300500810 (2023).
[56] S. Yaida. "Meta- MMtaa-Pinncipee Family of Hyperparameter Scaling Strategies". In: arXiv preprint arXiv:2210.44900 (2022).
[57] I. Goodfellow, Y. Bengio, and A. Courville. Deep Learning. http deeplearningbook. org. MIT Press, 2016.
www
[58] Zhu, Z. Wang, X. Wang, M. Zhou, and R. Ge. "Understanding -ttbbiit-tabilityy training dynamics with minimalist example". In: arXiv preprint arXiv:2210.0322 (2022).
[59] S. Arora, 7. Li, and A. Panigrahi. "Understanding gradient descent on the edge of stability in deep learning". . In: International Conference on Machine Learning. PMLR. 2022, pp. 948- 1024.
[60] K. Ahn, J. Zhang, and S. Sra. "Understanding the unstable convergence of gradient descent".
In: International Conference on Machine Learning. PMLR. 2022, pp. 247-257.
[61] J. M. Cohen, B. Ghorbani, S. Krishnan, N. Agarwal, S. Medapati, M. Badura, D. Suo, D. Cardoze, Z. Nado, G. E. Dahl, et al. " "Adaptive gradient methods the edge of stability". In: arXiv preprint arXiv:2207. S0t.22277.1484. 14484 (2022).
[62] G. Iyer, B. Hanin, and D. Rolnick. "Maximal initial learning rates in deep relu networks". In: International Conference on Machine Learning. PMLR. 2023, pp. 14500-14530.
[63] L. N. Smith and N. Topin. "Super-convergence: Very fast training of neural networks using large learning rates". In: Artificial intelligence and machine learning for multi-domain opera- tions upplications. Vol. 11006. SPIE. 2019, pp. 369-386.
[64] Y. Li, C. Wei, and T. Ma. "Towards explaining the regularization effect of initial large learning rate in training neural networks". In: Advances in Neural Information Processing Systems 32 (2019).
[65] D. S. Kalra and M. Barkeshli. "Phase diagram of early training dynamics in deep neural networks: effect of the learning rate, depth, and width". In: Thirty-seventh Conference on Neural Information Processing Systems. 2023.
[66] R. M. Neal. "Bayesian Learning for Neural Networks". PhD thesis. University of Toronto,
1995.
[67] J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, Pennington, and J. Sohl-Dickstein. "Deep neural networks as gaussian processes". In: arXiv preprint arXiv:1711.00165 (2017).
13

---

[68] A. Garriga-Alonso, C. E. Rasmussen, and L. Aitchison. "Deep convolutional networks as shallow gaussian processes". In: arXiv preprint arXiv: 05587 (2018).
[69] J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. "Infinite attention: NNGP and NTK for deep attention networks". In: International Conference on Machine Learning. PMLR. 2020,
pp. 4376 4386.
[70] G. Yang. "Tensor programs ii: Neural tangent kernel for any architecture". In: arXiv preprint arXiv:2006. 14548 (2020).
[71] L. Chizat, E. Oyallon, and F. Bach. "On lazy training in differentiable programming". In: Advances in neural information processing systems 32 (2019).
[72] L. Chizat and F. Bach. "On the global convergence of gradient descent for over-paramete pver-parameterized models using optimal transport". In: Advances in neural information processing systems 31 (2018).
[73] L. Chizat and Bach. "Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss". In: Conference on Learning Theory. PMLR. 2020, pp. 1305- 1338.
[74] S. Mei, T. Misiakiewicz, and A. Montanari. "Mean- --Menn-ieldd theory of two-layers neural networks: dimension-free free bounds and kernel limit". In: Conference on Learning Theory. PMLR. 2019,
pp. 2388-2464.
[75] S. Hayou, E. Clerico, B. He, G. Deligiannidis, A. Doucet, and J. Rousseau. "Stable resnet". In: International Conference on Artificial Intelligence and Statistics. PMLR. 2021, pp. 1324 1332.
[76] L. Noci, S. Anagnostidis, L. Biggio, A. Orvieto, S. P. Singh, and A. Lucchi. "Signal propaga- tion in transformers: Theoretical perspectives and the role of rank collapse". In: Advances in Neural Information Processing Systems 35 (2022), pp. 27 27198-27211
[77] S. Hayou and G. Yang. "Width and Depth Limits Commute in Residual Networks". In: arXiv preprint arXiv:2302.0044 00453 (2023).
[78] S. Hayou. "On the infinite-eppth limit of finite-width neural networks". In: arXiv preprint arXiv:2210.06 00688 (2022).
[79] B. Hanin and M. Nica. "Products of many large random matrices and gradients in deep neural networks". In: Communications in Mathematical Physics 376.1 (2020), pp. 287-322.
[80] L. Noci, G. Bachmann, K. Roth, S. Nowozin, and T. Hofmann. "Precise characterization of the prior predictive distribution of deep ReLU networks". In: Advances in Neural Information Processing Systems 34 (2021). pp. 20851-20862.
[81] M. Li, M. Nica, and D. Roy. "The future is log-Gaussian: ResNets and their infinite-depth- and-width limit at initialization". In: Advances in Neural Information Processing Systems 34 (2021), pp. 7852-7864.
[82] M. Li. M. Nica, and D. Roy. "The neural covariance SDE: Shaped infinite depth-and--widh networks initialization". In: Advances in Neural Information Processing Systems 35 (2022), pp. 10795-10808.
[83] L. Noci, C. Li, M. B. Li, B. He, T. Hofmann, C. Maddison, and D. M. Roy. "The shaped transformer: Attention models the infinite depth-and-width limit". In: arXiv preprint arXiv:2306. 17759 (2023).
[84] M. B. Li and M. Nica. "Differential Equation Scaling Limits of Shaped and Unshaped Neural Networks". In: arXiv preprini arXiv:2310. 12079 (2023).
[85] B. Hanin and M. Nica. "Finite depth and width corrections to the neural tangent kernel". In: arxiv preprint arXiv: Ntv::/999.0555 05989 (2019).
[86] S. Yaida. non-Gaussian processes and neural networks at finite widths". In: Mathematical and Scientific Machine Learning. PMLR. 2020, pp. 195-192.
[87] J. Zavatone- Veth, A. Canatar, B. Ruben, and C. Pehlevan. "Asymptotics of representation learning in finite Bayesian neural networks". In: Advances in neural information processing systems 34 (2021), pp. 24765-24777.
[88] I. Chizat and P. Netrapalli. "Steering Deep Feature Learning with Backward Aligned Feature Updates". In: Xixiv preprint arXiv:2311.187 (2023).
[89] A. Agarwala, F. Pedregosa, and J. Pennington. "Second-order regression models exhibit pro- gressive sharpening to the edge of stability". In: arXiv preprint arXiv:2210.04800 (2022).
14

---

[90] J. Gilmer, B. Ghorbani, A. Garg, S. Kudugunta, B. Neyshabur, D. Cardoze, G. E. Dahl, Z. Nado, and O. Firat. "A loss curvature perspective on training instabilities of deep learning models". In: International Conference on Learning Representations, 2021.
[91] R. Courant. "Über die Eigenwerte bei den Differentalggeichhnneen der mathematischen
Physik". In: Mathematische Zeitschrift 7.1 (1920), pp. 1-57.
[92] B. N. Parlett. The symmetric eigenvalue problem. SIAM, 1998.
[93] K. Osawa, S. Ishikawa, R. Yokota, S. Li, and T. Hoefler. "ASDL: A Unified Interface for Gradient Preconditioning PyTorch"..
arXiv preprint arXiv:230S. 04684 (2023).
[94] Z. Yao, A. Gholami, K. Keutzer, and M. W. Mahoney. "Pyhessian: Neural networks through the lens of the hessian". In: 2020 IEEE international conference on big data (Big data). IEEE.
2020, pp. 581-590.
[95] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner. M. De- hghani, M. Minderer, G. Heigold, S. Gelly, al. "An image is worth 16x16 words: Trans- formers for image recognition at scale". In: arXiv preprint arXiv:2010. /929 (2020).
[96] B. Bordelon, H. T. Chaudhry, and C. Pehlevan. "Infinite Limits of Multi-heat Transformer Dynamics".. In: arXiv preprint arXiv:2405. 15712 (2024).
15

---

Appendix
Table of Contents
A Related works
17
A.1 Learning Rate Transfer and Scaling Limits
18
A.2 Sharpness and Optimal Step Size
18
B Experiments in the absence a valid scaling limit
B.1 pPP withour 1/ /depth scaling of the residual branches
"rao.h..........................
18
B.2 Standard Parameterization (SP) experiments
19
B.3 μP- disabling residuals
19
B.4 μP experiments for full batch GD
20
Analysis of Two-Layer Linear Network
21
Relationship hetween sharpness and residual, Gauss-Newton
21
Proof Thm.5.1
23
C.3 Theory for Standard Parametrization (SP)
30
D Connection between the eigenvalues of the Hessian and NTK matrix
32
E Late-time dynamics and batch size ablations
33
E.1 Late-time dynamics
33
The effect Batch Size and Augmentation
F Large-Scale experiments, more Datasets and Optimizers
34
F.1 GPT-2 experiments WikiText..
34
F2 ConvNets Experiments on Larger Datasets and Adam(W)
35
G Time Evolution of other Spectral Quantities
36
Sharpness evolution in Random Feature Models
Directional sharpness
39
J Experimental details
39
J.1 Hessian Computation
40
J.2
GPT-2
41
J.3 Vision Transformers (ViTs)
................................
41
J.4 ResNet
................................
41
J.5 Coordinate check for yP
41
K Summary of Feature Learning Parametrizations
42
16

---

A Related works
Hyperparameter search. Hyperparameter tuning [40] has been paramount in order obtain good performance when training deep learning models. With the emergence of large language models, finding the optimal hyperparameters has become unfeasible in terms of computational resources. Classical approaches based on grid searching [41] over range of learning rates, even with improve- ments such as successive halving [42, 43] require training times on the order of weeks on large datasets, making them largely impractical. Bayesian optimization [44, 45] methods aim to reduce the search space for the optimal HPs by choosing what parameters to tune over in the next itera- tion, based on the previous iterations. A different approach for HP search involves formulating the problem as an optimization over the HP space and solving it through gradient descent [46-48].
Learning rate transfer. While meta-learning and neural architecture search (NAS) literature pro- vide methods for finding the optimal learning rate in neural networks, these are still dependent on the model size and can become costly for larger architectures. Parameter transfer methods have been studied in the literature, for learning rate transfer across datasets [49], and in the context of reusing previous hyperparameter optimizations for new tasks [50]. Perrone et al. [51] and Horváth et al. [52] proposed methods based on Bayesian optimization for hyperparameter transfer in various regimes. Yang and Hu [5], Yang et al. [6, 8], and Yang and Littwin [53] used the tensor programs framework to derive a model parameterization technique which leads to feature learning and learning rate trans- fer through width and depth, termed p.P. Bordelon et al. [7] used the DMFT framework [20, 54] to derive the Depth-uP limit, leading to optimal learning rate transfer across depth is models with residual connections. For MLPs, Jelassi et al. [55] the depth dependence of the jaP learning rates. Finally, conditions on the networks and its dynamics to achieve hyperparameter transfer have been analyzed in Yaida [56].
Training on the Edge of Stability. The choice of learning rate has been coined as one of the important aspects of training deep neural networks [57]. One phenomenon studied in the optimiza- tion literature is the fact that under gradient descent (GD), neural networks have sharpness close to 2/step size, termed the Edge of Stability (EoS) [16, 38, 58 60], with the extension to Adam being introduced by Cohen et al. [61]. Iyer et al. [62] study the maximal learning rate for ReLU networks and establish relationship between this learning rate and the width and depth of the model. Smith and Topin [63] and Y. Li et al. [64] study the effect of large initial learning rates on neural network training. Lewkowycz al. [15], Kalra and Barkeshli [65], and Kalra et al. [28] show empirically that the learning rate at initialization can lead the "catapult" phenomena. Song and Yun [36] analyze the trajectory of gradient descent in a two-layer fully connected network, showing that the initial- ization has a crucial role in controlling the evolution of the optimization. Finally, early evidence of Super Consistency was shown in Fig. 3 of Sagun et al. [11], where it was shown that at end of training the Hessian spectrum is similar across model sizes.
Scaling limits The study of scaling limits for neural network was pioneered by the seminal work of Neal [66] on the equivalence between infinite width networks and Gaussian Processes, and more recently extended in different settings and architectures [67-69] and under gradient descent training. leading to the Neural Tangent Kernel (NTK) [4, 19, 18, 70] or "lazy" limit [71]. The rich feature learning infinite width limit has been studied using different frameworks, either Tensor programs [5] or DMFT [20, 54]. The main motivation behind these works is to maximize feature learning as the width scaled up. In the two-layer case, the network's 's ffhnite-width dynamics have also been studied using other tools, such as optimal transport [72, 73] or mean-field theory [74]. The infinite depth analysis of 1/ depppth--salee residual networks was introduced in [75], and later applied to the Transformers (used here) in [76]. The infinite width-and-depth limit of this class of residual networks have been studied Hayou and Yang [77] and Hayou [78] at initialization and Bordelon et al. [7]. Yang et al. [8] for the training dynamics. Without the 1 1/v /depph-scallnng the joint limit has mainly been studied at initialization [79-84]. Deviations from the infinite dynamics can also be studied using perturbative approaches [85-87]. Finally, it is worth mentioning that method to control and measure feature learning have been recently proposed in Chizat and Netrapalli [88].
Scaling limits and Hyperparameter transfer It is worth mentioning that while for width limits feature learning a clear discriminant between NTK and mean-field limits, the issue becomes more subtle with depth limits. In fact, there exist a family of infinite depth limits that admit feature
17

---

learning (a E [1/2.1] 1] in Eq. 1, with an appropriate depth correction to the learning rate). [8] classifies the depth limits in terms of the feature diversity exponent, measure that quantifies the diversity between the features of different layers. With respect to this measure, = 1/2 is the one that maximizes it. Bordelon et al. [7] try to quantify the finite-size approximation to the infinite (continuous) model, arguing that hyperparameter transfer is achieved faster for models that have lower discretization error to the infinite model's dynamics.
A.1 Learning Rate Transfer and Scaling Limits
Our results on the early training dynamics complement the analysis of Vyass et al. [9] and Kalra et al. [28] on the consistency of the loss and the sharpness in the first few steps of training. However, we further extend it, showing how while the loss curves depart later in training (with "wider/deeper being better"), the consistency of the sharpness' dynamics is maintained longer in time and does not accumulate finite-size effects over time. Furthermore, our work explains some of the experimental results on the lack of progressive sharpening under the NTP parameterization [16] (Appendix H), by relating it to the lack of feature learning and the consequent absence of hyperparameter transfer. Our results on the role of feature learning are also compatible with [89]. where it is theoretically shown that the non linear dynamics exhibits progressive sharpening in simple non linear model. Our results are also in line with the recent experiments on the evolution of the sharpness for ReI.U MLPs trained with gradient descent under μP [28) (e.g. Figure 1). We extend these results to include the width (in) dependence behaviour of the convergent stability point, a crucial aspect for successful hyperparameter transfer.
Finally, It is worth mentioning that there exist a family of infinite depth limits that admit feature learning (a E [1/2, 1] in Eq. with an appropriate depth correction to the learning rate). Most of our experiments focus on models with a single layer per residual block, which exhibit transfer more consistently under the μP- /depth setting (i.e. = 1/2) and it is considered optimal in terms of feature diversity across blocks [8]. We adopt the same depth scaling with our experiments with Transformers which have multiple layers per block. Under this setting, both Bordelon al. [7] (e.g. Fig. 3) and Yang et al. [8] (Fig. 16) show good (albeit at time slightly worse) transfer across depth with Adam. More broadly, expect that scaling limit that admits learning rate transfer behaviour in the sharpness dynamics.
would find a corresponding
A.2 Sharpness and Optimal Step Size
One of the underlying assumptions to explain why super consistent sharpness causes hyperparameter transferis is that the sharpness influences the optimal learning rate. Although this is well-established in classic optimization [2]. Edge of Stability tells story of more intricate relationship: the choice of the learning rate also influences the sharpness. Also recent works on step size tuners show evidence of a more complex interaction in the joint dynamics of step size and sharpness [39], a relation that still has to be fully understood and could be leveraged to design better step size tuners. However, the sharpness is still arguably very good proxy for understanding the trainability of the model at large learning rates. For instance, Gilmer et al. [90] argue that maintaining a small sharpness in neural network optimization favors training at large learning rates. Interestingly, in Cohen et al. [16] (Appendix F) it is shown that choosing the step size as 1/X, at any step is suboptimal. The reason could be that gradient directions are not aligned with the largest curvature. Alternatively, one could claim there is another, more sophisticated, functional relationship between sharpness and the maximum step size allowed. Indeed, the "catapult mechanism") [[15] shows the maximum stable learning rate is Carch/ Ao. where Carch depends on the architecture. On related note, Gur-Ari et al. [17] shows that after the first few steps, the gradient direction is aligned with the top Hessian eigenvectors, underlying once again the importance of the sharpness and the first few eigenvalues in this phase of training, which we show here to be super consistent (Fig. 2).
B Experiments in the absence of a valid scaling limit
B.1 P without 1/Vdepth scaling of the residual branches
We first study the effect of increasing the depth in model without the 1 1/VT--caling of the residual branches, which results in exploding activations (and sharpness) as the depth increases. See Figure
18

---

6, where we first show that (top row) there is no learning rate transfer, and there is no consistent sharpness either. For instance. notice how the sharpness for the model of depth 24 quickly reaches its EoS value at its optimal learning rate of 0.068 (and larger models, e.g. depth 36 diverge). On the other hand, for the same learning rate the smaller depth models struggle to reach EoS at the same speed. This suggests that for the smaller-depth model, the learning rate can be safely increased. Indeed, the optimal learning rate for the smaller depth model is significantly larger, where the model reaches the EoS value very fast (optimal Ir: 0.53). The larger depth models are not trainable at these larger learning rates.
In the bottom row of Figure 6, we show that by adding linear warmup scheduler in the first phase of training. the network progressively more trainable at larger step-sizes, and the sharpness reaches its stability threshold at any width/depth that is trainable [90]. This suggests that the diverging initial sharpness can be counteracted with initial small learning rates. Also, we notice a better transfer, indicating that sustained period of depth-indeppenden sharpness can help learning rate transfer. On the other hand we do not observe good transfer after the first epoch, which correlates with the fact that the sharpness is not consistent due to a blow up with depth at initialization.
Epoch 20
2.5x 22
Depth
Depin
222
g 22.5 32"
1.7
377x32
10
Learning
Leamning
34p
Epoch
Epoch 20
Dupth
Depth
22"
30
222
23
20
8.60222
14"
20t1
10
Leamnngg 10-1
Learning
500p
Figure 6: juP without 1/ depth scaling of the residual branches for ConvNets trained using SGD on CIFAR10. (Top row): warmup. (Bottom row): 200 steps of linear warmup. Right column: sharpness dynamics in the first epoch for the optimal learning rates for the models of depth 3 and depth 24. Notice how in the first epoch (Left column) there is no learning rate transfer in both cases. This correlates with the fact that that the training starts at an increasing sharpness with depth, which causes no consistency the dynamics (and divergent behaviour large learning rates). Adding warmup alleviates this issue, allowing the model to reach edge of stability and improve learning rate transfer (middle column). One step in the sharpness plot corresponds batches. Parameters: batch size= 128, epochs= 20, using data augmentation.
B.2 Standard Parameterization (SP) experiments
Following the SP formulation introduced in [5], we analyze the sharpness and learning rate transfer under this regime in Figure 7. Note that our experiments use a fixed learning rate n). While Yang and Hu [5] use width scaled learning rate in order to parameterize SP as kernel limit in infinite width, the SP definition that we use does not have a well defined limit and thus diverges as the width is increased. For more details, we refer the reader to (Yang and Hu [5], Appendix J.3).
B.3 JP- disabling residuals
We also provide training runs of models parameterized with jP, while disabling the residuals. Note that the models quickly become untrainable under this regime when increasing the depth. This phenomenon is due to the vanishing curvature experienced by these models, as shown in Figure 8, which leads to vanishing signal.
19

---

Width
Width
Width
64
32
32
54
LOWS
64
32
128
128
128
10"
10P
8x 10-4
8x10-1
7x
10
Learning Rate
Learnnng Rata
Learting Rate
Width
107
Whdth
107
Whdth
52
52
52
84
10
107
236
10
Stop
Step
Stop
Figure 7: Standard Parameterization (SP) for ConvNets trained using SGD on CIFARIO, for varying number learning rate warmup steps. (Left column) No warmup. (Middle column) 1000 warmup steps and (Right column) 2000 warmup steps. Note that under SP, in the beginning the training starts from high curvature, which means that large step sizes would lead divergent behaviour. Adding warmup alleviates this issue, allowing the model to reach edge of stability and improve learning rate transfer. One step corresponds to 10 batches. Parameters: batch size=256, epochs=20, using data augmentation.
2 100
Depth
10a
Depth
3
10
Train Loss
6
109
1
12
24
10a
10
Ir
10-3
0.014678
0.5275
10
10
10"
10a
250
500
750
1000
Learning Rate
Step
Figure 8: μP parameterization on ConvNets with residuals disabled (r = 0) trained on CIFARIO. (Left) Learning transfer plot, showing that under this setting, the optimal learning rate transfers, but with increasingly larger shifts when increasing the depth due to the vanishing signal. (Right) Sharpness evolution during training, showing that the dynamics are following a depth independent trajectory for the optimal learning rate (0.5275). Note that deeper models become much harder train when residuals are disabled, motivated by the observation that the the depth 24 model suf- fers from vanishing curvature at larger learning rates. The spikes in the plot are due to the curva- ture approaching 0 in log-scale. One step corresponds 10 batches. Parameters: batch size=256. epochs-20, using data augmentation.
B.4 μP experiments for full batch GD
In this section, we investigate the effect of pP and on learning rate transfer and sharpness, when trained with full batch gradient descent. We subsampled 5000 sampled from CIFAR 10 in stratified fashion (i.e. 500 sampled from each of the 10 classes) and proceeded to train residual ConvNets under the fP parameterization with GD, following a similar procedure as [16]. Our findings show that under μP, when using a large enough learning rate, models are able to achieve edge of stability, well as have learning rate transfer. This is empirically shown in Figure 9, where we can see that while the sharpness has the typical oscillations around the FoS threshold studied by [16], still does maintain width andth-independent trend during training.
20

---

2x10
18x10P
Width
3x
Train Lass
1.6x109
L4x10
Tran
238
255
107
107
0.769773
Learning Rate
Eporr
Figure 9: Residual ConvNets trained using (full batch) GD on a 5000 sample subset of CIFARIO. (Left) Learning rate transfer plot, showing that the optimal learning rate transfers across different widths; the slight shift in the transfer plot is due to the oscillations around the EoS threshold seen in (middle) and (right). (Middle) Sharpness dynamics during training, showing consistent width independent dynamic throughout the whole training procedure; dashed line represents 2/n. (Right) Training loss dynamic during time for the optimal learning rate (0.76), showing the wider-is-better behaviour of the μP parameterization, as well as the oscillations induced by the EoS regime. Param- eters: batch size= 256, no warmup, using data augmentation.
C Analysis of a Two-Layer Linear Network
Recall the definition of our model:
L(
XEV Y
(6)
WVD
NTP loss
loss
16a
10
10-1
10-1
1000
10
20
10
Figure 10: Evolution of loss under μ.P and NTP for the toy example of Section 5: 1 VND
EV
w. I1, where u., 1 E E2, 100. This is minimal example of transfer captured by our theory: JPP trajectories align. Different linestyles correspond to different values of no (grid is different for J.P and NTP).
Under the previous assumptions regarding the data, we have that:
1 EVV
0g
W. VT
nNND
TVND
dy NDD
E EV NND
W..
C.1 Relationship between sharpness and residual, and Gauss-Newton
The following bound leverages a Gauss-Newton decomposition and leads analytical insights sup- porting See. 4.1. Proof of Lemma 5.2. which we restate here.
Lemma (GN bound). Let 2 [ G + R be Gauss-Newttr decomposition" (see Sec. 4.1) of the Hessian for the loss in Eq. 4, with G = K, where K E DND+++
Let denote the NTK matrix θ KKT E R DxD Then
O(E =
'Recall: [D] D in our simplified setting.
21

---

and
7 ND
Amax b2 Amax (m
w. 2-
The result above is actually more general: it indicates that the eigenvalues of the positive semidefinite portion the Hessian K K are fully² characterized by the eigenvalues a smaller matrix ④ = KK Furthermore, in model è has closed form that depends only the variables
Proof. The Hessian blocks become:
OEE
ID
&
VV
RNDXND
(7)
N2
1
F.
V
1
FVV th & IN NDNN
(8)
ND
WND
WND
Ovg
E
8V
V E w. & ERNND
(9)
ND 1 dvv 2ND
WND
TVND
E
EERNNN
(10)
Using these definitions, we can separate the Hessian into a sum of matrices, where one depends on the residual and one does not.
V=C(E.V)= = 1
I & VVT 2ND E T E
EOV
0 IN & (w W.
+ WNDD
IN @ (w w.)T 0
G(E.V)
R(E.V)
(11)
hence our quantity of interest:
()(((E ND
ID VVV E E T E
22
0
IN (w w.)
ND
IN X (w w. w
0
G(E.V)
R(E.V)
(12)
where w
T/ND" EV.
Study of G.
Note that
1
(ID@V)
G(E.V)=K' KT
(13)
VND
E
By SVD decomposition, it is easy that, the nonzero eigenvalues of & KT K are the same as the nonzero eigenvalues of θ K.
IDD 1
((E ND (IDOV E)
EE -
E
ND
ND
(14)
where e, are the quantities found in the main paper, and we therefore have
1
Amax (E V)] Amax (E,, V)]
Amax
EET
VTV.
(15)
ND
ND
'Simple application of the SVD decomposition.
22

---

Study of R and of the residual. Note that R(E, V) has both positive and negative eigenvalues, with spectrum symmetric along the real line. It easy to show that
72
Amnax (R, -min [R ND
1/ w
Using the fact that is Hermitian, of the sharpness from the maximum eigenvalue G in terms of the residual:
apply Weyl's inequality obtain bound the deviation
72 ND
72
Amax ( V)
w
w.ll2 2/2 S Amax b> 2 S Amax [O(
||w w.
W.112-
ND
(16)
Which finally yields:
72 ND
2 - Amax S
w w.112
(17)
C.2 Proof of Thm.5.1
We divide the proof into two parts. In the first part, study the evolution of the unnormalized quantities EE V and EV. Then, we study how normalization affects the dynamics.
Part one: dynamics in a smaller space. tions the beginning of this section.
We go step by step recalling the gradient descent equa-
Dynamics of EV. We have
E4 V+ = (F (V -(0v)
= -VV nggVV -nEdv + dggggdv
EV EVVTV. + To
170
W. VTV
To
EE
EV
TT8
EE w.
ND
VND
ND
VND
+ FVVT
yo
79/0
w. V
Yo ND
E
EV
77
F. w.
ND
VND
VND
EV
DO
EVV
MMo
W. V
No
EE EV
Mi VND W.V V
EE W
ND
VND
ND EVV E w.
V N
EVV E EV
EVEV
NOT
EVEV EV
72110 w. V E.w. w.
N D
NIDi
ND
Let us rename
N = EV,
V=VTV,
ē= EE
then the equation becomes more compact:
170
7To0
170
7770
w
vw
Dw.
Ew
EW..
ND
VND
ND
vND
Th 22
n
nr (w, )T +
22
(ww )>
(wn )w
(w, T ww
N D
N D
ND
(18)
Note that no quantities appear in the equation for + besides w, ú, ē. We will see that these do not appear also in the equations for ü, ē.
Dynamics of V V. Let write down here the equations for -njō and na for ease of reference:
-nOT
170
V
E
E
7oo
E
ND
AND
D
-10y
E
EV
77o0
E w.
ND
NND
23

---

we have
V+ -)/rr T ndv)
=VVV - -norv - 1/ av + oro
=VTV To
V EV + 7TO w EV
To ND
V EVEV
To V E. w. VND 70
ND
VND
(- To
V EEEE E+
7
W E
To ND
E' EV
EE w.
ND
VND
VND
VV
Ma
ETETEV
7
W
EV
Do
V E EV
EV
70
V E w.
ND
VND
ND
VND
E
W.
EE EV
N07
V E EEEEET w. +
72 w. EE w.
N2D
N
N
ND
Using our notation, equations get yet again simpler:
U+
2 110 ND
77 +
w.
TO
187 N D
wa ew + 22
w ēw..
NND
N
ND
(19)
Note that again no quantities appear in the equation for ot besides w, U, Dynamics of E E. For convenience, recall:
770 EVV
7To w. V VND Yi Vw,
ngg
ND
n
1JU VV E
ND
VND
we have
E. E = (E-ndmm) (E PTOT.TT
F.F.T
110g F.T FF.O r/ O
EET 9o0 EVV ET 77
E EVV ET
To ND
TOo
EVw!
ND
VND
VND
+
To0 EVVT + Mo
w. VT)
To ND
VV ET
Moo Vw
ND
VND
VND
EE
EVV ET + 77
W. E
EVV E
+
TTo
EVw! T
ND n
VND
ND
VND
FVV
E
Mi
2
W.W
VVVVVVV E
MO
22m EVV VwI +
N2 w. V VV...
N2 D
NID?
N
Using our notation, equations get yet again simpler:
/
700
7110 +
W. v0
7710
Wn.
ND
VND
VND
(20)
Tww + 16 Tw. w
N2 D2
N D
~ D
ND
Part two: Normalization.
Consider scalar reparameterizztions.. é
W
a
While we gave the form of these normalizers already in the paper, we keep it more general here real numbers and show that the right normalizers arise directly.
Reparameterization of EV. We have
24

---

w+ On.1i"
=
9/0
((w
7100 ow.
ND ((a
7/0
Ma0 Cw.
ND
VND
vND
+ N (ww (aw)))
V15
n
(am )w.
767
(w.w ((r
+
P (w- Jw..
N D
N D
ND
= =a ())
710
(ar)) + Moa (a,ū)ww-
100
(aro)))
Moaw (ack)
NDa,
VNDaa
NDoc
VNDae
n
mb
72
(a ww (a
(o Yw.
(o 2. ( (an +
(a v w
N o
N D aw
N
ND
V1o vw.
Now ew. VNDae
NDa
VNDaa
NDa
D2
7
w
N Daa.
'O7
(ww
)w,
N D&
I (w,w w yw +
ND
(w w
TwwTJ, )w..
We would like au a₂ to be such that on the right-hand side we have no width dependency. To do that we need (first and second line refer first and second lines the equation)
VN
1
VN
Ogy
X
aw
O ae
ag
x
N Ow OC ahe
N
1
Ohe X
N
Ohe
X
D
x N
where proportionality can depend on any factor (e.g. except of course N. Crucially note that the equations require 7 VN. So this can be done for pP but not for NTP (except if d = N). Further, we need
1
Orw Bv O
N
To get that w We. choose (as in the main paper)
1
1
Ow WND NVD
(for μP),
ae
0g ND
ND
So, get
110
ww + 700 vw.
110
700000
w NDo,
VNDa,
NDo
ew +
VN Dae
ew.
T ww T w N Dia aw
2216
(ww) T Jw., N D
(w.w) Jw +
(w. w T yw.
22
ND
W Too + Ww Thew + nnew
T
1 2
72 (w.w )w.
ww w ND
(ww) )w.
(w. w )w +
ND
ND)
ND
Reparameterization of T Recall that
n
2
T/o
7000
1e +
w
+
w 20 +
w, ow.
ND
vND
N2 2
N D
ND
This implies, after scaling
25

---

wt = =a,ūt
= (a,ū) - 10 TT 1 - 70 _. w
ND 10
VND
ah N ti èú 2
ya
221g0u
N w ēw +
w. ēw.
ND
10000 T w + ypy w w
NDa
VNDaw
nha T ew N N
1ona.
w I ew + ? 2 w. ew,
w.
N Da.
1 in addition with 7 = N gives independence of is easy to see that the choice aw de Ar x N the RHS width. Under our choices
1
aw NND NvD
(for μP),
ae ND
0g ND
get
?
78
v+
[-2w + 2w w]
+
[w T
2w + w ew.
ND
ND
Reparumeterizationn of EET. Let's substitute the scaled version in the equations
et =0ee
C -
www
+
Mo
w.wl T
Me Daw
ww.
N 2
VNDaw
www. T + 2 vw,w!
www
wwww
N2 D2
N Di00000
NED a
N Do,
With our choices
1
1
Ow WND
(for μP),
Oa ND
O. ND
NVD
we get
? 12ww T + w.w + ww!] +
no² ND
e+ =et
wwww" T
UW.W T wwww + Uw. w, T]
ND
C.2.1 Initialization
Proposition C.1. At initialization, as N + -0, 3 3ID and P 1 u=== = 1. Moreover, ND. Vi.j [D]. While for 7 1 (NTP) at initialization is in the limit Gaussian with elementwise variance D.D. for 7 /N (μP) we have w 4 w = 0, with elementwise variations scaling as VN' Elw, - 1 1.
+ in expectation: E/v-v errors from xc- initialization scale as VN
ND. Vi [D].
First, note that trivially
E Elury] = 0,
and
N
1
1
B
1
Σ
E V3Vv
2ND
D=2
jy'=1
26

---

Next:
1
Bleyy
E Bl Do.s.
ND
and
Eleiyl
E Ein B5x Bswl
N2D
k,,==1
For i j
1
B Ble
Σ
(i
N2
N D2
k,k'=1
For =
N
1
Ele
E 1ip
N2
k,k==1
N Σ
E E
E
N
N2 D2
k=1
N(N - 1)
3
N2D2
ND²
N+2
(i ===)) j).
ND
So
N + 2
2 ND
Varley]
ND2
D
Finally:
1 ND 1=1
1
B
B[V.V
D
and
N
1
N --1
3
N +
B[u²]
BIV.VVVV
Σ
EVV
B[V
N2 D
N2
N2 D
ND
N
ND2
1,'==
So
2
Varlu) ND2
C.2.2 Proof of Prop. 5.3
EoS Basic Derivation quadratic (global linear dynamics). In quadratic potentials } w Hw: we are at EOS if and only 7 = 2/\aax (H). The proof of this can be translated into Jacobian argument the gradient update map. The update map here is + Jacobian - n/H The eigenvalues of this map are in the range [1 [1- 11-Amax(HH)) where bounds are tight. For asymptotic stability (i.e. converge). we have the necessary condition Amax(H) (H) < 2/1,, which guarantees all eigenvalues are in the range (-1 To be at edge of stability means Amax (H) 2/ n, that is if and only if the Jacobian has eigenvalues in [-1, 1]. with an eigenvalue equal to -.
(I HH ye, which has Amin (H) mr],
Remark:
What changes here? We dealing with a nonlinear transition map where the Jacobian has explicit (17²)) terms. To be precise in this setting, we need carefully study the Jacobian and cannot assume its first-order approximation in %oo i.e. T - 7o 2 L,, is accurate.
27

---

Proof.
We need to compute the Jacobian for the dynamical system G (w, e, - (w* (w , Note that specifically at w.
Owt
Ow w==w" = In - mo(e + vf) + 1o 2
Ow+
Ow+
w. 8
w,,
0,
ND
de
w=w"
w=
(21)
de+
de+
de
Ow w==w
de
w=w
1p3,
dv
=
(22)
Ou+
Out
O0
dw
de
de
(23)
where we do not care about the values with a * because anyway the resulting matrix is lower- triangular:
[1-no(vdd + ex)
28 ND w. !
0 0
Cx, ex,
1 0
I
As in the quadratic setting, be the edge of stability, require this matrix have eigenvalues in [J, with at least an eigenvalue precisely equal to -1 Without loss in generality, we discuss eigenvalues in descending order for each one of matrices:
A A
AD.
Note that Ja Cx2 vx) is block lower-triangular, hence its eigenvalues are
XJa(Ww,,., e., v.)) = > X (1--
n
no(vvl +e.) ND w. Ut1
The necessary condition for being at the edge of stability is then
7 ND
Ap(T O
s w.
1. -1
where θ* == e.. We then ask the question: for which combination of O* and i/os can this
c.. happen? Note that for any eigenvalue index (eigenvalues sorted in descending order), 1Jo 2 0 and O" 0:
Au (I - ~00*) = =ToAD=x(""))
Thus:
Ap(T - m²0) (0"),
- = 1 (e*))
Let us set D = -- -a matrix of which we know the eigenvalues. The matrix we want to study n2 ND w. w.. Intuitively, this perturbation
rank-1 O(1²) perturbation: -" P where
matrix P does not modify much the cigenvalues.
To start, note that the CourantFischer Theorem [91] implies (see e.g. [92], exercise 10.2.2.):
An(T") S \D + P)
Next, recall Weyl's inequality for Hermitian (in this case, symmetric) matrices:
A+t-- + P) & ++(T)+
Apply this to = -1, D, we get
AD + P AD())) + AP()).
All in all, proved the bound
s +P) Ap((') (P))
28

---

Notice that this implies, given the structure of P.
An(TT) S An + +) S An(TT) +
ww /w,12 1
ND
For EoS, we require AD + P) = -1 Hence we necessarily need
Ti 2
AD(T*) -1
we 1
ND
Recall that (T AD(I - -2)0-))))--)) 1 ( hence the condition is
7677
l|,,11 $ - (00) -1.
ND
This implies
Mn 2
2
| -T0(((0)))) & -2.
W.
ND
Dividing everything by 17o.
2
7o
Iw,11 112 - =
170
ND
hence
S (00) S 70
2 2 + ND |w, 1|
To conclude, note that \1 coincides with the Hessian A1.
Note on marginal stability (sufficient conditions). At the same time, for stability, we need (sufficient condition)
-1 A +P) < ( P 1 The condition -1 S AD + P) is implied by -1 < \n()") (Courant-Fischer), This yields
-1 S 1 - m
that is
2
1/0 <
(
Next, need (I + PP) < 1. A sufficient condition for this is (again by Weyl)
XTTT*) + A2() 1
which implies
71
\T) $ 1
||w
ND
That is
72
1 - moA <
||w. 112
ND
i.e.
72
||w.112 S (0*),
ND
that is
No S \b(Or)ND
731110.11
All in all sufficient condition for marginal stability is
2 AD(Or)ND)
No min
(et))) ?2
The bound is not vacuous in the case of N 2 D and different data points.
29

---

Note. What we provided above is sufficient condition for marginal stability (all eigenvalues in [-1.1]. This derived just for interest of the reader but does not affect the result: indeed, sufficient condition can be not tight (as likely the case here). What we instead proved is that if we
are at EoS, than necessarily (i.e., in every instance of the problem) 2 - \1())) - 2 + ND || 1|?
C.3 Theory for Standard Parametrization (SP)
As in ypP and NTP, we consider the simplified loss
L(E.V) = EV w01
(24)
where E E RR DxN, VE RNx1, Compared to μP and NTP, normalizing factors appear in the initial-
M ization: Eij ~ 1/)), Vú NO, 1/N), Vi, Gradients are Öp = EVV w, V Oy
E EVE EVV-E Wx.
C.3.1 Dynamics equations, same form as μP and NTP
Let us rename w EV, v v= == V e= EET, then the dynamies gradient descent with stepsize 1) becomes, in (w, e, v) D+D space:
wt = ( Tn + e)( w.) + (ww w.w" "ww W.). et ut = c + -2w T w + Iw: w/ + n?2 [w TuTTen ew 2w ew + w ew.].
-2ww T + w.w" www. 17² [oww] vur w wwww uw. w.
Hence, under SP, (w, e, v) at different widths have width hidtt-inndepender evolution laws in D++ same happens under P. However, in contrast to uP,, initialization of this system does not converge as N 00. Hence, actual dynamics across widths are drastically different, as we will see next.
C.3.2 Initialization
B[E, First, note that trivially E[ws] = 0, and E[w:J² = = -,
1. This makes sense since the forward pass is normalized. However, at initialization:
D.
N
Σ EGE
deg
D
that already includes
dependency on the width N,
further:
Ble
Ex E|En Esg Ew }.
Hence
N Elegl
N
N
EIEi Eaw E[E
E
E EE
(i J),
D
N(N 1)
3N
N(N + +2)
Bleg]
Σ
E E 1e
Σ
ELERIEEE
EL
(i=
D2
D
D2
k
kek
A=1
So Varleyyl 0
Finally:
1
N
E
BVV. 1
1=1
N Σ
N
N(N 1)
3N
N(N +
El
EUGV Σ
E[V? + Σ B.V.
N2
N2
N2
=1
So Varlu] 0 (++z
30

---

C.3.3 Conclusion
While the laws (w, v) + (w+,et 2 are not width dependent in SP, has initialization of scale O(N). As such, while the forward pass (i.e. w) is O(1) at initialization, the trajectory of (w, e, u) under gradient descent starts at width-dependent point (O(1), ONN), O(1)); hence it is drastically different at different widths. Further, the NTK (c ID + e) is also O(N), and this controls the evolution of the forward pass w - 11 In + e w. + ( We therefore validate also in this simple setting the derivation by [5]: n = ( forward pass of SP blows up after one step of gradient descent.
31

---

D
Connection between the eigenvalues of the Hessian and NTK matrix
Recall that for MSE loss and one-dimensional output (x)), the Gauss-Newton decomposition of the Hessian H reads:
[D]
= + Σ Vof(=:)oo(zs) + + (
(25)
where G is the Gaussian- -Nwton matrix. This can be generalized to (1) different loss functions and (2) multidimensional output f(x) E RR, where k is the dimension of the logits (i.e. the number of classes classification problems). Here we notice that in (2) we exactly preserve the connection between GGN and NTK's spectra, while in (1) we have extra term in & that causes a deviation from the exact correspondence. However, we show that this deviation is largely negligible in practice, in the sense that G will have the same spectrum dominates R. as one would expect from the experiments in the main text (c.g Fig. 1). We begin by defining the Gauss-Newte matrix (GN) in the case of cross-entropy loss, following [12]:
the NTK as training progresses, and G still
P
g= Vof(a))
where now (af( R&P, and He E R is block-diagonal matrix where the A x k Hessian of the loss (H2) with respect to model output is repeated P times. Again, we can stack the Jacobian vectors (x;) for all the datapoints into E Ppxxp
thus:
G = K H&K
(26)
For MSE loss, Hc. is the identity, hence the correspondence to the NTK matrix is maintained (same sharpness). However, for the cross-entropy loss, the first derivative of the loss with respect to the model output A == V f(z) can be shown to be of(()) - y, where y is the one hot vector encoding the true classes, and (-)) denotes the softmax activation. Hence, for the Hessian, we have:
[Hc)i (f)))
(27)
which in general deviates from the identity. However, during training the model increase the proba- bility of correct predictions, and thus He gets closer to the identity, thus having an asymptotically negligible effect on the Gauss-Newton matrix g and its correspondence to the NTK.
0.95
Width
1.15
Width
0.90
1.05
0.85
128
1.00
256
0.80
0.95
0.75
0.70
10
10-1 Amax(G)
10a
10a
10a
Amax(C)
Figure 11: Norm of the residual and top eigenvalue of the GN matrix, where the vector field shows the evolution of these quantities during training for a fixed learning rate. Left: jP- note that all curves have a sharpening phase, after which the residual continues to decrease. Right: NTP Increasing the width reduces the sharpening, since it approaches the infinite width limit where the NTK matrix becomes asymptotically constant.
We now perform experiments to test whether & dominates the residual for cross-entropy loss, in order to support our claim on the connection between feature learning and optimization. We plot the evolution of the largest eigenvalue of the GN matrix and the residual norm through time in Figure 11 for μP and NTP. The largest eigenvalue of this matrix computed using power iteration algorithm based on the implementation provided in [93]. Note that while for μul we observe large amount of progressive sharpening during training, for NTP the sharpness becomes asymptotically constant. The architecture used for the plot is a convolutional neural network as described in J, trained on CIFAR-10 for 10 epochs using cross-entropy loss. These experiments confirm the results obtained for MSE loss in the main text (Fig. 5).
32

---

E Late-time dynamics and batch size ablations
E.1 Late-time dynamies
was noted in [16] (Appendix C) that with cross-entropy loss (adopted here), the sharpness de- creases towards the end of training. Here in Fig. 12, show that while the dynamics are remark- ably consistent during the first part of training, they diverge during the phase transition in which the sharpness begins to drop, with wider models starting to exhibit the sharpness drop earlier. Hence, this transition point is highly width-dependent, and it coincides with a slight shift the optimal learning rate. This late phase happens when the classifier maximizes the margin between classes [16], and can be largely prevented by using data augmentation techniques, as we exemplify in Sec.E.2.
0.32
00.3
6x107
256
512
1624
10P
1192
23
55
70
30
Epsch
Epoch 22
Spoch
Epom
1.0
337
Widh
255
16-1
15-2
2040
1056
$150
10"
30 30
10t1
18
12-
20th
100 10'
Learning
Leaming
Learning Rate
Figure 12: Late-time dynamics. We study the setting in Fig. 26 (for μP), under longer training time. Notice how when training longer, the sharpness decreases once the model gets closer convergence. In this phase, there is shift of the optimal learning rate, the bottom row shows.
E.2
The effect of Batch Size and Data Augmentation
Batch Size We test what happens to the sharpness and optimal learning rate when the batch size is increased. The sharpness well-known increase with the batch size, as it is shown also in Cohen et al. [16] and Jastrzkebski et al. [31].
WMd
54
-000
1524
12a
2043
0.2
Epoch
Epoch
Epech
0.60
2.45
B28
Witth
3.5
1.2
3.3% 12
54
356
23.4s 107
120
2.3
120
517
2.3 109
3.3%
2024
1.2% 107
1224
1.24
2045
1.1x 107
2342
32x
2342
4225
Learning
20a
30
Learning
1.2
32
Learning Rate
123
Figure 13: Batch size ablation three layer convolutional network. The red dashed line indicates the sharpness of 4/7u-- only shown for the largest learning rate where the sharpness rises above the EoS threshold. Parameters: dataset: CIFAR-10, 10, with data augmentation, epochs 50. The learning rate shown above are: (0.32, 0.88, 2.45)
Here we add that under jPP the batch size we tested (128, 256, 512) do not influence significantly the width-independence phenomenon of the sharpness, as summarize in Fig, 13. We observe good
33

---

learning rate transfer across all the tested batch sizes. We also observe that the optimal learning rate increases with the batch sizes by roughly factor of 2.
Data Augmentation We repeat the same experiment, varying the batch size, but this time we turn on data augmentation using standard transformations (random crops, horizontal flips, 10-degrees rotations). The results are in Fig. 14 (to compare with Fig. 13). Notice how data augmentation has stabilizing effect on the sharpness, delaying the late phase of training where the sharpness drops, as analyzed Sec. E.1 [16, 28]. On the other hand, Thus, under regularization techniques such as data augmentation, should expect better hyperparameter transfer.
9- 512
Weet
2348
10
10
4236
1r
0.33
10
0.25
Epoch
Epoch
1pach
2.45
B 64
B 256
D 512
309
Width
64.0
128.5
512.6
Toah
10
10
20-1
10
10t
30
30th
30
10
30"
301
Learning Rate
Learning Rate
Learning Rate
Figure 14: Batch size ablation. The red dashed line indicates the sharpness of 4/7jo, only shown for the largest learning rate where the sharpness rises above the EoS threshold. Parameters: dataset: CIFAR-10 10, without data augmentation, epochs = 50. The learning rate shown above are: (0.32, 0.88, 2.45)
F Large-Scale experiments, more Datasets and Optimizers
We provide empirical validation our findings and show that in realistic architectures, such as Transformers (ViTs, GPT-2) and ResNets, we achieve width (or depth, respectively) independent sharpness. These results empirically demonstrate that our findings extend different modalities and architectures.
F1
GPT-2 experiments on WikiText
Figure 15 shows the transfer of a GPT-2 model trained WikiText for 40 epochs with Depth- and Adam, without learning rate scaling. similar plot but for μP (without the depth parameterization) is presented in Figure 16. Note that the sharpness exhibits width independent behaviour.
4x
Depth
Train Loss
Dopth
10a
2
10º
1000
101
Learning Rale
0.00037
10
10
20 Epoch
Learning Rate
0.0009
Figure 15: GPT-2 trained on WikiText using Adam, with fixed learning rate, width 512. Note that due to the structure of the transformer block containing 2 layers per block, the transfer is imperfect and thus the sharpness is also not super consistent. Here, in contrast to Depth-uP's prescription in Table 1, we do not rescale the learning rate by I/vL. See also App.
34

---

Width
Width
128
10a
128
Train Loss
256
256
512
512
1024
1024
109
2048
2048
10a
Ir
0.0026
10 Learning Rate
10
10 Epoch
20
(a) juP Transfer
(b) Sharpness
Figure 16: Post-LNN Transformers (similar to GPT-2) trained with Adam on WikiText-2 parameter- ized with yP showing learning transfer in width (a) and super consistent sharpness evolution (b). HPs: layers, 2 heads, 20 epochs, batch size 512, 100 warmup steps, sequence length 35.
F22 ConvNets Experiments on Larger Datasets and Adam(W)
Figures and show residual ConvNets parameterized with yP and Depth-uP respectively trained on Tinylmagenet, and Figure 19 shows the evolution of a similar model parameterized with pP trained on Imagenet. Similarly, we study the evolution of residual ConvNets when trained with Adam on CIFAR-10 with 1 and 2 layers per block in Figure 20. Finally, we show the results of training residual ConvNets under uP with Adam and AdamW in Figure 21.
10:
4.6 1
10 10
Wieth
32
32.0
Tain Loss
10
64
128.0
4x
256.0
3.8
512
256
10
$12.0
0.008799
3.4 100
0.014678
0.024484
10
60
Learning
Figure 17: Residual convolutional networks (ResNets) trained Tiny-ImageNet with stochastic gradient descent. Left figure shows the learning rate transfers across width in ResNets parameterized with pP. Right figure shows that for fixed learning rate, the sharpness becomes width independent during training. Parameters: batch size 64, epochs 10.
10"
Depth
38"
10-1
6.0
10*
12
12.0
3.6x10'
24
10-1
35x
96
96.0
34x
303
0.018678
0.024484
10*
100
Leaming
Stop
Figure 18: Residual convolutional networks trained on Einy-mmaeeNe with stochastic gradient de- scent. Left figure shows the learning rate transfers across depth in ResNets parameterized with Depth Right figure shows that for fixed learning rate, the sharpness becomes depth indepen- dent during training. Parameters: batch size 64, epochs 10.
35

---

7.0
Width
6.5
16
Width
32
Train Loss
32
6.0
128
Sharpness
66
5.5
Leamiing Rate
0.0015e5
5.0
10
10
Learning Rate
Step
Figure 19: Residual convolutional networks (ResNets) trained ImageNet with stochastic gradient descent. Left figure shows the learning rate transfers across width in ResNets parameterized with j.P. Right figure shows that for fixed learning rate, the sharpness becomes width independent during training, Parameters: batch size 128, epochs 1.
1,5
1.5
1.45 x 10
Depth
1.4x
Depth
10
1.3x10
6
Train Los
103
12
1.2
103
24
Train
1.1x
24
x.25
48
1.2x10
10a
96
1.15
101
10-3 Learning Rate
10-4
Learning Rate
10a
Depth
Depth
6
10a
12
107
12
24
24
48
10e
10"
0.001233
0.000534
20
40
20
40
Epoch
Epoch
(a) Depth-ud with &
(b) Depth-uP with k = 2
Figure 20: ConvNets trained with Adam with fixed learning rate on CIFAR-IO, (a) One layer per block, showing that the learning rate transfers and have sharpness Super Consistency, (b) When training with k = 2 layers per block, notice that the transfer worsens. While the effect is subtle in this setting, translates the sharpness accumulating finite-size effects during training. Here, use the DP prescription for Adam reported in Table 1.
G Time Evolution of other Spectral Quantities
In this section we present convergence rates for various other spectral quantities of interest. In Fig- 22, show the same Super Consistency of the landscape in the of Depth-/uP, exhibited in the width case in Figure 2. In Figure 23, we present the convergence rate of the largest NTK eigenvalue, sharpness and loss respectively, as well as the evolution of the Hessian trace during time. Figure 24 (a) shows the loss curve evolution under jPP and NTP regime. Note in Figure 24 (b) and (c) how the sharpness achieves near EoS value in the case of μP,, whereas for NTP wider networks have a sharpness that remains closer to initialization. Finally, in Figure 25 we show that in the case of having multiple linear layers per block, this leads violations of Super Consistency, and thus imperfect learning rate transfer. Note that this is the same case as in the GPT-2 plots illustrated in Figure 155
36

---

Width
10*
Width
128
256
Train Loss
384
512
10a
512
Ir
0.0034
8 x 10
10
10 Epoch
20
Learning Rate
Width
10a
128
Width
128
256
Train Loss
256
384
109
<
10a
512
Ir
x 10
0.0034
x
10-2
0
10
15
20
Learning
Epoch
(c) jPP Transfer
(d) Sharpness
Figure 21: Convolutional networks trained with Adam (top) and AdamW (bottom, weight decay 0.001) on CIFAR-I0 parameterized with μP showing learning transfer in width (a) and super consis- tent sharpness evolution (b). HPs: 20 epochs, layers, no skips, batch size 256, 200 warmup steps.
2-121
30-
10
Opech
Qoach
10
33
-1
3.23
Epech
(a) Depth- pP ConvNet Depth-uP ConvNet
Depth-uP VIT
Figure 22: Convergence rates with respect to time for the losses and sharpnesses shown in Figure 4.
H Sharpness evolution in Random Feature Models
In this section, we compare the NTP and μP parameterizations to random feature model, i.e. model where we the weights of the intermediate layers are frozen to their value initialization, and only the final readout layer is trained. Crucially, this model does not learn features by construc- tion for any learning rate and any width. The results are shown in Fig. 26. Notice how the transfer in the random feature model is achieved only at very large widths compared to pPP However, the trans- fer is better than in NTP. This in line with our claim, as under random feature model increasing the learning rate does not induce more feature learning, is the case in NTP.
37

---

Epoch
Poode o -hhastuul
10e
Tor
10a
101
Width
Width
(a) Largest NTK eigenvalue Width /P
(b) Sharpness vs Width uP
10a
10a
(c(C) c(fe)
10
(r(H))
10a
64
32
128
10a
1824
0.47
0.77
0.65
0.64
2048
10
Width
Epech
(c) Loss vs Width JuP
(d) Hessian trace vs Epach uPP
Figure 23: (a) Convergence rate the largest NTK eigenvalue width at multiple steps during training. (b) Convergence rate of the sharpness in width at multiple steps during training. Note that the largest NTK eigenvalue starts width independent for both learning rates, but becomes width de- pendent during training, opposed the sharpness which maintains width independent dynamic throughout the whole optimization process (see Fig. 24). (c) Loss convergence rate in width; note that the loss accumulates finite-size effects during time, exhibiting wider is better effect during training. (d) Hessian trace evolution during training at various widths. Unlike the sharpness, the trace has width dependent period the beginning of training, but approaches width-indepedent threshold. Details: Residual ConvNet trained on CIFAR10 with cross- cnosspyn loss.
24
2x
18x18
Stes
15
Lex18
16418
1324
1324
2.5 10
5:4p
Stop
With
24x18
13-1
22xur
2-
23-4
224
13+
16x18
1326
1034
Shop
60mp
Witth
(a) Loss curves
(b) Sharpness dynamics
(c) Sharpness vs width
Figure 24: Early training dynamics in uP (top row) and NTP parameterization (bottom row). (a) Loss curves. (b): Sharpness dynamics. Notice how for yP progressive sharpening until EoS (black dashed line) is achieved at any width and with comparable speed, while for NTP the time to EoS progressively increase with width. Also, the loss curves start to depart from each other as training progresses, while À stays at EoS for selected time steps. For jPP, converges very fast in width, while in NTP diminishes. Other parameters: architecture: Three-layer convolutional network. Dataset: CIFAR-10, without data augmentation. B 128, epochs 1, learning rate To 0.8 for uP and 8 for NTP. The reported width the size of the readout layer.
sustained period ot time. (c) Sharpness vs width
38

---

10a
x
A(0)
10
A2
Amax
Azo
A
10a
A1o
10
40
20
40
60
Epoch
Epoch
Figure 25: Dynamics of some of the eigenvalues of the Hessian (left) and NTK (right) under Depth- μP scaling with & 2 layers per residual block. We observe violations on Super Consistency in both cases. The model is the same in Fig. 4 (center). This is compatible with absence of learning rate transfer.
24
22"
24:227
6-
126
10
1604
1624
2013
2003
2045
1104
8192
4095
9192
4096
12x107
10
Leaminn
12
Learning Rate
Learning Rate
with
Min
254
Strres
2004
2045
4006
4025
2x12
144
$7.20
Epoch
Tpech
(a) Random Features
(b) JP
(C) NTP
Figure 26: Learning rate transfer plot (top row) and sharpness dynamics (bottom row) for three- layer convolutional network and three different settings. (a) random feature model (only the readout layer is trained), (b) and (c) correspond to μP and NTP parameterizations, respectively. In random feature learning, the absence of feature learning prevents the sharpness' evolution at any width, thus learning rate transfer coincides with the convergence of the sharpness À at initialization. Also notice how for NTP, the progressive sharpening converges À 2/nu at much lower speed the width increases, epochs 20 for the μP/NTP models and 10 for the random feature model, dataset: CIFAR-10, without data augmentation.
line with the early dynamics reported in Fig. 24. Other parameters: B = 128,
I Directional sharpness
In Figure 27 we provide the evolution of the directional sharpness, which captures the curvature along the gradient direction during training under jP and Depth-uP respectively ConvNets. Note that, similar to the sharpness plots, the directional sharpness, defined as V.e HV.C also follows width independent trajectory during training. This measure has been used, for instance, in Gur-Ari et al. [17].
Experimental details
The experiments were ran on 1100 and H100 GPUs, with 80GB VRAM. Each experiment averaged less than 24 hours total execution time. Unless stated otherwise, we use data augmentations, where the random transformations are compositions of crops, horizontal flips, and 10-degree rotations. Additionally, we provide further details on the models used our experiments and the modifications we have introduced.
39

---

Width
10a1
16
10c
10-2
32
64
64
Train
128
10
x0-1
256
256
512
512
10-3
10
109
5
10 Epoch
15
20
0.5412
Learning Rate
(a) jP Transfer
(b) Gradient- Heessinn Alignment
Depth
Depth
1x
10
12
12
Train
6 10-1
24
24
10-1
48
5 10
0.4924
10
10
20
30
40
Learning Rate
Epach
(c) Depth-uP Transfer
(d) Gradient-Hessien Alignment
Figure 27: Convolutional networks trained with SGD on CIFAR- 10 parameterized with μPP (top) and Depth-pP (bottom) showing learning transfer (a) and super consistent Hessian Hessian-gradient alignment during training (b, d). HPs: (top) 100 warmup steps, batch size 256, 20 epochs. no residual connec- tions (bottom) 200 warmup steps, batch size 128, 40 epochs, VT. scaling (both) 6 layers, ReLU.
10
W.B.
Figure 28: Coordinate check p.P.
J.1 Hessian Computation
The implementations of our models are done in PyToorc.. For the Hessian measurements, we use PyHessian [94]. In particular, the library adopts the Power Iteration method for the top l: Hessian eigenvalues (i.e. including the sharpness) and the Hutchinson method for trace computation. Both methods adopts Hessian vector products avoid computing the whole Hessian. This reduces the time complexity from quadratic to linear in the number of parameters. In both algorithms, we fix the number of iterations and tolerance between consecutive eigenvalue computation the default values of 100 and 0.001, respectively. We measure the sharpness on the same fixed batch throughout training.
SGD Among the equivalent ways of parametrizing then network with ,P, we opt for the one that rescales the learning rate by the width, i.c. 17 = no No This effectively sets the EoS threshold the width-dependent value 2/ 2/(Nro). In our plot, we take this scaling difference into account by computing the eigenvalues of the scaled Hessian 22 = NH. With respect to learning rate transfer, such a rescaling makes intuitive sense, as is 7/o that transferring, and not 17.
Adam Adam updates are of the form 0++1 = 02 - r/PTt m4+1, where me a momentum vector computed in the exponential moving average fashion: m = 3 my + (1 )1 where gr are the gradient of the loss at time t. Pt+1 is diagonal preconditioner of the form
(1 88)
diag V -
el
35
40

---

where v := of this paper, we consider the preconditioned Hessian P- H, for which it is shown that largest eigenvalue converges to the EoS threshold of 2(1+81) in Cohen et al. [61]. Furthermore, we use n(1-8i)
where oi is the clement-wise squared gradients. In the experiments
the Adam μP parametrization in Table 8 of Yang et al. [6], which rescales the learning rates of the hidden layers (i.e. with both input and output dimensions that scale with N) by 1/N. To account for this, we further adjust the Hessian by computing H = DP-1H, where is diagonal matrix containing the learning rate each parameter. We always report spectral quantities of H in the experiments with Adam. With these modifications, we expect Amax ) (1-9.5*
2(1+&1) which is 38 for
the full-batch case and the default 31 = 0.9. Indeed, this is what we observe in our experiments. We stress that we expect Super Consistency of this preconditioned Hessian, where jP different layers may have a different dependence on the width [6]. Finally, we point out that in the Depth-uuP experiments, we reported the hessian of N/vLH. Although this produces results that deviate from the prescription of Cohen et al. [61], it is sufficient to capture Super Consistency depth, as each layer has the same depth-dependence.
J.2 GPT-2
The backbone architecture in the experiments presented Wikitext is the standard GPT-2 trans- former introduced by [33], with the Depth-uuP parameterization changes presented in [6, 8, 53]. Crucially, the following modifications are introduced by the μP parameterization:
The attention map is rescaled by 1 as opposed de
Vde
The residual branch downscaled by VI where is the number of layers the model
Our implementation based on the implementation provided by Yang al. [6], with the addition of the residual scaling. This uses different parametrization from the one reported in Table. but equivalent dynamics, obtainable using their "abc-rule". Similar the experiments performed on ViTs, the GPT-2 models are trained using Adam, where the base width is fixed and the depth is varied (and vice versa for the Depth-uu case). In addition, for the fixed width, increasing depth experiments, we place the layer normalization layer in front of the residual, following the Pre-LN architecture, unless stated otherwise, and we do not use any learning rate warmup. Note that this
setting we also train without the IvI scaling of the learning rate that the theory would prescribe (summarized in Table 1). This follows following the heuristic prescription of Bordelon et al. [7]. As in their case, we empirically observed that we did not get learning rate transfer if we used the scaling. In the fixed depth, increasing width case, we use a standard Post-LN GPT-2 style architecture.
linear rate warmup, with no decay, and
J.3 Vision Transformers (ViTs)
The ViT implementation is based on the work of [95] and follows the same tokenization and training protocol. In order to follow the Depth-uu parameterization, make the same modifications as J.2. The models are trained with Adam.
J.4 ResNet
We use convolutional neural networks with skip connections, with 3 x 3 kernels and stride 1. We apply pooling after every 3rd layer, followed by a subsequent convolutional layer. Following Bor- delon al. [7] and Yang et al. [8], we downscale the residual branches by I/VL. Note that in the Depth-pu setting we also scale the learning rate as 1 L in these experiments.
J.5 Coordinate check for ju
In Figure 28 we check the coordinatewise evolution of the activations at hidden layers within ConyNet. Note that the evolution is width independent, predicted by μP theory.
41

---

K
Summary of Feature Learning Parametrizations
Here summarize the parametrizations used. In Table I, we report the scaling of the learning rate, output multiplier depth-scaling exponents for the residual branches for Depth-uPP (both SGD and Adam) [7, 8, 96]. JuP is recovered as special case, where the depth dependence is ignored. Notice that this version of Depth-uP for Adam is obtained with simplification, setting Adam's 's parameter to zero, where Sign-GD is recovered.
?
a
Non residual block layers
SGD
1o N
2
To N 1 No
Adam
L
Table 1: Summary of Parametrizations. The non residual block layers represent those trainable vectors/matrices that are not in a residual block (typically, the first and last ones). For these layers, prescribe how to rescale both the weight variance and scaling multiplier.
42

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We provide extensive evaluations for the claims in the paper on Super Consis- tency and its relation to learning rate transfer.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We provide discussion on the limitations of this work. In particular, when it comes to the link between sharpness and optimal learning rate, and on the computational resources required to get Hessian, which prevents larger scale experiments.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
43

---

Answer: [Yes]
Justification: We explicitly state all the assumption and simplifications, as well as an orga- nized appendix with all the proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Answer: [Yes]
Justification: We provide detailed experimental details section, as well as formulas and choice of hyperparameters. We will release the code upone acceptance.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, if the contribution is specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
44

---

reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We will release the code upon acceptance
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
htte://uiieeeeedeeuubssseseooeee Polccyy)lll) for more details.
While we encourage the release of code and data, understand that this might be possible, so No an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips.. cc/
cc/ppclipcicccccccc public/ gtt///uw/e////e///ee /eee/oooduubmi wetiaubmiiiiiisii wwssssoollliieeeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have a dedicated experimental details Section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We have run the experiments with multiple seeds, and provided error bars in the plots (confidence intervals given by the Matplotlib Python library)
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" " if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
45

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: We provide this info.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't t make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttp://eeeip./eeeeeeeeeeeee ce httpppplcccchccccciieeeeeeeee
Answer: [Yes]
Justification:
We are conform with the code
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
No societal impact for this largely theoretical work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
46

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not pose such risks, and only use public data.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: provide info and citations for the packages that we use.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswi thcode com/ com/datasets has curated licenses for some datasets. Their licensing guide can help determine the cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
47

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification:
We do not release assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: Not applicable
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: not applicable.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
48


Reviews:
[{"review_id": "R2ynOx20wo", "reviewer": "Reviewer_f2zq", "summary": "The authors conduct a series of experiments in which they investigate for which attributes of the neural network the gap between its infinite-width (or infinite depth) value and the finite counterpart grows or shrinks during training. Among other attributes, they look at the training loss, largest loss Hessian eigenvalue, and largest NTK eigenvalue. They investigate how the results for different parameterisations (such as muP, neural tangent parameterisation, and some other unstable parameterisations with no well-defined infinite width behaviour).", "questions": "- Why are there two sets of lines for each width on Figure 1, one solid one dashed?\n- Is the Appendix G, which is meant to discuss “the effect of the lower order eigenvalues”, empty? \n- In Figure 2, the authors look at 4-th and 10-th largest eigenvalue. How come the authors decided to look at a fixed N-th largest eigenvalue? Given the scaling in width, wouldn't it be more interesting too look at say 4th and 10th percentile largest eigenvalue, given that the number of eigenvalues grows with depth/width?", "limitations": "- Definition 3.1 concerns asymptotic properties in *both* width/depth and training time, which are then doubly difficult to establish with certainty with finite width, depth, and training time experiments.", "rating": 7, "confidence": 5, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "The asymptotic scale properties of neural network training are an important area of research. This paper investigates an interesting question: how do the properties of the loss landscape, such as the curvature along the training trajectory, evolve as one scales up the neural network size.\n\nFurthermore, the experimental results are interesting, and well-presented in the figures. I liked Figure 3 in particular, which most convincingly illustrates the claim that the gap between finite-width and infinite-width attribute value shrinks/grows throughout training. In fact, I wish for every hypothesis investigated, a plot like that in Figure 3 was shown.\n\nLastly, it appears the authors did proper ablations, investigating multiple datasets and models, to verify their empirical conclusions.", "weaknesses": "The paper falls short in presentation and the formalism.\n\nOne of the largest issues is the definition of the term “super consistency”. Definition 3.1 has several issues:\n- The authors defines $S_N(t)$ as a function of the predictor $f_t(x)$. This is pretty vague, and not sure matches what authors are trying to do. If the predictor $f_t$ is interpreted to be a function $f_t:\\mathcal{X}\\to\\mathcal{Y}$ from some input space to some outputs space implied by the neural network architecture and weights at time $t$, then $S_N(t)$ cannot capture something like the curvature of the loss with respect to the weights. I'm convinced this hurdle can be overcome by carefully defining all objects, and what they are (e.g. a ‘predictor’).\n- If $S_N(t)$ depends on the weights (such as when considering the spectral norm of the loss Hessian), then it's a random variable. Hence, all the notions of distance and convergence need to be defined for random variables for the expressions in Definition 3.1 to make sense.\n- It's not clear the limit $S_\\infty(t)=\\lim_{N\\to\\infty}S_N(t)$ exists for many properties being considered. In fact, this already precludes the authors from talking about super-consistency of parameterisations that do not have well-defined infinite-width limit training dynamics (e.g. SP).\n- line 127 “if at any finite $N\\geq N_0$:” – did the authors mean to say ‘there exists some $N_0$ such that for all $N\\geq N_0$’? Otherwise this sentence doesn't make sense.\n- Does the symbol $\\sim$ here represent asymptotic equivalence as $t\\to\\infty$? This is something that should have been defined.\n\t- “$\\sim$ denotes the finite-time behaviour during training” is not a formal definition. I have no idea what it's meant to say.\n\t- Also, the whole condition of “$|S_N(t)-S_\\infty(t)|\\sim g(t)$ where $g(t)$ is a non-increasing function of time...” seems like it could have been equivalently stated much more simply as $|S_N(t)-S_\\infty(t)|=\\mathcal{O}(1)$ as $t\\to\\infty$.\n- This definition seems different from how super consistency was described in the abstract introduction: “certain [...] properties [... are largely independent] of the width and depth[...].\n\t- We name this property super consistency of the landscape”. The concept of super consistency, as defined in Definition 3.1, is simple enough to express in two sentences, that I don't see a reason it can't be explained in the introduction and/or abstract properly.\n\nLater in the paper, the authors proceed to use ‘super consistency’ in a sense completely disjoint from that in Definition 3.1. On lines 168-159 they say: “The optimal learning rate is preserved across widths/depths, indicating very fast convergence with respect to the scaling quantity (i.e. it's super consistency).” Definition 3.1. has seemingly nothing to do with the speed of convergence in the scaling quantity (width/depth).\n\nIn fact, I think the paper would have been stronger had the authors cut back on formalising things like “super consistency”, and just presented the empirical results for what they are. The term “super consistency” doesn't strictly seem necessary to convey the take-aways of the paper, and, at the moment, makes the paper more convoluted. Of course, a thorough formalism might be preferred, but it its current form it detracts from the paper.\n\nOthers:\n- I find the phrase “optimization happens along trajectories of super consistent top Hessian eigenvalues” quite confusing. It took me a while to realise the authors are saying that the ‘Hessian eigenvalues are super consistent along the optimisation trajectory’, which is the way I'd recommend they phrase it throughout.\n- There are several typos throughout the paper. Sometimes, seemingly parts of mathematical expressions are missing (e.g. line 161).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\nThe 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.\n\n**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:\n\n*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:*\n\n *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.*\n\n *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.*\n\n*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.*\n\nAnd a few lines later:\n\n*To give a quantitative measure to the finite-size accumulation property, we measure deviations over time by estimating the following quantity:\n$$g(t) := |S_{N}(t) - S_{\\infty}(t)|.$$\nWhen $g(t)$ is an increasing function (up to fluctuations), Super Consistency is violated.*\n\nThe 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).\nWe 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. \n\n\nOthers:\n1. “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.\n\n2. The typo in line 161 has been fixed.\n\n\n*Questions:*\n\n1. **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.\n\n2. **Appendix G**. We thank the reviewer for pointing this out. We have now completed the section.\n\n3. **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)\n\n**Other remarks:**\n\n1. **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. \n\n2. **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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "v2fEQ2tblk", "reviewer": "Reviewer_11Ly", "summary": "This paper proposes the concept of super consistency, which describes the stable properties of the loss landscape during training. By analyzing the maximum eigenvalue of the Hessian matrix, it is found that the sharpness under the μP and Depth-μP frameworks remains super-consistent and stable near the threshold. In contrast, NTK and other frameworks show significant differences.", "questions": "Is it feasible to extend the results of this article to other algorithms, such as AdamW?", "limitations": "The theoretical setting i simple. The author can consider trying to expand the two-layer linear network to two-layer ReLU or deep linear network.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper proposes a new concept of \"super consistency\", which provides a new perspective for understanding the behavior of models of different scales. Through a large number of experiments, the learning rate migration phenomenon under the μP and Depth-μP frameworks is verified, and the consistency of these frameworks on different tasks and datasets is demonstrated, including ResNets, Vision Transformers, and GPT-2.", "weaknesses": "This paper provides some theoretical analysis, mainly focusing on two-layer linear networks and failing to fully verify the theoretical applicability in nonlinear networks and complex structures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n1. **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. \n\n2. **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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vznlJ4sLCt", "reviewer": "Reviewer_47kG", "summary": "The authors argue that the top eigenvalues of the loss Hessian stabilize throughout training under width and depth muP scaling. This phenomenon is called the Super Consistency of the loss landscape. The authors provide theoretical convergence guarantees and empirical experiments supporting their claims. The learning rate transfer under muP scaling also correlates with the super consistent landscape, i.e. optimization follows the trajectories of sharpness (top eigenvalue of the loss Hessian). Under NTK or other suboptimal scaling, the super consistency is also violated and thus learning rate transfer failed. The authors further show that the dynamics of the weights are fully specified by a set of governing equations and thus one can derive a theoretical edge of stability result under scaling. The phenomenon of progressive sharpness towards stability also happens along with the NTK evolution with finite-size effects, suggesting other factors contributing to the super consistency of the sharpness.\n\n**[raising score from 6 to 7 after rebuttal]**", "questions": "1. How are you getting the eigenspectrum of the loss Hessian? It seems possible to get at least a few top eigenvalues using iterative Lanczos methods for models larger than 124 million parameters GPT-2 with A100-80GB.\n2. I don’t find enough evidence that the super consistent sharpness causally helps the learning rate transfer under muP scaling. In line 178, “optimization happens along trajectories of super consistent sharpness λ_max”. Is there evidence that the optimization does happen first along this direction, i.e. some kind of spectral bias during learning?", "limitations": "See weakness and questions", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The observation about the super consistency of the loss landscape and its relation to learning rate transfer under muP scaling is quite an important contribution to the community. I’m not an expert in this field, but this result seems novel to me. \n2. The paper is well-written and easy to follow.", "weaknesses": "The empirical evaluations are a bit limited. \n\n- First, a lot of the claims are made under the setting of ConvNet on CIFAR-10. It’s understandable from a compute perspective, but it’s not clear if super consistency scales to even larger models.\n- Second, the GPT-2 results on WikiText seem to break the super consistency, though it could be possible that it’s due to the Transformer block itself as the authors explained. So it’s not clear if the super consistency results hold across different models and modalities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging the importance and novelty of our results. Here we address the concerns:\n\n1. **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). \n\n2. **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.\n\n3. **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. \n\nQuestions:\n\n1. **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. \n\n2. **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.\n    + 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.\nFinally, 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).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MSrUOlqjNQ", "reviewer": "Reviewer_ssWb", "summary": "The paper investigates the loss landscape of the model with scaling width or depth,  through observing the largest eigenvalue of the Hessian marix and the NTK matrix. Authors show empirically that the loss Hessian evolve almost identically for different model sizes (which is named Super consistency), however, ntk accumulates finite-size effects over time. Authors also validate their empirical findings using theory in a two-layer NN with linear activations.", "questions": "- I am confused about Line 209, this line mentions we can decompose H into G + R, while G = K^T K and the dynamics are mainly driven by G. NTK and G share the same nonzero eigenvalues, doesn't this mean H and NTK should have similar behavior? Or feature learning happens in H which makes the difference between NTK and Hessian?", "limitations": "The authors have fully addressed the limitations in their paper", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper discusses an interesting phenomenon that the loss landscape gradually becomes stable, this might explain the transfer of learning rate\n- The paper shows NTK (lazy learning) has different behavior when comparing with Hessian of the loss, this suggests NTK is insufficient to explain the behavior\n- The paper gives an explicit evolution law under 2-layer linear NN setting,", "weaknesses": "- The authors do not conduct larger-scale environment since needing to compute Hessian", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for the strong overall score and for marking “excellent” our paper across all the three evaluation axes. \n\nOn the weakness:\n\n1. **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.\n\nQuestions: \n1. **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. \n\nWe hope that the new experiments and clarification will make the Reviewer more confident about their evaluation.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their initial reviews and interesting comments. In particular, we summarize that all the reviewers have acknowledged the validity and importance of Super Consistency as a novel and important property for understanding the loss landscape of neural networks at different scales, in particular with relation to the phenomenon of learning rate transfer. 

We did not find any common weakness shared across the majority of reviewers. The main concerns are: Super Consistency under very large models (Rev. 47kG and ssWb) and other optimizers such as AdamW (Rev. 11Ly) and the definition of the term Super Consistency (Rev. F2zq). 

1. **On larger scale experiments**. We would like to note that the main published work in this area that we take as reference (Cohen et al, 2021, Gur-Ari et al, 2018) train smaller models, and on a less diverse set of tasks and architectures. In fact, we have experimented across several tasks (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. Thus, we believe that our claims remain valid to a sufficient scale, that is compatible with previous work in this area. However, it remains interesting to see what would happen at an even larger scale. Thus, **we ran new experiments scaling up to about 300 million parameters**. See next:
2. **New experiments (Figures refer to the one-page pdf)**:
    + **Adam + Post-LN Transformers, width scaling (Figure 1)**. We have adapted the Hessian computation to the Adam case as in 
            Cohen et al. (https://arxiv.org/pdf/2207.14484) (Eq. 2), where the largest eigenvalue of the preconditioned Hessian is computed: 
            $\lambda_{\max}(P^{-1}H)$, where P is the diagonal Adam’s preconditioner. This is important for width scaling (in contrast to 
            depth scaling) because the preconditioner itself is width-dependent. Results are in the attached pdf (Figure 1). We show that 
            Super Consistency holds for a Post-LN Transformer architecture trained on wikitext, where the largest network has about 300 
            million parameters. 
    + **Alignment of gradient and top hessian eigenvector (Figure 3)**. Rev. 47Kg has raised a concern about how the Super 
            Consistency of the first few eigenvalues can explain learning rate transfer. The fact that only the first eigenvalues are important is 
            a finding of Gur-Ari et al, 2018, where it is shown that SGD happens in the space spanned by the top $k$ Hessian eigenvectors. 
            However, we have performed an experiment where we compute $g^T H g / ||g||^2$, which is an unnormalized measure of 
            alignment between gradient and Hessian, thus capturing the curvature along the gradient direction. We observe Super 
            Consistency in two experiments using residual networks, both under $\mu$P (width scaling) and Depth-$\mu$P (depth scaling). 
            This further strengthens the connection between the Super Consistency of the landscape and learning rate transfer and we 
            believe it resolves the reviewer’s concern. 
    + **Adam and AdamW experiments, width scaling (Figure 2)**.  We have performed new width-scaling experiments with Adam 
               and AdamW on 3 layers Convolutional Networks, extending Super Consistency to these settings. 

3. **Definition of Super Consistency**. Rev. F2zq has expressed concerns about Definition 3.1. We would like to point out that we found it overcomplicating to aim for an airtight mathematical description of our observation. To enhance clarity, we opted instead for an actionable definition that would reflect the actual measurements of e.g. Figure 3. However, we agree with the reviewer that some parts of Definition 3.1 can be potentially confusing, and can be interpreted as an attempt of rigorous formalization. We provide a revision below which considers all of the Reviewer's helpful comments. We are sorry the initial writing caused confusion, and we are positive our current description, now under no claim of "definition", better reflects our mental picture and resolves potential doubts. We have updated the paper by removing the Defintion environment and partially rephrasing some parts. Please see the rebuttal to Rev. F2zq for details.

Meta-review: The paper investigates the sharpness dynamics, the training loss dynamics and the learning rate transfer under $\mu$P, and explain why there is a consistency (named Super Consistency) of some properties across different model sizes.

The reviewers discussed the paper in detail (regarding the definition of the term "Super Consistency", large-scale experiments, experiments with Transformers, theoretical limitations beyond 2-layer NNs, etc.) and agreed upon positive ratings. I recommend accept.
I hope the authors incorporate the discussion with the reviewers into the camera-ready version.
Decision: Accept (poster)