Title: Super Consistency of Neural Network Landscapes and Learning Rate Transfer
Venue: NeurIPS 2024
Accepted: yes
Keywords: mup

deep learning theory

optimization theory

edge of stability

NTK
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

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 ( µ 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 loss Hessian, with a focus on its largest eigenvalue (i.e. the sharpness), and find that certain spectral properties under µ P are largely independent of the size of the network, and remain consistent as 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 dynamics? Through a connection between the Hessian's 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 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.

1 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, as models enlarge, so does the cost incurred in hyperparameter tuning which has led researchers to 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 network, not all of them facilitate learning rate transfer. For standard deep learning practices, such as networks parametrized with LeCun/Kaiming initializations [2, 3], a significant shift in 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 rate also varies as the width and depth of the network change. Alternatively, Yang and Hu [5] and Yang et al. [6] propose the µ P framework, designed to maximize the gradient update of the representations of the intermediate layers (i.e. feature learning) as the width increases. Under µ P scaling, and its depth extension for residual networks Depthµ P [7, 8], it has been empirically demonstrated that the learning

0

10

20

30

40

50

Epoch

10 0

10 1

Sharpness

P

0.19

0.53

1.47

0

10

20

30

40

50

Epoch

10 0

9×10 1

2×10 0

Train Loss

P

10

1

10 0

Learning Rate

P

64

128

256

512

1024

2048

4096

8192

0

20

40

Epoch

10

2

10

1

Sharpness

NTK

Learning Rate

6.81

52.75

0

20

40

Epoch

10 0

9×10 1

2×10 0

Train Loss

NTK

10 0

10 1

10 2

Learning Rate

NTK

rate transfers across both width and depth. In Vyas et al. [9] it 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 an example in Figure 1 (top center), where the training losses exhibits an increasing gap. The fact that the learning rate is 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 properties of the neural network's loss landscape that are preserved across training as a function 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 local curvature, and its structure for neural networks has been studied in several works [10-14]. Of great interest in optimization theory is the sharpness , i.e. the top Hessian eigenvalue, which for neural networks exhibit a rapid increase ( progressive sharpening ) towards a threshold called Edge of Stability ( EoS ) [15, 16]. However, although a few works have provided early insights [10, 16, 9], the scaling properties of the sharpness and Hessian's dynamics under different scaling limits remain unexplored. In this work we first present evidence of Super Consistency in the Hessian'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 µ P, NTP and other scaling limits. These results suggest that learning rate transfer happens in super consistent landscapes, as the geometry of the landscape does not significantly change with the network's size.

Then, we investigate the role of feature learning in the progressive sharpening phase, and argue that while in µ P feature learning causes progressive sharpening to reach a width-independent 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µ P it holds for the largest eigenvalues of the loss Hessian (Fig. 2), which converge to a largely width-independent 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 eigenvalues 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 µ P, Depthµ P, NTP and other parametrizations (Fig. 1, Fig. 4 and Sec. B). For µ P and Depthµ P, which do transfer, the sharpness stays super consistent, stabilizing to a 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µ P with multiple layers per residual block, the sharpness dynamics significantly separate during training for different widths, albeit in different ways. Also, here we 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 µ P. In Sec. 5 we provide intuition with a theoretical analysis on a two-layer linear network.

Finally, in Sec. 6 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 to the appendix (App. A).

2 Background and Definitions

We consider a neural network with residual connections, defined by the following recursive equations over the layer indexes ℓ ∈ [ L ] :

where N and L are the width and depth of the network, W ℓ ∈ R N × N for ℓ = 1 , . . . , L -1 , and τ is a factor that enables ( τ = 1 ) or disables ( τ = 0 ) the skip branch. We denote the output with f ( x ) = 1 γ W L ϕ ( h L ( x )) , where W L ∈ R 1 × N and γ scales the network output. Similarly, α has the role of interpolating between different depth limit regimes. At the first layer, we define h 1 ( x ) = 1 √ D W 0 x , where W 0 ∈ R N × D . All the weights θ = { W ℓ } L l =0 are initialized independently from N (0 , 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 = { ( x µ , y µ ) } |D| µ =1 of datapoints x µ ∈ R D and labels y µ ∈ R , we train the network with stochastic gradient descent (SGD) with batch size B and learning rate η ∈ R ,

where the loss L is a twice differentiable loss function. Defining f t := ( f t ( x µ )) µ ∈ [ |D| ] ∈ R |D| to be the vector of network's outputs at time t , if one considers continuous time, the corresponding gradient descent dynamics in function space df t /dt take the following form [18]: df t dt = -Θ( f t )∆( f t ) , where ∆( f t ) i := ∂ L ( f t ( x i )) /∂f t ( x i ) , i ∈ [ |D| ] is the vector of residuals, and Θ( f t ) ij := ⟨∇ θ f t ( x i ) , ∇ θ f t ( x j ) ⟩ for i, j ∈ [ |D| ] is the Neural Tangent Kernel (NTK).

Infinite Width. The parameters γ, α, η ∈ R determine the nature of the scaling limit. If γ = γ 0 , η = η 0 are O (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.e. Θ( f t ) = Θ( f 0 ) for all t ≥ 0 . Thus,

the network's dynamics become equivalent to a linear model trained on the first order term of the Taylor expansion of the model at 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 γ = γ 0 √ N , and η = η 0 γ 2 ( µ P, or mean-field parameterization), the features evolve in the limit (i.e. the NTK Θ( f t ) evolves), and the richer model's dynamics can be described using either Tensor Programs [5] or dynamical mean field theory [20]. Under µ P, Yang et al. [6] show that the learning rate η 0 as well as other hyperparameters transfer across width, in contrast to kernel limits, which we reproduce for our residual network in Fig. 1.

Infinite Depth. If on top of the µ P framework, the residual branches are scaled with α = 1 / 2 (Depthµ P), then Bordelon et al. [7] and Yang et al. [8] show that the infinite width dynamics also admit a feature-learning infinite depth limit. Under Depthµ P, the learning rate η 0 transfers with both width and depth. In 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µ P model. We summarize the feature learning parametrizations and report Depthµ P for Adam in Appendix K.

glyph[triangleinv]

∈

∞′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞

∞ ′

∞ ′

∞

§

∞ ′

glyph[owner]∈

W 〉glyph[ceilingleft]glyph[unionsq]〈

(a) Largest Hessian eigenvalues -

µ

P

∈

∞′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞ ′

(b) Largest NTK eigenvalues -

µ

P

′

glyph[triangle]

̸

∀

glyph[triangle]

′

∃

̸

′

glyph[triangle]

̸

∀

3 Super Consistency of the Optimization Landscape

In this work, we analyze the landscape through the lens of the preconditioned Hessian γ 2 H t , where H t := ∇ 2 θ L ( θ t ) := ∑ µ ∇ 2 θ L ( f t ( x µ )) ∈ R P × P , as θ t evolves with gradient descent. The Hessian is a key object in optimization theory [21], information geometry [22, 23], 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 1, 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 Super Consistency . The choice of the preconditioning factor γ 2 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 γ 2 set according to the corresponding parametrization.

More concretely, 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:

At realistically large N , S N ( t ) does not deviate significantly from its limit S ∞ ( t ) := lim N →∞ S N ( t ) . This is what is referred to as consistency in Vyas et al. [9].

S N ( t ) does not accumulate significant finite-size effects over time, i.e. the curves of S N ( t ) and S ∞ ( 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

′

glyph[arrowdblbothv]

glyph[triangle]

glyph[turnstileright]

⇐

glyph[triangle]

the other hand, the sharpness dynamics for µ P qualitatively exhibit little-to-no deviations. Also, notice that we assume existence of the limit lim N →∞ S N ( t ) . For those parametrizations (e.g. standard parametrization [6]) that do not have a well-defined limit, S N ( 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 µ P (and Depthµ P), Hessian eigenvalues are super consistent along the optimisation trajectory. Smaller eigenvalues have progressively different dynamics.

In Fig. 2 (a), we train a residual network on CIFAR-10 (a 10 classes image classification task) using cross-entropy loss, and show super consistent dynamics of three of the top k = 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 a small subspace where the gradient lies in the space spanned by top k Hessian eigenvectors (where k is the number of classes). Thus, our results show that the curvature along the training trajectory is preserved super consistently at different scales, thus suggesting that the geometry 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 trajectory regardless of the tiny-subspace assumption, in Appendix I we track the directional sharpness , that measures the curvature along the gradient direction.

To give a quantitative measure to the finite-size accumulation property, we measure deviations over time by estimating the following quantity:

When g ( t ) increases over time (up to fluctuations), Super Consistency is violated. We illustrate this in Fig. 3 (b, c), where we compute the left hand side of Eq. 3 for the loss L ( f t ) and the NTK's largest eigenvalue λ max (Θ) . To estimate the infinite width limit, we use a very large-width model as a proxy. Notice how under µ P the the loss dynamics progressively diverge from the infinite width model, indicating a finite-size accumulation over time. The same holds for λ max (Θ) . To study the rate of divergence g ( t ) , we fit a power law of the form y = at β to the observations. A larger β indicates a higher divergence rate. Notice how β > 0 . 6 for the loss, and β ≈ 2 for λ max (Θ) , indicating quadratic divergence. In comparison, in Fig. 3 (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 in fact width-independent . This might be due to the fact that the threshold is a stable attractor of the dynamics [28].

′

glyph[triangle]

∞ ′

∞

GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞

∞ ′

♣glyph[arrowdblbothv]

§

§

⇒ ♣

∞ ↙

∞ ↙∈∞

(a) Sharpness vs Time -

µ

P

∞ ′

∞

GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞

∈GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

♣⇐

⇒ ♣

↙ ̸ ∈

↙ ̸ ∈

(b) Loss vs Time -

µ

P

GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

♣glyph[arrowdblbothv]

§

§

⇒ ♣

W 〉glyph[ceilingleft]glyph[unionsq]〈

glyph[owner]∈

∞∈∀

glyph[triangleinv]∞∈

∈↙ ∈glyph[triangleinv]

∞ ↙

(c) NTK

λ

(Θ)

vs Time -

µ

P

glyph[triangle]

′

̸

′

↦

′

∀

′

∃

′

̸

′

↦

′

∀

′

∃

′

Figure 3: (a) Convergence rate of the sharpness at finite width N to the infinite limit proxy. Note that the distance approaches 0 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 is width 4096 , model is ConvNet, τ = 0 , η 0 = 0 . 7 .

4 Super Consistency and Learning Rate Transfer

Sharpness and Edge of Stability. We now focus on the sharpness λ := λ max ( γ 2 H ) , defined as the largest eigenvalue of the Hessian. In the theory of smooth convex [21], nonconvex [29], and stochastic [30] optimization, the sharpness plays a crucial role in in the guarantees of convergence

N

f

⇐

glyph[turnstileright]

glyph[arrowdblbothv]

f

⇐

glyph[turnstileright]

glyph[triangleinv]

↦

′

↙

∃

N

f

⇐

f

glyph[owner]

′

′

↙

∀

↦

′

N

⇐

glyph[turnstileright]

glyph[arrowdblbothv]

⇐

glyph[turnstileright]

glyph[triangleinv]

glyph[Ifractur]

glyph[Ifractur]

glyph[Ifractur]

∃

̸

glyph[owner]

of gradient methods and selection of the optimal step size. For instance, for a quadratic objective, λ t = λ is constant and gradient descent would diverge if the learning rate satisfies η 0 > 2 /λ , and training speed is maximized for η 0 = 1 /λ (LeCun et al. [2], page 28). Beyond this classical example, which assumes constant Hessian, the descent lemma [21] states that L ( θ t +1 ) ≤ L ( θ t ) if η ≤ 2 β where β := sup θ ∥∇ 2 L ( θ ) ∥ 2 , and ∥∇ 2 L ( θ ) ∥ 2 is the sharpness at θ . When it comes to deep neural networks, λ t is generally observed to increase during training ( progressive sharpening ): in the early phase of training it increases [31, 32] and then it decreases close to convergence [10]. Under full batch gradient descent training, the sharpness consistently rises above the EoS threshold of 2 /η 0 [16].

Conditions for hyperparameter transfer. The empirical success of hyperparameter transfer crucially relies on the following two observations, constituing a 'theoretical puzzle' [8].

The optimal learning rate is preserved across widths/depths, indicating very fast convergence with respect to the scaling quantity.

The models show consistent improvement in training speed with respect to the scaling quantity (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 λ max 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 µ P (and Depthµ P), the sharpness λ is super consistent along the training trajectory, while for NTP the sharpness decreases in width. This correlates with presence/absence of hyperparameter transfer.

In Fig. 1 we train a two-layer convolutional network under the µ P 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 λ . Notice how the sharpness' behaviour is qualitatively different in the two parameterizations: in µ P it reaches a width-independent value which is close to the EoS threshold of 2 /η 0 . On the other hand, in NTP we observe a progressive diminishing of the sharpness with width, as previously observed for Mean-Square-Error loss by Cohen et al. [16].

We then study the effect of depth under the Depthµ P model of Eq. 1. In Fig. 4 (left), we 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 parameterizations that do not exhibit learning rate transfer. In particular, we study the Depthµ P scaling of 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 are not super consistent, in that they accumulate finite-size effects over time. We study other parametrizations, including those without a stable limit in Appendix B, showing compatible results with those presented here. The observation that sharpness dynamics exhibit greater consistency compared to loss dynamics suggests that under µ P 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.

4.1 Feature Learning and Progressive Sharpening

Wenow study the effect of feature learning in the sharpness dynamics. Following the Gauss-Newton decomposition [25, 12], the Hessian can be decomposed as a sum of two matrices H = G + R , where

̸

glyph[triangleinv]

T

glyph[triangle]

glyph[triangleinv]

′

′

glyph[triangleinv]

′

glyph[triangleinv]

∞ ′

∞

∈GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∈

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

∞ ′

∞

∞ ′

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ′

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

∞∃∈

∞ ′

∞ ′

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ′

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

∞′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞

GLYPH<215>

∞ ′

∈↙ glyph[triangleinv]

GLYPH<215>

∞ ′

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

glyph[arrowbothv]∇

′↙∞∞

′↙∞∃

′↙glyph[owner]∈

(a) Depth-

µ

P ConvNet

k

= 1

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

∞∃∈

glyph[arrowbothv]∇

∞↙glyph[triangleinv]↦

(b) Depth-

µ

P ConvNet

k

= 2

∞′

∞glyph[triangleinv]

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞ ′

∞

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

glyph[arrowbothv]∇

(c) Depth-

µ

P ViT

glyph[owner]

′

glyph[triangle]

′

′

glyph[triangle]

′

̸

′

∀

′

′

G is the Gauss-Newton (GN) matrix and R depends on the Hessian of the model. For MSE loss,

where K ∈ R |D|× P is a matrix where each row is ∇ θ f ( x i ) (i.e. the Jacobian of f ( x i ) ), and y i ∈ R is the label. One can readily see that the NTK matrix can be written as Θ( f θ ) = KK ⊤ , thus the NTK and G share the same nonzero eigenvalues. In Figure 5, we show that under µ P 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 is fixed for NTP, while it evolves with time for µ 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 in 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 precisely in a simplified setting, providing further intuition on the sharpness dynamics and learning rate transfer.

P

N

T

P

′↙∈

∞↙′

∞ ′

∞ ′

∈

∞ ′

∞

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞′

∞′

∞ ′

∈

∞ ′

∞

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞′

∞′

⇐

′

↙

glyph[triangle]

′

↙

̸

′

↙

∀

′

↙

glyph[owner]

′

↙

glyph[owner]

glyph[triangleinv]

′

↙

glyph[triangle]

′

↙

glyph[triangle]

glyph[triangleinv]

⇐

Large scale 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-

⇐

glyph[owner]

glyph[triangle]

̸

glyph[triangle]

∀

∃

̸

̸

glyph[triangle]

∀

∃

̸

̸

T

glyph[triangle]

′

′

glyph[triangleinv]

′

glyph[triangleinv]

′

⇐

′

̸

glyph[triangle]

∀

∃

̸

glyph[owner]

̸

glyph[triangle]

∀

∃

̸

∃

∀

T

↦

̸

′

′

glyph[triangle]

glyph[triangleinv]

′

glyph[triangleinv]

′

glyph[owner]

̸

glyph[triangle]

∀

∃

̸

′

↙

glyph[triangle]

̸

̸

glyph[triangle]

∀

∃

̸

ImageNet, Wikitext), architectures (ViT, GPT-2 [33]), and optimizers (Adam [34] and AdamW [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 at 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 µ P this transition point is widthdependent, 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. 1 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, a close-to-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 L 2 loss. Our purpose is to characterize the dynamics of µP and NTP at the edge of stability through the lens of a simple example that shares a similar phenomenology with the more complex scenarios observed in the last section (see Fig. 10, App. C). In particular, the theory justifies the preconditioned Hessian γ 2 H ∝ NH as the right object of study when it comes to the sharpness computations (see Prop. 5.3). Also, it provides an intuition to the width-independent 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 a single datapoint or to vanishing targets 1 .

Notation and assumptions. Consider a dataset of |D| datapoints in D dimensions X ∈ R |D|× D ( |D| ≥ D ), and labels generated through a latent ground-truth vector w ∗ ∈ R D , that is Y = Xw ∗ . The neural network we use here is parametrized by weights W 0 ∈ R D × N , W 1 ∈ R N × 1 , where N is the width. To simplify the notation in our setting, we name E := W 0 and V := W 1 : f ( X ) = 1 γ √ ND XEV , L ( E,V ) = 1 2 ∥ f ( X ) -Y ∥ 2 . We initialize each entry of E,V i.i.d. Gaussian with mean zero and variance 1 . Recall that γ NTP = 1 , γ µP = √ N . We train with gradient descent (GD) with a learning rate η = η 0 γ 2 . Empirically, we observed (Fig. 10, App. C) that picking |D| = D and data X = I D ( I D 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 to the simple hypothesis class. The loss function reduces to:

Finally, we reparametrize the model by defining:

Note that using a learning rate η 0 γ 2 when optimizing L is equivalent to using a learning rate η 0 when optimizing γ 2 L . Next, we characterize how e, v evolve through time, and give conclusions for µ P.

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 + D 2 ) dimensions. This property is surprising because the original dynamical system described by GD on the variables E,V lives in N ( D + 1) dimensions. Concretely, this means we can 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 η = η 0 γ 2 on the loss of Eq. 4. The evolution of ( w,e, v ) is completely described by the following self-contained equation:

let the + denote updated quantities,

While the system above describes the evolution laws ( w k , e k , v k ) → ( w k +1 , e k +1 , v k +1 ) , the dynamics are influenced also by initialization. In Prop. C.1 in the Appendix, we 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 as the NTK Θ , which only depends on e and v .

Lemma 5.2 (GN bound) . Let γ 2 ∇ 2 L = G + R be Gauss-Newton decomposition 2 (see Sec. 4.1) of the Hessian for the loss in Eq. 4, with G = K ⊤ K , where K ∈ R D × ( ND + N ) . Let us denote the NTK matrix Θ = KK ⊤ ∈ R D × D . Then

and

We stress that this result implies that evolution of the NTK (i.e. 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.

Proposition 5.3 (EoS) . Let ( E,V ) evolve with GD with stepsize η = η 0 γ 2 on the loss of Eq. 4 towards a minimizer ( E ∗ , V ∗ ). Assume the corresponding solution in latent space ( w ∗ , e ∗ , v ∗ ) is marginally stable 3 . Then, λ max [ γ 2 ∇ 2 L ( E ∗ , V ∗ )] ∈ [ 2 η 0 , 2 η 0 + η 0 γ 2 ∥ w ∗ ∥ 2 ND ] .

Implications for NTP. Consider γ = 1 in Thm. 5.1. The dynamics of ( w,e, v ) are widthdependent . Let us take N → ∞ in the equation above to amplify this effect: the system becomes linear

While w evolves from w 0 as expected from standard NTK theory [4], e, v stay clamped at initialization. Applying Lemma 5.2 with γ = O (1) , the Hessian and the NTK have the same largest eigenvalue at large width (at rate O ( √ N ) ). This agrees with our intuition, as under NTP the predictor converges to a 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. 1 and throughout all our experiments). This derivation is also in line with our discussion in Sec. 4.1: we only have sharpening under feature learning, and for the same reason we cannot observe NTP at the edge of stability as N → ∞ (see stepsize dependency in Prop. 5.3), as 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.1.

Corollary 5.4. Consider µP ( γ = √ N ) and let ( E,V ) evolve with GD with stepsize η = η 0 γ 2 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, v ) is instead width-dependent, yet the error from N → ∞ case scales in expectation like 1 / √ N .

The corollary shows that µ P trajectories at different widths align in the latent space ( w,e, v ) , albeit with a vanishing perturbation in the initial condition (see Prop. C.1). While NTP's dynamics for e and v become slower as the width increases, for µ P their evolution laws are width independent. This implies that if the dynamics converge towards a minimizer for e and v , this will be at the sharpness value predicted by Prop. 5.3. Under µ P, where γ 2 ∝ N , this value will be width-independent , as Super 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 λ max (Θ) and the sharpness of the order of the residual term 1 / √ D ∥ w -w ∗ ∥ (Lemma 5.2). Finally, we stress that Prop. 5.3 prescribes the right scaling for the Hessian by including the preconditioning factor of γ 2 . Thus, we do not prove that at any finite time, the whole sharpness trajectory is width-independent, nor we are 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 transfer. 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 at 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 progressive sharpening and NTK evolution in the early phase of training. However, Figure 3 (b), shows how the NTK eigenvalues 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 are at play after progressive sharpening, such as Self-Stabilization [38]. In fact, progressive sharpening 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 a constant step size η 0 . However, an alternative is to use a step-size tuner, i.e. to automatically choose η 0 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 our results on 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 we 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 a larger scale than presented here. It would be interesting to see if Super Consistency still holds at an even larger scale.

References

W. X. Zhao, K. Zhou, J. Li, T. Tang, X. Wang, Y. Hou, Y. Min, B. Zhang, J. 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 arXiv:2303.18223 (2023). URL: http://arxiv.org/abs/2303.18223 .

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.

K. He, X. Zhang, S. Ren, and J. Sun. 'Delving deep into rectifiers: Surpassing human-level performance on imagenet classification'. In: Proceedings of the IEEE international conference on computer vision . 2015, pp. 1026-1034.

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

G. Yang and E. J. Hu. 'Tensor programs iv: Feature learning in infinite-width neural networks'. In: International Conference on Machine Learning . PMLR. 2021, pp. 11727-11737.

G. Yang, E. J. Hu, I. Babuschkin, S. Sidor, X. Liu, D. Farhi, N. Ryder, J. Pachocki, W. Chen, and J. Gao. 'Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer'. In: arXiv preprint arXiv:2203.03466 (2022).

B. Bordelon, L. Noci, M. B. Li, B. Hanin, and C. Pehlevan. 'Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit'. In: arXiv preprint arXiv:2309.16620 (2023).

G. Yang, D. Yu, C. Zhu, and S. Hayou. 'Tensor Programs VI: Feature Learning in InfiniteDepth Neural Networks'. In: arXiv preprint arXiv:2310.02244 (2023).

N. Vyas, A. Atanasov, B. Bordelon, D. Morwani, S. Sainathan, and C. Pehlevan. 'FeatureLearning Networks Are Consistent Across Widths At Realistic Scales'. In: arXiv preprint arXiv:2305.18411 (2023).

L. Sagun, L. Bottou, and Y. LeCun. 'Eigenvalues of the hessian in deep learning: Singularity and beyond'. In: arXiv preprint arXiv:1611.07476 (2016).

L. Sagun, U. Evci, V. U. Guney, Y. Dauphin, and L. Bottou. 'Empirical analysis of the hessian of over-parametrized neural networks'. In: arXiv preprint arXiv:1706.04454 (2017).

J. Martens. 'New insights and perspectives on the natural gradient method'. In: The Journal of Machine Learning Research 21.1 (2020), pp. 5776-5851.

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.

A. Orvieto, J. Kohler, D. Pavllo, T. Hofmann, and A. Lucchi. 'Vanishing curvature and the power of adaptive methods in randomly initialized deep networks'. In: arXiv preprint arXiv:2106.03763 (2021).

A. Lewkowycz, 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).

J. M. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. 'Gradient descent on neural networks typically occurs at the edge of stability'. In: arXiv preprint arXiv:2103.00065 (2021).

G. Gur-Ari, D. A. Roberts, and E. Dyer. 'Gradient descent happens in a tiny subspace'. In: arXiv preprint arXiv:1812.04754 (2018).

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

J. Lee, L. Xiao, S. 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).

B. Bordelon and C. Pehlevan. 'Self-consistent dynamical field theory of kernel evolution in wide neural networks'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 32240-32256.

Y. Nesterov. Introductory lectures on convex optimization: A basic course . Vol. 87. Springer Science & Business Media, 2013.

S. Amari. 'Information geometry'. In: Contemporary Mathematics 203 (1997), pp. 81-96.

S.-I. Amari. 'Natural gradient works efficiently in learning'. In: Neural computation 10.2 (1998), pp. 251-276.

B. Ghorbani, S. Krishnan, and Y. Xiao. 'An investigation into neural net optimization via hessian eigenvalue density'. In: International Conference on Machine Learning . PMLR. 2019, pp. 2232-2241.

J. Martens. Second-order optimization for neural networks . University of Toronto (Canada), 2016.

J. Martens and R. Grosse. 'Optimizing neural networks with kronecker-factored approximate curvature'. In: International conference on machine learning . PMLR. 2015, pp. 2408-2417.

A. Botev, H. Ritter, and D. Barber. 'Practical Gauss-Newton optimisation for deep learning'. In: International Conference on Machine Learning . PMLR. 2017, pp. 557-565.

D. S. Kalra, 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.02076 (2023).

G. Garrigos and R. M. Gower. 'Handbook of convergence theorems for (stochastic) gradient methods'. In: arXiv preprint arXiv:2301.11235 (2023).

R. M. Gower, N. Loizou, X. Qian, A. Sailanbayev, E. Shulgin, and P. Richtárik. 'SGD: General analysis and improved rates'. In: International conference on machine learning . PMLR. 2019, pp. 5200-5209.

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 arXiv:1807.05031 (2018).

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.09572 (2020).

A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. 'Language models are unsupervised multitask learners'. In: OpenAI blog 1.8 (2019), p. 9.

D. P. Kingma and J. Ba. 'Adam: A method for stochastic optimization'. In: arXiv preprint arXiv:1412.6980 (2014).

I. Loshchilov. 'Decoupled weight decay regularization'. In: arXiv preprint arXiv:1711.05101 (2017).

M. Song and C. Yun. 'Trajectory alignment: understanding the edge of stability phenomenon via bifurcation theory'. In: arXiv preprint arXiv:2307.04204 (2023).

E. Ott. Chaos in dynamical systems . Cambridge university press, 2002.

A. Damian, E. Nichani, and J. D. Lee. 'Self-stabilization: The implicit bias of gradient descent at the edge of stability'. In: arXiv preprint arXiv:2209.15594 (2022).

V. Roulet, A. Agarwala, and F. Pedregosa. 'On the Interplay Between Stepsize Tuning and Progressive Sharpening'. In: OPT 2023: Optimization for Machine Learning . 2023.

M. Claesen and B. D. Moor. Hyperparameter Search in Machine Learning . 2015. arXiv: 1502.02127 [cs.LG] .

J. Bergstra and Y. Bengio. 'Random search for hyper-parameter optimization.' In: Journal of machine learning research 13.2 (2012).

K. Jamieson and A. Talwalkar. 'Non-stochastic best arm identification and hyperparameter optimization'. In: Artificial intelligence and statistics . PMLR. 2016, pp. 240-248.

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

J. Snoek, H. Larochelle, and R. P. Adams. 'Practical bayesian optimization of machine learning algorithms'. In: Advances in neural information processing systems 25 (2012).

J. Snoek, 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: International conference on machine learning . PMLR. 2015, pp. 2171-2180.

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.

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.

D. Maclaurin, D. Duvenaud, and R. Adams. 'Gradient-based hyperparameter optimization through reversible learning'. In: International conference on machine learning . PMLR. 2015, pp. 2113-2122.

D. Yogatama and G. Mann. 'Efficient transfer learning method for automatic hyperparameter tuning'. In: Artificial intelligence and statistics . PMLR. 2014, pp. 1077-1085.

D. Stoll, J. K. Franke, D. Wagner, S. Selg, and F. Hutter. 'Hyperparameter transfer across developer adjustments'. In: arXiv preprint arXiv:2010.13117 (2020).

V. Perrone, R. Jenatton, M. W. Seeger, and C. Archambeau. 'Scalable hyperparameter transfer learning'. In: Advances in neural information processing systems 31 (2018).

S. Horváth, A. Klein, P. Richtárik, and C. Archambeau. 'Hyperparameter transfer learning with adaptive complexity'. In: International Conference on Artificial Intelligence and Statistics . PMLR. 2021, pp. 1378-1386.

G. Yang and E. Littwin. 'Tensor programs ivb: Adaptive optimization in the infinite-width limit'. In: arXiv preprint arXiv:2308.01814 (2023).

B. Bordelon and C. Pehlevan. 'Dynamics of Finite Width Kernel and Prediction Fluctuations in Mean Field Neural Networks'. In: arXiv preprint arXiv:2304.03408 (2023).

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 arXiv:2305.07810 (2023).

S. Yaida. 'Meta-Principled Family of Hyperparameter Scaling Strategies'. In: arXiv preprint arXiv:2210.04909 (2022).

I. Goodfellow, Y. Bengio, and A. Courville. Deep Learning . http : / / www . deeplearningbook.org . MIT Press, 2016.

X. Zhu, Z. Wang, X. Wang, M. Zhou, and R. Ge. 'Understanding edge-of-stability training dynamics with a minimalist example'. In: arXiv preprint arXiv:2210.03294 (2022).

S. Arora, Z. Li, and A. Panigrahi. 'Understanding gradient descent on the edge of stability in deep learning'. In: International Conference on Machine Learning . PMLR. 2022, pp. 9481024.

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.

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 at the edge of stability'. In: arXiv preprint arXiv:2207.14484 (2022).

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.

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 operations applications . Vol. 11006. SPIE. 2019, pp. 369-386.

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

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.

R. M. Neal. 'Bayesian Learning for Neural Networks'. PhD thesis. University of Toronto, 1995.

J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and J. Sohl-Dickstein. 'Deep neural networks as gaussian processes'. In: arXiv preprint arXiv:1711.00165 (2017).

A. Garriga-Alonso, C. E. Rasmussen, and L. Aitchison. 'Deep convolutional networks as shallow gaussian processes'. In: arXiv preprint arXiv:1808.05587 (2018).

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.

G. Yang. 'Tensor programs ii: Neural tangent kernel for any architecture'. In: arXiv preprint arXiv:2006.14548 (2020).

L. Chizat, E. Oyallon, and F. Bach. 'On lazy training in differentiable programming'. In: Advances in neural information processing systems 32 (2019).

L. Chizat and F. Bach. 'On the global convergence of gradient descent for over-parameterized models using optimal transport'. In: Advances in neural information processing systems 31 (2018).

L. Chizat and F. 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. 13051338.

S. Mei, T. Misiakiewicz, and A. Montanari. 'Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit'. In: Conference on Learning Theory . PMLR. 2019, pp. 2388-2464.

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

L. Noci, S. Anagnostidis, L. Biggio, A. Orvieto, S. P. Singh, and A. Lucchi. 'Signal propagation in transformers: Theoretical perspectives and the role of rank collapse'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 27198-27211.

S. Hayou and G. Yang. 'Width and Depth Limits Commute in Residual Networks'. In: arXiv preprint arXiv:2302.00453 (2023).

S. Hayou. 'On the infinite-depth limit of finite-width neural networks'. In: arXiv preprint arXiv:2210.00688 (2022).

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.

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.

M. Li, M. Nica, and D. Roy. 'The future is log-Gaussian: ResNets and their infinite-depthand-width limit at initialization'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 7852-7864.

M. Li, M. Nica, and D. Roy. 'The neural covariance SDE: Shaped infinite depth-and-width networks at initialization'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 10795-10808.

L. Noci, C. Li, M. B. Li, B. He, T. Hofmann, C. Maddison, and D. M. Roy. 'The shaped transformer: Attention models in the infinite depth-and-width limit'. In: arXiv preprint arXiv:2306.17759 (2023).

M. B. Li and M. Nica. 'Differential Equation Scaling Limits of Shaped and Unshaped Neural Networks'. In: arXiv preprint arXiv:2310.12079 (2023).

B. Hanin and M. Nica. 'Finite depth and width corrections to the neural tangent kernel'. In: arXiv preprint arXiv:1909.05989 (2019).

S. Yaida. 'Non-Gaussian processes and neural networks at finite widths'. In: Mathematical and Scientific Machine Learning . PMLR. 2020, pp. 165-192.

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.

L. Chizat and P. Netrapalli. 'Steering Deep Feature Learning with Backward Aligned Feature Updates'. In: arXiv preprint arXiv:2311.18718 (2023).

A. Agarwala, F. Pedregosa, and J. Pennington. 'Second-order regression models exhibit progressive sharpening to the edge of stability'. In: arXiv preprint arXiv:2210.04860 (2022).

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.

R. Courant. 'Über die Eigenwerte bei den Differentialgleichungen der mathematischen Physik'. In: Mathematische Zeitschrift 7.1 (1920), pp. 1-57.

B. N. Parlett. The symmetric eigenvalue problem . SIAM, 1998.

K. Osawa, S. Ishikawa, R. Yokota, S. Li, and T. Hoefler. 'ASDL: A Unified Interface for Gradient Preconditioning in PyTorch'. In: arXiv preprint arXiv:2305.04684 (2023).

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.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. 'An image is worth 16x16 words: Transformers for image recognition at scale'. In: arXiv preprint arXiv:2010.11929 (2020).

B. Bordelon, H. T. Chaudhry, and C. Pehlevan. 'Infinite Limits of Multi-head Transformer Dynamics'. In: arXiv preprint arXiv:2405.15712 (2024).

Appendix

Table of Contents

A Related works

Hyperparameter search. Hyperparameter tuning [40] has been paramount in order to 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 a range of learning rates, even with improvements 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 iteration, 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 provide 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 transfer through width and depth, termed µ P. Bordelon et al. [7] used the DMFT framework [20, 54] to derive the Depthµ P 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 µ P 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 optimization 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 a 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 et 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 initialization 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 is scaled up. In the two-layer case, the network's infinite-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 / √ depth-scaled 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 in Hayou and Yang [77] and Hayou [78] at initialization and in Bordelon et al. [7], Yang et al. [8] for the training dynamics. Without the 1 / √ depth-scaling, 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 a 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 is 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

learning ( α ∈ [1 / 2 , 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, a 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 Vyas 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 a simple non linear model. Our results are also in line with the recent experiments on the evolution of the sharpness for ReLU 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 ( α ∈ [1 / 2 , 1] in Eq. 1, 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 et 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, we expect that a scaling limit that admits learning rate transfer would find a corresponding width/depth-independent behaviour in the sharpness dynamics.

A.2 Sharpness and Optimal Step Size

One of the underlying assumptions to explain why super consistent sharpness causes hyperparameter transfer is that the sharpness influences the optimal learning rate. Although this is well-established in classic optimization [2], Edge of Stability tells a story of a 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 a 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 /λ t 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 c arch /λ 0 , where c arch depends on the architecture. On a 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 / √ depth scaling of the residual branches

We first study the effect of increasing the depth in a model without the 1 / √ L -scaling of the residual branches, which results in exploding activations (and sharpness) as the depth increases. See Figure

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 lr: 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 a linear warmup scheduler in the first phase of training, the network is 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 a sustained period of depth-independent 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.

10

2

10

1

10 0

Learning Rate

1.7 × 10 0

1.8 × 10 0

1.9 × 10 0

2×10 0

2.1 × 10 0

2.2 × 10 0

2.3 × 10 0

Train Loss

Depth

3

6

12

24

36

10

2

10

1

Learning Rate

10 0

Train Loss

Epoch 20

Depth

3

6

12

24

36

0

50

100

150

200

Step

10 0

10 1

10 2

10 3

Depth

3

6

12

24

lr

0.068129

0.5275

10

2

10

1

10 0

Learning Rate

1.8 × 10 0

1.9 × 10 0

2×10 0

2.1 × 10 0

2.2 × 10 0

2.3 × 10 0

Train Loss

Depth

3

6

12

24

36

10

2

10

1

Learning Rate

10 0

Train Loss

Epoch 20

Depth

3

6

12

24

36

0

50

100

150

200

Step

10 0

10 1

10 2

10 3

10 4

Depth

3

6

12

24

36

lr

0.068129

0.316228

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 η . While Yang and Hu [5] use a width scaled learning rate in order to parameterize SP as a 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 µ P - disabling residuals

We also provide training runs of models parameterized with µ P, 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.

10

2

10

1

Learning Rate

10 0

8×10 1

9×10 1

Train Loss

32

64

128

256

512

10

2

10

1

Learning Rate

10 0

7×10 1

8×10 1

9×10 1

Train Loss

32

64

128

256

512

10

2

10

1

Learning Rate

10 0

7×10 1

8×10 1

9×10 1

Train Loss

32

64

128

256

512

0

100

200

300

Step

10 2

10 3

32

64

128

256

512

lr

0.014678

0.024484

0.040842

0.068129

0

100

200

300

Step

10 2

10 3

32

64

128

256

512

lr

0.024484

0.040842

0.068129

0.113646

0

100

200

300

Step

10 1

10 2

10 3

32

64

128

256

512

lr

0.040842

0.068129

0.113646

0.189574

10

2

10

1

10 0

10 1

Learning Rate

10 0

2×10 0

Train Loss

Depth

3

6

12

24

0

250

500

750

1000

Step

10

4

10

3

10

2

10

1

10 0

10 1

10 2

Depth

3

6

12

24

lr

0.014678

0.5275

B.4 µ P experiments for full batch GD

In this section, we investigate the effect of µ P and on learning rate transfer and sharpness, when trained with full batch gradient descent. We subsampled 5000 sampled from CIFAR10 in a stratified fashion (i.e. 500 sampled from each of the 10 classes) and proceeded to train residual ConvNets under the µ P parameterization with GD, following a similar procedure as [16]. Our findings show that under µ P, when using a large enough learning rate, the models are able to achieve edge of stability, as 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 EoS threshold studied by [16], it still does maintain a width-independent trend during training.

T

glyph[triangle]

̸

∀

′

′

′

′

′

′

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ′

∞

∞ ′

∞ ′

∞ ↙

∈GLYPH<215>

∞ ′

∞ ↙

GLYPH<215>

∞ ′

∞ ↙

GLYPH<215>

∞ ′

∞ ↙

GLYPH<215>

∞ ′

∈GLYPH<215>

∞ ′

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞̸

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

∞′

∞ ′

∞

∞ ′

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞̸

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

glyph[arrowbothv]∇

∞′

∞ ′

∈GLYPH<215>

∞ ′

GLYPH<215>

∞ ′

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞̸

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

glyph[arrowbothv]∇

′

glyph[owner]

′

glyph[triangle]

′

E√glyph[wreathproduct]glyph[floorright]〈

′

glyph[owner]

′

glyph[triangle]

′

E√glyph[wreathproduct]glyph[floorright]〈

C Analysis of a Two-Layer Linear Network

Recall the definition of our model:

glyph[owner]

∞ ′

∞ ′

∞ ′

∞

∞ ′

∞

∞ ′

NTP←glyph[arrowbothv]glyph[wreathproduct]∫∫

N←glyph[Ifractur]∞′

∞ ′

∞ ′

∞ ′

∞

∞ ′

∞

∞ ′

∫∫

N←glyph[Ifractur]∞′

′

glyph[triangle]

′

〉glyph[unionsq]glyph[ceilingright]∇glyph[turnstileright]glyph[unionsq]〉glyph[wreathproduct]\

′

glyph[triangle]

′

P

glyph[arrowbothv]

glyph[wreathproduct]

〉glyph[unionsq]glyph[ceilingright]∇glyph[turnstileright]glyph[unionsq]〉glyph[wreathproduct]\

Under the previous assumptions regarding the data, we have that:

C.1 Relationship between sharpness and residual, and Gauss-Newton

The following bound leverages a Gauss-Newton decomposition and leads analytical insights supporting Sec. 4.1. Proof of Lemma 5.2, which we restate here.

Lemma (GN bound) . Let γ 2 ∇ 2 L = G + R be Gauss-Newton decomposition 4 (see Sec. 4.1) of the Hessian for the loss in Eq. 4, with G = K ⊤ K , where K ∈ R D × ( ND + N ) . Let us denote the NTK matrix Θ = KK ⊤ ∈ R D × D . Then

′

glyph[triangleinv]

glyph[owner]

glyph[owner]

̸

glyph[triangle]

N

glyph[Ifractur]

glyph[owner]

′

′

glyph[triangleinv]

glyph[owner]

N

glyph[Ifractur]

glyph[owner]

′

̸

glyph[triangle]

′

↙

↦

̸

∃

glyph[triangleinv]

glyph[owner]

T

′

′

′

̸

glyph[triangle]

′

↙

↦

̸

∃

glyph[triangleinv]

and

The result above is actually more general: it indicates that the eigenvalues of the positive semidefinite portion of the Hessian G = K ⊤ K are fully 5 characterized by the eigenvalues a smaller matrix Θ = KK ⊤ . Furthermore, in our model Θ has a closed form that depends only on the variables e, v .

Proof. The Hessian blocks become:

Using these definitions, we can separate the Hessian H into a sum of 2 matrices, where one depends on the residual and one does not.

hence our quantity of interest:

where w = 1 γ √ ND EV .

Study of G . Note that

By an SVD decomposition, it is easy to see that, the nonzero eigenvalues of G = K ⊤ K are the same as the nonzero eigenvalues of Θ = KK ⊤ :

where e, v are the quantities found in the main paper, and we therefore have

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 is easy to show that

Using the fact that G is Hermitian, we can apply Weyl's inequality to obtain a bound on the deviation of the sharpness from the maximum eigenvalue of G in terms of the residual:

Which finally yields:

C.2 Proof of Thm.5.1

We divide the proof into two parts. In the first part, we study the evolution of the unnormalized quantities EE ⊤ , V ⊤ V and EV . Then, we study how normalization affects the dynamics.

Part one: dynamics in a smaller space. We go step by step recalling the gradient descent equations at the beginning of this section.

Dynamics of EV . We have :

Let us rename

then the equation becomes more compact:

Note that no quantities appear in the equation for ˜ w + besides ˜ w, ˜ v, ˜ e . We will see that these do not appear also in the equations for ˜ v, ˜ e .

Dynamics of V ⊤ V . Let us write down here the equations for -η∂ V and -η∂ ⊤ V for ease of reference:

we have

Using our notation, equations get yet again simpler:

Note that again no quantities appear in the equation for ˜ v + besides ˜ w, ˜ v, ˜ e .

Dynamics of E ⊤ E . For convenience, recall:

we have

Using our notation, equations get yet again simpler:

Part two: Normalization. Consider scalar reparameterizations.

While we gave the form of these normalizers already in the paper, we keep it more general here to real numbers and show that the right normalizers arise directly.

Reparameterization of EV . We have

We would like α w , α e , α v to be such that on the right-hand side we have no width dependency. To do that we need (first and second line refer to first and second lines in the equation)

where proportionality can depend on any factor (e.g. d ) except of course N . Crucially note that the equations require γ ∝ √ N . So this can be done for µP but not for NTP (except if d ≃ N ) . Further, we need

To get that w → w ∗ , we choose (as in the main paper)

So, we get

Reparameterization of V ⊤ V . Recall that

This implies, after scaling

It is easy to see that the choice α w , α e , α v ∝ 1 N in addition with γ = √ N gives independence of the RHS to width.

Under our choices

we get

Reparameterization of EE ⊤ . Let's substitute the scaled version in the equations

With our choices

we get

C.2.1 Initialization

Proposition C.1. At initialization, as N → ∞ , e P → e ∞ := 1 D I D and v P → v ∞ := 1 D . Moreover, errors from ∞-initialization scale as 1 √ N in expectation: E | v -v ∞ | 2 , E | e ij -e ∞ i,j | 2 ≤ 2 ND , ∀ i, j ∈ [ D ] . While for γ = 1 (NTP) w at initialization is in the limit Gaussian with elementwise variance 1 /D , for γ = √ N ( µP ) we have w P → w ∞ := 0 , with elementwise variations scaling as 1 √ N : E | w i -w ∞ i | 2 = 1 ND , ∀ i ∈ [ D ] .

First, note that trivially and

Next:

and

For i = j

̸

For i = j

So

Finally:

and

̸

̸

̸

C.2.2 Proof of Prop. 5.3

EoS Basic Derivation on quadratic (global linear dynamics). In quadratic potentials 1 2 w ⊤ Hw , we are at EOS if and only if η = 2 /λ max ( H ) . The proof of this can be translated into a Jacobian argument on the gradient update map. The update map here is w + = ( I -ηH ) w , which has Jacobian I -ηH . The eigenvalues of this map are in the range [1 -λ max ( H ) η, 1 -λ min ( H ) η ] , where bounds are tight. For asymptotic stability (i.e. we converge), we have the necessary condition λ max ( H ) < 2 /η , which guarantees all eigenvalues are in the range ( -1 , 1) . To be at edge of stability means λ max ( H ) = 2 /η , that is if and only if the Jacobian has eigenvalues in [ -1 , 1] , with an eigenvalue equal to -1 .

Remark: What changes here? Weare dealing with a nonlinear transition map where the Jacobian has explicit O ( η 2 ) terms. To be precise in this setting, we need to carefully study the Jacobian and cannot assume its first-order approximation in η 0 , i.e. I -η 0 ∇ 2 L , is accurate.

Proof. We need to compute the Jacobian for the dynamical system G : ( w,e, v ) → ( w + , e + , v + ) . Note that - specifically at w = w ∗ ,

where we do not care about the values with a ⋆ because anyway the resulting matrix is lowertriangular:

As in the quadratic setting, to be at the edge of stability, we require this matrix to have eigenvalues in [ -1 , 1] , with at least an eigenvalue precisely equal to -1 . Without loss in generality, we discuss eigenvalues in descending order for each one of our matrices:

Note that J G ( w ∗ , e ∗ , v ∗ ) is block lower-triangular, hence its eigenvalues are

The necessary condition for being at the edge of stability is then

where Θ ∗ := v ∗ I + e ∗ . We then ask the question: for which combination of Θ ∗ and η 0 s can this happen? Note that for any eigenvalue index k (eigenvalues sorted in descending order), η 0 ≥ 0 and Θ ∗ ⪰ 0 :

Thus:

Let us set Γ ∗ = I -η 0 Θ ∗ - a matrix of which we know the eigenvalues. The matrix we want to study is a rank-1 O ( η 2 0 ) perturbation: Γ ∗ + P where P = η 2 0 γ 2 ND w ∗ ⊗ w ⊤ ∗ . Intuitively, this perturbation matrix P does not modify much the eigenvalues.

To start, note that the CourantFischer Theorem [91] implies (see e.g. [92], exercise 10.2.2.):

Next, recall Weyl's inequality for Hermitian (in this case, symmetric) matrices:

Apply this to j = 1 , i = D , we get

All in all, we proved the bound

Notice that this implies, given the structure of P ,

For EoS, we require λ D (Γ ∗ + P ) = -1 . Hence we necessarily need

Recall that λ D (Γ ∗ ) = λ D ( I -η 0 Θ ∗ ) = 1 -η 0 λ 1 (Θ ∗ ) , hence the condition is

This implies

Dividing everything by η 0 , hence

To conclude, note that λ 1 (Θ ∗ ) coincides with the Hessian λ 1 .

Note on marginal stability (sufficient conditions). At the same time, for stability, we need (sufficient condition)

The condition -1 ≤ λ D (Γ ∗ + P ) is implied by -1 ≤ λ D (Γ ∗ ) (Courant-Fischer). This yields

that is

Next, we need λ 1 (Γ ∗ + P ) ≤ 1 . A sufficient condition for this is (again by Weyl)

which implies

That is i.e.

that is

All in all a sufficient condition for marginal stability is

The bound is not vacuous in the case of N ≥ D and different data points.

Note. What we provided above is a sufficient condition for marginal stability (all eigenvalues in [ -1 , 1] ). This is derived just for interest of the reader but does not affect the result: indeed, a 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 η 0 ≤ λ 1 (Θ ∗ ) ≤ 2 η 0 + η 2 0 γ 2 ND ∥ w ∗ ∥ 2 .

C.3 Theory for Standard Parametrization (SP)

As in µ P and NTP, we consider the simplified loss

where E ∈ R D × N , V ∈ R N × 1 . Compared to µ P and NTP, normalizing factors appear in the initialization: E ij ∼ N (0 , 1 /D ) , V j ∼ N (0 , 1 /N ) , ∀ i, j . Gradients are ∂ E = EVV ⊤ -w ∗ V ⊤ , ∂ V = E ⊤ EV -E ⊤ w ∗ .

C.3.1 Dynamics equations, same form as µP and NTP

Let us rename w = EV, v = V ⊤ V, e = EE ⊤ , then the dynamics gradient descent with stepsize η becomes, in ( w,e, v ) ∈ R D + D 2 +1 space:

Hence, under SP, ( w,e, v ) at different widths have width-independent evolution laws in R D + D 2 +1 -same happens under µ P. However , in contrast to µ P, initialization of this system does not converge as N →∞ . Hence, actual dynamics across widths are drastically different, as we will see next.

C.3.2 Initialization

First, note that trivially E [ w i ] = 0 , and E [ w i ] 2 = ∑ N j,j ′ =1 E [ E ij E ij ′ V j V j ′ ] = ∑ N j =1 E [ E 2 ij V 2 j ] = 1 D . This makes sense since the forward pass is normalized. However, at initialization:

that already includes a dependency on the width N , further: E [ e 2 ij ] = ∑ N k,k ′ =1 E [ E ik E jk E ik ′ E jk ′ ] . Hence

̸

̸

So Var [ e ij ] = O ( N D 2 ) . Finally:

So Var [ v ] = O ( 1 N 2 ) .

̸

C.3.3 Conclusion

While the laws ( w,e, v ) → ( w + , e + , v + ) are not width dependent in SP, e 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, v ) under gradient descent starts at a width-dependent point ( O (1) , O ( N ) , O (1)) ; hence it is drastically different at different widths. Further, the NTK ( v · I D + e ) is also O ( N ) , and this controls the evolution of the forward pass w : w + = w -η ( v · I D + e )( w -w ∗ ) + O ( η 2 ) . We therefore validate also in this simple setting the derivation by [5]: if η = O (1) , forward pass of SP blows up after one step of gradient descent.

D Connection between the eigenvalues of the Hessian and NTK matrix

Recall that for MSE loss and one-dimensional output f ( x ) , the Gauss-Newton decomposition of the Hessian H reads:

where G is the Gaussian-Newton matrix. This can be generalized to (1) different loss functions and (2) multidimensional output f ( x ) ∈ R k , where k is the dimension of the logits (i.e. the number of classes in classification problems). Here we notice that in (2) we exactly preserve the connection between GGN and NTK's spectra, while in (1) we have an extra term in G 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 as the NTK as training progresses, and G still dominates R , as one would expect from the experiments in the main text (e.g Fig. 1). We begin by defining the Gauss-Newton matrix (GN) in the case of cross-entropy loss, following [12]:

where now ∇ θ f ( x i ) ∈ R kP , and ¯ H L ∈ R kP × kP is a block-diagonal matrix where the k × k Hessian of the loss ( H L ) with respect to model output is repeated P times. Again, we can stack the Jacobian vectors ∇ θ f ( x i ) for all the datapoints into K ∈ R |D|× kP , thus:

For MSE loss, ¯ H L 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 ∆ := ∇ f ( x ) L can be shown to be ∆ = σ ( f ( x )) -y , where y is the one hot vector encoding the true classes, and σ ( · ) denotes the softmax activation. Hence, for the Hessian, we have:

which in general deviates from the identity. However, during training the model increase the probability of correct predictions, and thus H L gets closer to the identity, thus having an asymptotically negligible effect on the Gauss-Newton matrix G and its correspondence to the NTK.

′

↙

∃

glyph[triangleinv]

̸

glyph[triangle]

̸

glyph[triangle]

∞ ′

∈

∞ ′

∞

∞ ′

∈

W 〉glyph[ceilingleft]glyph[unionsq]〈

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

∞ ′

∞

∞ ′

∈

∞↙′

∞↙′glyph[triangleinv]

∞↙∞′

∞↙∞glyph[triangleinv]

∈

W 〉glyph[ceilingleft]glyph[unionsq]〈

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

⇐

⇐

We now perform experiments to test whether G 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 is computed using a power iteration algorithm based on the implementation provided in [93]. Note that while for µ P we observe a 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).

′

↙

↦

′

↙

↦

glyph[triangleinv]

′

↙

∀

′

↙

∀

glyph[triangleinv]

′

↙

∃

′

′

↙

∃

′

↙

∃

glyph[triangleinv]

E Late-time dynamics and batch size ablations

E.1 Late-time dynamics

It was noted in [16] (Appendix C) that with cross-entropy loss (adopted here), the sharpness decreases towards the end of training. Here in Fig. 12, we show that while the dynamics are remarkably 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 in 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.

20

25

30

35

40

45

50

Epoch

10 1

2×10 0

3×10 0

4×10 0

6×10 0

0 =0.32

20

25

30

35

40

45

50

Epoch

10 0

0 =0.53

20

25

30

35

40

45

50

Epoch

10

1

10 0

0 =0.88

256

512

1024

2048

4096

8192

10

1

10 0

10 1

Learning Rate

10

3

10

2

10

1

10 0

Epoch 29

10

1

10 0

10 1

Learning Rate

10

3

10

2

10

1

10 0

Epoch 39

10

1

10 0

10 1

Learning Rate

10

4

10

3

10

2

10

1

10 0

256

512

1024

2048

4096

8192

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 is well-known to increase with the batch size, as it is shown also in Cohen et al. [16] and Jastrzkebski et al. [31].

0

10

20

30

40

50

Epoch

10 0

B=128

0

10

20

30

40

50

Epoch

10 0

B

=256

0

10

20

30

40

50

Epoch

10 0

B

=512

64

128

256

512

1024

2048

4096

lr

0.32

0.88

2.45

10

1

10 0

Learning Rate

10 0

2×10 0

Train Loss

B=128

64

128

256

512

1024

2048

4096

10

1

10 0

Learning Rate

10 0

1.1 × 10 0

1.2 × 10 0

1.3 × 10 0

1.4 × 10 0

1.5 × 10 0

1.6 × 10 0

1.7 × 10 0

Train Loss

B

=256

64

128

256

512

1024

2048

4096

10

1

10 0

Learning Rate

1.2 × 10 0

1.4 × 10 0

1.6 × 10 0

1.8 × 10 0

Train Loss

B

=512

64

128

256

512

1024

2048

4096

Here we add that under µ P, the batch size we tested ( 128 , 256 , 512 ) do not influence significantly the width-independence phenomenon of the sharpness, as we summarize in Fig. 13. We observe good

learning rate transfer across all the tested batch sizes. We also observe that the optimal learning rate increases with the batch sizes by roughly a 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 a stabilizing effect on the sharpness, delaying the late phase of training where the sharpness drops, as analyzed in Sec. E.1 [16, 28]. On the other hand, Thus, under regularization techniques such as data augmentation, we should expect better hyperparameter transfer.

0

10

20

30

40

50

Epoch

10 0

B

=64

0

10

20

30

40

50

Epoch

10 0

B

=256

0

10

20

30

40

50

Epoch

10 0

B

=512

64

128

512

1024

2048

4096

lr

0.32

0.88

2.45

10

1

10 0

10 1

Learning Rate

10

3

10

2

10

1

10 0

Train Loss

B

=64

10

1

10 0

10 1

Learning Rate

10

3

10

2

10

1

10 0

B

=256

10

1

10 0

10 1

Learning Rate

10

2

10

1

10 0

B

=512

64.0

128.0

512.0

1024.0

2048.0

4096.0

F Large-Scale experiments, more Datasets and Optimizers

We provide empirical validation of 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 to different modalities and architectures.

F.1 GPT-2 experiments on WikiText

Figure 15 shows the transfer of a GPT-2 model trained on WikiText for 40 epochs with Depthµ P and Adam, without learning rate scaling. A similar plot but for µ P (without the depth parameterization) is presented in Figure 16. Note that the sharpness exhibits a width independent behaviour.

10

4

10

3

Learning Rate

3×10 0

4×10 0

Train Loss

Depth

2

4

8

16

0

10

20

30

40

Epoch

10 1

10 2

10 3

Depth

2

4

8

16

Learning Rate

0.00037

0.0009

10

3

10

2

Learning Rate

10 0

Train Loss

128

256

512

1024

2048

(a)

µ

P Transfer

0

5

10

15

20

Epoch

10 1

10 2

10 3

128

256

512

1024

2048

lr

0.0026

(b) Sharpness

F.2 ConvNets Experiments on Larger Datasets and Adam(W)

Figures 17 and 18 show residual ConvNets parameterized with µ P and Depthµ P respectively trained on TinyImagenet, and Figure 19 shows the evolution of a similar model parameterized with µ P 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 µ P with Adam and AdamW in Figure 21.

10

2

Learning Rate

3.2 × 10 0

3.4 × 10 0

3.6 × 10 0

3.8 × 10 0

4×10 0

4.2 × 10 0

4.4 × 10 0

4.6 × 10 0

Train Loss

32

64

128

256

512

0

20

40

60

Step

10 0

10 1

Sharpness

32.0

64.0

128.0

256.0

512.0

Learning Rate

0.008799

0.014678

0.024484

10

2

Learning Rate

3.2 × 10 0

3.3 × 10 0

3.4 × 10 0

3.5 × 10 0

3.6 × 10 0

3.7 × 10 0

3.8 × 10 0

3.9 × 10 0

Train Loss

Depth

6

12

24

48

96

0

100

200

300

Step

10

1

3×10 2

4×10 2

6×10 2

Sharpness

Depth

6.0

12.0

24.0

48.0

96.0

Learning Rate

0.014678

0.024484

0.040842

10

3

Learning Rate

5.0

5.5

6.0

6.5

7.0

Train Loss

16

32

64

128

0

20

40

60

80

100

Step

10

1

10 0

10 1

Sharpness

16

32

64

128

Learning Rate

0.001565

0.002154

10

4

10

3

Learning Rate

1.1 × 10 0

1.15 × 10 0

1.2 × 10 0

1.25 × 10 0

1.3 × 10 0

1.35 × 10 0

1.4 × 10 0

1.45 × 10 0

1.5 × 10 0

Train Loss

Depth

6

12

24

48

96

10

4

10

3

Learning Rate

10 0

9×10 1

1.1 × 10 0

1.2 × 10 0

1.3 × 10 0

1.4 × 10 0

1.5 × 10 0

Train Loss

Depth

6

12

24

48

96

0

20

40

Epoch

10 0

10 1

Depth

6

12

24

48

96

lr

0.001233

(a) Depth-

µ

P with

k

= 1

0

20

40

Epoch

10 0

10 1

10 2

Depth

6

12

24

48

96

lr

0.000534

(b) Depth-

µ

P with

k

= 2

G Time Evolution of other Spectral Quantities

In this section we present convergence rates for various other spectral quantities of interest. In Figure 22, we show the same Super Consistency of the landscape in the case of Depthµ P, as 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 µ P and NTP regime. Note in Figure 24 (b) and (c) how the sharpness achieves a 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 to 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 15.

10

2

Learning Rate

10 0

8×10 1

9×10 1

Train Loss

128

256

384

512

0

5

10

15

20

Epoch

10 0

10 1

128

256

384

512

lr

0.0034

10

2

Learning Rate

10 0

8×10 1

9×10 1

Train Loss

128

256

384

512

(c)

µ

P Transfer

0

5

10

15

20

Epoch

10 0

10 1

128

256

384

512

lr

0.0034

(d) Sharpness

4

17

72

Epoch

10

3

10

2

10

1

|

(

f

)

(

fN

)|

0.2

0.21

0.42

10 1

2×10 1

Epoch

10

3

10

2

10

1

|

f

)

fN

)|

Depth

12

24

48

=1.09

=1.24

=2.11

10 1

6×10 0

2×10 1

Epoch

10

2

10

1

|

f

)

fN

)|

Depth

12

24

48

=0.42

=0.55

=0.69

10 1

4×10 0

6×10 0

Epoch

10

1

10 0

|

f

)

