Universal Sharpness Dynamics in Neural Network Training: Fixed Point Analysis, Edge of Stability, and Route to Chaos
Authors: Dayal Singh Kalra
Tianyu He
Maissam Barkeshli
Venue: ICLR 2025 Poster (2025)
Keywords: Sharpness Dynamics
Catapult Effect
Edge of Stability
Primary area: optimization

Abstract: In gradient descent dynamics of neural networks, the top eigenvalue of the Hessian of the loss (sharpness) displays a variety of robust phenomena throughout training. This includes early time regimes where the sharpness may decrease during early periods of training (sharpness reduction), and later time behavior such as progressive sharpening and edge of stability. We demonstrate that a simple 
-layer linear network (UV model) trained on a single training example exhibits all of the essential sharpness phenomenology observed in real-world scenarios. By analyzing the structure of dynamical fixed points in function space and the vector field of function updates, we uncover the underlying mechanisms behind these sharpness trends. Our analysis reveals (i) the mechanism behind early sharpness reduction and progressive sharpening, (ii) the required conditions for edge of stability, and (iii) a period-doubling route to chaos on the edge of stability manifold as learning rate is increased. Finally, we demonstrate that various predictions from this simplified model generalize to real-world scenarios and discuss its limitations.

Published as a conference paper at ICLR 2025
UNIVERSAL SHARPNESS DYNAMICS IN NEURAL NET- WORK TRAINING: FIXED POINT ANALYSIS, EDGE OF STABILITY, AND ROUTE TO CHAOS
Dayal Singh 2 Tianyu He & Maissam Barkeshli 1.3 (dayal, tianyuh, ma i ssam) ) @umd. edu
ABSTRACT
In gradient descent dynamics of neural networks, the top eigenvalue of the loss Hessian (sharpness) displays a variety of robust phenomena throughout training. This includes early time regimes where the sharpness may decrease during early periods of training (sharpness reduction), and later time behavior such as pro- gressive sharpening and edge of stability. We demonstrate that simple 2-layer linear network (UV model) trained on single training example exhibits all of the essential sharpness phenomenology observed in real-world scenarios. By analyzing the structure of dynamical fixed points in function space and the vector field of function updates, we uncover the underlying mechanisms behind these sharpness trends. Our analysis reveals (i) the mechanism behind early sharpness reduction and progressive sharpening, (ii) the required conditions for edge of stability. (iii) the erucial role of initialization and parameterization, and (iv) a period-doubling route to chaos on the edge of stability manifold as learning rate is increased. Finally, we demonstrate that various predictions from this simplified model generalize real-works scenarios and discuss its limitations.
INTRODUCTION
Over the last several years, it has been observed that the training dynamics of neural networks exhibits a rich and robust set of unexpected phenomena, stemming from the non-convexity of the loss landscape. These phenomena not only challenge our existing understanding loss landscapes but also open avenues for significantly enhancing model performance through improved optimization techniques. particular, the unexpected and robust phenomenology is mainly associated with the evolution of the Hessian of the loss function, which provides measure of the local curvature of the loss landscape and plays an important role in understanding generalization performance Keskar et al. (2016) Dziugaite & Roy (2017) Jiang et al. (2019). However, the relationship between sharpness and generalization has been called into question Dinh et al. (2017): Kaur et al. (2023).
On the one hand, it has been observed that at late training times, gradient descent (GD) typically exhibits "progressive sharpening," where the top eigenvalue of the loss Hessian A referred to as the sharpness, gradually increases with time, until it reaches roughly 2/1,, where 17 is the learning rate. Once the sharpness reaches roughly 2/1,, it stops increasing and typically oscillates near 2/1, a late-time training phenomenon referred to the "edge of stability (EoS)" Cohen et al. (2021) On the other hand, during early training, a decrease in sharpness is observed -relerred to as "sharpness reduction " Kalra & Barkeshli (2023) before hitting a temporary plateau.
For large enough learning rates, training temporarily destabilizes early on, and the network "catapults" out of its local basin, leading a temporary sudden increase in the loss in the first few steps, before eventually settling down in flatter region of the loss landscape characterized by lower sharpness Lewkowyez et al. (2020). Similar to the loss, sharpness may also spike within the first few steps of training and quickly decrease (sharpness catapult). A rich phase diagram as a function of network depth, width and learning rate summarizes the early training dynamics Kalra & Barkeshli (2023).
Depaarmment of Physics, University of Maryland, College Park DDpartment of Computer Science, University of Maryland, College Park Joint Quantum Institute, University of Maryland, College Park

---

Published as a conference paper at ICLR 2025
Te0.
10
100
10
12
10
(a)
(b)
120
0.4
100
80
0.2
12
117
10
10
10
(d)
(e)
(I)
Figure 1: Training loss and sharpness trajectories of ReLU FCNs trained on 5k subset of CIFAR-10 -0
examples using MSE loss and GD: (a, SP with o 0.5. (b, e) SP with o 2.0, (c, f) pP with a 2.0. The dashed lines in the sharpness figures show the 3/1 threshold.
The discovery of these intriguing sharpness phenomena has attracted significant attention, with emphasis on various toy models that exhibit similar phenomenology. Yet, the specific conditions and reasons why these phenomena occur still remain elusive. In this paper, analyze simple toy model, a 2- 2-layer linear network trained on one example, referred to as the UV model. We show that all of the phenomena deseribed above can be observed the UV model for appropriate choices learning rate, initialization, parameterization, and choice of training example. Through this exploration, we provide novel insights into the mechanisms at play and offer predictions that we validate in realistic architectures with both real and synthetic datasets.
Our Contributions. We revisit the four training regimes identified by Kalra & Barkeshli (2023) (early time transient, intermediate saturation, progressive sharpening, and late time EoS) in Section 3 focusing on the crucial role of initializations and parameterizations. Our findings reveal that models
in Standard Parameterization (SP) with large initializations do not exhibit EoS, even at late training
times. Moreover, we show that models in Maximal Update Parameterization (μP) Yang Hu (2021) do not experience early sharpness reduction. This result also holds for models in SP with small initializations.
We show the UV model exhibits all four training regimes and also captures the effect of initializations and parameterization discussed above. Through fixed-point analysis of the UV model in the function space, we analyze the origins of the various dynamical phenomena exhibited by the sharpness. Specifically, we demonstrate in Sections 4 and 5) (i) the emergence of various sharpness phenomena arising from the stability and position of the dynamical fixed points. (ii) critical learning rate ne- above which the model exhibits EoS on a sub-quadratic manifold, and (iii) a period-doubling route to chaos of sharpness fluctuations learning rate increased in the EoS regime.
In Section 6 we verify various non-trivial predictions from the UV model in realistic architectures with real and synthetic datasets. Our findings reveal: a teappness-weight norm correlation before the training enters the EoS regime, (ii) a phase diagram of EoS, revealing initializations and parameterizations that do not exhibit FoS, and (iii) a period-doubling route to chaos in real architectures trained synthetic datasets, while those trained real datasets exhibit long-range correlations at the EoS, with remnant of the period doubling route to chaos.
Related Works. Using the top eigenvalue of the Neural Tangent Kernel (NTK) X at initialization (t tt= 0), Lewkowyez et al. 2020) revealed a 'catapult phase', 2/25 < < 7/uax in which training converges despite initial spike in training loss. After the early training phase, sharpness continues to increase until reaches a break-even point (Jastrzebski et al.) 2020), beyond which GD dynamics typically enters the EoS regime (Cohen al. 2021). This has motivated various theoretical studies
2

---

Published as a conference paper at ICLR 2025
to understand GD dynamics at large learning rates Ma et al. (2022); Wang et al. (2022) Arora et al.
(2022); ; Damian et al. (2023); Rosca et al.] (2023); y Zhu et al. (2023); Wu et al. (2023); Chen & Bruna
(023 Ahn et al. 2022); Kreisler et al. (2023;; Song & Yun (2023); Chen et al. (2023). ). These
include an analysis by Wang et al. (2022). who show analyzed EoS in a 2-layer linear network using the norm of the last layer, restricting to cases that show progressive sharpening from initialization. Agarwala ct al. (2022) suggested that the UV model does not exhibit EoS behavior but showed that a modified model exhibits progressive sharpening and two-step oscillations EoS using NTK as the proxy. In contrast, we show that the UV model shows EoS behavior under the appropriate choice of parameterization and training example. Meanwhile Chen Bruna (2023) analyzed two-step gradient updates of single-neuron network and matrix factorization to gain insights into EoS. Furthermore, Song Yun (2023) analyzed 2-layer linear network under logistic loss and demonstrated that sharpness at late training times oscillates around 2f/nf", where f is the network output and l is the derivative of the loss. Another study by Chen et al. (2023) analyzed large learning rate dynamics of toy models which are characterized by one-dimensional cubic map. Our work, in contrast, delves into various sharpness phenomena occurring throughout the training trajectory and analyzes their origins. It is worth noting that a concurrent study by WWang al. (2023) also examines sharpness throughout training. Noci et al. (2024) relate sharpness dynamics to learning rate transfer in jP networks by showing that sharpness trajectories do not change appreciably when depth and width are varied. Given that our analysis spans the entire training trajectory. it relates to numerous studies. Hence, we defer a comprehensive discussion of related works to Appendix A
2 NOTATIONS AND PRELIMINARIES
This section describes the fundamental concepts and notations that form the basis of our analysis.
Dynamical Systems and Fixed Points: Consider a discrete dynamical system described by 0.+1 = (0.). A fixed point 0' of the dynamics satisfies M ((0)) = 0". The linear stability of a fixed point 0" is determined by analyzing the eigenvalues {>-"} } of the Jacobian (0') = M() le-e.. An eigendirection u!'" of a fixed point 0* is stable if |>'1| < 1 and unstable if | > 1 Out 2002)). The dynamics is captured by the vector field of updates ())) = M(0 - O. The corresponding unit vector is denoted G(0) = G(O)/IG(0)). Nullclines refer to curves where one of the variables, O,, remains invariant, i.e., f1: (0,).
Sharpness phenomena in neural networks are intrinsically tied network parameterization. Standard Parameterization (SP) /Soh- Dickststn et al. (2020) and Neural Tangent Parameterization (NTP) et al. (2018) are two commonly used parameterizations, which converge to kernel methods at infinite width. Yang & Huj (202T) proposed Maximal update Parameterization (nP), which allows for feature learning infinite width. For implementation details, see Appendix 1.
Parameterizations in Neural Networks:
The UV model refers to 2-layer linear network f: IR + R trained on a single example. We parameterize f as f 0) = 1 - Uz, where x E Rd is the input, n is the network width, and & E RK," E xd are trainable parameters, with each component drawn i.i.d. at initialization from a normal distribution (0, o²/n"") Here, E [0, 1] is a parameter that interpolates between NTP and jaP, and ner == n² 1P referred to as the effective width. We consider the network trained on a single training example (xx y) using MSE loss e (f(x;; 0). y) (f(=; 6) 2
UV Model:
REVIEW OF THE FOUR REGIMES OF TRAINING
Typical training trajectories of neural networks can be categorized into four training regimes Kalra & Barkeshli (003)), as shown in Figure Ta, d):
(TI) Early time transient: This corresponds to the first few steps of training. At small learning rates (1) < ),oss), loss and sharpness decrease monotonically. At larger learning rates (7 > Mess)) training catapults out of the initial basin, temporarily increasing the loss, and finally converges to a flatter region Lewkowyez et al. (2020). By the end of this regime, sharpness has decreased from initialization for all learning rates, and more substantially at larger learning rates.
3

---

