Phase diagram of early training dynamics in deep neural networks: effect of the learning rate, depth, and width
Authors: Dayal Singh Kalra
Maissam Barkeshli
Venue: NeurIPS 2023 poster (2023)
Keywords: Optimization dynamics
Phase diagrams
learning rate transition
Catapult effect

Abstract: We systematically analyze optimization dynamics in deep neural networks (DNNs) trained with stochastic gradient descent (SGD) and study the effect of learning rate $\eta$, depth $d$, and width $w$ of the neural network. By analyzing the maximum eigenvalue $\lambda^H_t$ of the Hessian of the loss, which is a measure of sharpness of the loss landscape, we find that the dynamics can show four distinct regimes: (i) an early time transient regime, (ii) an intermediate saturation regime, (iii) a progressive sharpening regime, and (iv) a late time "edge of stability" regime. The early and intermediate regimes (i) and (ii) exhibit a rich phase diagram depending on $\eta \equiv c / \lambda_0^H $, $d$, and $w$. We identify several critical values of $c$, which separate qualitatively distinct phenomena in the early time dynamics of training loss and sharpness. Notably, we discover the opening up of a "sharpness reduction" phase, where sharpness decreases at early times, as $d$ and $ 1/w$ are increased.

Phase diagram of early training dynamics in deep networks: effect of the learning rate, depth, and width
Dayal Singh Kalra *t dayal dayal@umd... @umd. edu
Maissam Barkeshli +1% maissundund, edu
Abstract
We systematically analyze optimization dynamics in deep neural networks (DNNs) trained with stochastic gradient descent (SGD) and study the effect of learning rate 17, depth d, and width w of the neural network. By analyzing the maximum eigenvalue X, of the Hessian of the loss, which a measure of sharpness of the loss landscape, (i) an early time transient regime, (ii) an intermediate saturation regime, (iii) a progressive sharpening regime, and (iv) late time "edge of stability" regime. The early and intermediate regimes (i) and (ii) exhibit rich phase diagram depending on 17 s/>f' , d, and w. We identify several critical values of c. which separate qualitatively distinct phenomena in the early time dynamics of training loss and sharpness. Notably, we discover the opening up of "sharpness reduction" phase, where sharpness decreases at early times, as d and 1/w are increased.
find that the dynamics can show four distinct regimes:
1 Introduction
The optimization dynamics of deep neural networks (DNNs) is rich problem that is of great interest. Basic questions about how to choose learning rates and their effect on generalization error and training speed remain intensely studied research problems. Classical intuition from convex optimization has lead to the often made suggestion that in stochastic gradient descent (SGD), the learning rate should satisfy 1) < 2/x", where A is the maximum eigenvalue of the Hessian H of the loss, in order to ensure that the network reaches a minimum. However several recent studies have suggested that is both possible and potentially preferable have the learning rate early in training reach v) 2/x" (66] 49 (2). The idea is that such a choice will induce temporary training instability, causing the network "catapult" out of local basin into flatter one with lower X where training stabilizes. Indeed, during the early training phase, the local curvature of the loss landscape changes rapidly
42 37 [6], and the learning rate plays erucial role in determining the convergence basin [37]. Flatter basins are believed to be preferable because they potentially lead to lower generalization eπror
39 and allow larger learning rates leading to potentially faster training.
From a different perspective, the major theme of deep learning is that it beneficial to increase the model size as much as possible. This has come into sharp focus with the discovery of scaling laws that show power law improvement in generalization error with model and dataset size [40]. This raises the fundamental question of how one can scale DNNs to arbitrarily large sizes while maintaining the ability to learn; in particular, how should initialization and optimization hyperparameters be chosen maintain similar quality of learning as the model size is taken infinity 33
48
69 58
70 68
"Condensed Matter Theory Center, University of Marylann, College Park I Institute for Physical Science and Technology. University of Maryland, College Park Department of Physics, University of Maryland, College Park 8 Joint Quantum Institute, University of Maryland, College Park
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