fN

)|

1.57

0.9

1.21

(a) Depth-

µ

P ConvNet

k

=

10 1

2×10 1

3×10 1 4 × 10 1

6×10 1

Epoch

10

5

10

4

10

3

10

2

|

f

)

fN

)|

Depth

12

24

48

=0.9

=0.91

=

0.02

(b) Depth-

µ

P ConvNet

k

= 2

10 1

6×10 0

2×10 1

Epoch

10

1

6×10 2

2×10 1

|

f

)

fN

)|

Depth

6

12

24

=0.18

=0.31

=0.36

(c) Depth-

µ

P ViT

1

H Sharpness evolution in Random Feature Models

In this section, we compare the NTP and µ P parameterizations to a random feature model, i.e. a model where we all the weights of the intermediate layers are frozen to their value at initialization, and only the final readout layer is trained. Crucially, this model does not learn features by construction 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 µ P. However, the transfer is better than in NTP. This is in line with our claim, as under a random feature model increasing the learning rate does not induce more feature learning, as is the case in NTP.

10 2

10 3

10 4

10 5

|

)

N

)|

Epoch

1

3

6

8

(a) Largest NTK eigenvalue vs Width -

µ

P

10 2

10 3

10

1

10 0

|

f

)

fN

)|

(b) Sharpness vs Width -