Published as a conference paper at ICLR 2025
(T2) Intermediate saturation: Following the initial transient regime, sharpness approximately plateaus before gradually increasing.
(T3) Progressive sharpening: In this regime, sharpness continues to increase until it reaches X x 2/n Dastrzebski et al. (2020); Cohen et al. (202T). At large effective widths or small learning rates, training may conclude before reaching this threshold.
(T4) Late-time dynamics (EoS): After progressive sharpening, for MSE loss, sharpness oscillates around 2/n. For cross-entropy loss, the sharpness oscillates when reaching approximately 2/v, while decreasing over longer time scales Cohen et al. (202T).
In this work, show that the sharpness dynamics heavily depends on the initialization and param- eterization of the network, and not every training trajectory shows all four regimes. For instance, Fig e) shows that FCNs in SP with large initialization (or large effective width) do not exhibit EoS, even when loss decreases to a value below 10-5, Following the early transient regime, sharpness monotonically decreases, with only nominal increase towards late training. In contrast, Figure f) shows that FCNs in pP (or small effective width) do not experience an initial sharpness reduction at small learning rates (1) < ). Rather, sharpness continues to increase until reaches 2/1, and then oscillates around it. At large learning rates (1) > Thhhrp), sharpness catapults and eventually settles into the same trend as above.
These different training regimes are generically observed for more complex architectures and datasets as we show in Appendix D.3 including CNNs and ResNets, trained on CIFAR-10 and Transformers trained on Wikitext-2.
In Appendix D.I we show that these trends remain consistent when NTP used instead of SP. Given this similarity in the training dynamics between SP and NTP, use NTP for theoretical analysis for clarity and SP in realistic experiments for implementation convenience.
demonstrates that the UV model displays all four training regimes. also captures the cases where sharpness reduction or EoS is not observed. Therefore, the simplified UV model can serve the sharpness dynamics. In the subsequent section, we perform fixed point analysis of the UV model and probe the origin of these complex phenomena in later sections.
Figure 2 (and Figure 7 in Appendix
effective model for understanding these universal behaviors in
FIXED POINT ANALYSIS OF THE UV MODEL
Under GD, the parameters of the UV model are updated as Ut+1 = Us - Afeer"
Ut+1
Vner where 7) is the learning rate and Af := f(x; 0c) y is the residual at training step 1. In function space, the dynamics can be completely described using the residual Aʃ, and trace of the loss Hessian À = Tr H 1 (x"UTUx + u u x æ). which is also the scalar neural tangent kernel in this case. The function space dynamics of the UV model can be fully described using two coupled nin-linear equations (for derivation, see Appendix C.I.
V. AfiUi
tue
Af+1 = Aft (1-12,
1i + 17 12122111 Af( v))
Teeff
17 ||00|
(Afe + 3)
A+11 A
Af? n/,,
(2)
Df.
Bell
with effectively three parameters: 17. ||/ll/ yourr and y. While similar equations have been considered in previous works Lewkowycz et al. (2020) Zhu et al. (2022); Agarwala et al. (002). the generalization to generic parameterizations novel and would be crucial observing different sharpness phenomena such as EoS. The y 0 case has been analyzed in prior works Lewkowycz et al. (2020); Kalra & Barkeshli (2023) for understanding catapult dynamics. Here, À can only decrease with time, as can be seen from Equation (2) with 17 7mma = 4/Ao (training diverges if n 7max)) As result, the model does not exhibit progressive sharpening and EoS. Below we focus on the case y > O, which allows for À to increase in time and consequently, much richer dynamics.
Equations (T) and (2) have four distinct fixed points/lines (referred to as I-IV) as detailed in Table ! of Appendix C.3 The fixed line I defines zero-loss line, meaning & = O for all points in I; the points
4

---

Published as a conference paper at ICLR 2025
m512. Meff m1. 200. v 0.20
m1, Roff
m1.
220, =0.23
512. R/// -512, =2.0 -0.50
20
11
15
2
20
-5
10
-5
0
50
50
A
Af
A
(a)
(b)
(c)
512, Reff -1. y 220,
Beft
y 200
512, Ref/ -512,
-2.00
12.5
16.0
10.0
7.5
7.0
3.0
2.54
25-
Pryyeair
0.04
0,0
A
A
(d)
(c)
(f)
Figure 2: Training trajectories of the UV model with ||>|| and y in the (Af, A) plane for different values of n, Netr and The columns show initializations with different n and Nefr, while the rows represent increasing learning rates for fixed initializations. The horizontal dash-dot line 1/> separates the stable (solid black vertical line) and unstable (dashed black vertical line) fixed points along the zero loss fixed line I. Forbidden regions, ||| + yl/ Tefff > A, (see Appendix C.2 are shaded gray. The nullclines ff++ Af, and Ac+1 = X are shown as orange and white dashed curves, respectively. Sharpness reduction, progressive sharpening. and divergent regions are colored green, yellow, and blue. The gray arrows indicate the local vector field G(Aʃ, , A), which is the direction of the updates. The training trajectories are depicted as black lines with arrows, with the star marking the initialization. In all cases, T)e = /nett/ 2 (introduced Section 5.2
in stable for 17 < 2 and unstable otherwise. Fixed point II at (-y, 0) corresponds to the origin in parameter space (U, v 0) and it is saddle point of the dynamics for convergent learning rates n. Both and are also fixed points of the GD optimization, i.e., critical points of the loss. The loss Hessian at I is positive definite, while fixed point II is saddle point in the loss landscape. The remaining two fixed points and IV are unstable and exist only in function space, representing non-trivial parameter space dynamics that leave the function space dynamics invariant.
Figure 2 shows the fixed points and the vector field Č(AS, >) determined by Equations I and 2) which Illustrates the direction of the updates at each point. Note that the stability of the fixed line (I) does not follow from & alone, as the magnitude G is required to determine stability. Figure also shows training trajectories for various parameter values. Using À as proxy for sharpness, see there are regions where increases (colored yellow) and decreases (colored green) along the flow, which refer to as progressive sharpening and sharpness reduction, respectively. It follows from Equation (2) that the condition y = 4(As + y) separates these regions. Importantly, the parameters 17. ||allTy /re and y influence the position of the fixed points. This, in turn, affects the extent of different regions and the vector field G, as illustrated Figure 2 In particular, on decreasing effective width nef increasing learning rate 77, fixed points III and IV move inward (see fixed point expressions in Table T). which relatively enlarges the progressive sharpening region while shrinking the overall convergent region. Overall, these illustrations demonstrate how the local stability and relative position of the fixed points collectively impact the dynamics. In the subsequent section, we discuss the dynamics in detail.
5

---

Published as a conference paper at ICLR 2025
5 UNDERSTANDING SHARPNESS DYNAMICS IN THE UV MODEL
In this section, we describe the origin of different robust phenomena in the dynamics of sharpness using the fixed point and linear stability analysis from the previous section. This explains the four training regimes observed in the UV model. We will discuss the influence of effective width and initializations, shedding light on the differences between NTP and jaP. For simplicity, we assume ||>|| = 1, while allowing Reff to vary continuously. Note that we use À == Tr H from the previous section as a proxy for sharpness; we have verified that the top eigenvalue of the Hessian the loss also follows À (see Appendix C.5) although it is more difficult to analyze analytically.
5.1 UNDERSTANDING SHARPNESS TRENDS THROUGHOUT TRAINING
Figure 2 shows that the training dynamics can exhibit different behavior depending on the initial region. Below we summarize these based on empirical observations.
(RI) Progressive sharpening region: As shown in Figure D d), initialization in this region experi- ences an upward push due to the flow originating from fixed point II, resulting in steady increase in X. Depending on relative to critical learning rate 7) (introduced in Section 5.2) different late-time dynamics arises. For 1/ < Te. training converges to stable fixed points on the zero-loss line (D, as shown in Figure 2la).
a). When 1) > T)c, all points along the zero-loss line (I) become unstable, as shown Figure 2 d). In this case, the network eventually converges to a line segment joining fixed points II and IV (the EoS manifold), where it continues oscillate indefinitely between these fixed points, leading to the EoS phenomena. This will be analyzed in more depth in the subsequent section.
(R2) Sharpness reduction region between fixed points Il and III: Fbg e) show that initializations in this region undergo a decrease in as the flow is towards saddle point II. On approaching this saddle point, the dynamics slows down, resulting in the intermediate saturation regime. Eventually, training moves away from this saddle and enters the progressive sharpening region. From here on, the dynamics becomes akin to the case (R1).
(R3) Sharpness reduction region h/w fixed line and point IV: Initializations in this region either converge to the nearby zero-loss solution for (1) < ne) or enter the progressive sharpening region for (1) > Me). In the latter case, the dynamics resembles those of case (RI).
So far, we have described the resultant dynamics when training is initialized in different regimes. Below, we describe the conditions that typically exhibit these regimes.
Neural Tangent Parameterization: In NTP, A and À follow normal distributions: fo W( o²) and A ~ N222 4z2/n). Hence, the model can be initialized in any of the three regions described above. Moreover, fixed points IIII and IV move outward with increasing width, affecting the local vector field G(Af, A). At large widths n > 1, G(A for An) at initialization points along [1 0j towards the zoro-loss line. For small learning rates (1) < 2/xu), training exponentially converges to the nearest loso-loss solution (see Figure Regardless of the initialization region, the change in À is minimal, receiving (1/xx) updates as per Equation (2) For large learning rates (1) > 3/Xo), the nearby zero-loss solution becomes unstable. Consequently, training catapults to a region with smaller X, while bouncing between fixed points III and IV. This is the catapult effect studied in Lewkowycz et al. (2020) and Figure ( demonstrates such a trajectory. By comparison, at small widths, the dynamics follows cases (RI-R3) discussed above.
Maximal Update (μP) Parameterizations: In contrast to NTP, the position of fixed points III- IV do not change with width n, and Auu follows the distribution: Afo ~ .N(yy, yy- 44/n), while An distribution remains unchanged. Consequently, at large widths, the model is initialized at (-y, 2) right above fixed point II in the progressive sharpening region (R1), satisfying the condition A < 4(A5o + y). FigureZZie d) shows such a trajectory. At small widths, fluctua- tions increase, making plausible for jP networks to start in the sharpness reduction regions. In this case, the dynamics follow case (R2) or (R3).
5.2 UNDERSTANDING EDGE OF STABILITY
This section analyzes the EoS behavior in the UV model, particularly from the fixed point perspective. As discussed in the previous section, the EoS behavior in the UV model arises when all fixed points

---

Published as a conference paper at ICLR 2025
6.00
0.50
9
7
(h)
Figure 3: (a) Bifurcation diagram depicting limiting values of À obtained by simulating Equation 33)
(b) Bifurcation diagram of the UV model. In both figures, || = 1, 2, nea = and T)c 0.5.
along the zero-loss line (I) become unstable wrt the learning rate. Yet, the gradient updates (shown as gray arrows in Figure 2) continue to point towards the zero loss line. As a result, training is trapped in this region, converging to the line segment that joins fixed points II and referreed as the EoS manifold -where it oscillates indefinitely.
EoS Manifold is Attractor: By examining the two-step dyanmics akin Agarwala al. (2022) Chen & Bruna (2023), we show in Appendix[C.. that training converges to the EoS manifold above critical learning rate Te-- For 17 < Te.. training converges the stable fixed points on the zero-loss line. By comparison, for > Tl. all points along the zero-loss line become unstable and the EoS manifold becomes dynamical attractor. The critical 7e for which all points on the zero-loss line become unstable thus gives a necessary condition for EoS:
Result 1. A necessary condition for the UV model to exhibit EoS is > ?e v l (see Appendix C.& for details). It is useful scale the learning rate as n7 = c/ in which case this condition becomes Ào < cllallaa//nen
For learning rates 17 > 2/A0. training can catapult to regions with Ar Ao In such cases. the condition \T < cl|/lla// f also applies.
Dynamics on the EoS Manifold and Route to Chaos:
The dynamics on the EoS manifold satisfies [|[=|[A/tv/// coupling Af and together. This yields the map Af++1 (Af)) describing the dynamics on the EoS manifold, with My defined as
(Af) AA
nAfe Noy
mmAe n-y
2) (Aft
Figure 3 shows the limiting values of À (i.e. the values of À that the network jumps between at late times) as a function of learning rate, obtained by simulating Equation (3). We refer this as the bifurcation diagram. As mentioned before, for Tee the zero-loss solution becomes unstable with oscillating around 2/n instead converging. These fluctuations exhibit a fractal structure, as the system undergoes series of period-doubling transitions with an increasing learning rate. This is the well-known the UV model for 2. The bifurcation diagram diagram extends up higher learning rates. This leads us to the following corollary of Result [
route to chaos Ott (2002). Figure 3(b) shows the bifurcation diagram of 0.8 before diverging
Corollary 5.1. Let Mmax be the maximum trainable learning rate for a given initialization. The bifurcation diagram is observed up to 1) <. Tmaxx If Tmax To the UV model does not exhibit EoS.
These results suggest that models with small λ and neir are more prone to show Eo$ behavior. As result, μP networks or those with small initial weight variance are more likely to exhibit EoS. On the other hand, large-width NTP networks may not show EoS behavior at all. In the next section, we will validate this prediction in real-world scenarios.
Ma et al. (2022) demonstrated that GD on sub-quadratic loss with large learning rates inherently results in EoS behavior. Here, we show that the loss on the EoS manifold of the UV model is sub-quadratic near its minimum. As noted above, the dynamics the EoS manifold satisfies À = |||=][(A++))) /ons. The loss on the EoS manifold is then given by
Connections to mub-quadratie loss:
(())) = 2 o (nc) 1)2 (see Appendix C.9 for derivation), where 0 denotes the parameters.