19
12
To acoseer
Texing 1e
10
63
Training trajectories of the (a) training loss, (b) sharpness, and (c) training accuracy of CNNs (d and w = 512) trained on CIFAR-10 with MSE loss using vanilla SGD with learning rates </>ă and batch size B = 512. Vertical dashed lines approximately separate the different training regimes. Horizontal dashed lines (b) denote the 2/n threshold for each learning rate.
Figure 1:
Motivated by these ideas, we perform a systematic analysis the training dynamics of SGD for DNNs as learning rate, depth, and width are tuned, across variety of architectures and datasets. We monitor both the loss and sharpness (A#) trajectories during early training, observing a number of qualitatively distinct phenomena summarized below.
1.1 Our contributions
We study SGD fully connected networks (FCNs) with the same number of hidden units (width) each layer, convolutional neural networks (CNNs), and ResNet architectures of varying width w and depth d with ReLU activation. For CNNs, the width corresponds to the number of channels.
We focus on networks parameterized in Neural Tangent Parameterization (NTP) [34], and Standard
Parameterization (SP) [62] initialized criticality [55 58) while other parameterizations and initializations may show different behavior. Further experimental details are provided in Appendix A We study both meaan- error (MSE) and cross-entropy loss functions and the datasets CIFAR- 10, MNIST, Fashion-MNIST. Our findings apply to networks with d/w S C, where C' depends on architecture class (e.g. for FCNs, C 1/16) and loss function, but independent of d, w. and r.. Above this ratio, the dynamics becomes doise-dominated, and separating the underlying deterministic dynamics from random fluctuations becomes challenging, as shown Appendix E We use sharpness to refer to X," the maximum eigenvalue of H at time-step f, and flatness refers to 1/x""
By monitoring the sharpness, we find four clearly separated, qualitatively distinct regimes throughout the training trajectory. Fig. [howw an example from a CNN architecture. The four observed regimes are: (i) early time transient regime where loss and sharpness may drastically change and eventually settle down, (ii) an intermediate saturation regime where the sharpness has lowered and remains relatively constant, (iii) a progressive sharpening regime where sharpness steadily rises, and finally, (iv) late time regime where the sharpness saturates around 2/n for MSE loss; whereas for cross- entropy loss, sharpness drops after reaching this maximum value while remaining less than 2/1 [8]. Note the log scale in Figure I highlights the early regimes (i) and (ii); in absolute terms these are much shorter time than regimes (iii) and (iv).
In this work, we focus on the early transient and intermediate saturation regimes. As learning rate, d and tuned, clear picture emerges, leading rich phase diagram, demonstrated in Section 2 Given the learning rate scaled as / = </>,,, we characterize four distinct behaviors in the training dynamics in the early transient regime (i):
Sharpness reduction phase (c < Closs): : Both the loss and the sharpness monotonically decrease during early training. There is particularly significant drop in sharpness in the regime Cerit < Closs which motivates us to refer to learning rates lower than Carit as sub-critical and larger than Ceris super-critical. We discuss Cerit in detail below. The regime Cerit significantly with increasing and 1/w, which is a new result of this work.
Closs opens up
Loss catapult phase (Cioss < Csharp) The first few gradient steps take training to a flatter region but with a higher loss. Training eventually settles down in the flatter region as the loss starts to decrease again. The shappness monotonically decreases from initialization in this early time transient regime.

---

Loss and sharpness catapult phase (Csharp < < Cmaa); ): In this regime both the loss and sharpness initially start to increase, effectively catapulting different point where loss and sharpness can start to decrease again. Training eventually exhibits a significant reduction in sharpness by the end of the early training. The report of a loss and sharpness catapult is also new to this work.
Divergent phase (c > Cmaa): The learning rate is too large for training and the loss diverges.
The critical values Closs. Caharp, Cmaa are random variables that depend on random initialization, SGD batch selection, and architecture. The averages of Closs Cshorps Gmaz shown in the phase diagrams show strong systematic dependence
depth and width. In order to better understand the cause of the sharpness reduction during early training we study the effect of network output at initialization by (1) centering the network, (2) setting last layer weights to zero, or (3) tuning the overall scale of the output layer. We also analyze the linear connectivity of the loss landscape in the early transient regime and show that for range of learning rates Class
Cbarvier, no barriers
exist from the initial state to the final point of the initial transient phase, even though training passes through regions with higher loss than initialization.
Next, we provide a quantitative analysis of the intermediate saturation regime. We find that sharpness during this time typically displays 3 distinct regimes as the learning rate is tuned, depicted in Fig. 5 By identifying an appropriate order parameter, we can extract a sharp peak corresponding to Cerit- For MSE loss Cerit R 2, whereas for crossentropy loss, 4 2 Cerit 2 2. For < Cerit, the network is effectively in a lazy training regime, with increasing fluctuations as and/or 1/u are increased.
Finally, we show that single hidden layer linear network the we model displays the same phenomena discussed above and we analyze the phase diagram in this minimal model.
1.2 Related works
A significant amount of research has identified various training regimes using diverse criteria, e.g., 33]. Here we focus on studies that characterize training regimes with sharpness and learning rates. Several studies have analyzed sharpness at different training times Ref. [8] studied sharpness at late training times and showed how large-batch gradient descent shows progressive sharpening followed by the edge of stability, which has motivated various theoretical studies 9 3. Ref. [37] studied the entire training trajectory of sharpness in models trained with SGD and cross-entropy loss and found that sharpness increases during the early stages of training, reaches a peak, and then decreases. In contrast, we find a tharpness-reducttin phase, < Closs which becomes more prominent with increasing d and 1/w, where sharpness only decreases during early training; this also occurs in the catapult phase Closs < c < Csharp- during which the loss initially increases before decreasing. This discrepancy is likely due to different initialization and learning rate scaling in their work [33
[
37
17
45
35
8.
37. 35 g 33
Ref. B5 examined the effect of hyperparameters studied the optimization dynamics of SGD with momentum using sharpness. Ref. (45) classify training into 2 different regimes using training loss, providing a significantly coarser description of training dynamics than provided here. Ref. (33] studied the scaling of the maximum learning rate with d and w during early training in FCNs and its relationship with sharpness at initialization.
sharpness at late training times. Ref. [20]
Refs. [2 present phase diagrams of shallow ReLU networks infinite width under gradient flow. Previous studies such as [41 69] show that 2/X is the maximum learning rate for convergence as w + oc. This limit results in the kernel regime as training time is restricted to 0(1) in the limit of infinite width, resulting in lazy training regime for learning rates less than 2/28T8 and divergent training for larger learning rates. In contrast, we analyze optimization dynamics at training timescales t. that grow with width Specifically, the end of the carly time transient period occurs at t, ~ log(w).
Ref. [49 analyzed the training dynamics at large widths and training times, using the top eigenvalue of the neural tangent kernel (NTK) as a proxy for sharpness. They demonstrated the existence of a new early training phase, which they dubbed the "catapult" phase, 2/XNTK < 1) < Mmaaa. in wide networks trained with MSE loss using SGD, in which training converges after an initial increase in training loss. The existence of this new training regime was further extended to quadratic models with large widths by 722 [53. Our work extends the above analysis by studying the combined effect of learning rate, depth, and width for both MSE and cross-entropy loss, demonstrating the opening
3

---

0.50
0.25
Teating 1e
1..5
0.20
2.00
Kas
0.30
0.8
0.30
0.00
622
6.6
1.0
((mp))
(stes)
(1)
(N)
(r)
2.0
1.8
1.0
0.3
(aloo
20
32
0.1
30
0.0
0.0
[(reS)
[mmp]
512) and (d, e, deep CNN (d = 10, w = 128) trained on CIFAR-I0 with MSE loss for t 10 steps using SGD for various learning rates y e/># and batch size B 512. (a, d) training loss, (b, e) sharpness, and (c, f) interpolated loss between the initial and final parameters after 10 steps for the respective models. For the shallow CNN, Closs 2.82, Csharp 5.65, Cmaz 17.14 and for the deep CNN,
Figure 2: Early training dynamics of (a, b, shallow (d
Gose 36.75, Caharp = 39.39, Cmaz = 48.50.
of a sharutess-reduction phase, the refinement of the catapult phase into two phases depending on whether the sharpness also catapults, analyzing the phase boundaries as d and 1/w is increased, analyzing linear mode connectivity in the catapult phase, examining different qualitative behaviors in the intermediate saturation regime (ii) mentioned above.
2 Phase diagram of early transient regime
For wide enough networks trained with MSE loss using SGD, training converges into a flatter region after initial increase in the training loss for learning rates
[49] Fig. 2]a, b) shows the first shallow (d = 3 and w = 512) CNN trained on
10 steps of the loss and sharpness trajectories the CIFAR-10 dataset with MSE loss using SGD. For learning rates, e 2.82, the loss catapults and training eventually converges into flatter region, as measured by sharpness. Additionally, we observe that sharpness may also spike initially, similar to the training loss (see Fig. (b)). However, this initial spike in sharpness occurs at relatively higher learning rates (c 2 5.65), which we will examine along with the loss catapult. We refer this spike sharpness as 'sharpness catapult.'
An important consideration is the degree to which this phenomenon changes with network depth and width. Interestingly, found that the training loss in deep networks average catapults at much larger learning rates than c = 2. Fig. dd. e) shows that for a deep (d = 10, w 128) CNN, the loss and sharpness may catapult only near the maximum trainable learning rate. In this section, characterize the properties of the early training dynamics of models with MSE loss. In Appendix show that similar picture emerges for cross-entropy loss, despite the dynamics being noisier.
2.1 Loss and sharpness catapult during early training
In this subsection, we characterize the effect of finite depth and width on the onset of the loss and sharpness catapult and training divergence. We begin by defining critical constants that correspond to the above phenomena.
Definition 1. (Closs: Caharps Cmaz) For learning rate s/xă let the training loss and sharpness step & be denoted by (c(c) and X (c). We define Cloes (Caharp) minimum learning rates constants

---

(a) 8
(b) d =
(c) d= 18
2
(d) w 512
(e) w 1024
(1) w 2048
Figure 3: Phase diagrams of early training of neural networks trained with MSE loss using SGD. Panels (a-c) show phase diagrams with width: (a) FCNs 8) trained on the MNIST dataset, (b) CNNs (d 7) trained on the Fashion-MNIST dataset, (c) ResNet (d 18) trained on the CIFAR-10 (without batch normalization). Panels (d-f) show phase diagrams with depth: FCNs trained on the Fashion-MNIST dataset for different widths. Each data point in the figure represents average of ten distinct initializations, and the solid lines represent a smooth curve fitted to the raw data points. The vertical dotted line shows curves for better visualization. For experimental details and additional results, see Appendices and respectively. The phase diagram of early training of FCNs with depth for three different widths trained on Mashion-MNISS with MSE loss using SGD.
for comparison, and various colors filled in between the various
such that the loss (sharpness) increases during the initial transient period:
Closs min(c|
max C.(c) Lo(c)), Csharp minfc (E(1.7)
(c) (c)). NE[1.T]
max
and Cmaz as the maximum learning rate constant such that the loss does not diverge during the initial transient period: Cmex maxxfe L((c) K, Yt E [1.Ti]}, where K a fixed large constant 3
Note that the definition of Cmaz allows for more flexibility than previous studies (33 in order to investigate a wider range of phenomena occurring near the maximum learning rate. Here, Closes Csharp- and Cmar random variables that depend the random initialization and the SGD batch sequence, and we denote the average over this randomness using (-)-
Fig. 3(a-c) illustrates the phase diagram of early training for three different architectures trained on various datasets with MSE loss using SGD. These phase diagrams show how the averaged values (Ctose). (Csharp), and (Cmaz) are affected by width. The results show that the averaged values of all the critical constants increase significantly with 1/2 (note the log scale). At large widths, the loss starts to catapult R 2. As 1/w increases, (Ctoss) increases and eventually converges to (Cmax) at large 1/w. By comparison, sharpness starts to catapult at relatively large learning rates at small 1/w, with (Caharp) continuing increase with 1/ur while remaining between (Closs) and (Cmaz)) Similar results are observed for different depths as demonstrated in Appendix [ Phase diagrams obtained by varying d are qualitatively similar to those obtained by varying 1/ shown in Figure d-f-). Comparatively, we observe that (Cmaz) may increase or decrease with 1/ur in different settings while consistently increasing with d, shown in Appendices and H
'We use K 10 to estimate Graz In all our experiments, La = O(1) (see Appendix A which justifies
the of a fixed value.
5

---

Cooe
2
24
2-
21
Glaar
(a) FCNs on MNIST
(b) CNNs Fashion-MNIST
(c) ResNeets CIFAR-10 10
Figure 4: The relationship between critical constants for (a) FCNs, (b) CNNs, and (c) ResNets. Each data point corresponds to a run with varying depth, width, and initialization. The dashed line represents the y = x line.
While we plotted the averaged quantities (Closs), (Csharp), (Gmaz) we have observed that their variance also increases significantly with d and 1/u: in Appendix C we show standard deviations about the averages for different random initializations. Nevertheless, we have found that the inequality Closs S Caharp S Craa typically holds, for any given initialization and batch sequences, except for some outliers due to high fluctuations when the averaged critical curves start merging at large d and 1/w. Fig. shows evidence of this claim. The setup the same as in Fig. 3 Appendix D presents extensive additional results across various architectures and datasets.
In Appendix F we show that cross-entropy loss shows similar results with some notable differences. The loss catapults at a relatively higher value (Class) Z 4 and (Cmex) consistently decreases with 1/us while still satisfying Closs
Csharp S Cmazz
2.2 Loss connectivity in the early transient period
In the previous subsection, we observed that training loss and sharpness might quickly increase before decreasing ("catapult") during early training for range of depths and widths. A logical next step is to analyze the region in the loss landscape that the training reaches after the catapult. Several works have analyzed loss connectivity along the training trajectory [21 [1] 64). Ref. [51] report that training traverses a barrier at large learning rates, aligning with the naive intuition of a barrier between the initial and final points of the loss catapult, the loss increases during carly training. In this section, we will test the credibility of this intuition in real-world models. Specifically, we linearly interpolate the loss between the initial and final point after the catapult and examine the effect of the learning rate, depth, and width. The linearly interpolated loss and barrier are defined as follows.
Definition 2. (C,na(s, c), Uc)) Let dg represent the initial set of parameters, and let Õr, represent the set of parameters at the end of the initial transient period, trained using learning rate constant c. Then, we define the linearly interpolated loss as Lime c) = [[ s) Gg + Gr,,, where E [0. 1] is the interpolation parameter. The interpolated loss barrier is defined as the maximum value of the interpolated loss over the range of s: U(c) max
Line(s) L(Ho).
Here we subtracted the loss's initial value such that positive value indicates a barrier to the final point from initialization. Using the interpolated loss barrier, we define Charrier follows.
Definition 3. (Charrier) Given the initial (0o) and final parameters (0r, we define Cbarrier as the minimum learning rate constant such that there exists a barrier from Go to 0r,:' Cbarrier minofe| (cc) > 0}.
Here, Cbarrier is also a random variable that depends on the initialization and SGD batch sequence. We denote the average over this randomness using (.) as before. Fig. 2ic, f) shows the interpolated loss of CNNs trained on the CIFAR- 10 dataset for 10 steps. The experimental setup is the same as in Section 2 For the network with larger width, we observe barrier emerging at Cbarrier = 5.65. while the loss starts to catapult at Closs 2.83. comparison, we do not observe any barrier from initialization to the final point at large d and 1/w. Fig. shows the relationship between (Charrier) and a/w for various models and datasets. We consistently observe that Csharp S Charriers suggesting that training traverses barrier only when sharpness starts to catapult during early training. Similar results

---

witth
0.3
0.4
03
9.
0.1
0.2
0.1
02
0.0
000
(a)
(b)
(c)
Figure 5: (a) Normalized sharpness measured c7 200 against the learning constant for 7- layer CNNs trained on the CIFAR-10 dataset, with varying widths. Each data point is an average over 5 initializations. where the shaded region depicts the standard deviation around the mean trend. (b, c) Smooth estimations of the first two derivatives, and X,, of the averaged normalized sharpness wrt the learning rate constant. The vertical lines denote Cerit estimated using the maximum of X,- For smoothening details, see Appendix 1.2
observed increasing instead of 1/w as shown in Appendix C phase diagram of early training using Cherrier as we did for other critical c's, as it is somewhat different in character than the other critical constants, which depend only on the sharpness and loss trajectories.
We chose not characterize the
These observations call into question the intuition catapulting out a basin for a range of learning rates in between Closs < c < Cbarrier- These results show that for these learning rates, the final point after the catapult already lies in the same basin as initialization, and even connected through a linear path, revealing inductive bias of the training process towards regions of higher loss during the early time transient regime.
Intermediate saturation regime
In the intermediate saturation regime, sharpness does not change appreciably and reflects the cumula- tive change that occurred during the initial transient period. This section analyzes sharpness in the in- termediate saturation regime by studying how it changes with the learning rate, depth, and width of the model. Here, we show results for MSE loss, whereas cross-entropy results are shown in Appendix F
We measure the sharpness A at time in the middle of the intermediate saturation regime. We choose that eT 2 2 For further details on sharpness measurement, see Appendix [. Fig. 5)a) illustrates the relationship between X and the learning rate for 7-layer deep CNNs trained on the CIFAR-10 dataset with varying widths. The results indicate that the dependence of X on learning rate can grouped into three distinct stages. (1) At small learning rates, A remains relatively constant, with fluctuations increasing and 1/w increase (c 2 in Fig. a)). (2) A
crossover regime where X is dropping significantly (2 < c 2 in Fig. aa). (3) saturation stage where X stays small and constant with learning rate 2$) in Fig. a)). In Appendix
show that these results are consistent across architectures and datasets for varying values of d and w. Additionally, the results reveal that in stage (1), where < 2 is sub-critien X decreases with increasing d and 1/w. other words, for small c and in the intermediate saturation regime, the loss is locally flatter d and 1/w increase.
We can precisely extract critical value of e that separates stages (1) and (2), which corresponds to the onset of an abrupt reduction of sharpness X To do this, we consider the averaged normalized sharpness over initializations and denote it by (>"/>")- The first two derivatives of the averaged normalized sharpness, Xx 2 (> / and X
8a (>* /xf') characterize the change in sharpness with learning rate. The extrema of X quantitatively define the boundaries between the three stages described above. In particular, using the maximum of X', we define (Cerit), which marks the beginning of the sharp decrease in X with the learning rate.
'time-step 7 = 200/c is in the middle of regime (ii) for the models studied. Normalizing by c callows proper comparison for different learning rates.
7

---

9-11
2
2-11
2-11
(a) vanilla
(b) centered
(c) last layer zero init
Figure 6: Phase diagrams of d = 8 layer FCNs trained on the CIFAR-10 dataset using MSE, demonstrating the effect of output scale at initialization: (a) vanilla network, (b) centered network. and (c) network initialized with the last layer set to zero.
Definition 4. ((Cers)) Given the averaged normalized sharpness (>#)>)) measured at T, we define Cerit to be the learning rate constant that minimizes its second derivative: (Cori) = arg max X
Here, we use (.) to denote that the critical constant is obtained from the averaged normalized sharpness. Fig. b, c) show Xr and X obtained from the results in Fig. 5(a). We observe similar results across various architectures and datasets, shown in Appendix I Our results show that (Cora) has slight fluctuations as d and 1/w are changed but generally stay in the vicinity of e = 2. The peak in X- becomes wider d and 1/w increase, indicating that the transition between stages (1) and (2) becomes smoother, presumably due to larger fluctuations in the properties of the Hessian H at initialization. In contrast to (Cerat), (Closs) increase with d and 1/u, implying the opening of the sharpness reduction phase (Corit) < c < (Closs) as d and 1/w increase. In Appendix F we show that cross-entropy loss shows qualitatively similar results, but with 2 x (Cerit) 4.
Effect of network output at initialization on early training
Here discuss the effect of network output [a 0.) initialization the early training dynamics. & is the input and Og denotes the set of parameters at time t. We consider setting the network output zero initialization, f(; 0n) 0, by either (I) considering the "centered" network: fo(z; 0) = f(x; 0) f(x; 0o), or (2) setting the last layer weights to zero at initialization (for details, see Appendix G). Remarkably, both (1) and (2) remove the opening up the sharpness reduction phase with 1/" w as shown in Figure 6. The average onset of the loss catapult, diagnosed by (Cioss). becomes independent of 1/w and
We also empirically study the impact of the output scale 19 on early training dynamics. Given network function / (a:; O), we define the scaled network ʃs (2; (x;; 0), where a is scalar, fixed throughout training. In Appendix H. we show that a large (resp. small) value of [>; 0p)l relative the one-hot encodings of the labels causes the sharpness decrease (resp. increase) during early training. Interestingly, we still observe an increase in (Cioss with d and 1/w, unlike the case of initializing network output to zero, highlighting the unique impact of output scale on the dynamics.
Insights from a simple model
the ue model, which shows much of the phenomena presented above. Define f(x) v ST u.r, with x, f(x) E R. Here, u, 2 E R are the
Here we analyze a two-layer linear network [56 60 49),
trainable parameters, initialized using the normal distribution, u, U; ~ NN 1) for [L...., . w). The model is trained with MSE loss on single training example (r,, y) = (1, which simplifies the loss to C(u, v) = f2/2, and which was also considered in Ref. [49 Our choice of y
motivated by the results of Sec. 4 which suggest that the empirical results of Sec. Sec.pare are intimately related the model having a large initial output scale [=; 0o)] relative to the output labels. We minimize the loss using gradient descent (GD) with learning rate 7). The early time phase diagram also shows similar features to those described in preceding sections (compare Fig. 7[a) and Fig. (3) Below develop understanding of this early time phase diagram in the we model.
8

---

Hows
(a)
(b)
(c)
Figure 7: The phase diagram of the ue model trained with MSE loss using gradient descent with (a) the top eigenvalue of Hessian X (b) the trace of Hessian tr( and (c) the square of the Frobenius
norm tr [H) He) used measure of sharpness. In (a), the learning rate is scaled as r 9/x6' while in (b) and (c), the learning rate is scaled as (/or(Ha). The vertical dashed line shows c =
(k: 2) for reference. Each data point is an average over 500 random initializations.
The update equations of the we model in function space can be written in terms of the trace of the Hessian tr( H) f fo- (1-m + w
f2
nf? (n)r((H)) 4). (I)
tr(H++1) = tr(H))
w
From the above equations, it is natural to scale the learning rate as 7 */u(Ho), Note that c naal" Ax /tr(H))) Also, we denote the critical constants in this scaling as klsss htrace. kmaz and kerse, where the definitions follow from Definitions I Land 4 on replacing sharpness with trace and use (.) to denote an average over initialization. Figure 7b) shows the phase diagram of early training. with tr( H+)) replaced with X as the measure of sharpness and with the learning rate scaled as 1) -= === (o)) Similar to Figure a),, we observe a new phase (kerit) k < (Kloss) opening up at small width. However, we do not observe the loss-sharpness catapult phase as tr(H) does not increase during training (see Equation Db. We also observe (kmaz) = 4, independent width.
we show that the critical value of k for which (Cx/ca) > 1 increases with 1/u, which explains why (kioss) increases with 1/w. Combined with (kerit) 2, this implies the opening up of the sharpness reduction phase as w is decreased.
In Appendix B.3
To understand the loss loss-sharpness catapult phase, we require some other measure tr (H) does not increase for 0 < k < 4. As X is difficult to analyze, we consider the Frobenius norm ||]||
y Vt(H"H) H) as a proxy for sharpness. We define krr as the minimum learning rate such that HH increases during early training. Figure ( shows the phase diagram of the uv model, with |H, as the measure of sharpness, while the learning rate scaled as 1) k/tr(Ho) We observe the loss-sharpness catapult phase at small widths. In Appendix B.4,, we show that the critical value of k for which I| |o ) > O increases from (kloss) as 1/w increases. This explains the opening up of the loss catapult phase at small w in Fig. 7 (c).
Fig. 8 shows the training trajectories of the uv model with large (w = 512) and small (w 2) widths in two-dimensional slice of parameters defined by tr tr(H) and weight correlation [wllo|l The above figure reveals that the first few training steps of the small-width network take the system in a flatter direction (as measured by (r( H)) as compared to the wider network. This means that the small-width network needs relatively larger learning rate get point of increased loss (loss catapult). We thus have the opening up of a new regime (kerit) < k < (kloss). ; in which the loss and sharpness monotonically decrease during early training.
The loss landscape of the ur model shown in Fig. reveals interesting insights into the loss landscape connectivity results in Section 2nd and the presence of Cbarrier. Fig. shows how even when there is loss catapult, as long as the learning rate is not too large, the final point after the catapult can be reached from initialization by linear path without increasing the loss and passing through barrier. However the learning rate becomes large enough, then the final point after the catapult may correspond to a region of large weight correlation, and there will be barrier in the loss upon linear interpolation.
The ue model trained on an example (ar,y) with y * 0 provides insights into the effect of network output at initialization observed in Section 4 In Appendix G we show that setting fo 0 and

---

age
0J woghh condationn
000 wiight comclatonn
0.5
woigd comeaaiion
wogltt cordatton
(s)
(d)
Figure 8: Training trajectories of the model trained (x. (1, 0), with (a, b) large and d) small width, in a two-dimensional slice of the parameters defined by the trace of Hessian tr(!f) and weight correlation, trained with (a, small (c 0.5) and (b, large (c 2.5) learning rates. The colors correspond to the training loss L, with darker colors representing smaller loss.
y * in the dynamical equations results loss catapult at k implying (Mitnas) ss (hierit) Rs irrespective of w.
6 Discussion
We have studied the effect of learning rate, depth, and width on the early training dynamics in DNNs trained using SGD) with learning rate scaled 1j -/xž"- We analyzed the early transient and intermediate saturation regimes and presented a rich phase diagram of early training with learning rate, depth, and width. We report new phases, sharpness reduction and loss-sharpness catapult, which have not been reported previously. Furthermore, we empirically investigated the underlying cause of sharpness reduction during early training. Our findings show that setting the network output to zero at initialization effectively leads to the vanishing of sharpness reduction phase at supercritical learning rates. We further studied loss connectivity the early transient regime and demonstrated the existence of regime (Ctoes) < < (Cborrier), in which the final point after the catapult lies in the same basin initialization, connected through linear path. Finally, we study these phenomena in 2-layer linear network (ur model), gaining insights into the opening of the sharpness reduction phase.
We performed preliminary analysis the effect of batch size the presented results in Appendix
The sharpness trajectories of models trained with a smaller batch size (B 32 vs. B = 512) show similar early training dynamics. In the early transient regime, we observe qualitatively similar phase diagram. In the intermediate saturation regime, the effect of reducing the batch size is to broaden the transition around Cerit-
In Section 2 we noted that for cross- cross-entropy loss, the loss starts to catapult around c =s 4 at large widths, as compared to Ctoss for MSE loss. Previous work, such [50), analyzed the catapult dynamics for the ur model with logistic loss and demonstrated that the loss catapult occurs above Masss = 4/X We summarize the main intuition about their analysis in Appendix B.9 However, a complete understanding of the catapult phenomenon the context of cross-entropy loss requires a more detailed examination.
The early training dynamics is sensitive to the initialization scheme and optimization algorithm used, and we leave it to future work to explore this dependence and its implications. In this work, we focused on models initialized at criticality [55] as it allows for proper gradient flow through ReLU networks at initialization 23 58). and studied vanilla SGD for simplicity. However, other initializations [46), parameterizations [69] (70), and optimization procedures [2 may show dissimilarities with the reported phase diagram of early training.
10

---

Acknowledgments
We thank Andrey Gromov, Tianyu He, and Shubham Jain for discussions, and Paolo Glorioso, Sho Yaida, Daniel Roberts, and Darshil Doshi for detailed comments on the manuscript. We also express our gratitude to anonymous reviewers for their valuable feedback for improving the manuscript. This work is supported by an NSF CAREER grant (DMR (731273240) and the Laboratory for Physical Sciences through the Condensed Matter Theory Center.
References
[1] Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep networks. International Conference Learning Representations, 2019.
[2] Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability, ArXiv, abs/2210.04866 2022.
[3] Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. In International Conference on Machine Learning, 2022.
[4] Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. In The Fleventh International Conference on Learning Representations, 2023.
[5] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[6] Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctua- tions in mean field neural networks. ArXiv, hads/234.03440, 2023.
[7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018.
[8] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representationss 2021.
[9] Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022.
[10] Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012.
[11] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[12] Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Gal Elidan, Kristian Kersting, and Alexander Ihler, editors, Proceedings of the Phirty-Third Conference on Uncertainty Artificial Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017. AUAI Press, 2017.
[13] Dumitru Erhan. Yoshua Bengio, Aaron Courville, Bierre-Antoin, Manzagol, Pascal Vincent, and Samy Bengio. Why does unsupervised pre-training help deep learning? Journal of Machine Learning Research, ((119)625-66/) 2010.
[14] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2021.

---

[15] Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy. and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In NeurIPS, 2020.
[16] Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. ArXiv, abs/1910.05922. 2019.
[17] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Linear mode connectivity and the lottery ticket hypothesis. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. JMLR.org. 2020.
[18] Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Perspective: A phase diagram for deep learning unifying jamming, feature learning and lazy training, 2020.
[19] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanicx: Theory and Experiment, 2020(11): 13301, nov 2020.
[20] Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Cardoze, George Edward Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Representations, 2022.
[21] lan J. Goodfellow and Oriol Vinyals. Qualitatively characterizing neural network optimization problems. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015.
[22] Priya Goyal, Piotr Dolkár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour. ArXiv, 706.0266777 2017.
[23] Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? In Neural Information Processing Systems, 2018
[24] Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference Learning Representations, 2020.
[25] Boris Hanin and David Rolnick. How to start training: The effect of initialization and architec- ture. In Neural Information Processing Systems, 2018.
[26] Soufianc Hayou, A. Doucet, and Judith Rousscau. On the selection of initialization and activation function for deep neural networks. ArXiv, abs/ 885608266, 2018.
[27] Soufiane Hayou, A. Doucet, and Judith Rousseau. Exact convergence rates of the neural tangent kernel the large depth limit. 2019.
[28] Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. 2015 IEEE International Conference Computer Vision (ICCV), pages 1026-1034, 2015.
[29] Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference Computer Vision and Pattern Recognition (CVPR), pages 770-778, 2016.
[30] 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.
[31] Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in neural information processing systems, 7, 1994.
[32] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):144,, 1997.
[33] Gaurav Iyer, Boris Hanin, and David Rolnick, Maximal initial learning rates deep relu networks. ArXiy, abs/2212.07295, 2022.
12

---

[34] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kemel: convergence and generalization in neural networks (invited paper). Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, 2018.
[35] 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 Learning Representations, 2020.
[36] Stanislaw Jastrzebski, Zac Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Amos Storkey, and Yoshua Bengio. Three factors influencing minima in SGD, 2018.
[37] Stanislaw Jastrzæbski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019.
[38] Stanislaw Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019.
[39] Yiding Jiang*, Behnam Neyshabur*, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020.
[40] Jared Kaplan. Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836,, 2020.
[41] Ryo Karakida, Shotaro Akaho, and Shun ichi Amari. Universal statisties of fisher information in deep neural networks: Mean field approach, 2019.
[42] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy. and Ping Tak Peter Tang. On large-batch batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[43] Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. In NIPS, pages 972-981, 2017.
[44] Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.
[45] Guillaume Leclere and Aleksander Madry. The two regimes of deep network training. ArXiv, abs/2002.10376, 2020.
[46] Yann LeCun, Léon Bottou, Genevieve B. Orr, and Neural Networks, 2012.
Müller. Efficient backprop. In
[47] Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018
[48] Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Narain Sohl-Dickstein, and Jeffrey S. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Journal of Statistical Mechanics: Theory and Experiment, 2020, 2019.
[49] 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, 0bs/2003.002118 88, 2020.
[50] Chunrui Liu, Wei Huang, and Richard Yi Da Xu. Implicit bias of deep learning in the large learning rate phase: A data separability perspective. Applied Sciences, 13(6), 2023.
[51] James Lucas, Juhan Bae, Michael R Zhang, Stanislav Fort, Richard Zemel, and Roger Grosse. Analyzing monotonic linear interpolation in neural network loss landscapes. arXiv preprint arX1v:210411104,, 2021.
13

---

[52] Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang, Phase diagram for two-layer relu neural networks at indtnite-width limit, 2020.
[53] David Meltzer and Junyu Liu. Catapult dynamics and phase transitions in quadratic nets, ArXiv, 00//2301.07737, 2023.
[54] Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl- Dickstein, and Samuel Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations, 2020.
[55] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In NIPS, 2016.
[56] Jasson D. M. Rennie and Nathan Srebro. Fast maximum margin matrix factorization for collaborative prediction. In Proceedings of the 22nd International Conference Machine Learning, ICML '05, page 713-719, New York, NY, USA, 2005. Association for Computing Machinery.
[57] Jasson D. M. Rennie and Nathan Srebro. Fast maximum margin matrix factorization for collaborative prediction. In Proceedings of the 22nd International Conference on Machine Learning, ICML '05, page 713-719, New York, NY, USA, 2005, Association for Computing Machinery.
[58] Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory. Cambridge University Press, 2022. https://deepleeminnggeeeey com
[59] Abraham. Savitzky and M. J. E. Golay. Smoothing and diffcrentiation of data by simplified least squares procedures. Analytical Chemistry, 36(8):6627-1639 1964.
[60] Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. CoRR, abs/ 121.61122.10. 2014.
[61] Vaishaal Shankar, Alexander W. Fang, Wenshuo Guo, Sara Fridovich-Ked Ludwig Schmidt, Jonathan Bagan-Kelley, and Benjamin Recht. Neural kernels without tangents. In ICML, 2020.
[62] Jascha Narain Sohl-Dickstein, Roman Novak, Samuel S. Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with standard parameterization. ArXiv, abs/2001.073001, 2020.
[63] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, CJ Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SeiPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272 2020.
[64] Tiffany J. Vlaar and Jonathan Frankle. What can linear interpolation of neural network loss landscapes tell us? In International Conference on Machine Learning, 2021.
[65] Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of stability. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[66] Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over parameterizeed learning: A dynamical stability perspective. In S. Bengio, H. Wallach, H. Larochelle, K. Grau- man, N. Cesa- Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
[67] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017. cite
Dataset is Benchmark is available at
freely available at https://gitthubb https//github.com/aaaaaado http:///aaadddeeeecc///ss/sooooo
http://fashion wwwtssiiiioomiiisssm... w-wiiittt-wbbbtte.... www.ste..cuccmmtrallaaaaeeeeeeae Law.
14

---

[68] Sho Yaida. Meta- Metaa-rincipped family of hyperparameter scaling strategies. ArXiv, abs/2210.04409,, 2022.
[69] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 Proceedings of Machine Learning Research, pages 11727-11737. PMLR, 18-24 Jul 2021.
[70] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub W. Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In Neural Information Processing Systems, 2022.
[71] Hanxu Zhou, Qixu Zhou, Zhenyuan Jin, Tao Luo, Yaoyu Zhang, and Zhi thi-Qin John Xu. Empirical phase diagram for three-layer neural networks with infinite width. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[72] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishman, and Mikhail Belkin. Quadratic models for understanding neural network dynamics. ArXiv, 1iss/2055.11887 2022.
15

---

A Experimental details
Datasets: We considered the MNIST [O]. Fashion-MNIST [67, and CIFAR-10 44 datasets. We standardized the images and used one-hot encoding for the labels.
Models: We considered fully connected networks (FCNs), Myrtle family CNNs 61 and ResNets (version 1) [29) trained using the JAX [7], and Flax libraries [30]. We use d and wo to denote the depth and width of the network. Below, we provide additional details of the models and clarify what width corresponds to for CNNs and ResNets.
1. FCNs: We considered ReLU FCNs with constant width w in Neural Tangent Parameteriza- tion (NTP) Standard Parameterization (SP), initialized criticality (55). The models do not include bias or normalization. The forward pass of the pre-activations from layer to 1 ++ 1 given by
W
Σ W! (h)))
where ((.) the ReLU activation and 7 is a constant. For NTP, 9 2/v6 t and the weights
W' are initialized using normal distribution, i.e., W ~/ 1). For SP, 7" = 1 and the
weights W' are initialized as W ~ (0, 3/w). For the last layer, we have y 1/vl for
NTP and W ~ (O 1/w) for SP.
For 4/w Z 1/16, the dynamics is noisier, and becomes challenging to separate the underly- ing deterministic dynamics from random fluctuations (see Appendix Eb.
2. CNNs: We considered Myrtle family ReLU CNNs [61] without any bias or normalization in Standard Parameterization (SP), initialized using He initialization (29]. The above model uses a fixed number of channels in each layer, which we refer to as the width of the network.
In this ease, the forward pass equations for the pre-activations from layer to layer + are given by
(a)
Σ Σ (8)o())aa B)),
(3)
BEker
where at 3 label the spacial location. The weights are initialized as ~ (10, where k is the filter size.
3. ResNets: We considered version 1 ResNet 29 implementations from Flax examples without Batch Norm or regularization. For ResNets, width corresponds to the number of channels in the first block. For example, the standard ResNet-18 18 has four blocks with widths (w, 2w, 4w, 8w), with
64. We refer to w as the width or the widening factor. We considered ResNet-18 and ResNet-34.
All the models are trained with the average loss over the batch DB = {(Epp Mp) i.e., L(x, )D,) = 1/B 1/B >i= E( Mn), where {(.) is the loss function. This normalization, along with initialization, ensures that the loss is O(1) at initialization.
Bias: Throughout this work, we have primarily focused on models without any bias for simplicity. In we demonstrate that bias does not have an appreciable impact on the results.
Appendix K
Batch size: We use a batch size of 512 and scale the learning rate as = =</x" in all our experiments, unless specified. Appendix prpendixxshowr shows results for smaller batch size B 322
Learning rate: We scale the learning rate constant as c = 2", with x E (-1.0,... . } in steps of 0.1. Here, Fmar is related to the maximum learning rate constant Craar 2
Sharpness measurement: We measure sharpness using the power iteration method with 20 iterations. We found that 20 iterations suffice both for MSE and cross-entropy loss. For MSE loss, we use 2048 randomly selected training examples for evaluating sharpness at each step. In comparison, we found that cross-entropy requires large number of training examples to obtain good approximation
16

---

of sharpness, Given the computational constraints, we use 4096 training examples to approximate sharpness for eross-entropy loss.
Averages over initialization and SGD runs: initializations and the SGD runs. In experiments, found that the fluctuations from initial- ization at large d/w outweigh the randomness coming from different SGD runs. Thus, we focus on initialization averages in all our experiments.
All the critical constants depend on both the random
A.1 Compute usage
We utilized different computational resources depending the task complexity. For less demanding tasks, we performed computation for a total of 2800 hours, utilizing a seventh of an NVIDIA A100 GPU. For more computationally intensive tasks, utilized full NVIDIA A100 GPU for a total 300 hours.
A.2 Reproducibility
The main results of this paper can be reproduced using the associated GitHub reposi- tory https///ghhbbboommaayaalllall/aa
A.3 Details of Figures in the main text:
Figure I A shallow CNN (d = 5, = 128) in SP trained on the CIFAR-10 dataset with MSE loss for 1000 epochs using SGD with learning rates 7 =/>0" and batch size = 512. We measure sharpness at every step for the first epoch, every epoch between 10 and 100 epochs, and every 10 epochs beyond 100.
Figure 2)
(top panel) A wide (d = 5, w 512) and (bottom panel) deep CNN (d = 10, w = 128) in SP trained the CIFAR-10 dataset with MSE loss for 10 steps using vanilla SGD with learning rates = c/>0 and batch size B = 512.
Figure 3 Phase diagrams of early training of neural networks trained with MSE loss using SGD. Panels (a-c) show phase diagrams with width: (a) FCNs (d = 8) trained on the MNIST dataset, (b) CNNs (d = 7) trained on the Fashion-MNIST dataset, (c) ResNet (d 18) trained on the CIFAR-10 (without batch normalization). Panels (d-f) show phase diagrams with depth: FCNs trained on the MishionnMNIST dataset for different widths. Each data point in the figure represents average of ten distinct initializations, and the solid lines represent a two-degree polynomial y = a + br + cx2 fitted to the raw data points. Here, where 1/w, and y can take on of three values: Closs; Csharp and Cmaan-
Figure 4 (a) FCNs in SP with d E (4,8, 361 and E {256, 512, 1024, 2048) (0r8iir on the MNIST
dataset, (b) CNNs in SP with d E (5, 7, 10} and w E {64, 128, 25,,256,55 512} trained on the Fashion- MNIST dataset, (c) ResNet in SP with d E {18, 34) and w E {32, 64, 128) trained on the CIFAR-10 dataset (without batch normalization).
Figure 6 Phase diagrams of d = 8 layer FCNs trained on the CIFAR-10 dataset using MSE, demonstrating the effect of output scale initialization: (a) vanilla network, (b) centered network, and (c) network initialized with the last layer set to zero. The values of widths are the same as in Figure 3
200 against the learning rate constant for 7-layer CNNs in SP trained on the CIFAR-10 dataset, with w E { (128, 256, 512]. Each data point is an average over five random initialization. Smoothening details are provided in Appendix I.2
Figure 5 Normalized sharpness measured
Figure 7 The phase diagram of the ue model trained with MSE loss using gradient descent with (a) the top eigenvalue of Hessian A/ (b) the trace of Hessian tr ((HH)) and (c) the square of the Frobenius norm tr (H, H₂) used as a measure of sharpness. In (a), the learning rate is scaled as r = -/>/"
while in (b) and (c), the learning rate is scaled r (/tr(H,). The vertical dashed line shows
(k = 2) for reference. Each data point is an average over 500 random initializations.
Figure 8: Training trajectories of the model with (a, b) large (w 512) and (c. d) small (w = 2) width, trained for / = 10 training steps on a single example (x. (1,0) with MSE loss using vanilla gradient descent with learning rates (a. c) 0.5 and (b, d) c 2.50.
17

---

10
0.5
Tran Ims
04
0.2
16
30
(shmp)
[step)
1.0
1.0
0.6
Taste
9.
0.0
0.4
0.2
0.0
(step)
[(ste)
(step)
(b)
Figure 9: Training trajectories of the ur model with (a) large (wr 512) and (v) a small (w 2) widths trained for & = 10 training steps on a single example (ar, y) = (1, 0). For the wide network, Goss 2.1, Cakarp 2.6, Cmea = 4.0, and for the narrow network, Closs 3.74, Csharp = 4.63, Cmaz = 4.93.
B Additional results for the we model
B.1 Details of the model
Consider two-layer linear network in (NTP) with unit input-output dimensions
f(x)
nT uc, Vww
where x, f(x) E RR. Here, u, E R are trainable parameters, with each element initialized using the normal distribution, uy (0, 1) for E (1....,w) w. The model is trained using MSE loss on a single training example (xr y) = (1, (), which simplifies the loss to
C(u, v) = 3.
(5)
The trace of the Hessian tr(H) has a simple expression in terms of the norms of the weight vectors
22 (||u|| + [|u|1²),
(r(H)
(6)
which is equivalent to the NTK for this model. The Frobenius norm of the Hessian |||| can be written in terms of the loss L and tr(H))
2 2) 1+
tr(H)2 46 1+ 2
| - tr(2² 2
w
The gradient descent updates of the model trained using MSE loss on single training example (z, y) (1,0) 0) are given by
18

---

witth
32.0
0.6
512
0.4
0.2
0.0
-0
(a)
(b)
(c)
Figure 10: (a) Normalized sharpness measured at 100 steps against the learning rate constant for the uv model trained on (z, y) (1, 0), with varying widths. Each data point is an average over 500 initializations, where the shaded region depicts the standard deviation around the mean trend. (b, c) Smooth estimations of the first two derivatives, Xr and Xr, of the, averaged normalized sharpness wrt the learning rate constant. The vertical dashed lines denote Cerit estimated for each width, using the maximum of X-. Here, we have removed the points beyond c = 3.5 for the calculation of derivatives to avoid large fluctuations the divergent phase. Smoothening details are described in Appendix [.2
U++1 = U;
Ux
(8)
V@
1
U++1
ii nft
Ur
(9)
W
The update equations in function space can be written terms of the trace the Hessian (r(H).
fi11 = fi (1 ( (1 n n7
(10)
nf?
tr(H++11) = tr( + (ntr(H))) -4).
Figure shows the training trajectories of the ue model trained on (z, y) = (1,0) 0) using MSE loss for 10 training steps. The model shows similar dynamics those presented in Section is worth mentioning that the above equations have been analyzed in 49 at large width. In the following subsections, extend their analysis by incorporating the higher-order terms analyze the effect of finite width.
B.2
The intermediate saturation regime
The model trained on (æ, y) (1, does not show the progressive sharpening and late-time regimes (iii) and (iv) described in Section Hence, we can measure sharpness at the end training analyze how it is reduced upon increasing the learning rate and compare it with the intermediate saturation regime results in Section 3
Figure 10 a) shows the normalized sharpness measured behavior reproduces the results observed the intermediate saturation regime in Section 3 In particular,
100 steps for various widths. This
stages (1) and (2), where > /> starts off fairly independent of learning rate constant and then dramatically reduces when > 2; stage (3), where x²" plateaus at small value as a function of is too close to the divergent phase this model be clearly observed. The corresponding derivatives of the averaged normalized sharpness, Xr. and X², are shown in Figure b0 e). The vertical dashed lines denote Carit estimated for each width, using the maximum of X,- We observe that Corit 2 for all widths.
19

---

4/
42
4.4
45
48
44 (TV(
(a)
(b)
(c)
Figure 11: (a, b) The averaged loss at the first step (C//co) against the learning rate constant & for varying widths obtained from (a) inequality [ō]and (b) numerical experiments. The intersection of (Ci/Co) with the horizontal line endpoints of kioss small and large widths. The shaded region in (b) shows the standard deviation around the mean trend. (c) The scaling of X and tr( tr(Ho) with width.
depicts kloss. The two vertical lines k = 2 and A = mark the
B.3 Opening of the sharpness reduction phase in the uo model
This section shows that O(Yw) terms Equation I0 effectively lead the opening of the sharpness reduction phase with 1/w in the ue model. In Appendix B.2 we demonstrated that for the model, Ceri for all values of widths. Hence, it suffices to show that Closs increases from the value as 1/u increases. We do so by finding the smallest k such that the averaged loss over initializations increases during early training.
follows from Equation 10 that the averaged loss increases in the first training step the following holds
C1 Eo
+ 22 -(1-
1,
(II)
where (.) denotes the average initializations. On scaling the learning rate with trace 7 H/u(Ho), we have
4
22
8f
Co
=
+
w
tr(Ho)²
(12)
L1
22
13
k4
ft tr((o))
)) > 1.
Lu
=
(1 - k:)2 + 2(1
k)
w
tr
(T(H))?
+
w²
(13)
The required two averages have the following expressions as shown Appendix B.8
(14)
(2(
4(w + 1)
få
3(w + 2)
T(w)
(15)
Tr(Ho)
16
I(w +
Inserting the above expressions in Equation 13 on average the loss increases in the very first step the following inequality holds
L1 Co
-(1-k)²
((1
1.2 - k)
314
+
1
(16)
2(w 1)
16(w 3)( + 1)
20

---

The graphical representation of the above inequality shown in Figure [T is in excellent agreement with the experimental results presented in Figure IT(bb..
Let us denote Kosss as the minimum learning rate constant such that the average loss increases in the constant if the loss increases in the first 10 steps. Then, Koss increases from the value 2 as 1/w increases as shown in Figure I1 a). By comparison, the initialization, this condition becomes k < kloss both the loss and trace monotonically decrease in the first training step. These arguments can be extended to later training steps, revealing that the loss and trace will continue to decrease for k < KoDEs.
first step. Similarly, let kloss denote the learning
trace reduces at any step if 1/ tr( tr(
4. Hence, for
Next, let Mloos denote the learning rate corresponding to Closs Then, we have Moss
Stoss A0
Kioas (x(
implying
Closs
kloss tr(Ho)
(17)
Figure [1fc) shows that X 2 tr Ho) for all widths, implying Closs 2 Kloss Hence, Closs increases with 1/w as observed in Figure 7 a)) In Appendix B.2 we demonstrated that for the ue model, Cerit = 2 for all values of widths. Incorporating this with Closs increases with 1/e, we have sharpness reduction phase opening up as l/u increases.
B.4 Opening of the loss catapult phase at finite width
In this section, we use the Frobenius norm of the Hessian || as a proxy for the sharpness and demonstrate the emergence of the loss-sharpness catapult phase at finite width. In particular, We analyze the expectation value (tr(H" (H"HH))))))))) H)) after the first training step near k = kioss and show that ktoss < kfrrb, with the difference increasing with 1 1/w. First, write tr(H) H, H,) in terms of La and tr(H,)
tr(H: H.) = tr(H,) + (1+ 2 L.
(18)
Next, using Equations write down the change in tr(H, Hi) after the first training step in terms of (r(H)) and Lo
2
A Atr(H1f (C([(( tr(H1] Ho) = tr(HH1)2 - tx(())²2 +4 4
1
(Ci - Lo)
w
(n)r(HH)) - 4) T
2
Atr(Hf
(n)r(Ho) - 4) + 2(r(HH))
+ 4
[
(Li - Lo)
(19)
Next, we substitute = k/ tr() to obtain the above equation as a function of k
A Atr(HTHH) H1) k(k - 4)
k(k 4)
fd ir(Ho)²
+ (1+ 2
+ +2ʃ0
(Ci - LLo)
(20)
Finally, we calculate the expectation value of (Atr (r(Hf H1))
A At(HH H1))
k:l 4)
[kkk- 4)
fd
+((58)
+4 + 2 2) (C1 Lo),
(r((H))
w
w w
(21)
21

---

0
ApPlteletttteeee ffgh begrng
Caoer
Cassa
4
Figure The relationship between the critical constants for the un model trained on a single training examples (z, y) = () with MSE loss using gradient descent. Each data point corresponds to random initialization
for
by estimating (2((o))²
using the approach demonstrated in the previous section
3w
(22)
br br(Ho)²
4(w 3)
Inserting
fg [n((H))²
in Equation 21 along with (f6) = 1, we have
k(k 4)
3k(k 4)
2 w
(Air(HFHH)) H.))
+
+4 1+
(Li -Lo)
(23)
4(w + +3)
[(k;
At infinite width, the above equation reduces (A tr( (r(H H1) = ( Lo), and hence, kyrob kloss For any finite width, I(k, w) 0 for < k < 4. At S Kioss L1 Lo S 0. and therefore (A tr(HF H.)) < 0. In order for the sharpness to catapult, we require (A ir( H1)) > 0 and therefore frob > kloss. As 1/w increases (kk,, w)l also increases, which means a higher value of L1 Lo is required to reach a point where (A (r( H H1)) 2 0. Thus kro kloss increases with 1/w.
B.5 The early training trajectories
Figure showws the early training trajectories of the wv model with large (w 512) and small (w = 2) widths. The dynamics depicted show several similarities with early training dynamics of real-world models shown in Figure 2 At small widths, the loss catapults at relatively higher learning rates (specifically, Closs 3.74, which is significantly higher than the critical value of Ceris 2).
B.6 Relationship between critical constants
Figure 12 shows the relationship between various critical constants for the we model. The data show that the inqquality Closs S Csharp I Cmaz holds for every random initialization of the ue model.
B.7 Phase diagrams with error bars
This section shows the variation in the phase diagram boundaries of the uv model shown in Figure Za b). Figure [3 shows these phase diagrams. Each data point an average of over 500 initializations. The horizontal bars around each data point indicate the region between 25% and 75% quantile.
B.8 Derivation of the expectation values
S and [Tr(Hof²)
Here, we provide the detailed derivation of the averages
We begin by
Tr(Ha)²
12
finding the average
TT(
22

---

2-1
2-1
Nowel
2-1
Vione
2-2
2-2
2-
2-
2.a
2-4
2-5
2-4
(a)
(b)
(c)
Figure 13: The phase diagram of the ue model trained with MSE loss using gradient descent with (a) sharpness X (b) trace of Hessian trØ and (c) the square of the Frobenius norm tr HH, H₂) used as measure of sharpness. rate is scaled as H/tr(Ho). Each data point denotes an average of over 500 initialization, and the smooth curve represents 2-degree polynomial fitted to the raw data. The horizontal bars around the average data point indicate the region between 25% and 75% quantile.
(a), the learning rate is scaled as 7/ - s/Aff while in (b) and (c), the learning
O0 Tr(Ho)²
du;du,
||u|1
1|-l13)
Σj UyUzurt ([]u|12 [|u112)2
exp
(24)
where ||-|| denotes the norm of the vectors.
The above integral is non-zero only if = k. Hence, it is a sum of w identical integrals. Without any loss of generality, we solve this integral for and multiply by obtain the final result, i.e.
000
duadus
||u||2 + ||0||
uf (||u + 21
w
II
exp
(25)
Tr(Ho)²
2
1=1
Consider transformation of u, E IR" into dimensional spherical coordinates such that
U1 Tu COS Pu1
U1 Ty CO8 (Pu
(26)
which yields,
f8
2 (2
/ dr drudr an wn 1 exp
- + r
r cos² 1urr cos? Pui
w-1
2r(Ho)2
(r + r3)²
(27)
13
(2 / drudroexpy exp
w²
r + r (-
(rw+1 + rott) /
r2r3
du dw cos² Pay cos? pm
Tr(Ho)2
2
(28)
2 (2x))z / druddreep exp
2 + r
r2
dl cos? 9p
Tr(Ho)²
2
20-+1 (r
+ r
2
1.
(29)
where dl denotes the w dimensional solid angle element. Here, we denote the radial and angular integrals by Ir and Lp respectively. The radial integral 1, is
23

---

1, (0 dr dr
Ta --1t
r + r
exp
(30)
(r + r3)
2
Let Ta cos 1 and r, Rsin 0 with R E (0, 00) and 0 E E[-3
2: r].
then we have
| dR R e- for
T/2
1,
d0 cos #+1 0sin "+1
(31)
VT [
I, 23
(32)
12+3
where r(.) denotes the Gamma function. The angular integral 1p is
/ dip dipidip- ... d sin" 6p cos? 11 sin w-3 42 sin f
1
(33)
1. - / /dpidipp- ... diw sin 2-2 1p sin
sin w-3 92... sin p So dip sin "-2 10 cos? 401
(34)
So dpi sin w-2 401
71//2
I.
(35)
(()=2)
Plugging in Equations
32 and 35 into Equation 29
obtain very simple expression
2
W
(36)
Tr(Ho)
2
4(w + 1)*
The other integral
fd T+(Ho))
can be obtained by generalizing the above approach as described below
fd Tr(Ho))
10
du;du;
||u||2 + l0ol²) E Ejkdmm ur
2
exp
2m
2
([]u||2 + ||u112)2
(37)
The integral is zero if either j I and 1 = m or = k = = which we consider separately. Without loss of generality, we find the following integrals
fd
doddu;
||u
||011
uu
Tr(Ho)
=
2
2
exp
2
(1|u 12 + 110122)
(38)
22
-00
fd
du;du;
||u112 + ||0/13
uf
exp
(39)
Tr(Ho)*
2T
2
[||/| 1 + |10
which have the following expressions
f4
w²
I(w)
(40)
Tr(Ho)
16
T(w 4)
22
fd
9w2
I(w)
(41)
Tr(Ho)
16
l'w + 4)
24

---

Pawi
2-1
2-1
2-10
(a) d
d
d 16
Figure 14: Phase diagrams of FCNs in NTP with varying depths trained on the MNIST dataset using MSE loss.
where r(.)) denotes the gamma function. On combining the expressions with their multiplicities, we obtain the final result
ff
fo Tr(Ho)
ff (r(Ho))
3w0( - 1)
+ w
(42)
Tr(Ho)
fd Tr(Ho)*
3(w + 22
[(w)
(43)
16
F(w + +
B.9 Insights into the catapult effect in crossentropy loss using we model
In this section, we summarize the main intuition behind the discrepancy in the values of Closs for cross-entropy loss at large widths. We consider the ur model trained on classification task using logistic loss, as presented in [50].
Consider the ue model trained on binary classification task using the logistic loss on two training
examples (Ir1, v1) (1, 1) and (I22, 32) (1, -1). Then, the total loss is L(f) &/2 log(2 2 cosh(f)). Hence, the loss grows monotonically as the output function |ʃ] increases. The update
equation of the function is given by:
Je+1 Se. (1 (T((H)))()) + -C
(44)
St.
where v is the learning rate and C.). is the derivative of the loss. At large width, if the condition holds, then output function continues to decrease. Given that f'ff) < 1/2 in the above case, this decrease persists for tr(H) < 4. This result provides some intuition behind the discrepancy.
[1 ((())
C Phase diagrams of early training
This section describes experimental details and shows additional phase diagrams of early training. The results include (1) FCNs in NTP trained on MNIST, MNSIS, and CIFAR-10 datasets, (2) CNNs in SP trained on Busshon-MNIST and CIFAR-10, and (3) ResNets in SP trained on CIFAR- 10 datasets using MSE loss. Figures 14 to 19 show these results. The depths and widths are the same as specified in Appendix A Each data point is average over 10 initializations. The horizontal bars around the average data point indicate the region between 25% and 75% quantile. Phase diagrams for cross-entropy results shown in Appendix F
Additional experimental details : train each model for 10 steps using SGD with learning rates / '/>g' and batch size of 512, where 2 with E (0.0. ... Tmmax in steps of 0.1. Here, Zmaz relatd the maximum trainable learning rate constant as Cmaz 25max We have
considered 10 random initializations for each model. As mentioned in Appendix A do not
25

---

2-1
2-1
2-10
2-1
(a) d
(b) d 8
(c) d 16
Figure 15: Phase diagrams of FCNs in NTP with varying depths trained on the Fashion-MNIST dataset.
2-1
2-1
2-1
2-1-
2
(a) d = 4
(b) d = 8
(c) d = 16
Figure 16: Phase diagrams of FCNs in NTP with varying depths trained on the CIFAR- 10 dataset.
Tounl Iikaol Buai
loudl
2-2
2-
2-
1-
2-
2-3
2-3
2-2
2-1
2
2-1
23
2
(a) d
(b)
(c) 10
Figure 17: Phase diagrams of Convolutional Neural Networks (CNNs) in SP with varying depths trained on the Fashion-MNIST dataset.
(tun)
Piwrr!
2-2
2-
2-4
2-4
2
2-2
2
27
21
2
2'
2
27
(a) d
(b)
(c) 10
Figure 18: Phase diagrams of Convolutional Neural Networks (CNNs) in SP with varying depths trained on the CIFAR-10 dataset.
26

---

Grued Goshap (fmee)
o (Fuee)
2
22
2
22
(a) d 18
(b) d 34
Figure 19: Phase diagrams of Resnets in SP with different depths trained on the CIFAR-10 dataset.
Vhec? Outerg
Irue Istep Doai
dogth
death
2
(a) 512
(b) 1024
(c) w 2048
Figure 20: Phase diagrams of FCNs in NTP with varying widths trained the CIFAR-10 dataset.
consider averages over SGD runs as the randomness from initialization outweighs it. Hence, we obtain 10 values for each of the critical values in the following results. For each initialization, we compute the critical constants using Definitions 1 and 3 To avoid random increase in loss and sharpness due to fluctuations, we round off the values of x// /x/" and Ci/Co to their second decimal places before comparing with 1. We denote the average values using data points and variation using horizontal bars around the average data points, which indicate the region between 25% and 75% quantile. The smooth curves are obtained by fitting two-degree polynomial x 1/w and y can take on one of three values: Gloss: Csharp and Cyan.
+ ba + cx with
Phase diagrams with depth Figure 20 shows the phase diagrams with depth for FCNs in NTP trained on the CIFAR-10 dataset. The phase diagrams look qualitatively similar compared to the 1/v phase diagrams.
D Relationship between various critical constants
Figure 21 illustrates the relationship between the early training critical constants for models and datasets. The experimental setup is the same as in Appendix C Typically, find that Closs S Csharp < Ceharp holds true. However, there are some exceptions, which are observed at high values of 4/w (d, e)), where the trends of the critical constants converge, and large fluctuations can cause deviations from the inequality.
E The effect of d/w on the noise in dynamics
In this section, we demonstrate that FCNs with d/ar 2 1/16, the dynamics becomes noise- dominated. This aspect makes it challenging to distringuish the underlying deterministic dynamics from random fluctuations. demonstrate this, we consider FCNs trained on CIFAR-10 using MSE and cross-entropy loss and use 4096 training examples for estimating sharpness.
27

---

to
Gue
Glus
Canz
Cana
2
Conz
21
22
21
21
Colorp
Calorn
Calery
(a) FCNs on MNIST
(b) FCNs on Dushion-MNIST
(c) FCNs on CIFAR-10
Class
Class
Class
Came
Canez
Cans
24
21
21
22
21
21
2
Calorp
Cabare
Cbbay
(d) CNNs on Fashion-MNIST
(c) CNNs on CIFAR-10
(f) ResNets CIFAR-10
Figure 21: The relationship between various critical constants for various models and datasets. Each data point corresponds to model with random initialization. The dashed line denotes the values where x = y-
losc. with 1024 desth: 16
with 256 repil:
42
1012
with
cepth
Iser mad,, witth desth:
mnt, withi: depth
022
0.09
600
Figure 22: Training loss trajectories of ReLU FCNs with d = trained on the CIFAR- 10 dataset with MSE loss using SGD with learning rate o/>g' and batch size B = 512.
Figures 22 and 23 show the training loss and sharpness of FCNs with d 16 and varying widths, trained on CIFAR- 10 using MSE loss. We observe that the sharpness dynamics becomes noisier for S 64.
Figures 24 and 25 shows the training dynamics with loss switched to cross-entropy, while keeping the initialization and SGD batch sequence the same in the MSE loss case. In comparison to MSE
28

---

lost mes withh 1004 Feeth 18
los mse, witte: 532, depti:
les ras witte: 251, depti:
1.
12
022
22
10100
112
05-
0.4
06
Figure 23: Sharpness trajectories of ReLU FCNs with d 16 trained on the CIFAR-10 dataset with MSE loss using SGD with learning rate 17 = c/>² and batch size B = 512.
lose wer, softh: 517 depth 14
loss seet addtt 255 depth 14
Soos: semt, wedtt 138, desti: 16
wont, dapth:
kest andt wilks 32, dagly 15
330
035
999
032
280
mp
Figure 24: Training loss trajectories of ReLU FCNs with d = 16 trained on the CIFAR- 10 dataset with cross-entropy loss using SGD with learning rate n </>." and batch size 512.
loss, the training loss and sharpness dynamics show higher level of noise, especially for S 256. As result, becomes difficult to characterize the training dynamics for d/u- Z 1/16.
F
Crossentropy
In this section, we provide additional results for models trained with cross-entropy (xent) loss and compare them with MSE results. Broadly speaking, models trained with cross-entir coossenttrppy loss show similar characterstics to those trained with MSE loss, such as, (i) sharpness reduction during early training, (ii) an increase in critical constants Cooss Csharp with d and 1/m, (iii) Closs S Csharp S Cmazz However, the dynamics of models trained with cross-entropy loss is noisier compared to MSE as shown in the previous section, and characterizing these dynamics can be more complex. In the
29

---

loss oest, widia: 1824, degth:
wer, with deycls
koo weet, with 250, darch
1.
1ee 200t, widts desth 25
1.2
022
123
635
04
02
Figure 25: Sharpness trajectories of ReLU FCNs with d 16 trained on the CIFAR-10 dataset with cross-entropy loss using SGD with learning rate 17 = 5/>0" and batch size B 512.
14
2-1
23
(a) MSE, d = 4
(b) MSF, = 8
(c) MSF, = 16
2-9
Hell
2-13
2-30
2
2
(d) xent,
(e) xent, d
(1) xent, 16
Figure 26: The phase diagrams of early training of FCNs trained the CIFAR-10 dataset using (a. b, c) MSE and (d, c. f) cross-entropy loss. Each data point is an average over 10 initializations, and solid lines represent smooth curve fitted to raw data points. The horizontal bars around the averaged data point indicates the region between 25% and 75% quantile. For cross-entropy phase diagrams, the line is shown for reference only and does not relate to Cerit
following experiments, consider models trained on the CIFAR-10 dataset and used 4096 training examples to estimate sharpness.
F11 Phase diagrams
Figure 26 compares the phase diagrams of FCNs in SP trained on the CIFAR-10 dataset, using both MSE and cross-entropy loss. The estimated critical constants for cross-entropy loss generally
30

---

loss: mse
loss: xent
Risce
21
20
2
24
Oue
24
Gwar
Takars
(a)
(b)
Figure 27: Comparison of the relationship between critical constants for FCNs in SP trained on CIFAR-10 using MSE and cross-entropy loss. Each data point corresponds to a randomly initialized model with depths and widths mentioned in Appendix A
more noisy, as quantified by the confidence intervals. comparison to phase diagrams of models trained with MSE loss, observe a few notable differences. First, the loss starts to catapult at value appreciably larger than c = at large widths. Primarly, S Closs X 8. Additionally, Cmaz generally decreases with 1/w. This decreasing trend becomes less sharp at large depths.
Despite these differences, the phase diagrams for both loss functions share various similarities. First, we observe sharpness reduces during early training for с < Csharp (see the first row of Figure 25) Next, we observe that the inequality Closs S Csharp S Gmaz generally holds for both loss functions as demonstrated in Figure 27 barring some exceptions.
Figure 28 shows the phase diagrams for CNNs and ResNets trained on the CIFAR-10 dataset using cross-entropy loss. The observed critical constants are much noisier as quantified by the confidence intervals. Nevertheless, the phase diagram shows similar trends as mentioned above. For large 1/w models, we found that progressive sharpening begins after 5 - 10 training steps. For these cases, we only use the first steps measure sharpness to avoid progressive sharpening. For CNNs, we observed that the dynamics becomes difficult to characterize for w Z 32 and d Z 10, due to large fluctuations. Consequently, we've opted not to include these particular results.
27
2
23
(a) CNN, 5
(b) CNN,
(c) ResNet, 18
Figure 28: Phase diagrams of (a, b) CNNs and (c) ResNets trained the CIFAR-10 dataset with cross-entropy loss using SGD with 7 = s/x" and = 512.
F.2 Intemediate saturation regime
Figure 29 shows the normalized sharpness measured at 100 for FCNs trained on CIFAR-10 using cross-entropy loss. ? Similar to MSE loss, we observe an abrupt drop in sharpness at large learning rates. However, this abrupt drop occurs at 2 S Cerit S 4. The estimated sharpness is noisier (compare with Figure 38), which hinders a reliable estimation of Cerst- We speculate that we require a
"The time step r = 101/2 is in the middle of the intermediate saturation regime for most of the models. For further details estimating sharpness, Appendix I.I
31

---

width
178
widt
180
250
512
512
103
1033
122
0.75
0.4
0.30
0.2
112
0.25
2
2
(a) d
(b) d
(c) d = 16
Figure 29: Sharpness measured cT 100 against the learning rate constant for FCNs trained on the CIFAR-IE 10 dataset using cross-entropy loss, with varying depths and widths. Each curve is an average over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical dashed line shows for reference.
large number of averages for reliable estimation of Cerit for cross cross-entropy loss. We leave the precise characterization of Cerit for cross-entropy loss for future work.
G The effect of setting model output to zero at initialization
In this section, demonstrate the effect of network output f(; 06) at initialization the early training dynamics. In particular, we set the network output to zero at initialization, fa; 86) = by (1) "centering' the network by its initial value fe(2; 0.) [(x; 0) f(;; 0o) or (2) setting the last layer weights to zero at initialization. We show that both (1) and (2) remove the opening of the
sharpness reduction phase with 1/w. Resultantly, the average onset of loss catapult occurs at Class
independent of depth and width.
Throughout this section, use 'vanilla' networks to refer to networks initialized in the standard way. For simplicity, we train FCNs using full batch gradient descent with MSE loss using a subset consisting of 4096 examples of the CIFAR-10 dataset.
G.1 The effect of centering networks
Given network function (x; 0c), we define the centered network fe(æ; 0t) fe(z; 0,) = fx;; (f; f(r;00),
(45)
where (x; 0o) the network output at intialization. By construction, network output is zero at initialization. It is noteworthy that centering network is an unusual way of training deep networks as it doubles the cost training because of two forward passes.
Figure 30 compares the training loss and sharpness dynamics of vanilla networks and centered networks. Unlike vanilla networks, we do not observe decrease in sharpness for c < Closs during early training. Rather, we observe a slight increase in sharpness. To distinguish this slight increase from sharpness catapult, we introduce a threshold c, comparing normalized sharpness X #//"" with 1 1+ G, to define sharpness catapuit 8 As demonstrated in Appendix G.3 the uv model trained on single training example (x, y) with y # 0 sheds lights on this initial increase in sharpness.
Interestingly, irrespective of depth and width, we observe that loss catapults at Ctass ~ 2, as demonstrated the phase diagrams in Figure33] a, c). These findings suggest strong correlation between large network output at initialization [x;; fo) and the opening of the sharpness reduction phase discussed in Section 2
G.2 The effect of setting the last layer to zero
An alternative way to train networks with f(x; 0,) is by setting the last layer to zero at initialization. The principle of criticality at initialization [5 58 68 does put any constraints
In experiments,
0.05. We use the same threshold for zero-init networks.
32

---

vanilla nets
mse, width: 512, depth:
vanilla nets
loss: mac. wicth: 512, depth: B
Training 6
1.5
1.0
0.5
0.0
10
15
20
10
15
20
stop
(a)
(b)
centered
centered nets
loss: mse, width: 512, depth:
loss: mse. wieth: 512, depth: 8
2.5
1.0
X
1.0
0.5
0.5
0.0
3
10
15
20
0.0g
10
15
90
stop
(c)
(d)
zero-init nets
zero-init nets
mie, width: depth:
loss: mse: width: 512, depth: 8
1.5
1.0
0.5
0.5
5
20
15
20
5
201
15
20
step
step
(c)
(f)
Figure 30: Comparison of the early training dynamics of (a, b) vanilla, (c, d) centered, and (e, zero- initialized FCNs (with depth- = 8 and width = 512), trained on the CIFAR-10 dataset with MSE loss using gradient descent for 20 steps.
the last layer weights. Hence, setting the last layer to zero does not affect signal/gradient propagation at initialization. Yet, setting the last layer to zero results initialization in a flat curvature region at initialization, resulting in access to larger learning rates. We refer to these networks 'zero-init' networks.
Figure 30 compares the training dynamics of zero-ini init networks with vanilla and centered networks. We observe that the dynamics is quite similar to the centered networks: (i) sharpness does not reduce for small learning rates and (ii) loss catapults Closs ~ 2, irrespective of depth and width. Figure [T] f) show the phase diagrams
networks with zero-initialized networks. Like centered networks, the critical constants do not scale with depth and width. Again, suggesting that large network output initialization (x; (o) related the opening of the sharpness reduction phase in the early training results shown in Section 2
33

---

Voun Vinwog
2-1
2-1
2-1
(a) centered, 4
(b) centered, 8
(c) centered, d = 16
1
2-1
9-1
9-1 30
91 97 9-1
(d) zero-init init, d = 4
(e) zero-init, d = 8
(1) zero-init, d = 16
Figure 31: The phase diagrams of early training dynamics (a, b, c) centered and (d, e, f) zero-init networks trained on CIFAR-I0 using MSE using gradient descent. Each data point is an average over 10 initializations. The horizontal bars around the average data point indicate the region between 25% and 75% quantile.
G.3 Insights from we model trained on (x, y)
In this section, gain insights into the effect of setting network output at initialization using wv model trained on an example (I, y). In particular, we show that loss catapults at kloss = 2 and sharpness increases during early training.
Consider the we model trained on a single training example (x, y) with +
f(x) = V Σ
u&u; x.
This simplifies the loss function
C - (f(x)
(1()))
y)² -1) A,
(46)
where Af the residual. The trace of the Hessian tr(H) is
x² (1|1) + ||u|12)-
tr tr(H)
(47)
The Frobeinus norm can be written in terms of the trace and the network output
2 )
H - =2
1+
(48)
wAf)
The function and residual updates are given by
Note that for y 0, the network is already at minimum for fo 0.
34

---

0.1
0.1
lass: width: 512, depth:
losz: width: 512, depth: 8
0.000
(a)
(b)
5
1.0
loss: mse. width: 512, depth:
loss: width: 512, depth: 8
0.5
1.5
0.4
Trainig
1.0
0.22
(c)
(d)
10:0
loes: mse, width: 512, depth:
loss: mie, width: 512, depth:
BERES
1.5
0.5
0.0
10
stop
stop
Figure 32: The early training dynamics of FCNs with a fixed output scale trained on the CIFAR-10 dataset with MSE loss using gradient descent.
fe-1 -t- ()r((((H. + T222 S.A
(49)
AS+1 Asi (1-71((1))) + 1222
(50)
10
Similarly, we can obtain the trace update equations
AAf w
vr(H+41)= (r(H))
()((H))
(51)
Af.
Let us analyze them for the networks with zero output at initialization. The loss at the first step increases if
35

---

T1 Lo
-
-(11-
21 + npx? fosfi) )
(52)
(53)
Setting fo and scaling the learning rate as 17 k/ (r(H)), we see that the loss increases the first step if k > 2.
L1 To
= ((1 ((1 k)?) >
(54)
Next, analyze the change in trace during the first training step. Setting fo 0, we observe that the trace increases for all learning rates
tr(H1) tr(Ho) + 172x2
tr(o).
(55)
modulated by the learning rate and width. Finally, analyze the change in Frobenius norm in the first training step at k = kloss, which implies AfP ASo
(A|H11P)) = (Ir(H1,) (r(H,)2 2nd (Aʃ? Af7))
(56)
As (H(H) increases in the first training step, | || و also increases the first training step.
H The effect of output scale on the training dynamics
Given neural network function (ax with depth d and width w, define the scaled network as f.(x) = o (x), where a is referred to as the output scale. In this section, we empirically study the impact of the output scale the early training dynamics. In particular, small) value of [=; 0o) relative to the one-hot encodings of the labels causes the sharpness to decrease (resp. increase) during early training. Interestingly, we still observe increase in (Ctoss) with d and 1/w, unlike the case of initializing network output to zero, highlighting the unique impact of output scale on the dynamics. For simplicity, we train FCNs using gradient descent with MSE loss using a subset consisting of 4096 examples of the CIFAR- 10 dataset, as in the previous section.
show that large (resp.
H.1 The effect fixed output scale at initialization
In this section, we study the training dynamics of models trained with fixed output scale at initialization. Given a network output function f(0), we define the 'scaled network' as
fa(0) sf(()
(57)
[lf()))]l
where s is scalar, fixed throughout training. By construction, the network output norm [ff)ll))]
equals For standard initialization, (l/ (VV)), where k are the number of classes.
Figure 32 shows the training dynamics FCNs for three different values of the output scale s. The training dynamics of networks with = 1.0 and
networks initialized with smaller output scale 10.0 share qualitative similarities. In contrast, 0.1) exhibit distinctly different dynamics. In
particular, we observe that for large output scales (s (..) sharpness decreases during early training. while sharpness increases for small output scales Furthermore, the training dynamics tends to
We empirically observed that sharpness reduces for output scales as small as s 0.5, which relatively small compared to Vk.
36

---

1.0 lass
10.0 depth:
loss: mis, depth
depth: 4
2-1
148
1e4
H
1/20
2
2
(c)
1.0
10.0
lass: mis, depth:
depth: 8
depth:
2-5
0 Beod
2-1
Fel TM
14
2
2
2-1
(f)
Figure 33: The phase diagrams of early training dynamics for ReLU FCNs with fixed output scale trained on a subset the CIFAR-10 10 dataset using MSE loss using gradient descent. Each data point an average over 10 initializations. The horizontal bars around the average data point indicate the region between 25% and 75% quantile.
be noisier at small output scales, making it difficult to characterize catapult dynamics amidst these fluctuations. In summary, the training dynamics of networks with small output scale deviate from the training dynamics discussed in the main text, particularly the sharpness quickly increases during early training.
Figure 33 shows the trends of various critical constants with width for FCNs for three different values of s. Similar to vanilla networks, we observe that Closs increases with d and 1/w. In comparison, sharpness decreases (increases) for large (small) values of These experiments suggest that the output scale primarly influences the increase/decrease sharpness during early training and does not affect the scaling of Closs with depth and width.
Note that we do not generate phase diagrams for these experiments as the training dynamics of networks with small output scales at initialization deviate from the training dynamics disucssed in the main text.
H.2 Scaling the output scale with width
In this section, we study the training dynamics of models with an output scale scaled with width as = w- which is commonly used in the literature [19
We consider three distinct values (-0.5, 0.O, 0.5), where o -0.5 represents the lazy regime, or = 0.5 corresponds to feature learning (rich) regime and 0.0 correponds to standard (vanilla) initialization.
6 4
Figure shows the training loss and sharpness trajectories of FCNs trained on for different values. We observe that the training trajectories in the lazy regime look identical to standard initialization. In comparison, the training trajectories in the feature learning regime distinctly different. We observe that in the standard and lazy regimes, sharpness decreases during early training, whereas sharpness tends to increase in the feature learning regime and eventually oscillates around the edge of stability regime. Moreover, we observe that sharpness can catapult before the training loss in the feature learning regime (compare catapult peaks [4[e, f)). These results are parallel to the fixed output scale networks studied in the pervious section.
37

---

-0.5
IT:
-0.5
width: 512, depth:
width: 512, depth:
1.50
3000
1.20
Tranny bee
1900
1000
0.75
0.30
500
0.25
I0
step
step
(h)
d:
0.0
loss:
width: 512, depth:
mie width: 512, depth
1.50
L.25
1.00
Traiiing
0.30
0.25
15
21
0.00
10
20
stop
step
(c)
(d)
d: 0.5
0.5
depth:
depth:
2,5
0.000
2.1
Traning bue
1.5
0.043
1.0
0.000
stop
(f)
Figure 34: The early training dynamics of FCNs with output scale w- --taineed on the CIFAR- 10 dataset with MSE loss using gradient descent.
summarizes the early training dynamics FCNs with different o values. We observe similar results in the previous section. The output scale affects the initial increase/decrease of sharpness but does not affect the scaling trend of Closs with depth and width. Moreover, we observe a systematic pattern of Crar scaling with width. In the lazy regime, observe that Cmaa increases with 1/w, while Cmaz decreases with 1/w in the feature learning regime.
Figure 35
I Sharpness curves in the intermediate saturation regime
This section shows additional results for Section B for MSE loss. Cross- Cross-entropy results are shown in Appendix
Figures 36 to 40 show the normalized sharpness curves for different depths and widths.
L.1 Estimating the sharpness
This paragraph describes the procedure for measuring the sharpness to study the effect of the learning rate, depth, and width in the intermediate saturation regime. We measure the sharpness X a time in the middle of the intermediate saturation regime. We choose so that x 200, for learning rates c 22, where [-1.0, 4.0) in steps of 0. The value 200 is chosen such that J is in the middle of the intermediate saturation regime. Next, we measure sharpness over range of steps
38

---

-0.5
0.0 lass
0.5 depth:
lass: depth:
depth: 4
2-1
141
1,8
11
(c)
-0.5
0.0
0.5 depth:
lass:
depth:
depth: 8
2-1
M
Fell F--
2.1
231
(f)
Figure 35: The phase diagrams of early training dynamics for ReLU FCNs with varying depths and output scale.
t EE IT -- 5 5] and average over i to reduce fluctuations. We repeat this process for various initializations and obtain the average sharpness.
0.8
width
width
withh
2048
2043
2043
0.2
0.2
0.0
0.0
0.0
(a) d
(b)
(c) 16
Figure 36: Sharpness measured at ct = 200 against the learning rate constant for FCNs trained on the MNIST dataset, with varying depths and widths. Each curve average ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote Cerit estimated using the maximum of X
I.2 Estimating the critical constant Carit
This subsection explains how to estimate Cerit from sharpness measured time т. First, we normalize the sharpness with its initial value, and then average over random initializations. Next, we estimate the critical point Cerit using the second derivative of the order parameter curve. Even if the obtained averaged normalized sharpness curve is somewhat smooth, the second derivative may become extremely noisy as minor fluctuations amplify on taking derivatives. This can cause difficulties in obtaining Cerit- We resolve this issue by estimating the smooth derivatives of the averaged order parameter with the SavltzkyGoolay filter [59] using its scipy implementation
63 The estimated
Cori is shown by vertical lines in the sharpness curves in Figures 36 to 40
39

---

0.5
width
512
512
512 1024
1034
1026
01.4
0.2
0.2
0.0
2 0.0
2
2
2
2
(a) d =
(b) d 8
(c) d = 16
Figure 37: Sharpness measured at cr 200 against the learning rate constant for FCNs trained on the Fashion-MNIST dataset, with varying depths and widths. Each curve is an average over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote Corit estimated using the maximum X,
width
width
0.5
width
0.6
2043
2043
2043
0.4
0.3
0.0
0.0
0.0
(a) d
(b)
(c) 16
Figure 38: Sharpness measured at cT = 200 against the learning rate constant for FCNs trained on the CIFAR-10 dataset, with varying depths and widths. Each curve is average initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote Cerit estimated using the maximum of X--
wisth
width
0.5
width
128
11.5
128
120
0.6
106
0.4
0.2
012
0.0
0.0
0.0
(a) d
(b) d
(c) 10
Figure 39: Sharpness measured CT 200 against the learning rate constant for Myrtle-CNNs trained on the CIFAR- 10 dataset, with varying depths and widths. Each curve is an average of over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote Cerit estimated using the maximum of X
The effect of batch size on the reported results
J.11 The early transient regime
Figure 41 shows the phase diagrams of early training dynamics of FCNs with d = trained on CIFAR-10 dataset using two different batch sizes. The phase diagram obtained is consistent with the findings presented Section 2 except for one key difference. Specifically, we observe that when d/w is small and small batch sizes are used for training, sharpness may increase from initialization
40

---

withh
0.5
178
188
120
0.4
022
0.0
0.0
0.0
2
2 0
2
2
2
2
21
(a) d = 18
(b) d = 34
(c) d = 50
Figure 40: Sharpness measured at cT = 200 against the learning rate constant for ResNets trained on the CIFAR-I0 dataset, with varying depths and widths. Each curve an average of over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote Corit estimated using the maximum X,.
relatively smaller values of c. This is reflected in Fig. by (Csherp) moving the left as B is reduced from 512 to 128. However, this initial increase in sharpness is small compared to the sharpness catapult observed larger batch sizes. We found that this increase small batch sizes is due to fluctuations in gradient estimation that can cause sharpness to increase above its initial value by chance.
Boow?
2-1
2
(a) 128
(b) 512
Figure 41: The phase diagram of early training for FCNs with d 4 trained the CIFAR-10 dataset with MSE loss using SGD with different batch sizes.
J.2
The intermediate saturation regime
1653
256
0.6
87
264
1034
512
2003
1004
X 800
520
0.2
512.
24
21
(a)
(b)
(c)
200 against the learning rate constant for FCNs with trained on the CIFAR-IO dataset, with varying widths. Each data point is an average over 10 initializations, where the shaded region depicts the standard deviation around the mean trend. (b, c) Smooth estimations of the first two derivatives, X+ and X the averaged normalized sharpness wrt the learning rate constant.
Figure 42: (a) Normalized sharpness measured
Figure shows the normalized sharpness, measured at et 200, and derivatives for various widths and batch sizes. The results consistent with those in Section with lowering in the peak
41

---

heights of the derivatives and X' at small batch sizes. The lowering of the peak heights means the full width at half maximum increases, which implies a broadening of the transition around Cerit. at smaller batch sizes.
K
The effect of bias on the reported results
In this section, we show that FCNs with bias show similar results as presented in the main text. We considered FCNs in SP initialized with He initialization [29].
Figure 43 shows the phase diagrams of early training for FCNs with bias trained on the CIFAR-10 dataset. We observe similar phase diagram compared to the no-bias case (compare with Figure 26).
fownl
2-1
2-1
(a)
(b)
Figure 43: The phase diagram of early training for FCNs with bias trained on the CIFAR-10 dataset with MSE loss using SGD with different depths.
42


Reviews:
[{"review_id": "XI7GpF88WL", "reviewer": "Reviewer_bXcV", "summary": "This work mainly investigates the learning dynamics depending on the learning rate and categorizes its qualitative behaviors in some phases. In particular, the authors identify new phases; sharpness reduction, and loss-sharpness catapult. They investigate the dependence of the critical learning rates (mainly, $c_{class}, c_{sharp}, c_{max}$ and $c_{barrier}$) on the width, depth & loss type, and try to identify a universal relationship between them. \nTheir empirical observations in some deep neural networks are also confirmed in a simple linear network and the obtained phase diagram is expected to hold in various situations. ", "questions": "**Connection to the infinite width limit**\n\nWhile the current work focuses on finite-sized networks, some previous work investigated the infinite-width limit [64]. In particular, [64] claims that for the infinite-width limit, the learning dynamics *must be* the kernel regime for SP initialization (nevertheless, the learning dynamics does not progress or explodes).  How is this infinite-width study related to your results? Since the critical learning rate of the kernel regime is given by $2/\\lambda_{0}^{NTK}$, does this suggest that all  $c_{loss,sharp,max}$ converge to 2 in the infinite width limit? \n\n**Phase diagram in previous work**\n\nRelated to [64], there are several studies on the phase diagram of convergence in the infinite width case (e.g. https://arxiv.org/abs/2012.15110 and https://arxiv.org/abs/2007.07497). It would be better to mention them for giving a rich overview of the phase-diagram studies in the literature.  More related to the current work, the phase diagram for width v.s. learning rate is also investigated in the literature (https://arxiv.org/abs/1806.01316). This work implies that 2/$\\lambda_{0}^{NTK}$ works as a critical learning rate of the convergence in the infinite width. \n\n**The effect of early learning dynamics on eventual minima**\n\nI understand that the main focus of the current work is on the early training regime, but it is quite curious how the sharpness reduction, loss catapult, and loss-sharpness catapult phases determine the eventual generalization performance and sharpness of minima in the trained models. If this point would be clarified, it would increase the significance of this paper much more. ", "limitations": "**No normalization layers**\n\nThe architectures in this work have no batch/layer normalization layer. Since the batch norm is known to strongly affect the sharpness (e.g. https://arxiv.org/abs/1901.10159), it is curious how the normalization layer potentially changes the results.  \n\n**Difference between MSE and cross-entropy cases**\n\nIt seems that the cross-entropy loss shows a different behavior of the critical learning rates. In particular, in most cases, $c_{loss}$ and $c_{sharp}$ decrease close to 2 as $1/w$ decreases in the MSE case while they keep far from 2 in the cross-ent case. In the revised version, I believe it would be better to give a more elaborated explanation of the difference in the main text.\n\n\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This work attacks the challenging problem to find any universal behavior of early learning dynamics in finite-size networks (I guess that the terminology of \"phase diagram\" comes from physics or chemistry but such a dynamic and finite-sized diagram seems not known). Despite this difficulty, the authors verified new phases (characterized by $c_{loss},c_{sharp},c_{max}$) in various settings and these new phases seem to have the originality. Their careful and elaborated experiments are persuasive and enhance the quality of the paper. \n", "weaknesses": "(a) The whole findings are empirical. Even in a simple model of $uvx$, the authors did not solve the dynamics in an analytical form. It is quite hard to obtain any intuition on why the three phases characterized by $c_{loss}<c_{sharp}<c_{max}$ universally appear. \n\n (b)They are several unclear points that the authors should more clearly mention. See below.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to read our manuscript and for providing detailed comments. Below are our responses to the comments provided:\n\n> Connection to the infinite width limit:\n\nIn our study, we are looking at the behavior of the optimization dynamics at training timescales $t_\\star$ that grow with width $w$. Specifically, the end of the early time transient period occurs at $t_\\star \\sim \\log(w)$. However, in the usual infinite width studies in SP initialization [64], the training dynamics is restricted to $O(1)$ times in the limit of infinite width; this leads to the kernel regime, for which the system is in a lazy training regime for learning rates less than $2/\\lambda_0^{NTK}$ and divergent training for larger learning rates. Our results for critical constants in the limit of large width, therefore, do not converge to $2$ as would be expected from the analysis of [64] (except for $c_{loss}$, which converges to $2$ for MSE loss). We will provide a more detailed explanation in the updated version of the paper to clarify this distinction.\n\n> Phase diagram in previous work:\n\nWe will incorporate the above suggestion by including a paragraph discussing the phase diagram from previous studies in the updated version of our manuscript.\n\n> The effect of early learning dynamics on eventual minima:\n\nWe agree that analyzing generalization properties through the lens of the phase diagram could lead to a more comprehensive understanding and enhance the importance of the study. However, this is a non-trivial task due to the extensive hyperparameter search space (including depth, width, learning rates, initializations, and batch size) and the significantly large training time required for each workload (model and training task). Moreover, defining the success measure, such as finding the best generalization error for a fixed compute versus the shortest training time for a given target error, adds another layer of complexity. A comprehensive analysis would result from a cumulation of studies, with our study contributing to this effort.\n\n> Difference between MSE and cross-entropy cases:\n\nThe disparity in the limiting value of $c_{loss}$ at large width is mostly due to the non-constant Hessian of the cross-entropy loss, as discussed in [62]. Previous work, such as [47], analyzed the catapult dynamics for the $uv$ model with logistic loss and demonstrated that the loss catapult occurs above $\\eta_{loss} = 4 / \\lambda_0^{NTK}$. \n Their argument is as follows.\n \n Consider the $uv$ model trained on a binary classification task using the logistic loss on two training examples $ (x_1, y_1) = (1, 1) $ and $ (x_2, y_2) = (1, -1) $. Then, the total loss is $ \\mathcal{L}(f) = \\frac{1}{2} \\log (2 + 2 \\cosh(f)) $. Hence, the loss grows monotonically as the output function $ |f| $ increases.\nThe update equation of the function is given by:\n\n\\begin{align}\n    f_{t+1} = f_t \\left( 1 -  \\frac{\\eta \\lambda^{NTK}_t \\mathcal{L}'(f)}{f_t} + \\frac{\\eta^2 \\mathcal{L}'(f_t)^2}{w}  \\right),\n\\end{align}\nwhere $\\eta$ is the learning rate, $w$ is the width, and $\\mathcal{L}'(.)$ is the derivative of the loss. \nAt large width $w$, if the condition $ | 1 - \\eta \\lambda^{NTK} \\mathcal{L}'(f_t) / f| < 1 $ holds, then the output function continues to decrease. Given that $ \\mathcal{L}'(f) / f \\leq 1/2 $ in the above case, this decrease persists for $ \\eta \\lambda^{NTK} < 4 $.\n\nThis result provides some intuition behind the discrepancy. We can describe the above analysis in more detail in the updated version of the manuscript. However, a complete understanding of the catapult phenomenon in the context of cross-entropy loss requires a more detailed examination.\n\n \n\n\n"}, {"review_id": "D0G58FNQ9i", "reviewer": "Reviewer_qjGV", "summary": "In this paper, the authors studied the effects of learning rate, width and depth of neural networks for early training dynamics. Specifically, the authors focused on the value of loss and the maximum eigenvalue of Hessian of loss and found that there are 4 different possible training dynamics in the early training. These different regimes are determined by the value of learning rate and are related with depth and width. Experiment results are provided to support the claim. A simple example with the same phenomena is also given and analyzed.", "questions": "1.\tI was wondering for these $c$ values ($c_{loss}$, $c_{sharp}$, $c_{max}$), how many times do the experiments repeat?\n2.\tFor the definition of these $c$ values, I was wondering how time $T_1$ is chosen.\n3.\tIn Definition 4, should $\\chi_\\tau$ be $\\chi_\\tau^\\prime$?\n4.\tIn Figure 4, I was wondering if the networks have the same depth and width, or different point represents network that may have different depth and width.\n", "limitations": "The limitations are discussed in the paper. This is a theoretical work and therefore does not seem to have negative societal impact.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.\tUnderstanding the effect of large learning rate and its effect on the nonlinear training dynamics is an important problem.\n2.\tThe paper presents many experiment results to support their findings.\n3.\tThe finding that there exist different regimes depending on the learning rate ($c_{loss}$, $c_{sharp}$, $c_{max}$) in the early training seems to be interesting and new. \n", "weaknesses": "1.\tThe current paper seems to have lots of results and experiments in the main text. As a reader, it is not very easy for me to get the main conclusion for each section/set of experiments. It might be good to highlight the conclusions so that the readers can understand the point easier.\n2.\tThe current paper tries to study the effect of depth and width on those $c$ values ($c_{loss}$, $c_{sharp}$, $c_{max}$) related with stepsize. However, it seems to me that only a few choices of width $w$ are tried in the experiments (Figure 3) and no results showing the effect of depth $d$ in the main text. It might be good to include more results on different width and depth so that the fitted curves could use more than only 3 or 4 data points and the support of the conclusion could be stronger.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and careful reading. Below are our responses to the questions and comments:\n\n> The current paper seems to have lots of results and experiments in the main text. As a reader, it is not very easy for me to get the main conclusion for each section/set of experiments. It might be good to highlight the conclusions so that the readers can understand the point easier.\n\nWe can highlight/summarize the conclusion of each experiment in each section if it helps readers to grasp the main points easily. This requires minor edits, and we can implement them in the updated version of the manuscript.\n\n> The current paper tries to study the effect of depth and width on those values ($c_{loss}, c_{sharp}, c_{max}$) related with stepsize. However, it seems to me that only a few choices of width $w$ are tried in the experiments (Figure 3) and no results showing the effect of depth $d$ in the main text. It might be good to include more results on different width and depth so that the fitted curves could use more than only 3 or 4 data points and the support of the conclusion could be stronger.\n\nFor each architecture, we explored various values of depth and width, considering $10$ distinct initializations for each. Specifically, we tested $4$ values of widths and $3$ values of depths for FCNs and $3$ values of widths for CNNs. We concur that adding more values could strengthen our conclusions. Figure 1 of the PDF attached to the global response shows the phase diagrams of three neural networks with ~10 width values: (a) 8-layer FCNs trained on MNIST trained with MSE, (b) 7-layer CNNs trained on Fashion-MNIST using MSE, and (c) 16-layer FCNs trained on CIFAR-10 using MSE. We will make these adjustments in the revised version. We also agree that including the effect of depth in the main text would be beneficial and will make these enhancements in the updated version of the manuscript. \n\n> I was wondering for these $c$ values ($c_{loss}, c_{sharp}, c_{max}$), how many times do the experiments repeat?\n\nAs mentioned in line 4 of Figure 3 caption, each experiment is repeated for $10$ distinct initializations for each depth and width to obtain the $c$ values.\n\n> For the definition of these $c$ values, I was wondering how time $T_1$is chosen.\n\nWe choose $T_1$ to be the smallest value of step $t$ that contains the entire duration of the catapult effect. For the widths and depths considered, the catapult typically lasts at most $\\sim 10$ steps. Thus, for computational efficiency, we've set $T_1 = 10$.\n\n> In Definition 4, should $\\chi_\\tau$ be $\\chi_\\tau'$?\n \nWe thank the reviewer for pointing this out. We have fixed this in the updated version of the manuscript.\n\n> In Figure 4, I was wondering if the networks have the same depth and width, or different point represents network that may have different depth and width. \n\n In Figure 4, each data point can correspond to a run with varying depth, width, and initialization. With $n_w$ values of width, $n_d$ values of depth, and $10$ initializations, there are a total of $n_w \\times n_d \\times 10$ data points in this Figure, which illustrates that the inequality $c_{loss} \\leq c_{sharp} \\leq c_{max}$ holds regardless of depth, width, and initialization.  We thank the reviewer for pointing this out. We will update the Figure caption to include this information in the revised version of the manuscript.\n\n"}, {"review_id": "jO00fub6rz", "reviewer": "Reviewer_BcD3", "summary": "The paper identifies new phenomena involving the dynamics of the sharpness (largest eigenvalue of the Hessian) in the early stages of neural net training, for ReLU nets of various architectures that are initialized using He initialization.   Previously, it was known that for very wide networks, there are two 'phases' in hyperparameter space: on the one hand, if the step size exceeds 2/(initial sharpness), then a \"catapult\" occurs (https://arxiv.org/abs/2003.02218) and the sharpness drops; on the other hand, if the step size is below this threshold, then the sharpness remains at its initial value.  This paper fills in the picture for _narrow_ nets.   For these nets, the paper shows that the sharpness decreases at the beginning of training even when the step size is less than 2/(initial sharpness).   As a consequence, the train loss does not spike even when the step size is bigger than 2/(initial sharpness), since the sharpness drops first.   Only for larger learning rates does the train loss initially spike.  Then, for still larger learning rates, the paper shows the _sharpness_ also spikes (which the paper calls a \"sharpness catapult\"), which has not been observed previously.  For several architectures, the paper draws phase diagrams depicting which learning rates cause which behavior.  The paper then investigates connections to the loss connectivity literature, finding that only very large learning rates result in a barrier in the loss landscape (and that not every catapult results in a barrier).  The paper then investigates the role of initial network output, finding that initial sharpness reduction goes away if the network is made to output zero at initialization.  The paper finally reproduces all of these phenomena analytically in a simplified model (the u-v model from the catapult paper).", "questions": "- Have you looked at loss connectivity between the final parameters (at the end of training) and the initialization? ", "limitations": "discussed above", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "- The paper identifies and carefully studies a novel phenomenon involving the dynamics of the sharpness at the early stage of training.\n- The experiments are very thorough (though limited to ReLU nets initialized with He initialization)\n- There is an analytical explanation in a simplified model\n- The findings shed some light on mode connectivity\n", "weaknesses": "- The paper doesn't explain the early-stage sharpness reduction for general neural nets (i.e. beyond the simple u-v model)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging comments. Here is our response to the question on loss connectivity:\n\n> Have you looked at loss connectivity between the final parameters (at the end of training) and the initialization? \n\nAs the main focus of this work is the early training dynamics, we did not analyze the loss connectivity between the initial and final parameters. However, we refer to a couple of studies that linearly interpolate the loss between the initial and final parameters.\n\nGoodfellow et al. 2014 (https://arxiv.org/abs/1412.6544) linearly interpolate the loss between the initial and final points of training and show that training does not encounter any barriers. However, they do not analyze the effect of large learning rates. A more recent study by Lucas et al. 2021 (https://arxiv.org/abs/2104.11044) analyzed the effect of large learning rates in Section 4.1 of their paper. They report that training traverses a barrier at large learning rates, which is in agreement with the naive intuition of a barrier between the initial and final points of the loss catapult. However, they do not distinguish their result as a function of different learning rate phases. We believe that analyzing such results through the lens of our phase diagram could lead to a more comprehensive understanding of these phenomena.\n \n"}, {"review_id": "K7gdPgNqts", "reviewer": "Reviewer_bhX5", "summary": "The paper studies the effect of depth, width and learning rate on the early training dynamics of DNNs. The authors first describe four phases throughout, and then focus on the first two phases.\n\nIn the first phase, they identify three ranges of learning rates, according to how they affect the evolution of the training loss and the sharpness of the loss in the first few training steps:\n- for small learning rates both loss and sharpness decrease monotonically. Surprisingly, this regime extends to learning rates that are larger than the traditional threshold of $2 / \\lambda_{max}$ which guarantees a close match of GD to GF.\n- For larger learning rates, the loss is non-monotone but the sharpness still decreases.\n- For yet larger learning rates, both the loss and sharpness are non-monotone.\nAny larger learning rates lead to divergence.\nIn general larger depths $d$ and smaller widths $w$ shifts these regime to larger values of the learning rate.\n\nIn the second phase, the sharpness settles to some value. For small learning rates $\\eta$, this value is roughly constant w.r.t. $\\eta$ (but which depends on the width), but for large $\\eta$ the sharpness decreases and settles to a value which seems to be independent of the width.\n\nThe simple model of a shallow linear network with one datapoint and one dimensional inputs and outputs is studied, and the qualitative features described above are recovered.", "questions": "It is shortly mentioned that setting the outputs of the network at initialization to zero changes the ranges of learning rates corresponding to the regimes of phase 1. In the appendix, it appears that the transitions between ranges become almost independent of the width in this case, and the smallest range (where both loss and sharpness decrease) stops much closer to the expected value of $2 / \\lambda_{max}$. This suggests that the dependence of these values on the width $w$ which is quite thoroughly studied in the main paper, might only be related to the variance size of the outputs at initialization, and not some deeper property of the network. Ideally I think that the plots of the appendix describing this phenomenon should be put in the main, as it greatly impacted my personal interpretation of the results of the paper. \n\nDid the authors also check the behavior of the output function norm in the first few timesteps? I was wondering whether the non-monotonic behavior at the begining of training might be related to the function norm having a similar non-monotonic behavior (scaling up then down, possibly to some value close to zero), since the function norm has an obvious impact on both the loss and sharpness of the loss.\n", "limitations": "The limitations of the results presented are discussed well.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Understanding the different regimes of training as a function of width, depth and learning rates is an important problem, and any work that studies these questions either mathematically or empirically is welcome. The analysis and discussions are clear and insightful. A range of different architectures are tested. The regimes proposed and studied are to my knowledge new.", "weaknesses": "The authors define a large number of regimes acoording to some rather subtle changes in behavior sometimes, and it is not always clear whether these regimes are also correlated with other quantities that we actually care about (e.g. generalization/feature learning). The training phase of neural networks features many strange/non-monotonic behavior and it is not clear that all of them have to be identified, since they might have a rather small impact.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging comments. Below are our responses to the comments:\n\n> The authors define a large number of regimes acoording to some rather subtle changes in behavior sometimes, and it is not always clear whether these regimes are also correlated with other quantities that we actually care about (e.g. generalization/feature learning). The training phase of neural networks features many strange/non-monotonic behavior and it is not clear that all of them have to be identified, since they might have a rather small impact.\n\nWe believe that analyzing the generalization and feature learning properties through the lens of our phase diagram could lead to a more comprehensive understanding. Various studies, including [46, 66] and a recent paper (https://arxiv.org/abs/2306.04815), claim that large learning rates lead to better performance. \nIn particular, [46] conjectures that optimal performance in wide networks is obtained for learning rates between $c_{loss} = 2$ and $c_{max}$. We believe that our finer characterization of the phase diagram, particularly at finite depth and width, could help in narrowing down the range of learning rates with optimal performance.\n\nFrom a broader perspective, we believe that any systematic behavior in optimization dynamics that occurs repeatedly throughout architectures and datasets is worth documenting and trying to understand. While we sympathize with the referee's comment that it is unclear what will be useful for quantities we care about, we do think that systematic studies that shed light on the nature of the loss landscape and the behavior of nonconvex optimization are worthwhile as it is difficult to predict which discoveries will eventually be practically useful.  \n\n\n> It is shortly mentioned that setting the outputs of the network at initialization to zero changes the ranges of learning rates corresponding to the regimes of phase 1. In the appendix, it appears that the transitions between ranges become almost independent of the width in this case, and the smallest range (where both loss and sharpness decrease) stops much closer to the expected value of $2 / \\lambda_{max}$. This suggests that the dependence of these values on the width $w$ which is quite thoroughly studied in the main paper, might only be related to the variance size of the outputs at initialization, and not some deeper property of the network. Ideally I think that the plots of the appendix describing this phenomenon should be put in the main, as it greatly impacted my personal interpretation of the results of the paper.\n\nWe understand that this aspect may deserve more attention in the main text. Accordingly, we can move the relevant phase diagrams (for example, Figure 30) from the appendix to the main part of the paper to highlight this observation. \n\n> Did the authors also check the behavior of the output function norm in the first few timesteps? I was wondering whether the non-monotonic behavior at the begining of training might be related to the function norm having a similar non-monotonic behavior (scaling up then down, possibly to some value close to zero), since the function norm has an obvious impact on both the loss and sharpness of the loss.\n\nBased on the reviewer's feedback, we performed a preliminary analysis of the output function norm during early training in FCNs. For models trained with MSE loss, we found that the critical learning rate constant $c_{output}$ at which the output norm increases from initialization aligns with $c_{loss}$ (see table below). This is expected because the output norm is closely related to the loss. However, for cross-entropy loss, we found that the critical constant for the output norm is smaller than $c_{loss}$ as shown in the table below, specifically, $c_{output} < c_{loss}$. A thorough study would be needed to fully understand the relationship between the output function norm and cross-entropy loss.\n\n**Table: Average values of different critical constants of 4 layer FCNs trained with MSE loss using SGD with batch size $512$. Each value is an average over 10 distinct initializations.**\n\n| $d$ | $w$ | $c_{loss}$ | $c_{sharp}$ | $c_{output}$ | $c_{max}$ |\n|-----|-----|------------|-------------|--------------|----------|\n| 4   | 256  | 3.05      | 3.85        | 3.04         | 7.69     |\n| 4   | 512  | 2.52      | 3.34        | 2.55         | 8.41     |\n| 4   | 1024 | 2.30      | 3.35        | 2.30         | 9.86     |\n| 4   | 2048 | 2.24      | 3.39        | 2.24         | 10.22    |\n\n---\n\n**Table: Average values of different critical constants of 4 layer FCNs trained with cross-entropy loss using SGD with batch size $512$. Each value is an average over 10 distinct initializations.**\n\n| $d$ | $w$ | $c_{loss}$ | $c_{sharp}$ | $c_{output}$ | $c_{max}$ |\n|-----|-----|------------|-------------|--------------|----------|\n| 4 | 256   | 7.70       | 10.94       | 4.79         | 21.68    |\n| 4 | 512   | 7.03       | 13.23       | 4.59         | 37.52    |\n| 4| 1024   | 6.66       | 22.32       | 4.14         | 69.42    |\n| 4| 2048   | 5.51       | 5.83        | 3.64         | 101.80   |\n\n\n"}]

Author rebuttal:
We would like to express our gratitude to all the reviewers for taking the time to review our paper and providing detailed comments. Attached below is the PDF containing the figures required for some responses. 


Meta-review: The paper studies empirically nonlinear training dynamics under different learning rate, network width and height. It is an important topic and reviewers agree that the paper is well-written and the experiments are relatively thorough. Therefore I vote for acceptance.
Decision: Accept (poster)