µ

P

10 1

10 2

10

3

10

2

10

1

10 0

|

(

f

)

(

fN

)|

0.47

0.77

0.65

0.64

(c) Loss vs Width -

µ

P

0

2

4

6

8

10

Epoch

10 0

10 1

10 2

tr

(

H

)

32

64

128

256

512

1024

2048

(d) Hessian trace vs Epoch -

µ

P

0

50

100

150

200

Step

1.4 × 10 0

1.6 × 10 0

1.8 × 10 0

2×10 0

2.2 × 10 0

2.4 × 10 0

Train Loss

P

8

16

32

64

128

256

512

1024

2048

0

50

100

150

200

Step

10

1

10 0

P

8

16

32

64

128

256

512

1024

2048

10 1

10 2

10 3

10 0

2×10 0

5×10 1

2.5 × 10

1

P

Step

16

32

64

128

0

50

100

150

200

Step

1.6 × 10 0

1.8 × 10 0

2×10 0

2.2 × 10 0

2.4 × 10 0

Train Loss

NTK

8

16

32

64

128

256

512

1024

2048

(a) Loss curves

0

50

100

150

200

Step

10

1

2×10 1

5×10 2

4×10 1

NTK

8

16

32

64

128

256

512

1024

2048

(b) Sharpness dynamics