---

Published as a conference paper at ICLR 2025
Test acturacy
31.
2.5
100
52000
1.33
1050
20
0.5
0.5
4
<2
(a)
(b)
RA /2
Test accuracy
3.5
13.5
2.5
2,5
40.0
2.0
37.5
0.75
0.50
16 0:25
1000
0.50
(c)
(d)
Figure 4: (a, b) Heatmap of nA /2 and test accuracy of ReLU FCNs in SP trained on 5k subset of CIFAR- 10 until 99% training accuracy is achieved, with the weight variance o and learning rate multiplier nAy axes. As the color varies from blue white, n/" increases. (b, d) Same heatmaps with fixed o² = 2.0, but varying 8 continuously.
Since À ~ 3) the loss is of the form [()) ~ (al00) - 6)2 and is sub-quadratic near its minimum. The GD dynamics near the minimum is given by cubic map. which is known to show the period-doubling route to chaos Rogers & Whitley (193)). Chen et al. (203)) showed similar route to chaos by considering a two-layer network with quadratic activation, with the last layer vector fixed through training and each entry set to one. In this model, the loss is sub-quadratic by construction ([|2 - -²²² and the dynamics given by a cubic map.
PREDICTIONS AND VERIFICATIONS IN REAL-WORLD SCENARIOS
The preceding analysis offers broader insights and predictions for optimization in real-world models. In this section, study realistic architectures with real and synthetic datasets and examine the extent to which insights from the UV model generalize.
Experimental Setup: Consider network / [a;; f), with trainable parameters 0, initialized using normal distribution with zero mean and variance o.2 in appropriate parameterization. In this section, we use the interpolating parameterization with with 0 are equivalent to networks in SP width n goes to infinity and those with
[0, 1] (detailed in Appendix [B.2II. where networks
are in μ... The network trained on a dataset with P examples using MSE loss and GD. The learning rate scaled as </>". where c is the learning rate constant, and X is the sharpness at initialization. Additional details provided figure captions and Appendix B.2
6.1 IMPLICATIONS OF INITIALIZATION AND PARAMETERIZATION FOR REAL-WORLD MODELS
The analysis in Section 5.1 unveils crucial insights into the implicit biases of parameterization in real-world networks. Figure 2 a, d) shows that μuP networks begin training flat region of the landscape, where gradients point towards increasing sharpness, and approach the zero loss line while maintaining a low sharpness bias. In contrast, networks NTP (or equivalently SP), characterized as large initializations, experience sharpness reduction during early training and might not approach with a minimal sharpness bias. The agreement of these observations with networks trained on real-world datasets (Figure TD suggests that these inherent biases hold in practical scenarios.
8

---

Published as a conference paper at ICLR 2025
2.
110000000
20-
92.
1,5
1000
10000
0.
step
(a)
(b)
(c)
6.722
1.5
0.51
0.0
5000 D100
1000
05
step
(d)
(f)
Figure 5: 2-layer linear FCNs trained on (first row) 5, 000 iid random examples with unit output dimension and (second row) 5, 000 CIFAR-I0 examples. Different columns correspond to the bifurcation diagram, late-time sharpness trajectories, and the power spectrum of sharpness trajectories. The power spectrum computed using the last 1000 steps of the trajectories.
6.2 SHARPNESS & WEIGHT-NORN CORRELATION AND THE ORIGIN OF FOUR REGIMES
Section 5. revealed that, for a wide variety initializations, at early times trajectories move closer to the saddle point II, resulting in interim decrease in À (also proportional the weight norm in this case), before eventually increasing. This critical point where all parameters are zero also exists in real-world models. We thus anticipate that in real-world models, the origin of the four training regimes may be related to a similar mechanism. This would predict a decrease in weight norm as training passes near the saddle point, followed by an eventual increase. In Appendix E we validate this hypothesis. During the sharpness reduction and intermediate saturation regimes, we see decrease the weight norm, followed by an increase in the weight norm as the network undergoes progressive sharpening, following the prediction from the UV model.
6.3
THE PHASE DIAGRAM OF EDGE STABILITY
Result presents a necessary condition for EoS to occur in the UV model: Ào <lj=lly// /aa real-world models, the initial sharpness A can be controlled using the initial variance of the weights Therefore, this result predicts that real-world models with (i) small initial weight variance of (ii) large interpolating parameter 8, or (iii) large learning constant e are more likely to exhibit EoS behavior. Figure 4 a, show the phase diagram of EoS, validating these predictions. Additional phase diagrams in Appendix Findicate an enhanced tendency for CNNs and ResNets to exhibit EoS. Furthermore, Figure 4 b. show the corresponding test accuracy heatmaps for the above experiments, revealing a positive correlation between observing EoS and improved model performance.
v In
6.4 ROUTE CHAOS AND BIFURCATION DIAGRAMS
The analysis in Section 5.2 unveiled structured fluctuations in À at the EoS, with a period- doubling route to chaos observed as the learning rate tuned. This motivates to analyze fluctuations the EoS in real-world models trained on realistic and synthetic datasets. Figure shows the bifurcation diagram, late-time sharpness trajectories, and power spectrum of sharpness trajectories for a 2-layer linear FCN. In the first row, the model is trained on random synthetic data with 5, 000 iid examples with unit output dimension, whereas, in the second row, on 5, 000 example subset of CIFAR-10 10. Similar to the UV model, FCNs trained random data exhibit shown in Figure a). By comparison, FCNs trained on CIFAR-10 only show dense bands in the sharpness rather than exhibiting a clear drubbing route to chaos.
-eobbiddoubling route to chaos, as
On analyzing the sharpness trajectories at EoS, we observe range correlations in time in real datasets, with fluctuations increasing with the learning rate (see Figure e))) By comparison,

---

Published as a conference paper at ICLR 2025
sharpness trajectories of models trained random datasets exhibit short-period oscillations (see Figure(S[b)). The power spectrum of these sharpness trajectories further quantifies these observations, as shown in Figure[5jc. f). In the random dataset case, high-frequency modes corresponding to the period-doubling route to chaos emerge at EoS as shown in Figure (tc. In contrast, real datasets exhibit low-frequency modes at small learning rates. As the learning rate is increased, nigh-frequency modes, reminiscent of the period-doublin Appendix G.I we demonstrate that CNNs and ResNets trained image datasets show dense bands of sharpness similar to those in FCNs.
route to chaos, start emerging (see Figure [J()). In
To understand when the period-dowd -rriodd route to chaos arises, we perform further analysis in Appendix G A key determining feature appears to be whether the singular value spectrum of the input-input and output-input covariance matrices are flat or have power-law decay. In Appendix G.2 we show that a 2-layer FCNs trained on random dataset with power-law singular value spectrum in the input exhibits dense sharpness bands. Appendix G.3 we show that linear FCNs trained on synthetic datasets with random inputs, such as teacher-student settings and generative settings (details in Appendix B.I). exhibit the period-doubling route to chaos. In contrast, non-linear networks trained on these tasks exhibit dense sharpness bands observed in real datasets. These observations shed some light on the nature of EoS observed in realistic settings. Nevertheless, a complete understanding of sharpness fluctuations at EoS requires separate detailed examination.
7 DISCUSSION
In this work, analyzed the crucial effect of initializations and parameterizations on the sharpness dynamics of neural networks and characterized conditions under which sharpness phenomena, such as sharpness reduction, progressive sharpening, and EoS, occur during training.
To develop a deeper understanding of these sharpness phenomena, we analyzed the UV model, which exhibits these sharpness trends. Through fixed point analysis, we uncover the underlying mechanisms behind these complex sharpness phenomena, such as (i) the mechanism behind early sharpness reduction and progressive sharpening, (ii) the required conditions for the edge of stability, (iii) the crucial role of initialization and parameterization, and (iv) a derioi-doubling route to chaos on the edge of stability manifold as the learning rate is increased. Finally, we demonstrated that various predictions from this simplified model generalize to real-world scenarios.
While the UV model does capture major sharpness phenomena, does have some limitations. The UV model does not capture the mon-oonoo decrease in the training loss at EoS. This is because the UV model dynamics is effectively described by two variables and the EoS condition À = 2/11 puts another constraint. As result, the loss also oscillates along with sharpness at EoS. In comparison, real-world neural networks have a large number of degrees of freedom, and the training can still converge in stable cigendirections. Furthermore, the UV model does not exhibit loss catapult behavior for small effective widths (nei = 1). In such cases, the threshold for loss catapult (1) = 2/xo) is comparable to the maximum trainable learning rate. Finally, the UV model does not capture long-range correlations in sharpness dynamics observed in realistic datasets. Through extensive experiments, we identified that long-range correlations in the input covariance matrix results these long-range correlations.
The applicability of the fixed point analysis extends well beyond the UV model and can be employed in settings involving complex architectures and adaptive optimizers. A prerequisite for applying this method is the closure of the dynamical equations describing the model. While this closure requirement is manageable for deep linear networks through additional function space variables. becomes challenging for non-linear DNNs where such closure with a few variables is difficult. By analyzing the fixed points of such equations in broader classes of models where closure is achievable, we can gain significant insights into their training dynamics, thereby advancing our understanding of non-convex optimization in neural networks.
Various results such as the phase diagram of EoS, the bifurcation diagram, and the late-time sharpness analysis depend on the training time. Nevertheless, we found that training the models longer does not impact the conclusions presented. In Appendix [.2] we show results are robust for reasonably small batch sizes (B x 512). For even smaller batch sizes, the dynamics becomes noise-dominated, and separating the inherent dynamics from noise becomes challenging.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
D.S.K. and M.B. are supported by an NSF CAREER awards (DMR1753222 and (MR-2345644) and T.H. supported by the NSF CAREER award (DMR-2045181), The authors acknowledge the University of Maryland supercomputing resources httpp: http//ppc.... umd. edu) made available for conducting the research reported in this paper. We thank the anonymous reviewers for their valuable feedback in improving this work. We also thank Darshil Doshi, Andrey Gromov, Jeremy Cohen, Alex Damian, Alex Atanasov, and Lorenzo Noci for insightful discussions. D.S.K. performed part of this work at the Institute for Physical Science and Technology at the University of Maryland, College Park.
REPRODUCIBILITY STATEMENT
We have provided extensive details for reproducing various theoretical and experimental results presented in the paper. Appendix details the derivations of the theoretical results for the UV model, while Appendix [B] provides extensive details for reproducing the presented experiments. Furthermore, the Supplementary materials include the JAX code required reproduce key results.
REFERENCES
Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. second-order regression model shows edge of stability behavior. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop). 2022. URL ptttps::/openreviiee
Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the "edge of stability". ArXiv, abs/2212.07466, 2022.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 10241024. PMLR, 17-23 Jul 2022. URL https: http:////occeeeeng..e.c .mw..ngg....r.....m............. wwtpp/eewweee.ereeeeeeeeeeeeeeee wwwpp//wwwwa....aaaaa html
James Bradbury. Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miln and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
Lei Chen and Joan Bruna. Beyond the edge of stability via two-step gradient updates. arXiv, 2206.04172, 2023.
Xuxing Chen, Krishnakumar Balasubramanian, Promit Ghosal, and Bhavya Agrawalla. From stability to chaos: Analyzing gradient descent dynamics in quadratic regression. 2310.01687, 2023.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2021. URL /oattps://opexxeeeee
net for.m/oormmmmmm -ht.:e/eooooommmed------eee-eeee rherihttrree
Alex Damian, Eshaan Nichani, and Jason D. Lee. stabiiztaii The implicit bias of gra- dient descent the edge of stability. The Eleventh International Conference on Learning Representations, 2023. URL httthss:t///peneeee wwt.eeeewwwwweeeeeeeeeeeeeeeeeee
net/ fett:/oormmmmmdnhiia555
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, (9(6): 141-142, 2012.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1019-1028. PMLR, 06-11 Aug 2017. URL dinh1 7b. html
mlr. prwss/v70. wwtp::sswwwwlpp..seeeeeemm

---

Published as a conference paper at ICLR 2025
Darshil Doshi, Tianyu He, and Andrey Gromov. Critical initialization of wide and deep neural networks through partial jacobians: General theory and applications. arXiv, 2111.12143, 2023.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv, 1703.1 1008, 2017.
Kaiming He, X. Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL
httt:://githuee com goog e flax
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and general- ization in neural networks. S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018. URL https: http:://ccccceceeeeeeeeeo neur ips paper fi e 5a fa
2222 Paper.
Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho and Krzysztof Geras*. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2020. URL https:/7
openreview www.:eee/eeeeee mett//oommmdddrl13/1
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv, 1912.02178, 2019.
Dayal Singh Kalra and Maissam 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. URL BRLLhttps:/// /ttpps::/openzeview http:/n.e//////m forum?
id=A1 10GKK
Simran Kaur, Jeremy Cohen, and Zachary Chase Lipton. On the maximum hessian eigenvalue and generalization. In Javier Antorán, Arno Blaas, Fan Feng, Sahra Ghalebikesabi, Ian Mason, Mclanie F. Pradier, David Rohde. Francisco R. Ruiz, and Aaron Schein (eds.), Proceedings on Can't Believe It's Not Better! Understanding Deep Learning Through Empirical Falsification" at NeurlPS 2022 Workshops, volume 187 of Proceedings of Machine Learning Research, pp. 51- 65. PMLR, 03 Dec 2023. httts:/::/ http://pooceeeenneeeeeee
/ ml prw.s poress/vi8//// v187/kaur23a
html
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ArXiv, 1609.04836, 2016.
Lingkai Kong and Molei Tao. Stochasticity of deterministic gradient descent: Large learning rate for multiscale objective function. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33. pp. 2625- 2638. Curran Associates, Inc., 2020. UR./iitpsss///rrceeeiiggg neuripsss.......so.ee...... papeaprr fi hllles/pape/22220/f555 w1t.0/00011/e//// http://00000.0...0..oooooooaaaa0
pdf Itai Kreisler, Mor Shpigel Nacson, Daniel Soudry, and Yair Carmon. Gradient descent monotoni- cally decreases the sharpness of gradient flow solutions in scalar networks and beyond. ArXiv, 13064, 2023.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. ArXiv, 2003.02218, 18, 2020.
ChaoKunin Ma, Lie Wu, and Lexing Ying. Beyond the quadratic approximation: The multiscale structure of neural network loss landscapes. Journal of Machine Learning, 1())2477267, 2022 ISSN 2790-2048, doi: https://doi.orr/ 2200///..10020///.2.22222242
Intro/articl / hmm
12

---

Published as a conference paper at ICLR 2025
Lorenzo Noci, Alexandru Meterez, Thomas Hofmann, and Antonio Orvieto. Why do learning rates
transfer? reconciling optimization and scaling limits for deep learning, 2024. URL https
Edward Ott. Chaos in Dynamical Systems. Cambridge University Press, 2 edition, 2002. doi: 10.1017/CB0978051 1803260.
Thomas D. Rogers and David C. Whitley. Chaos in the cubic mapping. Mathematical Modelling 4(1): 9-25, 1983. ISSN 0270-0255. doi: http://bbaooooo/0000000000/s3530 URL https
www www.sciiecediic.c recte com/ httpa///c........ooooooooooooooa
Mihaela Rosca, Yan Wu, Chongli Qin, and Benoit Dherin. On a continuous time model of gradient descent dynamics and instability in deep learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Vaishaal Shankar, Alexander W. Fang, Wenshuo Guo, Sara Fridovich-Keil, Ludwig Schmidt, Jonathan Sagan-Kelley, and Benjamin Recht. Neural kernels without tangents. In ICML, 2020.
Jascha Sohl-Dickstein, Roman Novak, Samuel S. Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with standard parameterization. arXiv, 2001.07301, 2020.
Minhak Song and Chulhee Yun. Trajectory alignment: Understanding the edge of stability phe- nomenon via bifurcation theory. arXiv, 2307. 3307.04204, 2023.
Yuqing Wang, Zhenghao Xu, Tuo Zhao, and Molei Tao. Good regularity creates large learning rate implicit biases: edge stability, balancing, and catapult. arXiv, 2310.17083 2023.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and
sharpening and edge of stability. Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL htttps://opentt openreviewwwee.ee/op/peee..o.ww. wwwee//oormmiiommm toot://ooooii----ooott0
Jingfeng Wu, Vladimir Braverman, and Jason D. Lee. Implicit bias of gradient descent for logistic regression at the edge of stability. arXiv, 2305. 13088.1788, 2023.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mman
a novel image dataset for benchmarking machine learning algorithms, 2017. URL http: org/
Stss/1708.07749 cite arxiv: 708.7007777 1ee.11778.777777 077777Comment Dataset is freely available at mnist Benchmark is available http://fashhonnmmiiss.---
http:///httbb..oommmmmmmdmeeeeee websitc.cu-ccentral-I htw.://ccccccttlllamammooooooooo wwwaaaaww..comm
Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite- neural networks In Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 11727-11737. PMLR, 18-24 Jul 2021. URL
ml wwwpiiinssmm.rrrrre....e htt.a/acccccccaa...eeerrss/////a
html
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for understanding neural network dynamics. ArXiv, 2205.11781 2022.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of- stability training dynamics with a minimalist example. In The Eleventh International Confer- ence on Learning Representations, 2023. URL hattp:://pperrett net wtt.:/ouuttteteieeeoeemm
/EagBsMABO
A FURTHER DISCUSSION ON RELATED WORKS
Lewkowycz et al. 2020) examined the early training dynamics of wide networks at large learning rates. Using the top eigenvalue of the Neural Tangent Kernel (NTK) A at initialization = 0), they revealed a 'catapult phase', 2/1g < 17 < 7mmax: in which training converges despite an initial spike in training loss. Kalra & Barkeshli 2023) analyzed early training dynamics for arbitrary depths and
13