10 2

10 3

10

1

2×10 1

5×10 2

4×10 1

NTK

Step

16

32

64

128

(c) Sharpness vs width

T

glyph[triangle]

glyph[ceilingright]

glyph[owner]

glyph[turnstileright]

〈

S

glyph[triangle]

̸

∀

⇐

′

′

′

′

′

′

′

′

glyph[triangleinv]

′

̸

′

↦

′

′

glyph[owner]

′

E√glyph[wreathproduct]glyph[floorright]〈

glyph[triangle]

′

glyph[triangleinv]

′

̸

′

↦

′

∞′

∞ ′

∈

∈

∞ ′

∞′

∞ ′

∈

∞ ′

∞ ′

glyph[owner]

′

E√glyph[wreathproduct]glyph[floorright]〈

glyph[triangle]

′

P

N

T

K

∞ ′

∈

∞ ′

∞

∞ ′

∞ ′

∞

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ↙

∈GLYPH<215>

∞ ′

∞ ↙

GLYPH<215>

∞ ′

∞ ↙

GLYPH<215>

∞ ′

∞ ↙

GLYPH<215>

∞ ′

∈GLYPH<215>

∞ ′

∈↙ ∈GLYPH<215>

∞ ′

∈↙ glyph[triangle]

GLYPH<215>