---

Published as a conference paper at ICLR 2025
width and revealed "sharpness reduction phase', 2/>8 < 2m///0, which opens up significantly as Closs increases with depth and 1/width.
Beyond early training, sharpness continues to increase, until it reaches break-even point @Jastzzeskii
ict al. 2020), beyond which GD dynamics typically enters the EoS regime (Cohen et al. (021). This
has motivated various theoretical studies to understand GD dynamics at large learning rates: Arora
eet al.
(2022); Rosca et al.] (2023); Zhu et al.] (2023); ; Wu et al. (2023): : Chen & Bruna (2023): :: Ahn
al. )
; Kreisler et al. (2023); y |Son & Yun (023); Chen al. (003)). In particular, Ma et
al.
Ma et al. 2022) showed that loss functions with sub-quadratic growth exhibit EoS behavior.
Arora et al. (2022) show that normalized gradient descent reaches the EoS regime. [Damian al.)
(2023) analyze the dynamics of the cubic approximation of the loss. Assuming a negative correlation
between the gradient direction and the top eigenvector of Hessian, they show that gradient descent dynamics enters a stable cycle in the EoS regime. (hhn et al. (2022) analyze EoS in a single-neuron 2-layer network and a simplified three-parameter ReLU network assuming the existence of a 'forward invariant subset' near the minima. Kreisler et al. (2023) analyzed scalar linear networks to show that the sharpness attained by the gradient flow dynamics monotonically decreases in the EoS regime. Wu et al. (2023) demonstrate that gradient descent, with any learning rate in the EoS regime, optimizes logistic regression with linearly separated data over large time scales. Below, we discuss closely related works in detail and clarify the distinction with our work.
Wang et al. 2022) analyze EoS in a 2-layer linear network using the norm of the last layer. They solely focus on cases that exhibit progressive sharpening right from initializations by considering assumptions (refer to Assumptions 4.1 and 4.2 of their paper) on the training dataset. Contrary to these assumptions, Figure ee) demonstrates that such assumptions are invalid in many realistic settings, where progressive sharpening is not observed at all.
Agarwala et al. (2022) showed that a modified model exhibits progressive sharpening and two-step oscillations at EoS using NTK as the proxy. They state that the UV model does not exhibit EoS behavior (see Section 3.2.1 of the referenced paper). This is because their analysis is restricted to the Standard Parmaeterization corresponding to Figure Ze, f) (c.f. Figure 8 of AAgawala et al. (2022). In contrast, we show that the UV model exhibits EoS behavior under the appropriate choice of parameterization and training example.
Zhu et al. 2023) proved EoS convergence for the loss { - 1)2, where æ; y E RR Additionally, they empirically demonstrated a bifurcation diagram in the space of abstract variables of and y. It is worth noting that while these bifurcations arise from the same underlying behavior, they contrast with our route to chaos bifurcation diagrams which quantify sharpness fluctuations with learning rate.
Chen & Bruna 2023) analyze two-step gradient updates of a single-neuron network and matrix factorization gain insights into EoS. Similar to work, they show bifurcation diagram of sharpness against the learning rate for the matrix factorization problem. While the scalar matrix factorization problem can be mapped to the UV model with a specific choice of || it is not //aar straightforward to apply their conclusions to the neural network setting, as it requires the correct choice of parameterization. In particular, the UV model under NTP parameterization, as shown in Figure cc, f), does not display EoS behavior at considerable widths, finding also noted by Agarwala et al. (0022). Observing EoS requires the correct choice of the parameterization (μPP) and training example. Furthermore, although the scalar matrix factorization in Chen & Bruna (2023) can be mapped to special case of the UV model considered in our work, we provide significant additional insights. In particular, with respect to the bifurcation phenomena in the UV model, we explain the existence of an attractor submanifold on which the EoS behavior occurs. We further show that on the EoS submanifold, the loss becomes subquadratic in nature and the gradient descent dynamics therefore become approximated by the cubic map, which is well-studied in the chaos literature. This makes clear the origin of the veriod-doubling route to chaos. Additionally, in Section 6, we extend the analysis of EoS beyond the UV model, comparing sharpness trajectories of synthetic and real datasets at FoS. In contrast to the synthetic datasets, sharpness trajectories of real datasets show long-range correlations in time. We take the first steps by attributing these long-range correlations to correlations in the dataset. Note that this setting cannot be mapped to the matrix factorization setting.
Song & Yun 2023) show that late-time trajectories oscillate around 2f/ne', where f is the network output and e" is the derivative of the loss. They refer to the term bifurcation diagram to describe these phenomena, contrasting with the sharpness versus learning rate bifurcation diagrams presented in our
14

---

Published as a conference paper at ICLR 2025
study. We quote Section 3 from their paper ...we plot the bifurcation diagram r(p) (p) and observe that GD trajectories tend to align with this curve.." Here, p and correspond to f and 2 in
nd
setting. They plot trajectories in the (p, 4) f 2 plane and their condition r(p) simply corresponds to the EoS condition À = for MSE loss. In contrast, our work presents bifurcation diagrams resulting from how the sharpness fluctuations vary with the learning rate. Therefore, the bifurcation diagrams from these works are not directly related the route-to-chaos bifurcation diagrams presented in our work.
Kong & Tao (2220) were the first ones to show that gradient descent dynamics becomes chaotic at large learning rates and converges to a statistical distribution instead of a minimum.
Chen al. (2023) analyzed large learning rate dynamics of toy models which are characterized by one-dimensional cubic map and demonstrated five different training phases: (a) monotonic, (b) catapult, (c) periodic, (d) chaotic, and (e) divergent. particular, they considered two-layer network with quadratic activation, where the last layer vector v is not trained and each entry is set to one. This model belongs to a family that effectively described by one variable Áf. In this model, the loss is sub- aobo-uudratic by construction (||Ux||12 y)² In contrast, the UV model that study is an effectively two-variable model and in these cases, training dynamically finds the attractive EoS manifold such that the loss has sub-quadratic nature this submanifold.
Concurrent work by Wang et al. (2023) categorizes training trajectories into three stages: (i) sharpness reduction, (ii) progressive sharpening, and (iii) edge of stability. They argue that different large learning rate behavior depends on the 'regularity' of the loss landscape. Specifically, they generalize toy landscapes from existing studies with parameters controlling the regularity. They show that models with good regularity first experience a decrease in sharpness and then progressive sharpening and enter the edge of stability.
Noci et al. (2024) examined the sharpness dynamics of networks with parameterization. They argue that the learning rate transfer property of μP correlated with consistent sharpness trajectories across varying depths and widths.
B EXPERIMENTAL DETAILS
B. DATASETS
Standard image datasets: We considered the MNIST Deng (2012). Fashion-MNIST Xiao et al.) (017)), and CIFAR-10 Krizhevsky (2009) datasets. The images are standardized to have zero mean and unit variance across the feature dimensions, and target labels are represented as one-hot encodings.
Random dataset: We construct a random dataset (X,Y) Y) {(x", .e"".""")) > with x ~ NN( I) and y ~ (0, both sampled independently. Note there no correlation between inputs and outputs.
Teacher-student dataset: Consider a teacher network f(;; with 6g initialized randomly as described in Appendix B.2 Then, we construct teacher-student dataset (X, {( } with nn (0, 7) and y f(x" 0p).
Random power-law dataset: Starting with the random dataset (X', Y'), we utilize the singular value decomposition of the input and output matrices
x'= P.
Y' Py
Next, we rescale the t singular value of Sy and Sy as
(Sz); (s
(Sy)x A,(S )ak
(5)
and re-construct input and output matrices as below
15

---

Published as a conference paper at ICLR 2025
X P=S
Y=
(6)
The variables Ags Bg, Ay, and By uniquely characterize the dataset.
Generative image dataset: Given a pre-trained network [(æ; 0) on standard image dataset listed above, we construct a generative image dataset (X,Y) = {(x") ) with x ~ ~N(00, and yd ((*";0).
B.2 MODELS
FCNs:
We considered ReLU FCNs without bias with uniform hidden layer width n.
We considered Myrtle family ReLU CNNs Shankar et al. 20200) without any bias with fixed number of channels in each layer, which refer to the width of the network.
CNNs:
ResNets: We adapted ResNet He et al. (016)) implementations from Flax examples Our imple- mentation uses Layer norm and inttialize the weights as (0, <./tama). For ResNets, refer the number of channels in the first block as the width.
We implemented all models using the JAX Bradbury et 2018), and Flax libraries Heek al. (2020).
B.2.1 DETAILS OF NETWORK PARAMETERIZATTIO
In this section, we describe different parameterizations used in the paper. For simplicity, we describe the parameterizations for FCNs. Nevertheless, these arguments generalize other architectures.
Standard Parameterization (SP): Consider neural network f Rd - Rddous with d layers and constant width Then, standard parameterization is defined follows:
h (x) w(),,
h (++1(()))) w((-)) (z))
f(x; 0) 1 ( (1(()11(xx))
where W(1) ol W ~ <2)n) for 1 < l < d, and W(d) ~ ((-) is the elementwise activation function. The input normalized such that |0| =00==aa din-
Neural Tangent Parameterization (NTP): Consider neural network f R + Rdaa with d
layers and constant width n. Then, the Neural Tangent Parameterization is defined as follows:
h (x)) =
Tw (1) x,
din
(1+1
(x)
w(+t) (x)). Vn
Tw
1
(()) (x((-1)(z))
f(x;0)
where w ~ (0, 1) for < 1 d and o6.) is the elementwise activation function. The input is normalized such that [x]] 2 = din-
Both SP and NTP closely related parameterizations SP with learning rate 17 @(1nn) learning rate becomes equivalent to NTP when the input dimension is equal to the width (d,, = n) Yang Hu (2021)
16

---

Published as a conference paper at ICLR 2025
Interpolating Parameterization: Consider a neural network f R + Roo with d layers and constant width n. Let W denote the weight matrix at layer I. Then, "interpolating parameterization" defined as follows:
(x) n n22
h 11110()))) (x) -((+)) (1))))(((())))) (x))
f(x;0)= n1/2 wid (n(d (z)d-))
(9)
Here, 8 is parameter that interpolates between standard-like parameterization and maximal update pa- rameterization. The weight matrices are sampled from Gaussian distributions: W()) ~ NN(0, ea/n"). W </))) for 1 < < d, and W(d) NNO, 1/n). We normalize the input such that ||||| 1.
Maximal update Parameterization (uP):: The maximal update parameterization corresponds to the = case in the above setting.
B.3 DETAILS FIGURES
Training loss and sharpness trajectories of 4-layer ReLU FCNs with = 512, trained on a subset of 5, 000 CIFAR-10 examples using MSE loss and GD: (a, d) SP with o = 0.5, (b, e) SP with oh 2.0, (c, f) pP with o = 2.0.
Figure 1
Figure 2 Training trajectories of the UV model with ||ж||| 1 and y 2 in the (Af, plane for different values of n, Rer and 17. The columns show initializations with different n and Retl, while the represent increasing learning rates for fixed initializations. The horizontal dash-dot line 1/ separates the stable (solid black vertical line) and unstable (dashed black vertical line) fixed points along the zero loss fixed line I. Forbidden regions, 2||=|||AA yl/ /neir X, (see Appendix C.2 are shaded gray. The nullclines A/ Af and A4+1 = X are shown as orange and white dashed curves, respectively. Sharpness reduction, progressive sharpening, and divergent regions are colored green, yellow, and blue. The gray arrows indicate the local vector field G(Af, f).,,),, A), which is the direction of the updates. The training trajectories are depicted as black lines with arrows, with the star marking the initialization. all cases, T)e = /iett/2 (introduced Section 5.2
UV model dynamics on the EoS manifold:(a) Bifurcation diagram depicting late-tima limiting values of À obtained by simulating Equation 33. (b) Bifurcation diagram of the UV model. In both figures, ||a:|| 1,y and Ner and Y)c 0.5.
Figure 3
Figure 4
Phase diagram of EoS: (a) Heatmap n /2 of 3-layer ReLU FCNs with 8 0 trained on subset of 5, 000 CIFAR-10 examples for 10k steps, with the weight variance o and learning rate multiplier e = ff as axes. X is obtained by averaging X over last 200 steps. As the color varies from blue to white, qス" increases, where the brightest white region indicates the EoS regime with /2 2 1. (b) Same heatmap with fixed o = 2.0, but varying s continuously.
EoS in synthetic vs rral-datasets: 2-layer lincar FCN trained on (first row) 5, 000 iid random examples with unit output dimension and (second row) 5, 000 CIFAR-10 examples. Different columns correspond to the bifurcation diagram, late-time sharpness trajectories, and the power spectrum of sharpness trajectories. Both models are trained for 10k steps using GD.
Figure 5)
Two-step phase portrait of UV model in (Af,9) phase plane: These plots equivalent to Figure d-f), but with training trajectory and local are plotted for every other step.
Figure 9
Sharpness and Weight Norm of 3-layer ReLU FCNs in SP with o = 1/3 and width 200, trained on subset of CIFAR-10 with 5, 000 examples using GD.
Figure 15:
17

---

Published as a conference paper at ICLR 2025
B.4 SHARPNESS MEASUREMENT
We measure sharpness using the power iteration method with m iterations. Typically, m 20 iterations ensure convergence. Exceptions requiring more iterates discussed separately.
B.5 POWER SPECTRUM ANALYSIS
For a given signal (t),, we standardize the signal
z(1) (1(1)
(10)
where is the mean and o² is the variance of the signal. Subtracting the mean removes the zero frequency component in the power spectrum. Next, consider the discrete Fourier transform F(w) of rtt):
1 T
F(w)
12m xtt),
(11)
Then, the power spectrum is P(w) = [F()) The normalization by in the Fourier transform ensures that the sum of the power spectrum equal to the variance of the signal, i.e., Ew, P(w) o².
B.6 ESTIMATION OF COMPUTATIONAL RESOURCES USED
Most of our experiments, aside from the phase diagrams, required minimal computational resources, estimated to be less than 50 A100 hours. In contrast, each phase diagram required 50 A100 hours, totaling 500 A100 hours for all phase diagrams. Including initial experiments, we expect our total usage be under 600 A100 hours.
C PROPERTIES OF THE UV MODEL
C.I DERIVATION OF THE FUNCTION SPACE DYNAMICS
Equations T and (2) can derived using the gradient descent update equations:
A
U2+1 = U₃ - 1)
(12)
Vieff
AfU
1t+1 = 1t -
(13)
Vneil
At step t + 1, the residual Af+4 can be written in terms the gradient updates of U and v:
Aftt1 Stt1 y
(14)
oit U1110-1 -
(15)
Tetf
A T
)
UU, (U;- Afwix"
v. - 11
x --y
(16)
/noi
nASe (a (NTU. Up u. + 17?
= Asi
v,
x
UTTT
(17)
Neff
Noft
= -Afi (1 1
1² || Af((Dft + +1))
(18)
Ball
18

---

Published as a conference paper at ICLR 2025
Here, Afi+1 only depends Af and X Similarly, write down the A4+1 using the gradient update equations:
Ac+1
UTTUUt1 + vo 11+1"++12"
(19)
Neeff
n||2=1?
Af((Aft + y) + 77 1/2||0||?3SF X
A
(20)
Nett
x 2
nd: - Aʃc
Al: +
A
A
(21)
Rell
Df
Equations 18 and 21 form closed system. This means that Aft+1 and Ac+1 are completely described using fe and A As a result, the complete dynamics of the UV model can fully
||al and y.
described using only these variables with three parameters effective parameters f);
C.2 FORBIDDEN REGIONS OF THE UV MODEL
In this section, utilize the non-negativity of À derive the condition for allowed regions within the phase plane for the UV model. Consider the function space equations written in terms of the pre-activation hax)) Uw:
1
f(x; 0)
v" h(x)
(22)
Reff
1
A
(||0)) [||0||||<0|) + |1h11²).
(23)
Naaf
Let cos(h, v) denote the cosine similarity between v and h. Then, the network output is bounded as
Melt |
As yl
cos(h, (h, v)
(24)
Next, using [||||||=||| |(b1)) 2 0, we can bound the product
||
|//||
using
2 2||2=||
A + +
(25)
The derived inequality describes the allowed phase plane regions for the UV model.
C.3 FIXED POINTS AND LINE
To identify the fixed points, we set Af++ Aft and A111 X in Equations and (2) This yields the dynamical fixed points the UV model. Table lists these fixed points along with their stability. Additionally, it provides the eigenvalues and eigenvectors of the Jacobian for the update maps described by Equations ( and (2), evaluated at the fixed points.
C.4 THE MAXIMUM LEARNING RATE TUPPER
In Section 4 stated that for 7) nupper 27c. training diverges for all initializations except for those at the fixed points. Here, we justify this claim.
First, shows that 7 approaches Yupper- fixed point III merges with fixed point II, reducing the convergence region to the EoS manifold. At this learning rate, the stability fixed point changes
from saddle to unstable as corresponding eigenvalue (1 ))
surpasses Consequently,
any initialization outside the EoS manifold results in divergence. Next, Figure Fa) shows that on the EoS manifold training diverges for y) ?upper- This corroborates initial claim.
19

---

Published as a conference paper at ICLR 2025
(A/")*))
eigenvalues
eigenvectors
Linear stability
Ratt 4
stable unstahle
(0, A) for A
1.1 1.1-
viai
1y >
(-M.0)
silelxy
) (1 allel
saddle
Tel
21=11 V
9,5
unstable
Mall
Iaily
1 3
9,5 2nliellv
P|[m|
unstable
Table 1: Fixed line (I) and points (IIIV) and corresponding eigenvalues and eigenvectors of the Jacobian of the update map in Equations (T) and (2). The stability is determined for < Nupper = 2n, = 2//ac//lar llaly Above this threshold, training diverges for all initializations except for those at the fixed points, as demonstrated in Appendix C.4
512, ny -1, -200 0.50
512, Ruy -. y -20 n 0.75
512. Nu17 -1, 1 -2.0, -1.00
10
Af
A,
Af
(a)
(b)
(c)
Figure 6: Phase portrait of the UV model for different learning rates 1j. The critical learning rate is
ne 0.5 and the maximum learning rate is Tupper
1.0.
C.5 SHARPNESS VERSUS THE TRACE OF HESSIAN
RAP
Training lass
Training loss
Training loss
18
30
(a)
(b)
(c)
2.19
2.0
2.00
AP TA
1.0
X
11F
10
step
step
(e)
Figure 7: Training trajectories of the UV model trained on a single example with ||>||| = 1 and using MSE loss and GD: (a. d) NTP with o² 0.5. (b, e) NTP with 512, o² 1.0, and (c, f) pP with 512, o² = 1.0.
In this section, show that the trace of the Hessian À (which is also the scalar NTK in this case), is an adequate proxy for sharpness. Figure shows training trajectories of the UV model, with X as proxy for sharpness and learning rate scaled as v) R/xo- These trajectories show similar trends to
20

---

Published as a conference paper at ICLR 2025
1/4
100
Training las
Training loss
Training
102
100
000
(a)
(b)
Ag
Au/A0
177
10
02
10
10
17
10
step
stop
(d)
(e)
(I)
Figure 8: UV model shows four training regimes: Training trajectories of the UV model trained on single example (x, y) with ||ж|| = 1 and y 2 using MSE loss and gradient: (a, d) NTP with n and o² 0.5 (b, e) NTP with n = 512 and o = 1.0, and (c, f) uP with n = and o² = 1.0.
those of H observed in Figure 7 with one key difference: during early training, À does not catapult during early training large widths (compare Figure ()) and Figure e)). Otherwise, effectively captures other qualitative behavior of A
C.6 THE DISTRIBUTION OF RESIDUAL AND NTK AT INITIALIZATION
In this section, we compute the distribution of A and À for the UV model initialization. Consider the UV model written in terms of the pre-activation h(x) = Uæ,
1 v- h(x)
f(x;;) 0) = Vnl-p"
(26)
1
A n-ーp (||0|/2|<0111 + |1<]1))
(27)
with Unj ~ N(0 oa/p)))))) Then, each pre-activation h; is normally distributed at initialization with zero mean and variance
du
dm
o
2 | 1
Es [hil = (UyUn)ayjR
(28)
np
nP
J,K=1
J,k=1
Hence, each pre-activation is distributed h, N(0 oB ||>||| /n²)) It follows that the network output also normally distributed at initialization with zero mean and variance
n
x
2
t |
Eg [
(hah)) 1
Σ
(29)
n1
n1 1-
np
np
np
Hence, the residual at initialization is distributed as Afo ~ N (-y, of 1/17). Similarly, we can also compute the distribution of À at initialization. The mean value of X given by
21