∞ ′

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

Rglyph[turnstileright]\glyph[ceilingleft]glyph[wreathproduct]glyph[arrowdblbothv]←F

glyph[ceilingright]glyph[unionsq]glyph[intersectionsq]∇glyph[ceilingright]←Mglyph[wreathproduct]glyph[ceilingleft]glyph[arrowbothv]

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞∈∀

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

∞′∈glyph[triangle]

∈′glyph[triangle]∀

∀∞∃∈

∞ ′

∞

∞ ′

∞ ′

∞

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ′

∞

∞ ′

W 〉glyph[ceilingleft]glyph[unionsq]〈

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

∞′∈glyph[triangle]

∈′glyph[triangle]∀

∀∞∃∈

∞ ′

∞

∞ ′

∈

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ′

∞

∞ ′

W 〉glyph[ceilingleft]glyph[unionsq]〈

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

∞′∈glyph[triangle]

∈′glyph[triangle]∀

∀∞∃∈

∈

E√glyph[wreathproduct]glyph[floorright]〈

∈GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

∇ √ \

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞∈∀

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

∞′∈glyph[triangle]

∈′glyph[triangle]∀

∀∞∃∈

glyph[arrowbothv]∇

∞↙glyph[triangle]↦

(a) Random Features

∞′

∞glyph[triangleinv]

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞ ′