---

Published as a conference paper at ICLR 2025
1
Egg Aol
n1- (11=1] (1|)) <1]A]19) 20 ||:
(30)
where we have used (I|=|1?) 1 and (||1 o 11 Using similar computations. the second moment of À is given by:
4(n +
Ea [ 12-2p
(||=||"<<|e|| + <I|A11 + + 2|| +||==|((10((|1)111}})]
(31)
n
Hence, the À at initialization is distributed as λo N (20 || | 1==l"//)//)
C.7 EoS MANIFOLD IS A DYNAMICAL ATTRACTOR
To demonstrate that late time trajectories for > Te converge to the EoS manifold, we define Y X (Af + y). 3 lies the direction orthogonal to the EoS manifold, such that 0 corresponds the manifold itself, while is forbidden. Under this transformation, updates as 3 = B1 + =||= j It follows that 3 = 0 stays invariant under the dynamics and defines a nullcline.
Vier
Due to oscillations in Af near convergence, is instructive examine the two-step dynamics Agar wala al. (202)); Chen & Bruna (203), compactly denoted (Af++2, A++2) M(2) (A) Ae).
Figure 9 shows the two-step trajectories and the corresponding vector field ( B) in the
(Af, 3) plane.
We observe that there exists critical 7e such that for 1 < Me. (A B) points towards the stable zero-loss line (see Figure Pl)). By comparison, for 1j > T/., all points along the zero-loss line become unstable and the vector field directs towards points the B = 0 line, shown in Figure (b). The critical Ne for which all points on the zero-loss line become unstable thus gives a necessary condition for EoS:
512, Beff =1,
2.0,, v =0 0.20
=512, Dory =1, =2.0, rj =0.00000 0.60
Sieegnt
Doenget
15
85
5
TV
-5
Af
Af
Figure 9: These plots are equivalent to Figure 2 a. d) (n. = 0.5). but with training trajectory and local vector field plotted for every other step in (A 8) plane. The tilted dash-dotted line indicates the y 2 line.
C.8 CRITICAL LEARNING RATE FOR EDGE OF STABILITY
In this section, we estimate the required condition on the learning rate for the UV model to exhibit EoS. We specifically focus on the case with y >> 0 as for y = 0, can only decrease. As a result, the model does not exhibit progressive sharpening and EoS. In Section occurs as the zero-loss minima with the smallest À becomes unstable. From Equation (25) it follows that the smallest À with zero loss is
observed that the EoS
22

---

Published as a conference paper at ICLR 2025
2| 1u
Amin
(32)
/Tefff
This minimum becomes unstable if the learning rate 7/ exceeds critical value Tee given by
Beff
(33)
x yyy
It is worth noting that this is necessary condition for À to oscillate around 2/n. Otherwise, training converges to the zero-loss minimum with Amin for v) < 7mmax.
We can also derive the exact same result by analyzing the dynamics the EoS manifold. As discussed in Section 5.2 the dynamics on the EoS manifold is given by the map Aft41 M (A))), where
M(Af) Af (1-
21|| /
|1|æ=| Vmeet
2
Af(Af+)))
(Af
y)
(34)
As demonstrated Section 5.2 EoS in the UV model follows the period doubling route to chaos, with the period two cycle marking the onset. Hence, the conditions required for the emergence of the period two cycle are also necessary conditions for EoS. Consider the two-step dynamics on the EoS manifold given by the map (Af) = (M()))). This map has six fixed points (excluding three fixed points of the map M) summarized below
Af" Tēō(1 mpy 1 ( 1) npy)
(35)
211²
3 h(n, + V2
5th(n
wayl nuy k(n.F.w))
A
(36)
4m
Here ã = |||||| and h( F, v)
/-7 Tyn(2 + Zyn). For the fixed points to exist, we require the expressions inside the square root to non-negative, i.c.
var
nll Tett ml T|| lly
- 1)
glllelly +3 Met
Tell
20
72 1/1
(37)
x lly
//|l=|ly Man
(v32 32 2) /neil
+
- 1
72
(38)
As 17 < 72, the necessary condition for the period two cycle to emerge is > 7) //aa//la which coincides with the condition obtained earlier in this section.
C.9 SUB-QUADRATIC LOSS
In this section, we detail how the UV model naturally leads to sub-quadratic loss on the EoS manifold demonstrated Ma al. (2022).
23

---