∞

W 〉glyph[ceilingleft]glyph[unionsq]〈

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

∞′∈glyph[triangle]

∈′glyph[triangle]∀

∀∞∃∈

glyph[arrowbothv]∇

′↙glyph[owner]∈

∞↙glyph[triangle]↦

(b)

µ

P

∞′

∞glyph[triangleinv]

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∈

∞ ′

∞

W 〉glyph[ceilingleft]glyph[unionsq]〈

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

∞′∈glyph[triangle]

∈′glyph[triangle]∀

∀∞∃∈

glyph[arrowbothv]∇

glyph[owner]∞↙̸∈

(c) NTP

′

glyph[triangleinv]

′

glyph[triangleinv]

glyph[triangle]

↙

′

∀

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 µ P and Depthµ P respectively in ConvNets. Note that, similar to the sharpness plots, the directional sharpness, defined as ∇ θ L ⊤ H ∇ θ L ∥∇ θ L∥ 2 also follows a width independent trajectory during training. This measure has been used, for instance, in Gur-Ari et al. [17].

J Experimental details

The experiments were ran on A100 and H100 GPUs, with 80 GBVRAM.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 in our experiments and the modifications we have introduced.

′

̸

glyph[triangle]

glyph[triangle]

′

⇐

glyph[owner]