Published as a conference paper at ICLR 2025
First, recall that the EoS manifold that connects fixed points II and IV satisfies the relation À [|]=/(((//t// //ar Then:
21Ix + v)
1. Starting Relation:
(39)
Vieff
\vnnft
2. Solving for Af:
2/ æ+l
y
(40)
X \
3. Factoring Out y
A
(41)
2112:11y
Veel
4. Defining T)c
?e
|| lly
(42)
nad
5. Expressing A with Tec
A y
-1
(43)
1²2
ncd
6. Computing the Loss: L =
AF²
(44)
2
As we mentioned in the main text, since ~ 0(1 the loss sub-quadratic near minimum.
D ADDITIONAL RESULTS ON SHARPNESS DYNAMICS
D.1 SHARPNESS DYNAMICS OF NTP NETWORKS
17
(a.
(b)
120
6:00
11
(c)
(d)
Figure 10: Training loss and sharpness trajectories of ReLU FCNs in NTP trained on subset of 10 examples using MSE loss and GD: (a, c) o 0.5, (b, d) o² 2.0.
Figure T0 shows that the sharpness dynamics of FCNs in NTP aligns with the behavior of FCNs in SP demonstrated in Figure
D.2 THE EFFECT OF BATCH SIZE ON THE FOUR TRAINING REGIMES
In this section, we examine the effect of batch size B on the results presented in the main text. We find that our conclusions are robust for reasonable batch sizes around B z 512. For even smaller batch sizes, the dynamics becomes domsn-dom,n and separating the inherent dynamics from noise becomes challenging. This observation further supports the use SGD to reduce the computational cost of experiments in the subsequent sections involving CNNs and ResNets.
24

---

Published as a conference paper at ICLR 2025
0.86
Training scouraty
Training
0:01-
AP
6022
127
32
-2
00
10
30
10
107
10
00
10
Step
step
step
(a)
(b)
(c)
0.08
Training acurcay
Training las
6:00
500
6.07
6:00
32
szep
stop
(d)
(e)
(I)
1.0
Training accuracy
Training lass
6:00
8.0u
0.02
5tep
step
step
(g)
(h)
Figure 11: Comparison of SGD trajectories with their GD counterpart for 3-layer FCNs in SP with o2 = 0.5 trained on a subset of CIFAR-10 consisting 5, 000 training examples with MSE loss. The learning rate is scaled 1 s/>Ø and batch sizes (a-c) B 512, (d-f) B 128, (g-i) B 32 are considered. GD trajectories are plotted using solid lines with transparency.
Figure |I shows that SGD trajectories of FCNs in SP begin to deviate from their GD counterpart significantly for batch sizes around B z 128. In contrast, for jPP networks this deviation begins at larger batch size of B zv 512 as shown in Figure 12
D.3 GENERIZABILITE THE FOUR TRAINING REGIMES
This section shows that the four regimes are generically observed for different architectures, loss functions, and datasets. Figure 13 shows training trajectories of CNNs and ResNets trained SGD with batch size = 512. Figure 14 show the training trajectories Transformers trained on the WikiText-2 2 dataset using the cross-entropy loss. These results further exemplify that four regimes of training are generically observed for complex architectures, datasets, and loss functions.
E SHARPNESS-WEIGHE NORM CORRELATION
Section 5.1 reveals that several aspects of the training dynamics are controlled by the fact that, for wide variety initializations, at early times trajectories move closer to the saddle point II, resulting in interim decrease in À (also proportional to the weight norm in this case), before eventually increasing. This critical point where all parameters are zero also exists in real-world models. We thus anticipate that in real- world models, the origin of the four training regimes may be related to similar mechanism. This would predict a decrease in weight norm as training passes near the saddle point, followed by eventual increase.
25

---

Published as a conference paper at ICLR 2025
6.0
20
Training scouraty
312
Training
90/00
6.77
6:11
0.00
12
32
60
10
10
10
00
10
Step
step
step
(a)
(b)
(c)
6.00
500
Training ccurccy
Training
6.2
6/31
12
step
stop
(d)
(I)
0.5
6106
Training accuracy
0.2
5tep
step
step
(h)
Figure 12: Same setting Figure 11 except we used 3-layer FCNs in pP with o = 2.0.
Figure 15 validates this hypothesis. During the sharpness reduction and intermediate saturation as the
decrease in the weight norm, followed by an increase in the weight network undergoes progressive sharpening, following the prediction from the UV model. Similar correlations between the last layer weight norm and sharpness are utilized by Wang al. (2022) to analyze the EoS phase. By comparison, we focus on the correlation between sharpness and weight norm during early training attribute the emergence of four regimes to the critical point corresponding to all parameters being zero. In Appendix E we provide further evidence for this correlation between sharpness and weight norm, extending this relationship to CNNs and ResNets.
regimes,
This section presents additional results for Section sharpness and weight norm during training-
6 further supporting the relationship between
Figure Tolis is an extended version of Figure 15 where we plotted the whole training trajectories and measured Pearson correlation
Et=n (0r - 0:) (x (2
(45)
Eten (0v 0.)² 2 - (X - (A*))
(1*
Here t 12 2 and 0, (1)=11
Figure 17 shows the weight norm of each layer separately for the experiment in Figure 15 This result shows high correlation between weight norm and sharpness through training.
We also confirm these correlations between weight norm and sharpness in CNNs for the experiment in Figure 13 a, b).
26

---

Published as a conference paper at ICLR 2025
Training
130
250
00
07
16
10
10
stop
stop
(a)
(b)
IA
10.
10
100
10"
10
step
stop
(c)
(d)
Figure 13: Training trajectories of (a, b) a 5-layer CNN in SP with n 64, and (c, d) ResNet- 18 with LayerNorm in SP, also with 64. Both models are trained the CIFAR-10 10 dataset with MSE loss using SGD. The learning rate scaled as r = 5/>," and batch size is = 512. In panel (d), X becomes negative during early training. This is due to the power iteration method returning the largest eigenvalue by magnitude.
62
65
05
2.50
(b)
.12
1.8
28
(e)
Figure 14: Training loss, sharpness and test loss trajectories of Transformers trained on Wikitext-2 using cross-entropy loss: (a, b, c) Pre-LN Transformers and (d, e, f) Pre-LN Transformer without last LayerNorm.
E.1 SETTING THE LAST LAYER TO ZERO ELIMINATES SHARPNESS REDUCTION
Figure 19 shows setting the last layer zero eliminates sharpness reduction during early training.
27

---

Published as a conference paper at ICLR 2025
88.00
$0.09
5002
5300
step
step
(a)
(b)
Figure 15: Sharpness and Weight Norm of 3-layer ReLU FCNs in SP with o 1/3, trained on subset of CIFAR-10 with 5, 000 examples using GD.
1.00
03.03
88.04
33:03
0.80
0.77
step
step
step
(b)
Figure 16: Sharpness and Weight Norm of 3- 3-layer ReLU FCNs in SP with o 1/3, trained subset of CIFAR-10 with 5, 000 examples using GD.
a2
9.25
AM
100
20
8.70
32.0
10
8.00
107
10
10
10
stop
stop
(b)
Figure 17: Weight Norm of each layer in 3-layer ReLU FCNs (same experiments as Figure 15 (a. c) SP with of = 1/3. All results obtained by training subset of CIFAR- 10 with 5, 000 examples using GD.
F ADDITIONAL PHASE DIAGRAMS OF EOS
This section demonstrates additional phase diagrams of EoS and quantifies the effect of batch size in the EoS regime. Figure 20 shows phase diagrams of EoS for FCNs trained CIFAR-I0 with MSE loss using SGD for 10, 000 steps for three different batch sizes. We observe that as the batch size decreases, X oscillates at value different from 2/1 depending on of and s. For large o and small H favors smaller value for smaller batch size, which is in agreement with the observation in Cohen et 2021). In contrast, H can be larger than 2/n for small o² and large s late training times.
Figures 21 and 22 show the phase diagrams of EoS for CNNs and ResNets trained on the CIFAR-10 10 dataset with MSE loss using SGD for 10, 000 steps with learning rate 17 c/2 and batch size B = 512. In contrast to the FCN phase diagrams, these architectures exhibit EoS behavior at smaller values of and larger values of o indicating their implicit bias towards EoS. Moreover, we observe in ResNets, that EoS is less sensitive to change o² likely due to combination LayerNorm and residual connections Doshi al. 2023
28

---

Published as a conference paper at ICLR 2025
255
11.
115
282
75.4
22.5
112
10
11
10 stop
10
107
stop
(a)
(b)
(c)
11.8
2.34
11.2
11.7
11.00
11.5
34
11.3
022
10
10
107
step
(d)
(e)
(I)
Figure 18: Weight Norm of each layer for 5-layer CNNs in SP (same experiments as Figure 13 a. b): (a) Total weight norm; (b-f) Weight norm of each layer. We see that for 16, the initial catapult in sharpness X (Figure b)) is accompanied by catapult in total weight norm. Notably, the total weight norm and per-layer weight norm, whether catapults (a, c-e) or (b, f), show a decreasing trend during the early sharpness decreasing stage, followed by an eventual increase.
0.30
0/70-
3
0000
(a)
0.0
0.27
10
507
(d)
(c)
(f)
Figure 19: Sharpness trajectories of ReLU FCNs trained 5k subset of CIFAR-10 10 examples using
MSE loss and GD. (a, b, c) SP with o² 0.5, SP with o = 2.0 and μP with o 2.0. (d, e, f) the bottom row shows the effect of setting the last layer to zero initialization. The dashed lines in the sharpness figures show the 2/n threshold.
It is worth noting that EoS boundaries in these phase diagrams are time-dependent. For instance, models close to the EoS boundary may eventually reach EoS on training longer (see Figure (3(b) = 1.0 for example), causing shift in the EoS boundary. Nevertheless, models with small learning rates, large of and small may never show EoS behavior, regardless of training duration, as predicted by the UV model and seen in Figure e)).
29

---

Published as a conference paper at ICLR 2025
30
30
25
23
20
20-
-14
1.3
10
10
03
01
25
10
15
75
3.0
0.0
02
04
0.6
0.5
10
(a)
(b)
2.1
20
11
1.5
10
12
6.1
65
02
81
25
10
1.0
(c)
(d)
1.5
9
75
(c)
(f)
Figure 20: Phase diagram of EoS for 3-layer FCNs trained on CIFAR-10 with MSE loss using SGD with three different batch sizes: (a. b) B = 512, (c, d), = 128, and (e, f) B 32. The color indicates the value of 7y" /2, where X is obtained by averaging X over the last 200 steps. Except for the batch size, all settings are identical to Figure ! Black dash-dotted lines indicate the phase boundary n //2 1. For elarity, these lines generated from data smoothed with Gaussian kernel.
40
011
30
34
20
93
100
1.2
0.1
100
07
04
0.6
0%
00
<2
(a)
(b)
Figure 21: Phase diagram of EoS for 5-layer CNNs SP with width 64 trained with MSE loss using SGD for 10, 000 steps with learning rate 17 = 7/x and batch size B 512.
G ROUTE TO CHAOS
G.1 ROUTE TO EOS IN REAL DATASETS
This section presents additional bifurcation diagrams for different architectures and datasets. These results show the reminiscent the period-doubling route to chaos observed different architectures and datasets. In all figures, choose the smallest and largest learning rate exhibiting FoS for
30

---

Published as a conference paper at ICLR 2025
40
60
301
30
20
23
10
12
01
10
70 30 00 5.0
02
0. 0.5
10
(a)
(b)
Figure 22: Phase diagram of EoS for ResNet-18 SP with width n trained with MSE loss using SGD for 10, 000 steps with learning rate e/x5 and batch size B = 512. For 1, the average eigenvalue X observed to less than 2/n. Upon detailed investigation of the trajectories, found that X," oscillates around value lower than 2/n. We leave this anomalous behavior as observation.
plotting the trajectories and power spectrum. The structured route to chaos in realistic experiments can be disrupted due to variety of reasons. Below, we discuss a few of them.
Measurement of only the top eigenvalue of Hessian: In our experiments, we only measured the top eigenvalue of the Hessian. However, when multiple eigenvalues of Hessian enter EoS, plotting only the top eigenvalue of Hessian is projection that could obscure all the structured routes to chaos that the system may exhibit.
The effect of correlations in real-world datasets: Real-world datasets inherently contain correlations between different samples (x, y). These correlations can be quantified using the input-input covariance matrix Exx XX" E Rddn and output-input covariance matrices Σyx YX" E igdauxXay In Appendix G we find that key determining factor in observing route-to-chaos is whether the power spectrum of Σxx is flat exhibits power law decay. We show that power-law decay in the singular values of the Σxx results in long-range correlations in time and dense sharpness bands observed in real datasets.
1.0-
(.903
0.8-
0L985
0.2
0,0
0.6
e-1S 7830
0.2
0.4-
0.655
0.60
0.0
05
1.0
155
0.0
00.
0.2
0.3
0.4
0.5
(a)
(b)
(c)
Figure 23: 4-layer FCN in SP with width n 512 trained on a subset 5000 examples of MNIST with MSE loss using GD. Both power spectrums computed using the last 1000 steps of the corresponding trajectories.
G.2 THE EFFECT OF POWER-LAW TRENDS IN DATA ON SHARPNESS TRAJECTORIES
In this section, analyze 2-layer lincar FCN trained the power law dataset described in Appendix B.I to understand the origin of long-range correlations sharpness trajectories and dense sharpness bands in realistic datasets.
Figure 28 shows the bifurcation diagram, late time trajectories, and the associated power spectrum
of the network trained the power-law dataset with the same A, 1.0 and Ag = 1.0, for four different combinations pawer-law exponents: (i) B, = 0.0, B, = 0.0, (ii) B, = 1.0, B, = 0.0, (iii) By 0.0, By 1.0, and (iv) B, 1.0, By 1.0. We observe that power-law trend to the
31

---

Published as a conference paper at ICLR 2025
1.0
100
0.05
0.81
0.66
01002
3/44
083
2.0
0600
0550
0.0
0.1
0.3
0.1
0.5
(a)
(b)
(c)
Figure 24: 4-layer FCN in SP with width n 512 trained on subset of 5000 examples of Fashion- MNIST with MSE loss using GD. Both power spectrums are computed using the last 1000 steps of corresponding trajectories.
1535
C- 1606
1000
0.8
0.2
1.005
0.0
0.4.
n 1-33
0.655
0.2-
0002
0003
2
0.0
0.11
0.3
0.1
0.5
(a)
(b)
(c)
Figure 4-layer FCN in SP with width n 512 trained on subset of 5000 examples of CIFAR- 10 with MSE loss using GD. Both power spectrums are computed using the last 1000 steps of the corresponding trajectories.
-6.6699
C- 0E42
1.00
1,00
0.70
0.01
0.2
105
0.22
1000
0.00
0.00
0.70
1006
0.38
0.10
8000
0.0
011
0.2
03
0.1
05
(a)
(b)
(c)
Figure 26: 5-layer CNN in SP with width n = 32 trained on a subset of 1000 examples of CIFAR-10
with MSE loss using GD.
singular values of the input matrix results in dense sharpness bands observed in real datasets. is worth noting that this is one way to obtain dense sharpness bands and in general, there can be many other methods.
G.3 ROUTE TO CHAOSS IN SYNTHETIC DATASETS
In this section, we analyze the route to chaos in synthetic datasets to gain insights into the dense sharpness bands in realistic datasets. We considered datasets, defined as follows:
Teacher-student dataset: Consider a teacher FCN f dd - Rd with din = 3072, dou = 10, depth d, and width 512 in Standard Paramaterization. Then, construct teacher-student dataset (X, Y) consisting of P = 5000 examples with x ~ NO, I) and y" = f(x"; 9p). Next, train student FCN with the same depth and depth the teacher FCN this dataset.
32

---

Published as a conference paper at ICLR 2025
00033
01120001
LIZZ
0.10
0,51
1.00-
1108
0000
0.2
1002
0.50
8000
0.0
011
0.2
0.3
0.1
0.5
(a)
(b)
(c)
Figure 27: ResNet-18 18 in SP with width n 32 trained on a subset of 1000 examples of CIFAR-10
with MSE loss using GD.
Figures 29 and 30 show the bifurcation diagram, late time sharpness trajectories and the associated power spectrum of linear and ReLU FCNs trained the teacher-student task. These figures show that while linear FCN shows the period-doubling route to chaos, ReLU FCN shows long-range correlations observed in real datasets.
Consider 5-layer CNN fæx ()) in SP with n = 64, trained on the CIFAR-10 dataset with MSE loss using SGD with learning rate 1j = 12/x" and momentum m for 100k steps. This model achieves test accuracy of 76.9%. Then,
Generative dataset:
construct generative image dataset (X, Y) consisting of P = 5000 examples with x ~ A00 I) and y f(*"; 0). Next, we train an FCN in SP with depth d, width and weight variance of 0.5 on the generated dataset.
Figures 31 and 32 show the bifurcation diagram, late time trajectories and the associated power spectrum of a 4- layyer ReLU FCN with linear and ReLU activations, trained on the generative CIFAR 10 dataset. We observe that while the linear network shows period-doubling route to chaos, the ReLU shows long-range correlations observed in real-datasets.
33

---

Published as a conference paper at ICLR 2025
1.4
1.3-
1.001
0.5
1.2
1.3
-06
3.00
0.0-
83.00
5000
0.0
step
(a)
(b)
(c)
2002
2.0
1.00-
0.2
1.04-
0..3
12
0.01
1.0
1.5
0.5
1.0
1.0
5000
9200
2400
5000
5800
10000
0.0
0.1
0.2
0.4
0.5
step
(d)
(e)
(l)
15112005
1.5
15
0.5
1.0
1.0
0.00
2.0
0.5
10
20
2600)
0920
1403)
1900
1681
10003
0.0
0.1
0.2
0.3
0.1
05
stop
(h)
1024
1.05-
2.00
0.00
1.03
0.05
0.0
C-11
1.2
0.02
0.01
10.5
1.0
1.5
1.0
9000
9.001
3400
5003
5800
1300
0.0
0.1
0.2
0.1
0.5
(k)
(I)
Figure 28: Bifurcation diagrams, late-time sharpness trajectories, and power spectrum of a 2-layer linear network trained the power-law dataset for different parameter values: (a-c) B2 0.0, By 0.0, (d-f) BN = 1.0, By = 0.0, (g-i) B, = 0.0, Bg = 1.0, and (i-l) B, = 1.0, B, = 1.0. All power spectrums are computed using the last 1000 steps of the corresponding trajectories.
34

---