glyph[triangle]

̸

glyph[triangle]

′

glyph[triangle]

̸

∀

∃

̸

glyph[triangle]

′

∃

̸

′

↙

glyph[triangleinv]

glyph[owner]

′

′

′

glyph[triangle]

′

∃

̸

glyph[triangle]

′

∃

̸

′

glyph[triangle]

′

∃

̸

glyph[triangle]

′

∃

̸

∀

↦

↙

∃

′

↙

′

↙

W 〉glyph[ceilingleft]glyph[unionsq]〈

J.1 Hessian Computation

The implementations of our models are done in PyTorch. For the Hessian measurements, we use PyHessian [94]. In particular, the library adopts the Power Iteration method for the top k Hessian eigenvalues (i.e. including the sharpness) and the Hutchinson method for trace computation. Both methods adopts Hessian vector products to 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 to 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.e. η = η 0 γ 2 = η 0 N . This effectively sets the EoS threshold to the width-dependent value of 2 / ( Nη 0 ) . In our plot, we take this scaling difference into account by computing the eigenvalues of the scaled Hessian γ 2 H = NH . With respect to learning rate transfer, such a rescaling makes intuitive sense, as it is η 0 that is transferring, and not η .

Adam Adam updates are of the form θ t +1 = θ t -ηP -1 t +1 m t +1 , where m t is a momentum vector computed in the exponential moving average fashion: m t +1 = β 1 m t +(1 -β 1 ) g t +1 , where g t are the gradient of the loss at time t . P t +1 is a diagonal preconditioner of the form