Published as a conference paper at ICLR 2025
12-25 2859
0.4
17/4 (2
0.00
1.25
0.2
1.0-
1,204
1.15
0.2.0
01805
0.150
(1875
(1560)
0.0
00.
0.5
(a)
(b)
(c)
Figure 29: 2-layer linear FCN in yP trained on the teacher-student task. Both power spectrums are computed using the last 1000 steps of the corresponding trajectories.
1.0
1990
7542
3,9885
3,980
0.2
0.51
445673
0.01
c- EDST
C-4
0.2
0.6
0.9000
0.
0.5
1.01
1.5
2.0
25
0130
1800
50.01
S140)
500
19400
01.1
stop
(a)
(b)
(c)
Figure 30: 4-layer ReLU FCNs in j&P trained on the teacher-student task. Both power spectrums are computed using the last 1000 steps of the corresponding trajectories.
0.504
0.4
3/4 /2
0,010-
0,000
000
1.0
155
2.0
0320
step
(a)
(b)
(c)
Figure 31: 4-layer linear FCNs in SP with oi with MSE loss using GD. Both power spectrums are computed using the last 1000 steps of the corresponding trajectories.
0.5 trained on the generative CIFAR-10 task
1.0.
[.000
0.8-
0.6
(110
3
1.
0.00
0.1
0.55
00
0.1
0.2
033
0.4
05
0750
93/0
0050
9900
1000
0.9
00
(a)
(b)
(c)
Figure 32: 4-layer ReLU FCNs in SP with o² = 0.5 trained on the generative CIFAR-10 10 task with MSE loss using GD. Both power spectrums are computed using the last 1000 steps of the corresponding trajectories.
35


Reviews:
[{"review_id": "4NhIKdEmvr", "reviewer": "Reviewer_Knqc", "summary": "The paper studies nonlinear gradient descent dynamics in a two-layer linear model trained on a single example. The model in question (a \"UV\" model) extends the model analyzed in Lewkowycz et al. (2020) by including a nonzero target scalar and incorporating different parameterizations (standard / neural tangent / $\\mu P$) via an effective width parameter. The resulting gradient descent dynamics has a rich phase portrait with several stable/unstable fixed points/lines and different evolution patterns depending on initial conditions and hyperparameters. In particular, the authors describe the absence or presence of a catapult regime, sharpness reduction, progressive sharpening, edge of stability oscillations. For the edge of stability (EoS) regime, it is shown to occur on a particular line in the parameter space at learning rates above a particular critical value, and its origin is related to the well-known period-doubling fractal mechanism. After that, the paper discusses a number of predictions one can make for realistic models based on the obtained phase picture of the toy model: a correlation of the parameterization with sharpness, temporary decrease of the weight norm, conditions of EoS and sharpness fluctuations.", "questions": "Line 813: \"constant width networks in SP with learning rate $\\eta=\\Theta(1/n)$ are equivalent to those in NTP Yang & Hu (2021).\" - Where exactly is this shown in Yang & Hu (2021)? Table 1 and Figure 2 there actually show SP and NTP as different parameterizations.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is easily readable and clear. A substantial background and literature review are provided. A large number of appendices provide substantial additional details of the methodology, computations, and experiments. \n\nThe toy model proposed in the paper is simple, but sufficiently rich to describe multiple complex effects observed in realistic models (catapult regime, sharpness reduction, progressive sharpening, edge of stability oscillations). The paper gives a systematic discussion of how these effects depend on the initialization, parameterization, learning rate.\n\nThe paper uses the toy model to conjecture several qualitative predictions on the learning behavior in realistic models/data. These are illustrated by multiple experiments with different networks and data.", "weaknesses": "I generally like the paper, but I'm not convinced in the sufficient novelty and significance of its results. The UV model studied in the paper is a toy model and is only a slight generalization of the well-known model from Lewkowycz et al. (2020). The study of this model is a fairly straightforward analysis of the 2D phase portrait of a simple $(\\Delta f, \\lambda)$ dynamics with respect to several hyperparameters. The paper does not develop any new analytic methods or prove any new general theorems. All the regimes discussed in the paper - the catapult mechanism, edge of stability, etc. - have already been demonstrated repeatedly in very similar models in multiple previous works (and I should give the credit to the authors for carefully listing them in two related work sections). The added value of the present paper - detailed analysis of the extended toy model and some predictions for realistic models - can certainly be of interest to a part of ML community, but just looks relatively marginal to me on the ICLR scale.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W8ZpIX4d0c", "reviewer": "Reviewer_AyGJ", "summary": "This paper covers a tremendous amount of theoretical work, presenting a simple model which the authors analyze to study the role of initialization and parameterization on initial movements in sharpness during training; the role of initialization, parameterization, and learning rate on the existence of the Edge of Stability; and finally drawing analogy between the EoS and period-doubling bifurcations.", "questions": "1) Under light investigation, the mathematical derivations for Eqs. 1 and 2 are sound. However, one might expect that the dynamics of the theoretical model might be independent to the sign of $\\delta f_t$. Can you provide an intuition to why this might not be the case?\n2) What is the justification for setting $\\eta = c/\\lambda_0$?\n3) What is a 'low/minimal sharpness bias', as in Section 6.1?\n4) What can I do to load this PDF more quickly?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "The Edge of Stability (EoS) is an important phenomenon observed in recent research, spurring questions about the understanding of dynamics for non-convex optimisation with deep neural networks. This paper studies the early dynamics of these networks, providing insights on the early movements in sharpness and the existence of EoS later on in the training trajectory.\n\nThe theoretical analysis is clear, and the resulting phase diagram in Section 4 appears novel, and is insightful and easy to understand. \n\nThe applicability of the analysis appears general to a number of common parameterizations and can be used to draw insights to deep neural networks in general.", "weaknesses": "While this work presents impressive analysis, revealing insights toward important phenomena in optimization, the empirical support for individual claims are less apparent, possibly as a combination of limited empirical evidence and poor presentation. Additionally, the theoretical approach is hugely similar to existing work (Lewkowyz, 2020), though many additional insights revealed, especially concerning early dynamics. Finally, the connection of this work to common motivations for optimization analysis (e.g. generalization, feature learning) is indirect - only the suggestion that the existence of EoS can improve generalization error is presented in Fig 4 and discussed sparingly. \n\nThe areas for improvement in empirical support are detailed as follows:\n1. Section 6.1 claims that $\\mu P$ networks begin training in a flat region of the landscape where gradients point to increased sharpness, while NTP/SP networks experience a reduction, the evidence in Fig. 1 suggests that this distinction is not so clear. \n2. Section 6.2 suggests that real-world deep neural networks (sometimes) undergo an interim decrease in $\\lambda$ and the weight norm, but the real-world models presented in Fig. 1 and App. E exhibit a decrease in weight-norm. As real-world initializations typically use 0-mean weight initializations, one might expect that at least some of the models can be initialized in the R1 regime and immediately enter EoS. Therefore, the evidence that initialization in R1 can exist for real-world networks is limited, questioning the extent to which the theoretical model can apply to real-world models.\n3. In general, the oscillation of EoS is normally observed around $2/\\eta$, which is also stated in the paper in Sections e.g. 1 and 3. However, the point of oscillation predicted by the UV model (as seen in the phase diagrams) do not necessary have to be near 2/\\eta. This distinction does not appear to be sufficiently addressed. \n4. Some of the empirical evidence discusses the effects of SGD, but it is an open question how and whether EoS phenomena translates to SGD. The methodology at which EoS is computed for SGD experiments in the Appendix are extremely unclear. An important scaling for the learning rate as $\\eta = c/\\lambda_0$ is stated, but not justified or explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4hQtjUKmLH", "reviewer": "Reviewer_yW4m", "summary": "This work investigates training dynamics of a simple UV model and shows that it replicates several phenomenological properties that have been observed in larger networks. In particular, the same 4 \"phases\" of training are observed. The authors use the tractability of the UV model to identify why this occurs and when/why some of these phases can be \"missed\". The authors then show that some predictions from the UV model are found in larger models.", "questions": "Q1. How truly \"similar\" are the dynamics in the UV model and the FCNs used to motivate this paper? \n\nQ2. Does the analysis performed by the authors demonstrate \"Universal\" properties?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written and motivated. \n\n2. The numerical experiments are well designed and insightful. \n\n3. The main conclusions of the analytical work are well explained and relatively easy to grasp. \n\n4. The goal of understanding training dynamics is an important one that is well aligned with the NeurIPS community.", "weaknesses": "1. The use of simplified models is important and I appreciate that the authors demonstrated that the UV model, despite being very simple, has some of the same qualitative properties as what has been observed in other models. While the authors provide evidence that the features of their analysis can be used to understand training dynamics in larger models, I think there are a few weaknesses in these results: \n    a. First, I think it is worth noting that (at least as far as I could see), comparing Fig. 1 and Fig. 7 (FCN vs. UV model), there seem to be some differences in behavior of the loss and $\\lambda$. In particular, there are oscillations in the training loss that do not appear in the FCN, and there is less of a spike in the training loss and $\\lambda$ for say the model with $\\mu$P (Fig. 7C, F). This makes it unclear how \"similar\" the phenomenology are between the UV model and the training phases that have previously been studied. One way to try to quantify the \"similarity\" would be to use metrics for estimating topological conjugacy between training dynamics (arXiv:2302.09160). \n    b. The authors show that $c$, $\\sigma_W$, and $s$, which are predicted in the UV model analysis to play an important role in determining EoS are predictive of EoS in larger networks (e.g., Fig. 4). However, it seems that some parameters are much more important than others, e.g. for large enough value of $s$ any value of $c$ suffices to see EoS. This was not discussed in the main text and it was unclear to me why this would be. Similarly, the authors say there is a correlation between EoS and test accuracy, which is in my opinion weakly shown in Fig. 4. Maybe this is due to the fact that the color is saturated at 1? Or perhaps the way of estimating EoS ($\\eta \\lambda^H / 2$ ) is not the best way of estimating EoS? \n     c. The lack of period doubling in large networks is a very interesting point (and the idea that the data could be suppressing it somehow is really cool!), but to me this suggests again the possibility that there are fundamental differences in the dynamics between the UV model (trained with one example) and modern architectures that makes it unclear how generally valid the analysis performed is. Providing more discussion on this would be helpful. \n\n2. I did not follow how the authors got that the loss was $\\mathcal{L}(\\theta) = 1/2 \\Delta f^2 = y^2/2 (\\eta_c \\lambda / 2 - 1)^2$ in Sec. 5 (\"Connections to sub-quadratic loss\"). Providing more detail would be useful. \n\n3. The title of the paper \"Universal sharpness\" is perhaps too strong. I did not feel like the authors ever really made a claim of universality, nor did they provide much analysis of it, in the main text. \n\nMINOR POINTS: \n1. I was under the impression that the role of loss landscape sharpness and generalization was called into question by Dinh et al. (2017). Is there a reason this was not cited in the Introduction when discussing loss landscape and generalization? \n\n2. In Appendix E, the examples used to validate the hypothesis about the weight norms is a 3-layer FCN, while Fig 5 has a 2-layer FCN. Is there a reason for these differences? Were the weight norms consistent with the hypothesis in 2-layer FCNs?\n\n3. The authors mention that fixed point analysis is a generally powerful approach (Sec. 7). I think it would be good to discuss it's limitations more, particularly for nonlinear DNNs. \n\n4. Very minor, but I think putting references in parantheses would make the paper a little easier to read. I.e., \"the early training dynamics [Kalra and Barkeshli (2023)]\" (page 1).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4X17JWMtkB", "reviewer": "Reviewer_GRmS", "summary": "The paper analyzes four phases of training NNs (early time transient, intermediate saturation, progressive sharpening, edge of stability) from the perspective of the top eigenvalue of the Hessian of the loss. The authors show that these dynamics depend strongly on the parameterization of the network (width; standard vs. maximum parameterization) and the optimizer settings (large vs. small learning rate). This is not only done experimentally, but also in a linear toy model, the UV model. For this model, a fixed point analysis of the gradient dynamics can be performed, which shows that even this simple model can exhibit all four stages.", "questions": "- What actionable insights can be derived from the analysis?\n- Does Section 6 rely on the trace of the Hessian, or does it analyze its top eigenvalue? If it relies on the trace, can it be ensured that the qualitative picture is the same (i.e., do the results from Appendix C.5 for the UV model carry over to more complex architectures)?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "The paper is mathematically dense and, to the best of my knowledge, rigorous. The analysis of the edge of stability manifold as a chaotic attractor is indeed interesting, and so is the observation that EoS is correlated with improved performance. The analytical model for the UV system is valid and accessible. To the best of my (very) limited knowledge in this regard, this is the first analysis of this kind. Being not familiar with the related work, however, I am not very confident in my assessment.\n\nThe paper contributes to a better understanding of training neural networks by carefully studying a simple toy system and by showing that fully-connected NNs also exhibit similar properties. Via this connection, one can argue that the theoretical analysis indeed sheds light on the internal workings of NN training.", "weaknesses": "Among the weaknesses I would count (and please excuse my ignorance in this field):\n- The analysis in Sec. 5.2 is quite simple and appears obvious or at least little surprising. While the observation that the EoS manifold is an attractor is quite relevant, the remaining statements (Result 1, Corollary 5.1, bifurcation diagrams) appear straightforward consequences of this observation. This is not a weakness of the paper in itself, but this observation could be used to shorten this section.\n- It is not immediately clear how the observations made in this paper will affect deep learning engineering practices. In other words, some actionable insights would be appreciated, linking the insights from the paper to recommendations regarding learning rate, parameterization, and architectural choices.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)