′

↙

glyph[owner]

T

̸

↦

̸

T

glyph[triangleinv]

′

{

glyph[floorright]

′

↙

glyph[owner]

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞ ′

∞

∞ ′

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

∞ ′

GLYPH<215>

∞ ′

∞

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞̸

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

(a)

µ

P Transfer

∞′

∞glyph[triangleinv]

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞ ′

∞ ′

∞ ′

∈

∞ ′

∞

∈

W 〉glyph[ceilingleft]glyph[unionsq]〈

∞̸

glyph[owner]∈

∞∈∀

∈glyph[triangleinv]̸

glyph[triangleinv]∞∈

glyph[arrowbothv]∇

′↙glyph[triangleinv]glyph[triangle]∞∈

(b) Gradient-Hessian Alignment

∞ ′

∞

∞ ′

Lglyph[ceilingright]glyph[turnstileright]∇ \〉}←R glyph[turnstileright]glyph[unionsq]glyph[ceilingright]

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

GLYPH<215>

∞ ′

∞

∇glyph[turnstileright]〉\←Lglyph[wreathproduct]∫∫

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

(c) Depth-

µ

P Transfer

∞′

E√glyph[wreathproduct]glyph[floorright]〈

∞ ′

∞ ′

∈

∈

Dglyph[ceilingright]√glyph[unionsq]〈

∞∈

∈glyph[triangle]

glyph[arrowbothv]∇

′↙glyph[triangle]∃∈glyph[triangle]

(d) Gradient-Hessian Alignment

′

′

glyph[owner]

′

glyph[triangle]

′

′

↙

glyph[triangle]

′

↙

glyph[triangle]

′

↙

glyph[triangle]

′

↙

glyph[triangle]

∞ ′

∈

′↙∞

′↙∈

glyph[arrowbothv]∞

glyph[unionsq]glyph[Ifractur]∞

Lglyph[turnstileright]†glyph[ceilingright]∇

Cglyph[wreathproduct]\glyph[subsetsqequal]′∞

Cglyph[wreathproduct]\glyph[subsetsqequal]′∈

Cglyph[wreathproduct]\glyph[subsetsqequal]′glyph[owner]

∞ ′

∈

′↙∞

′↙∈

glyph[arrowbothv]∞

glyph[unionsq]glyph[Ifractur]∈

Lglyph[turnstileright]†glyph[ceilingright]∇

Cglyph[wreathproduct]\glyph[subsetsqequal]′∞

Cglyph[wreathproduct]\glyph[subsetsqequal]′∈

Cglyph[wreathproduct]\glyph[subsetsqequal]′glyph[owner]

∞ ′

∈

′↙∞

′↙∈

glyph[arrowbothv]∞

glyph[unionsq]glyph[Ifractur]glyph[owner]

Lglyph[turnstileright]†glyph[ceilingright]∇

Cglyph[wreathproduct]\glyph[subsetsqequal]′∞

Cglyph[wreathproduct]\glyph[subsetsqequal]′∈

Cglyph[wreathproduct]\glyph[subsetsqequal]′glyph[owner]

∞ ′

∈

′↙∞

′↙∈

glyph[arrowbothv]∞

glyph[unionsq]glyph[Ifractur]glyph[triangle]

Lglyph[turnstileright]†glyph[ceilingright]∇

Cglyph[wreathproduct]\glyph[subsetsqequal]′∞

Cglyph[wreathproduct]\glyph[subsetsqequal]′∈

Cglyph[wreathproduct]\glyph[subsetsqequal]′glyph[owner]

W 〉glyph[ceilingleft]glyph[unionsq]〈

′

{

glyph[floorright]

̸

glyph[triangle]

glyph[owner]

̸

glyph[triangle]

∀

g

H

g

g

H

g

′

↙

′

↙

glyph[owner]

g

g

glyph[triangleinv]

glyph[triangle]

glyph[owner]

glyph[owner]

′

glyph[triangleinv]

{

glyph[floorright]

′

↙

′

↙

glyph[owner]

W 〉glyph[ceilingleft]glyph[unionsq]〈

̸

glyph[triangle]

glyph[owner]

̸

glyph[triangle]

∀

{

glyph[floorright]

where ν t := β 2 ν t -1 +(1 -β 2 ) g 2 t , where g 2 t is the element-wise squared gradients. In the experiments of this paper, we consider the preconditioned Hessian P -1 H , for which it is shown that its largest eigenvalue converges to the EoS threshold of 2(1+ β 1 ) η (1 -β 1 ) in Cohen et al. [61]. Furthermore, we use 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 -1 H , where D is a diagonal matrix containing the learning rate for each parameter. We always report spectral quantities of ˜ H in the experiments with Adam. With these modifications, we expect λ max ( ˜ H ) = 2(1+ β 1 ) (1 -β 1 ) , which is 38 for the full-batch case and the default β 1 = 0 . 9 . Indeed, this is what we observe in our experiments. We stress that we expect Super Consistency of this preconditioned Hessian, where in µ P different layers may have a different dependence on the width [6]. Finally, we point out that in the Depthµ P experiments, we reported the hessian of N/ √ LH . Although this produces results that deviate from the prescription of Cohen et al. [61], it is sufficient to capture Super Consistency in depth, as each layer has the same depth-dependence.

J.2 GPT-2

The backbone architecture in the experiments presented on Wikitext is the standard GPT-2 transformer introduced by [33], with the Depthµ P parameterization changes presented in [6, 8, 53]. Crucially, the following modifications are introduced by the µ P parameterization:

The attention map is rescaled by 1 d Q , as opposed to 1 √ d Q

The residual branch is downscaled by 1 √ L , where L is the number of layers in the model

Our implementation is based on the implementation provided by Yang et al. [6], with the addition of the residual scaling. This uses a different parametrization from the one reported in Table. 1 but equivalent dynamics, obtainable using their 'abc-rule". Similar to 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µ P 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 in this setting we also train without the 1 √ L 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 linear rate warmup, with no decay, and we use a standard Post-LN GPT-2 style architecture.

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µ P parameterization, we make the same modifications as in J.2. The models are trained with Adam.

J.4 ResNet

We use convolutional neural networks with skip connections, with 3 × 3 kernels and stride 1 . We apply pooling after every 3 rd layer, followed by a subsequent convolutional layer. Following Bordelon et al. [7] and Yang et al. [8], we downscale the residual branches by 1 / √ L . Note that in the Depthµ P setting we also scale the learning rate as 1 √ L in these experiments.

J.5 Coordinate check for µ P

In Figure 28 we check the coordinatewise evolution of the activations at hidden layers within a ConvNet. Note that the evolution is width independent, as predicted by µ P theory.

K Summary of Feature Learning Parametrizations

Here we summarize the parametrizations used. In Table 1, we report the scaling of the learning rate, output multiplier γ , depth-scaling exponents α for the residual branches for Depthµ P (both SGD and Adam) [7, 8, 96]. µ P is recovered as a special case, where the depth dependence is ignored. Notice that this version of Depthµ P for Adam is obtained with a simplification, setting Adam's ϵ parameter to zero, where Sign-GD is recovered.

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 Consistency 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 in 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 of the work performed by the authors?

Answer: [Yes]

Justification: We provide a 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 a 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 authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on 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 approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or 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 it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

Justification: We explicitly state all the assumption and simplifications, as well as an organized 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide a 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 to 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 to 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, or if the contribution is a specific model and empirical evaluation, it 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 to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce 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 authors 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as 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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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 a subset of experiments are reproducible, they 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 to 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, hyperparameters, 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 to a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).

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

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably 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 to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to 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, or 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 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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We are conform with the code

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not pose such risks, and only use public data.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to 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 not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: 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?

Answer: [Yes]

Justification: we provide info and citations for the packages that we use.

Guidelines:

The answer NA means that the paper does not use 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 a URL.

The name of the license (e.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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

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 documentation provided alongside the assets?

Answer: [NA]

Justification: We do not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was 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 paper include the full text of instructions given to participants and screenshots, if applicable, as 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in 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 on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: not applicable.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If 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 NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

STRENGTHS:
- The paper discusses an interesting phenomenon that the loss landscape gradually becomes stable, this might explain the transfer of learning rate
- 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
- 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

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
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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]**\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nSee weakness and questions", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

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

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}