Title: Phase diagram of early training dynamics in deep neural networks: effect of the learning rate, depth, and width
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

We systematically analyze optimization dynamics in deep neural networks (DNNs) trained with stochastic gradient descent (SGD) and study the effect of learning rate η , depth d , and width w of the neural network. By analyzing the maximum eigenvalue λ 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 η ≡ c / λ H 0 , 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.

1 Introduction

The optimization dynamics of deep neural networks (DNNs) is a 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 η < 2 / λ H , where λ H 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 it is both possible and potentially preferable to have the learning rate early in training reach η > 2 / λ H [66, 49, 72]. The idea is that such a choice will induce a temporary training instability, causing the network to 'catapult' out of a local basin into a flatter one with lower λ H where training stabilizes. Indeed, during the early training phase, the local curvature of the loss landscape changes rapidly [42, 1, 37, 16], and the learning rate plays a crucial role in determining the convergence basin [37]. Flatter basins are believed to be preferable because they potentially lead to lower generalization error [31, 32, 42, 12, 39, 14] and allow larger learning rates leading to potentially faster training.

From a different perspective, the major theme of deep learning is that it is 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 to maintain a similar quality of learning as the model size is taken to infinity [34, 47, 48, 11, 69, 58, 70, 68]?

10 0

10 1

10 2

10 3

10 4

10 5

t

(step)

10

-

1

10 0

Training loss

(a)

(i)

(ii)

(iii)

(iv)

10 0

10 1

10 2

10 3

10 4

10 5

t

(step)

0

1

2

3

λ H

t

/λ H

0

(b)

(i)

(ii)

(iii)

(iv)

10 0

10 1

10 2

10 3

10 4

10 5

t

(step)

0

.

2

0

.

4

0

.

6

0

.

8

Training accuracy

(c)

(i)

(ii)

(iii)

(iv)

c

1.0

1.4

2.0

2.8

4.0

5.7

8.0

11.3

16.0

22.6

Motivated by these ideas, we perform a systematic analysis of the training dynamics of SGD for DNNs as learning rate, depth, and width are tuned, across a variety of architectures and datasets. We monitor both the loss and sharpness ( λ H ) trajectories during early training, observing a number of qualitatively distinct phenomena summarized below.

1.1 Our contributions

We study SGD on fully connected networks (FCNs) with the same number of hidden units (width) in 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 at criticality [55, 58], while other parameterizations and initializations may show different behavior. Further experimental details are provided in Appendix A. We study both mean-squared error (MSE) and cross-entropy loss functions and the datasets CIFAR10, MNIST, Fashion-MNIST. Our findings apply to networks with d / w ≲ C , where C depends on architecture class (e.g. for FCNs, C ≈ 1 / 16 ) and loss function, but is independent of d , w , and η . Above this ratio, the dynamics becomes noise-dominated, and separating the underlying deterministic dynamics from random fluctuations becomes challenging, as shown in Appendix E. We use sharpness to refer to λ H t , the maximum eigenvalue of H at time-step t , and flatness refers to 1 / λ H t .

By monitoring the sharpness, we find four clearly separated, qualitatively distinct regimes throughout the training trajectory. Fig. 1 shows an example from a CNN architecture. The four observed regimes are: (i) an 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) a late time regime where the sharpness saturates around 2 / η for MSE loss; whereas for crossentropy loss, sharpness drops after reaching this maximum value while remaining less than 2 / η [8]. Note the log scale in Figure 1 highlights the early regimes (i) and (ii); in absolute terms these are much shorter in time than regimes (iii) and (iv).

In this work, we focus on the early transient and intermediate saturation regimes. As learning rate, d and w are tuned, a clear picture emerges, leading to a rich phase diagram, as demonstrated in Section 2. Given the learning rate scaled as η = c / λ H 0 , we characterize four distinct behaviors in the training dynamics in the early transient regime (i):

Sharpness reduction phase ( c < c loss ) : Both the loss and the sharpness monotonically decrease during early training. There is a particularly significant drop in sharpness in the regime c crit < c < c loss , which motivates us to refer to learning rates lower than c crit as sub-critical and larger than c crit as super-critical. We discuss c crit in detail below. The regime c crit < c < c loss opens up significantly with increasing d and 1 / w , which is a new result of this work.

Loss catapult phase ( c loss < c < c sharp ) : 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 sharpness monotonically decreases from initialization in this early time transient regime.

Loss and sharpness catapult phase ( c sharp < c < c max ): In this regime both the loss and sharpness initially start to increase, effectively catapulting to a 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 > c max ): The learning rate is too large for training and the loss diverges.

The critical values c loss , c sharp , c max are random variables that depend on random initialization, SGD batch selection, and architecture. The averages of c loss , c sharp , c max shown in the phase diagrams show strong systematic dependence on 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 a range of learning rates c loss < c < c barrier , 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 c crit . For MSE loss c crit ≈ 2 , whereas for crossentropy loss, 4 ≳ c crit ≳ 2 . For c ≪ c crit , the network is effectively in a lazy training regime, with increasing fluctuations as d and/or 1 / w are increased.

Finally, we show that a single hidden layer linear network - the uv 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., [13, 1, 15, 37, 17, 45, 35, 8, 33]. Here we focus on studies that characterize training regimes with sharpness and learning rates. Several studies have analyzed sharpness at different training times [37, 16, 35, 8, 33]. 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, 2, 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 sharpness-reduction phase, c < c loss which becomes more prominent with increasing d and 1 / w , where sharpness only decreases during early training; this also occurs in the catapult phase c loss < c < c sharp , during which the loss initially increases before decreasing. This discrepancy is likely due to different initialization and learning rate scaling in their work [33].

Ref. [35] examined the effect of hyperparameters on sharpness at late training times. Ref. [20] 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.

Refs. [52, 71] present phase diagrams of shallow ReLU networks at infinite width under gradient flow. Previous studies such as [41, 69] show that 2 / λ NTK 0 is the maximum learning rate for convergence as w → ∞ . This limit results in the kernel regime as training time is restricted to O (1) in the limit of infinite width, resulting in a lazy training regime for learning rates less than 2 / λ NTK 0 and divergent training for larger learning rates. In contrast, we analyze optimization dynamics at training timescales t ⋆ that grow with width w . Specifically, the end of the early 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 / λ NTK 0 < η < η max , 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 [72, 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

0

2

4

6

8

10

t

(step)

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Training loss

(a)

0

2

4

6

8

10

t

(step)

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

λ H

t

/λ H

0

(b)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

s

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

L

int

(

s, c

)

(c)

c

2.0

3.0

4.0

5.7

8.0

0

2

4

6

8

10

t

(step)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Training loss

(d)

0

2

4

6

8

10

t

(step)

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

λ H

t

/λ H

0

(e)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

s

0

.

1

0

.

2

0

.

3

L

int

(

s, c

)

(f)

c

2.0

5.7

11.3

22.6

45.3

of a sharpness-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 an initial increase in the training loss for learning rates c > 2 [49]. Fig. 2(a, b) shows the first 10 steps of the loss and sharpness trajectories of a shallow ( d = 5 and w = 512 ) CNN trained on the CIFAR-10 dataset with MSE loss using SGD. For learning rates, c ≥ 2 . 82 , the loss catapults and training eventually converges into a flatter region, as measured by sharpness. Additionally, we observe that sharpness may also spike initially, similar to the training loss (see Fig. 2 (b)). However, this initial spike in sharpness occurs at relatively higher learning rates ( c ≥ 5 . 65 ), which we will examine along with the loss catapult. We refer to this spike in sharpness as 'sharpness catapult.'

An important consideration is the degree to which this phenomenon changes with network depth and width. Interestingly, we found that the training loss in deep networks on average catapults at much larger learning rates than c = 2 . Fig. 2(d, 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, we characterize the properties of the early training dynamics of models with MSE loss. In Appendix F, we show that a 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. ( c loss , c sharp , c max ) For learning rate η = c / λ H 0 , let the training loss and sharpness at step t be denoted by L t ( c ) and λ H t ( c ) . We define c loss ( c sharp ) as minimum learning rates constants

2

-

1

2 0

2 1

2 2

2 3

2 4

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

〈

c

barrier

〉

(a)

d

= 8

2

-

1

2 0

2 1

2 2

2 3

2 4

2 5

2 6

c

2

-

8

2

-

7

2

-

6

2

-

5

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

〈

c

barrier

〉

(b)

d

= 7

2

-

1

2 0

2 1

2 2

2 3

2 4

2 5

c

2

-

7

2

-

6

2

-

5

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

〈

c

barrier

〉

(c)

d

= 18

2

-

1

2 0

2 1

2 2

2 3

2 4

2 5

c

2 2

2 3

2 4

d

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

〈

c

barrier

〉

(d)

w

= 512

2

-

1

2 0

2 1

2 2

2 3

2 4

c

2 2

2 3

2 4

d

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

〈

c

barrier

〉

(e)

w

= 1024

2

-

1

2 0

2 1

2 2

2 3

2 4

c

2 2

2 3

2 4

d

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

〈

c

barrier

〉

(f)

w

= 2048

such that the loss (sharpness) increases during the initial transient period:

and c max as the maximum learning rate constant such that the loss does not diverge during the initial transient period: c max = max c { c | L t ( c ) < K, ∀ t ∈ [1 , T 1 ] } , where K is a fixed large constant. 5 .

Note that the definition of c max allows for more flexibility than previous studies [33] in order to investigate a wider range of phenomena occurring near the maximum learning rate. Here, c loss , c sharp , and c max are random variables that depend on 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 ⟨ c loss ⟩ , ⟨ c sharp ⟩ , and ⟨ c max ⟩ are affected by width. The results show that the averaged values of all the critical constants increase significantly with 1 / w (note the log scale). At large widths, the loss starts to catapult at c ≈ 2 . As 1 / w increases, ⟨ c loss ⟩ increases and eventually converges to ⟨ c max ⟩ at large 1 / w . By comparison, sharpness starts to catapult at relatively large learning rates at small 1 / w , with ⟨ c sharp ⟩ continuing to increase with 1 / w while remaining between ⟨ c loss ⟩ and ⟨ c max ⟩ . Similar results are observed for different depths as demonstrated in Appendix C. Phase diagrams obtained by varying d are qualitatively similar to those obtained by varying 1 /w , as shown in Figure 3(d-f). Comparatively, we observe that ⟨ c max ⟩ may increase or decrease with 1 / w in different settings while consistently increasing with d , as shown in Appendices F and H.

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(a) FCNs on MNIST

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(b) CNNs on Fashion-MNIST

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(c) ResNets on CIFAR-10

While we plotted the averaged quantities ⟨ c loss ⟩ , ⟨ c sharp ⟩ , ⟨ c max ⟩ , we have observed that their variance also increases significantly with d and 1 / w ; in Appendix C we show standard deviations about the averages for different random initializations. Nevertheless, we have found that the inequality c loss ≤ c sharp ≤ c max 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. 4 shows evidence of this claim. The setup is 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 ⟨ c loss ⟩ ≳ 4 and ⟨ c max ⟩ consistently decreases with 1 / w , while still satisfying c loss ≤ c sharp ≤ c max .

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 a 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, 51, 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, as the loss increases during early 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. ( L int ( s, c ) , U ( c )) Let θ 0 represent the initial set of parameters, and let θ T 1 represent the set of parameters at the end of the initial transient period, trained using a learning rate constant c . Then, we define the linearly interpolated loss as L int ( s, c ) = L [(1 -s ) θ 0 + s θ T 1 ] , where s ∈ [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 s ∈ [0 , 1] L int ( s ) -L ( θ 0 ) .

Here we subtracted the loss's initial value such that a positive value indicates a barrier to the final point from initialization. Using the interpolated loss barrier, we define c barrier as follows.

Definition 3. ( c barrier ) Given the initial ( θ 0 ) and final parameters ( θ T 1 ), we define c barrier as the minimum learning rate constant such that there exists a barrier from θ 0 to θ T 1 : c barrier = min c { c | U ( c ) > 0 } .

Here, c barrier 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. 2(c, f) shows the interpolated loss of CNNs trained on the CIFAR-10 dataset for t = 10 steps. The experimental setup is the same as in Section 2. For the network with larger width, we observe a barrier emerging at c barrier = 5 . 65 , while the loss starts to catapult at c loss = 2 . 83 . In comparison, we do not observe any barrier from initialization to the final point at large d and 1 / w . Fig. 3 shows the relationship between ⟨ c barrier ⟩ and 1 / w for various models and datasets. We consistently observe that c sharp ≤ c barrier , suggesting that training traverses a barrier only when sharpness starts to catapult during early training. Similar results

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

128

256

512

(a)

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

χτ

width

128

256

512

(b)

2

-

1

2 0

2 1

2 2

2 3

2 4

c

-

0

.

3

-

0

.

2

-

0

.

1

0

.

0

0

.

1

0

.

2

0

.

3

χ

′

τ

width

128

256

512

(c)

were observed on increasing d instead of 1 / w as shown in Appendix C. We chose not to characterize the phase diagram of early training using c barrier 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.

These observations call into question the intuition of catapulting out of a basin for a range of learning rates in between c loss < c < c barrier . 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 an inductive bias of the training process towards regions of higher loss during the early time transient regime.

3 Intermediate saturation regime

In the intermediate saturation regime, sharpness does not change appreciably and reflects the cumulative change that occurred during the initial transient period. This section analyzes sharpness in the intermediate 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 λ H τ at a time τ in the middle of the intermediate saturation regime. We choose τ so that cτ ≈ 200 . 6 For further details on sharpness measurement, see Appendix I.1. Fig. 5(a) illustrates the relationship between λ H τ 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 λ H τ on learning rate can be grouped into three distinct stages. (1) At small learning rates, λ H τ remains relatively constant, with fluctuations increasing as d and 1 / w increase ( c < 2 in Fig. 5(a)). (2) A crossover regime where λ H τ is dropping significantly ( 2 < c < 2 3 in Fig. 5(a)). (3) A saturation stage where λ H τ stays small and constant with learning rate ( c > 2 3 ) in Fig. 5(a)). In Appendix I, we 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 c < 2 is sub-critical, λ H τ decreases with increasing d and 1 / w . In other words, for small c and in the intermediate saturation regime, the loss is locally flatter as d and 1 / w increase.

We can precisely extract a critical value of c that separates stages (1) and (2), which corresponds to the onset of an abrupt reduction of sharpness λ H τ . To do this, we consider the averaged normalized sharpness over initializations and denote it by ⟨ λ H τ / λ H 0 ⟩ . The first two derivatives of the averaged normalized sharpness, χ τ = -∂ ∂c ⟨ λ H τ / λ H 0 ⟩ and χ ′ τ = -∂ 2 ∂c 2 ⟨ λ H τ / λ H 0 ⟩ , characterize the change in sharpness with learning rate. The extrema of χ ′ τ quantitatively define the boundaries between the three stages described above. In particular, using the maximum of χ ′ τ , we define ⟨ c crit ⟩ , which marks the beginning of the sharp decrease in λ H τ with the learning rate.

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(a) vanilla

2

-

1

2 0

2 1

2 2

2 3

2 4

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(b) centered

2

-

1

2 0

2 1

2 2

2 3

2 4

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(c) last layer zero at init

Definition 4. ( ⟨ c crit ⟩ ) Given the averaged normalized sharpness ⟨ λ H τ / λ H 0 ⟩ measured at τ , we define c crit to be the learning rate constant that minimizes its second derivative: ⟨ c crit ⟩ = arg max c χ ′ τ .

Here, we use ⟨ . ⟩ to denote that the critical constant is obtained from the averaged normalized sharpness. Fig. 5(b, c) show χ τ and χ ′ τ obtained from the results in Fig. 5(a). We observe similar results across various architectures and datasets, as shown in Appendix I. Our results show that ⟨ c crit ⟩ has slight fluctuations as d and 1 / w are changed but generally stay in the vicinity of c = 2 . The peak in χ ′ τ becomes wider as 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 ⟨ c crit ⟩ , ⟨ c loss ⟩ increase with d and 1 / w , implying the opening of the sharpness reduction phase ⟨ c crit ⟩ < c < ⟨ c loss ⟩ as d and 1 / w increase. In Appendix F, we show that cross-entropy loss shows qualitatively similar results, but with 2 ≲ ⟨ c crit ⟩ ≲ 4 .

4 Effect of network output at initialization on early training

Here we discuss the effect of network output f ( x ; θ t ) at initialization on the early training dynamics. x is the input and θ t denotes the set of parameters at time t . We consider setting the network output to zero at initialization, f ( x ; θ 0 ) = 0 , by either (1) considering the 'centered" network: f c ( x ; θ ) = f ( x ; θ ) -f ( x ; θ 0 ) , 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 of the sharpness reduction phase with 1 /w as shown in Figure 6. The average onset of the loss catapult, diagnosed by ⟨ c loss ⟩ , becomes independent of 1 /w and d .

We also empirically study the impact of the output scale [19, 5, 4] on early training dynamics. Given a network function f ( x ; θ ) , we define the scaled network as f s ( x ; θ ) = αf ( x ; θ ) , where α is a scalar, fixed throughout training. In Appendix H, we show that a large (resp. small) value of ∥ f ( x ; θ 0 ) ∥ relative to the one-hot encodings of the labels causes the sharpness to decrease (resp. increase) during early training. Interestingly, we still observe an increase in ⟨ c loss ⟩ with d and 1 / w , unlike the case of initializing network output to zero, highlighting the unique impact of output scale on the dynamics.

5 Insights from a simple model

Here we analyze a two-layer linear network [56, 60, 49], the uv model, which shows much of the phenomena presented above. Define f ( x ) = 1 √ w v T ux , with x, f ( x ) ∈ R . Here, u, v ∈ R w are the trainable parameters, initialized using the normal distribution, u i , v i ∼ N (0 , 1) for i ∈ { 1 , . . . , w } . The model is trained with MSE loss on a single training example ( x, y ) = (1 , 0) , which simplifies the loss to L ( u, v ) = f 2 / 2 , and which was also considered in Ref. [49]. Our choice of y = 0 is motivated by the results of Sec. 4, which suggest that the empirical results of Sec. 2 are intimately related to the model having a large initial output scale ∥ f ( x ; θ 0 ) ∥ relative to the output labels. We minimize the loss using gradient descent (GD) with learning rate η . The early time phase diagram also shows similar features to those described in preceding sections (compare Fig. 7(a) and Fig. 3). Below we develop an understanding of this early time phase diagram in the uv model.

1

2

3

4

c

2

-

6

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

1

2

3

4

k

2

-

6

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

1

/w

〈

kloss

〉

〈

ktrace

〉

〈

kmax

〉

(b)

1

2

3

4

k

2

-

6

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

1

/w

〈

kloss

〉

〈

kfrob

〉

〈

kmax

〉

(c)

The update equations of the uv model in function space can be written in terms of the trace of the Hessian tr( H )

From the above equations, it is natural to scale the learning rate as η = k / tr( H 0 ) . Note that c = ηλ H 0 = kλ H 0 / tr( H 0 ) . Also, we denote the critical constants in this scaling as k loss , k trace , k max and k crit , where the definitions follow from Definitions 1 and 4 on replacing sharpness with trace and use ⟨ . ⟩ to denote an average over initialization. Figure 7(b) shows the phase diagram of early training, with tr( H t ) replaced with λ H t as the measure of sharpness and with the learning rate scaled as η = k / tr( H 0 ) . Similar to Figure 7(a), we observe a new phase ⟨ k crit ⟩ < k < ⟨ k loss ⟩ 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 1). We also observe ⟨ k max ⟩ = 4 , independent of width.

In Appendix B.3, we show that the critical value of k for which ⟨ L 1 / L 0 ⟩ > 1 increases with 1 / w , which explains why ⟨ k loss ⟩ increases with 1 / w . Combined with ⟨ k crit ⟩ ≈ 2 , this implies the opening up of the sharpness reduction phase as w is decreased.

To understand the loss-sharpness catapult phase, we require some other measure as tr( H ) does not increase for 0 < k < 4 . As λ H t is difficult to analyze, we consider the Frobenius norm ∥ H ∥ F = √ tr( H T H ) as a proxy for sharpness. We define k frob as the minimum learning rate such that || H t || 2 F increases during early training. Figure 7(c) shows the phase diagram of the uv model, with || H t || 2 F as the measure of sharpness, while the learning rate is scaled as η = k / tr( H 0 ) . We observe the loss-sharpness catapult phase at small widths. In Appendix B.4, we show that the critical value of k for which 〈 || H 1 || 2 F -|| H 0 || 2 F 〉 > 0 increases from ⟨ k loss ⟩ 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 a two-dimensional slice of parameters defined by tr( H ) and weight correlation ⟨ v,u ⟩ / || u |||| v || . 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 tr( H ) ) as compared to the wider network. This means that the small-width network needs a relatively larger learning rate to get to a point of increased loss (loss catapult). We thus have the opening up of a new regime ⟨ k crit ⟩ < k < ⟨ k loss ⟩ , in which the loss and sharpness monotonically decrease during early training.

The loss landscape of the uv model shown in Fig. 8 reveals interesting insights into the loss landscape connectivity results in Section 2.2 and the presence of c barrier . Fig. 8 shows how even when there is a loss catapult, as long as the learning rate is not too large, the final point after the catapult can be reached from initialization by a linear path without increasing the loss and passing through a barrier. However if 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 a barrier in the loss upon linear interpolation.

̸

The uv model trained on an example ( x, 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 f 0 = 0 and

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

weight correlation

0

1

2

Tr

(

H

)

(a)

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

weight correlation

(b)

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

weight correlation

(c)

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

weight correlation

(d)

10

-

2

10

-

1

10 0

10 1

10 2

10 3

̸

y = 0 in the dynamical equations results in loss catapult at k = 2 , implying ⟨ k loss ⟩ ≈ ⟨ k crit ⟩ ≈ 2 , 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 as η = c / λ H 0 . 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 two 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 in the early transient regime and demonstrated the existence of a regime ⟨ c loss ⟩ < c < ⟨ c barrier ⟩ , in which the final point after the catapult lies in the same basin as initialization, connected through a linear path. Finally, we study these phenomena in a 2-layer linear network ( uv model), gaining insights into the opening of the sharpness reduction phase.

We performed a preliminary analysis on the effect of batch size on the presented results in Appendix J. 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 a qualitatively similar phase diagram. In the intermediate saturation regime, the effect of reducing the batch size is to broaden the transition around c crit .

In Section 2, we noted that for cross-entropy loss, the loss starts to catapult around c ≈ 4 at large widths, as compared to c loss = 2 for MSE loss. Previous work, such as [50], analyzed the catapult dynamics for the uv model with logistic loss and demonstrated that the loss catapult occurs above η loss = 4 /λ NTK 0 . We summarize the main intuition about their analysis in Appendix B.9. However, a complete understanding of the catapult phenomenon in 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 [22] may show dissimilarities with the reported phase diagram of early training.

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 (DMR1753240) and the Laboratory for Physical Sciences through the Condensed Matter Theory Center.

References

Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep networks. In International Conference on Learning Representations , 2019.

Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. ArXiv , abs/2210.04860, 2022.

Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. In International Conference on Machine Learning , 2022.

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 Eleventh International Conference on Learning Representations , 2023.

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.

Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. ArXiv , abs/2304.03408, 2023.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

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 on Learning Representations , 2021.

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.

Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine , 29(6):141-142, 2012.

Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations , 2020.

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 Thirty-Third Conference on Uncertainty in Artificial Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017 . AUAI Press, 2017.

Dumitru Erhan, Yoshua Bengio, Aaron Courville, Pierre-Antoine Manzagol, Pascal Vincent, and Samy Bengio. Why does unsupervised pre-training help deep learning? Journal of Machine Learning Research , 11(19):625-660, 2010.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations , 2021.

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.

Stanislav Fort and Surya Ganguli. Emergent properties of the local geometry of neural loss landscapes. ArXiv , abs/1910.05929, 2019.

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.

Mario Geiger, Leonardo Petrini, and Matthieu Wyart. Perspective: A phase diagram for deep learning unifying jamming, feature learning and lazy training, 2020.

Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment , 2020(11):113301, nov 2020.

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.

Ian J. Goodfellow and Oriol Vinyals. Qualitatively characterizing neural network optimization problems. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings , 2015.

Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. ArXiv , abs/1706.02677, 2017.

Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? In Neural Information Processing Systems , 2018.

Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations , 2020.

Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. In Neural Information Processing Systems , 2018.

Soufiane Hayou, A. Doucet, and Judith Rousseau. On the selection of initialization and activation function for deep neural networks. ArXiv , abs/1805.08266, 2018.

Soufiane Hayou, A. Doucet, and Judith Rousseau. Exact convergence rates of the neural tangent kernel in the large depth limit. 2019.

Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. 2015 IEEE International Conference on Computer Vision (ICCV) , pages 1026-1034, 2015.

Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 770-778, 2016.

Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020.

Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in neural information processing systems , 7, 1994.

Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation , 9(1):1-42, 1997.

Gaurav Iyer, Boris Hanin, and David Rolnick. Maximal initial learning rates in deep relu networks. ArXiv , abs/2212.07295, 2022.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural networks (invited paper). Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing , 2018.

Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho*, and Krzysztof Geras*. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations , 2020.

Stanisław Jastrz˛ ebski, Zac Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Amos Storkey, and Yoshua Bengio. Three factors influencing minima in SGD, 2018.

Stanisław Jastrz˛ ebski, 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.

Stanisław Jastrz˛ ebski, 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.

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.

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.08361 , 2020.

Ryo Karakida, Shotaro Akaho, and Shun ichi Amari. Universal statistics of fisher information in deep neural networks: Mean field approach, 2019.

Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations , 2017.

Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. In NIPS , pages 972-981, 2017.

Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009.

Guillaume Leclerc and Aleksander Madry. The two regimes of deep network training. ArXiv , abs/2002.10376, 2020.

Yann LeCun, Léon Bottou, Genevieve B. Orr, and Klaus-Robert Müller. Efficient backprop. In Neural Networks , 2012.

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.

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.

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 , abs/2003.02218, 2020.

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.

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 arXiv:2104.11044 , 2021.

Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at infinite-width limit, 2020.

David Meltzer and Junyu Liu. Catapult dynamics and phase transitions in quadratic nets. ArXiv , abs/2301.07737, 2023.

Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha SohlDickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations , 2020.

Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In NIPS , 2016.

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.

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.

Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory . Cambridge University Press, 2022. https://deeplearningtheory.com .

Abraham. Savitzky and M. J. E. Golay. Smoothing and differentiation of data by simplified least squares procedures. Analytical Chemistry , 36(8):1627-1639, 1964.

Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. CoRR , abs/1312.6120, 2014.

Vaishaal Shankar, Alexander W. Fang, Wenshuo Guo, Sara Fridovich-Keil, Ludwig Schmidt, Jonathan Ragan-Kelley, and Benjamin Recht. Neural kernels without tangents. In ICML , 2020.

Jascha Narain Sohl-Dickstein, Roman Novak, Samuel S. Schoenholz, and Jaehoon Lee. On the infinite width limit of neural networks with a standard parameterization. ArXiv , abs/2001.07301, 2020.

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, C J 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. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods , 17:261-272, 2020.

Tiffany J. Vlaar and Jonathan Frankle. What can linear interpolation of neural network loss landscapes tell us? In International Conference on Machine Learning , 2021.

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.

Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017. cite arxiv:1708.07747Comment: Dataset is freely available at https://github.com/zalandoresearch/fashion-mnist Benchmark is available at http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/.

Sho Yaida. Meta-principled family of hyperparameter scaling strategies. ArXiv , abs/2210.04909, 2022.

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 of Proceedings of Machine Learning Research , pages 11727-11737. PMLR, 18-24 Jul 2021.

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.

Hanxu Zhou, Qixu Zhou, Zhenyuan Jin, Tao Luo, Yaoyu Zhang, and Zhi-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.

Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for understanding neural network dynamics. ArXiv , abs/2205.11787, 2022.

A Experimental details

Datasets: We considered the MNIST [10], 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 w 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.

FCNs: We considered ReLU FCNs with constant width w in Neural Tangent Parameterization (NTP) / Standard Parameterization (SP), initialized at criticality [55]. The models do not include bias or normalization. The forward pass of the pre-activations from layer l to l +1 is given by

where ϕ ( . ) is the ReLU activation and γ l is a constant. For NTP, γ l = 2 / √ w and the weights W l are initialized using normal distribution, i.e., W l ij ∼ N (0 , 1) . For SP, γ l = 1 and the weights W l are initialized as W l ij ∼ N (0 , 2 / w ) . For the last layer, we have γ L = 1 / √ w for NTP and W L ij ∼ N (0 , 1 / w ) for SP.

For d / w ≳ 1 / 16 , the dynamics is noisier, and it becomes challenging to separate the underlying deterministic dynamics from random fluctuations (see Appendix E).

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 case, the forward pass equations for the pre-activations from layer l to layer l +1 are given by

where α, β label the spacial location. The weights are initialized as W l ij ( β ) ∼ N (0 , 2 / k 2 w ) , where k is the filter size.

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 has four blocks with widths [ w, 2 w, 4 w, 8 w ] , with w = 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 D B = { ( x µ , y µ ) } B µ =1 , i.e., L ( x, y D B ) = 1 / B ∑ B µ =1 ℓ ( x µ , y µ ) , 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 Appendix K, we demonstrate that bias does not have an appreciable impact on the results.

Batch size: We use a batch size of 512 and scale the learning rate as η = c / λ H 0 in all our experiments, unless specified. Appendix J shows results for a smaller batch size B = 32 .

Learning rate: We scale the learning rate constant as c = 2 x , with x ∈ {-1 . 0 , . . . x max } in steps of 0 . 1 . Here, x max is related to the maximum learning rate constant as c max = 2 x max .

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 m = 2048 randomly selected training examples for evaluating sharpness at each step. In comparison, we found that cross-entropy requires a large number of training examples to obtain a good approximation

of sharpness. Given the computational constraints, we use 4096 training examples to approximate sharpness for cross-entropy loss.

Averages over initialization and SGD runs: All the critical constants depend on both the random initializations and the SGD runs. In our experiments, we found that the fluctuations from initialization at large d / w outweigh the randomness coming from different SGD runs. Thus, we focus on initialization averages in all our experiments.

A.1 Compute usage

We utilized different computational resources depending on 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, we utilized a 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 repository:https://github.com/dayal-kalra/early-training.

A.3 Details of Figures in the main text:

Figure 1: A shallow CNN ( d = 5 , w = 128 ) in SP trained on the CIFAR-10 dataset with MSE loss for 1000 epochs using SGD with learning rates η = c / λ H 0 and batch size B = 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) a deep CNN ( d = 10 , w = 128 ) in SP trained on the CIFAR-10 dataset with MSE loss for t = 10 steps using vanilla SGD with learning rates η = c / λ H 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 Fashion-MNIST dataset for different widths. Each data point in the figure represents an average of ten distinct initializations, and the solid lines represent a two-degree polynomial y = a + bx + cx 2 fitted to the raw data points. Here, where x = 1 / w , and y can take on one of three values: c loss , c sharp and c max .

Figure 4: (a) FCNs in SP with d ∈ { 4 , 8 , 16 } and w ∈ { 256 , 512 , 1024 , 2048 } trained on the MNIST dataset, (b) CNNs in SP with d ∈ { 5 , 7 , 10 } and w ∈ { 64 , 128 , 256 , 512 } trained on the FashionMNIST dataset, (c) ResNet in SP with d ∈ { 18 , 34 } and w ∈ { 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 at 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.

Figure 5: Normalized sharpness measured at cτ = 200 against the learning rate constant for 7 -layer CNNs in SP trained on the CIFAR-10 dataset, with w ∈ { 128 , 256 , 512 } . Each data point is an average over five random initialization. Smoothening details are provided in Appendix I.2.

Figure 7: The phase diagram of the uv model trained with MSE loss using gradient descent with (a) the top eigenvalue of Hessian λ H t , (b) the trace of Hessian tr( H t ) and (c) the square of the Frobenius norm tr( H T t H t ) used as a measure of sharpness. In (a), the learning rate is scaled as η = c / λ H 0 , while in (b) and (c), the learning rate is scaled as η = k / tr( H 0 ) . The vertical dashed line shows c = 2 ( k = 2 ) for reference. Each data point is an average over 500 random initializations.

Figure 8: Training trajectories of the uv model with (a, b) large ( w = 512 ) and (c, d) small ( w = 2 ) width, trained for t = 10 training steps on a single example ( x, y ) = (1 , 0) with MSE loss using vanilla gradient descent with learning rates (a, c) c = 0 . 5 and (b, d) c = 2 . 50 .

0

2

4

6

8

10

t

(step)

0

20

40

60

80

Training loss

0

2

4

6

8

10

t

(step)

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

0

2

4

6

8

10

t

(step)

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Tr H

t

/Tr H

0

c

1.0

1.5

2.0

2.5

3.0

(a)

0

2

4

6

8

10

t

(step)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

Training loss

0

2

4

6

8

10

t

(step)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

0

2

4

6

8

10

t

(step)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Tr H

t

/Tr H

0

c

1.5

2.2

3.0

3.7

4.5

(b)

B Additional results for the uv model

B.1 Details of the model

Consider a two-layer linear network in (NTP) with unit input-output dimensions

where x, f ( x ) ∈ R . Here, u, v ∈ R w are trainable parameters, with each element initialized using the normal distribution, u i , v i ∼ N (0 , 1) for i ∈ { 1 , . . . , w } . The model is trained using MSE loss on a single training example ( x, y ) = (1 , 0) , which simplifies the loss to

The trace of the Hessian tr( H ) has a simple expression in terms of the norms of the weight vectors

which is equivalent to the NTK for this model. The Frobenius norm of the Hessian ∥ H ∥ F can be written in terms of the loss L and tr( H )

The gradient descent updates of the model trained using MSE loss on a single training example ( x, y ) = (1 , 0) are given by

0

1

2

3

4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

τ

/λ H

0

width

2.0

8.0

32.0

512.0

(a)

0

1

2

3

4

c

-

0

.

2

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

χτ

width

2

8

32

512

(b)

0

1

2

3

4

c

-

3

-

2

-

1

0

1

2

χ

′

τ

width

2

8

32

512

(c)

The update equations in function space can be written in terms of the trace of the Hessian tr( H ) .

Figure 9 shows the training trajectories of the uv model trained on ( x, y ) = (1 , 0) using MSE loss for 10 training steps. The model shows similar dynamics to those presented in Section 2. It is worth mentioning that the above equations have been analyzed in [49] at large width. In the following subsections, we extend their analysis by incorporating the higher-order terms to analyze the effect of finite width.

B.2 The intermediate saturation regime

The uv model trained on ( x, y ) = (1 , 0) does not show the progressive sharpening and late-time regimes (iii) and (iv) described in Section 1. Hence, we can measure sharpness at the end of training to analyze how it is reduced upon increasing the learning rate and to compare it with the intermediate saturation regime results in Section 3.

Figure 10(a) shows the normalized sharpness measured at τ = 100 steps for various widths. This behavior reproduces the results observed in the intermediate saturation regime in Section 3. In particular, we can see stages (1) and (2), where λ H τ / λ H 0 starts off fairly independent of learning rate constant c , and then dramatically reduces when c > 2 ; stage (3), where λ H τ / λ H 0 plateaus at a small value as a function of c is too close to the divergent phase in this model to be clearly observed. The corresponding derivatives of the averaged normalized sharpness, χ τ , and χ ′ τ , are shown in Figure 10(b, c). The vertical dashed lines denote c crit estimated for each width, using the maximum of χ ′ τ . We observe that c crit = 2 for all widths.

0

1

2

3

4

k

0

2

4

6

8

L

1

L

0

〉

w

2

4

8

16

32

64

(a)

0

1

2

3

4

k

0

2

4

6

8

L

1

L

0

〉

w

2

4

8

16

32

64

(b)

4

.

2

4

.

4

4

.

6

4

.

8

4

λ H

0

/Tr

(

H

0

)

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

2 0

d/w

(c)

B.3 Opening of the sharpness reduction phase in the uv model

This section shows that O ( 1 / w ) terms in Equation (10) effectively lead to the opening of the sharpness reduction phase with 1 / w in the uv model. In Appendix B.2, we demonstrated that for the uv model, c crit = 2 for all values of widths. Hence, it suffices to show that c loss increases from the value 2 as 1 / w increases. We do so by finding the smallest k such that the averaged loss over initializations increases during early training.

It follows from Equation 10 that the averaged loss increases in the first training step if the following holds

where ⟨ . ⟩ denotes the average over initializations. On scaling the learning rate with trace as η = k / tr( H 0 ) , we have

The required two averages have the following expressions as shown in Appendix B.8.

Inserting the above expressions in Equation 13, on average the loss increases in the very first step if the following inequality holds

The graphical representation of the above inequality shown in Figure 11(a) is in excellent agreement with the experimental results presented in Figure 11(b).

Let us denote k ′ loss as the minimum learning rate constant such that the average loss increases in the first step. Similarly, let k loss denote the learning rate constant if the loss increases in the first 10 steps. Then, k ′ loss increases from the value 2 as 1 / w increases as shown in Figure 11(a). By comparison, the trace reduces at any step if η tr( H t ) < 4 . At initialization, this condition becomes k < 4 . Hence, for k < k ′ loss , 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 < k ′ loss .

Next, let η loss denote the learning rate corresponding to c loss . Then, we have η loss = c loss λ H 0 = k loss tr( H 0 ) , implying

Figure 11(c) shows that λ H 0 ≥ tr( H 0 ) for all widths, implying c loss ≥ k loss . Hence, c loss increases with 1 / w as observed in Figure 7(a). In Appendix B.2, we demonstrated that for the uv model, c crit = 2 for all values of widths. Incorporating this with c loss increases with 1 / w , we have sharpness reduction phase opening up as 1 / w increases.

B.4 Opening of the loss catapult phase at finite width

In this section, we use the Frobenius norm of the Hessian ∥ H ∥ F 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 T H ) ⟩ after the first training step near k = k loss and show that k loss ≤ k frob , with the difference increasing with 1 /w . First, we write tr( H T t H t ) in terms of L t and tr( H t )

Next, using Equations 1, we write down the change in tr( H T t H t ) after the first training step in terms of tr( H 0 ) and L 0

Next, we substitute η = k/ tr( H 0 ) to obtain the above equation as a function of k

Finally, we calculate the expectation value of ⟨ ∆tr( H T 1 H 1 ) ⟩

2

3

4

5

6

csharp

2

3

4

5

6

closs

cmax

by estimating 〈 f 4 0 tr( H 0 ) 2 〉 using the approach demonstrated in the previous section

Inserting 〈 f 4 0 tr( H 0 ) 2 〉 in Equation 21 along with ⟨ f 2 0 ⟩ = 1 , we have

At infinite width, the above equation reduces to 〈 ∆tr( H T 1 H 1 ) 〉 = 4 ⟨L 1 -L 0 ⟩ , and hence, k frob = k loss . For any finite width, I ( k, w ) < 0 for 0 < k < 4 . At k ≤ k loss , L 1 -L 0 ≤ 0 , and therefore 〈 ∆tr( H T 1 H 1 ) 〉 < 0 . In order for the sharpness to catapult, we require 〈 ∆tr( H T 1 H 1 ) 〉 > 0 and therefore k frob > k loss . As 1 /w increases | I ( k, w ) | also increases, which means a higher value of L 1 -L 0 is required to reach a point where 〈 ∆tr( H T 1 H 1 ) 〉 ≥ 0 . Thus k frob -k loss increases with 1 /w .

B.5 The early training trajectories

Figure 9 shows the early training trajectories of the uv 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, at c loss = 3 . 74 , which is significantly higher than the critical value of c crit = 2 ).

B.6 Relationship between critical constants

Figure 12 shows the relationship between various critical constants for the uv model. The data show that the inequality c loss ≤ c sharp ≤ c max holds for every random initialization of the uv 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 7(a, b). Figure 13 shows these phase diagrams. Each data point is 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

Here, we provide the detailed derivation of the averages 〈 f 2 0 Tr ( H 0 ) 2 〉 and 〈 f 4 0 Tr ( H 0 ) 4 〉 finding the average 〈 f 2 0 Tr ( H 0 ) 2 〉

1

2

3

4

c

2

-

6

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

1

2

3

4

k

2

-

6

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

1

/w

〈

kloss

〉

〈

ktrace

〉

〈

kmax

〉

(b)

1

2

3

4

k

2

-

6

2

-

5

2

-

4

2

-

3

2

-

2

2

-

1

1

/w

〈

kloss

〉

〈

kfrob

〉

〈

kmax

〉

(c)

where ∥ . ∥ denotes the norm of the vectors.

The above integral is non-zero only if j = k . Hence, it is a sum of w identical integrals. Without any loss of generality, we solve this integral for j = 1 and multiply by w to obtain the final result, i.e.,

Consider a transformation of u, v ∈ R w into w dimensional spherical coordinates such that

which yields,

where d Ω denotes the w dimensional solid angle element. Here, we denote the radial and angular integrals by I r and I φ respectively. The radial integral I r is

Let r u = R cos θ and r v = R sin θ with R ∈ [0 , ∞ ) and θ ∈ [ -π 2 , π 2 ] , then we have

where Γ( . ) denotes the Gamma function. The angular integral I φ is

Plugging in Equations 32 and 35 into Equation 29, we obtain a very simple expression

The other integral 〈 f 4 0 Tr ( H 0 ) 4 〉 can be obtained by generalizing the above approach as described below

The integral is zero if either j = k and l = m or j = k = l = m , which we consider separately. Without loss of generality, we find the following integrals

which have the following expressions

2

-

1

2 0

2 1

2 2

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 4

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 8

2

-

1

2 1

2 3

2 5

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c)

d

= 16

where Γ( . ) denotes the gamma function. On combining the expressions with their multiplicities, we obtain the final result

B.9 Insights into the catapult effect in crossentropy loss using uv model

In this section, we summarize the main intuition behind the discrepancy in the values of c loss for cross-entropy loss at large widths. We consider the uv model trained on a classification task using logistic loss, as presented in [50].

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 L ( f ) = 1 2 log(2 + 2 cosh( f )) . Hence, the loss grows monotonically as the output function | f | increases. The update equation of the function is given by:

where η is the learning rate and L ′ ( . ) is the derivative of the loss. At large width, if the condition | 1 -η tr( H ) L ′ ( f ) /f | < 1 holds, then output function continues to decrease. Given that L ′ ( f ) /f ≤ 1 / 2 in the above case, this decrease persists for η tr( H ) < 4 . This result provides some intuition behind the discrepancy.

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, Fashion-MNIST, and CIFAR-10 datasets, (2) CNNs in SP trained on Fashion-MNIST and CIFAR-10, and (3) ResNets in SP trained on CIFAR10 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 an 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 are shown in Appendix F.

Additional experimental details : We train each model for t = 10 steps using SGD with learning rates η = c / λ H 0 and batch size of 512 , where c = 2 x with x ∈ { 0 . 0 , . . . x max } in steps of 0 . 1 . Here, x max is relatd to the maximum trainable learning rate constant as c max = 2 x max . We have considered 10 random initializations for each model. As mentioned in Appendix A, we do not

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 4

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 8

2

-

1

2 1

2 3

2 5

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c)

d

= 16

2

-

1

2 0

2 1

2 2

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 4

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 8

2

-

1

2 1

2 3

2 5

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c)

d

= 16

2

-

1

2 1

2 3

c

2

-

9

2

-

8

2

-

7

2

-

6

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 5

2

-

1

2 1

2 3

2 5

c

2

-

9

2

-

8

2

-

7

2

-

6

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 7

2

-

1

2 1

2 3

2 5

c

2

-

9

2

-

8

2

-

7

2

-

6

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c)

d

= 10

2

-

1

2 1

2 3

c

2

-

9

2

-

8

2

-

7

2

-

6

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 5

2

-

1

2 1

2 3

2 5

c

2

-

9

2

-

8

2

-

7

2

-

6

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 7

2

-

1

2 1

2 3

2 5

c

2

-

9

2

-

8

2

-

7

2

-

6

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c)

d

= 10

2

-

1

2 1

2 3

2 5

c

2

-

7

2

-

6

2

-

5

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 18

2

-

1

2 1

2 3

2 5

c

2

-

7

2

-

6

2

-

5

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 34

2

-

1

2 1

2 3

2 5

c

2 2

2 3

2 4

depth

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

w

= 512

2

-

1

2 1

2 3

c

2 2

2 3

2 4

depth

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

w

= 1024

2

-

1

2 1

2 3

c

2 2

2 3

2 4

depth

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c)

w

= 2048

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 a random increase in loss and sharpness due to fluctuations, we round off the values of λ H t / λ H 0 and L t / L 0 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 a two-degree polynomial y = a + bx + cx 2 with x = 1 / w and y can take on one of three values: c loss , c sharp and c max .

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 / w 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, we find that c loss ≤ c sharp ≤ c sharp holds true. However, there are some exceptions, which are observed at high values of d / w (see 21 (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 for FCNs with d / w ≳ 1 / 16 , the dynamics becomes noisedominated. This aspect makes it challenging to distringuish the underlying deterministic dynamics from random fluctuations. To demonstrate this, we consider FCNs trained on CIFAR-10 using MSE and cross-entropy loss and use 4096 training examples for estimating sharpness.

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(a) FCNs on MNIST

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(b) FCNs on Fashion-MNIST

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(c) FCNs on CIFAR-10

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(d) CNNs on Fashion-MNIST

2 2

2 4

2 6

csharp

2 1

2 2

2 3

2 4

2 5

2 6

closs

cmax

(e) CNNs on CIFAR-10

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

closs

cmax

(f) ResNets on CIFAR-10

0

5

10

15

20

step

0

2

4

6

8

10

Training loss

loss:

mse, width: 1024, depth: 16

c

3

6

9

12

15

0

5

10

15

20

step

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

Training loss

loss:

mse, width:

512, depth: 16

c

1

2

3

4

5

0

5

10

15

20

step

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

Training loss

loss:

mse, width:

256, depth: 16

c

2

4

6

8

10

0

5

10

15

20

step

0

.

1

0

.

2

0

.

3

0

.

4

Training loss

loss:

mse, width:

128, depth: 16

c

2

4

6

8

10

0

5

10

15

20

step

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

Training loss

loss:

mse, width:

64, depth: 16

c

2

4

6

8

10

0

5

10

15

20

step

0

.

06

0

.

08

0

.

10

0

.

12

0

.

14

0

.

16

0

.

18

Training loss

loss:

mse, width:

32, depth: 16

c

1

2

3

4

5

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 w ≲ 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 as in the MSE loss case. In comparison to MSE

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

λ H

t

/λ H

0

loss:

mse, width: 1024, depth: 16

c

3

6

9

12

15

0

5

10

15

20

step

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

loss:

mse, width:

512, depth: 16

c

1

2

3

4

5

0

5

10

15

20

step

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

loss:

mse, width:

256, depth: 16

c

2

4

6

8

10

0

5

10

15

20

step

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

loss:

mse, width:

128, depth: 16

c

2

4

6

8

10

0

5

10

15

20

step

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

loss:

mse, width:

64, depth: 16

c

2

4

6

8

10

0

5

10

15

20

step

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

loss:

mse, width:

32, depth: 16

c

1

2

3

4

5

0

5

10

15

20

step

2

4

6

8

10

12

14

Training loss

loss:

xent, width: 1024, depth: 16

c

15

30

45

60

0

5

10

15

20

step

2

3

4

5

6

7

8

Training loss

loss:

xent, width:

512, depth: 16

c

8

16

24

32

40

0

5

10

15

20

step

2

3

4

5

6

Training loss

loss:

xent, width:

256, depth: 16

c

15

30

45

60

0

5

10

15

20

step

2

.

00

2

.

25

2

.

50

2

.

75

3

.

00

3

.

25

3

.

50

3

.

75

Training loss

loss:

xent, width:

128, depth: 16

c

15

30

45

60

0

5

10

15

20

step

2

.

1

2

.

2

2

.

3

2

.

4

2

.

5

2

.

6

2

.

7

Training loss

loss:

xent, width:

64, depth: 16

c

4

8

12

16

20

0

5

10

15

20

step

2

.

15

2

.

20

2

.

25

2

.

30

2

.

35

2

.

40

Training loss

loss:

xent, width:

32, depth: 16

c

4

8

12

16

20

loss, the training loss and sharpness dynamics show a higher level of noise, especially for w ≲ 256 . As a result, it becomes difficult to characterize the training dynamics for d / w ≳ 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-entropy loss show similar characterstics to those trained with MSE loss, such as, (i) sharpness reduction during early training, (ii) an increase in critical constants c loss , c sharp with d and 1 / w , (iii) c loss ≤ c sharp ≤ c max . 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

0

5

10

15

20

step

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

1

.

6

λ H

t

/λ H

0

loss:

xent, width: 1024, depth: 16

c

15

30

45

60

0

5

10

15

20

step

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

1

.

4

1

.

6

λ H

t

/λ H

0

loss:

xent, width:

512, depth: 16

c

8

16

24

32

40

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

λ H

t

/λ H

0

loss:

xent, width:

256, depth: 16

c

15

30

45

60

0

5

10

15

20

step

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

t

/λ H

0

loss:

xent, width:

128, depth: 16

c

15

30

45

60

0

5

10

15

20

step

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

λ H

t

/λ H

0

loss:

xent, width:

64, depth: 16

c

4

8

12

16

20

0

5

10

15

20

step

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

λ H

t

/λ H

0

loss:

xent, width:

32, depth: 16

c

4

8

12

16

20

2

-

1

2 0

2 1

2 2

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(a) MSE,

d

= 4

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(b) MSE,

d

= 8

2

-

1

2 1

2 3

2 5

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(c) MSE,

d

= 16

2

-

1

2 1

2 3

2 5

2 7

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(d) xent,

d

= 4

2

-

1

2 1

2 3

2 5

2 7

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(e) xent,

d

= 8

2

-

1

2 1

2 3

2 5

2 7

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(f) xent,

d

= 16

following experiments, we consider models trained on the CIFAR-10 dataset and used 4096 training examples to estimate sharpness.

F.1 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 are generally

2 1

2 2

2 3

2 4

2 5

csharp

2 1

2 2

2 3

2 4

2 5

loss: mse

closs

cmax

(a)

2 2

2 4

2 6

csharp

2 2

2 4

2 6

loss:

xent

closs

cmax

(b)

more noisy, as quantified by the confidence intervals. In comparison to phase diagrams of models trained with MSE loss, we observe a few notable differences. First, the loss starts to catapult at a value appreciably larger than c = 2 at large widths. Primarly, 4 ≲ c loss ≲ 8 . Additionally, c max 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 c < c sharp (see the first row of Figure 25). Next, we observe that the inequality c loss ≤ c sharp ≤ c max 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 5 steps to measure sharpness to avoid progressive sharpening. For CNNs, we observed that the dynamics becomes difficult to characterize for w ≲ 32 and d ≳ 10 , due to large fluctuations. Consequently, we've opted not to include these particular results.

2

-

1

2 1

2 3

2 5

c

2

-

8

2

-

7

2

-

6

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(a) CNN,

d

= 5

2

-

1

2 1

2 3

2 5

c

2

-

8

2

-

7

2

-

6

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(b) CNN,

d

= 7

2

-

1

2 1

2 3

2 5

c

2

-

7

2

-

6

2

-

5

1

/w

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(c) ResNet,

d

= 18

F.2 Intemediate saturation regime

Figure 29 shows the normalized sharpness measured at cτ = 100 for FCNs trained on CIFAR-10 using cross-entropy loss. 7 Similar to MSE loss, we observe an abrupt drop in sharpness at large learning rates. However, this abrupt drop occurs at 2 ≲ c crit ≲ 4 . The estimated sharpness is noisier (compare with Figure 38), which hinders a reliable estimation of c crit . We speculate that we require a

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

λ H

τ

/λ H

0

width

256

512

1024

2048

(a)

d

= 4

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

256

512

1024

2048

(b)

d

= 8

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

λ H

τ

/λ H

0

width

256

512

1024

2048

(c)

d

= 16

large number of averages for a reliable estimation of c crit for cross-entropy loss. We leave the precise characterization of c crit for cross-entropy loss for future work.

G The effect of setting model output to zero at initialization

In this section, we demonstrate the effect of network output f ( x ; θ t ) at initialization on the early training dynamics. In particular, we set the network output to zero at initialization, f ( x ; θ 0 ) = 0 , by (1) 'centering' the network by its initial value f c ( x ; θ t ) = f ( x ; θ ) -f ( x ; θ 0 ) 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 c loss ≈ 2 , independent of depth and width.

Throughout this section, we 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 a network function f ( x ; θ t ) , we define the centered network f c ( x ; θ t ) as

where f ( x ; θ 0 ) is the network output at intialization. By construction, the network output is zero at initialization. It is noteworthy that centering a network is an unusual way of training deep networks as it doubles the cost of 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 a decrease in sharpness for c < c loss during early training. Rather, we observe a slight increase in sharpness. To distinguish this slight increase from sharpness catapult, we introduce a threshold ϵ , comparing normalized sharpness λ H t / λ H 0 with 1 + ϵ , to define a sharpness catapult. 8 As demonstrated in Appendix G.3, the uv model trained on a 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 c loss ≈ 2 , as demonstrated in the phase diagrams in Figure 31(a, b, c). These findings suggest a strong correlation between a large network output at initialization ∥ f ( x ; θ 0 ) ∥ 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 ) = 0 is by setting the last layer to zero at initialization. The principle of criticality at initialization [55, 58, 68] does not put any constraints on

0

5

10

15

20

step

0

1

2

3

4

Training loss

vanilla nets

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(a)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

λ H

t

/λ H

0

vanilla nets

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(b)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

Training loss

centered nets

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(c)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

λ H

t

/λ H

0

centered nets

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(d)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

Training loss

zero-init nets

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(e)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

λ H

t

/λ H

0

zero-init nets

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(f)

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 in initialization in a flat curvature region at initialization, resulting in access to larger learning rates. We refer to these networks as 'zero-init' networks.

Figure 30 compares the training dynamics of zero-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 c loss ≈ 2 , irrespective of depth and width. Figure 31(d, e, f) show the phase diagrams of networks with zero-initialized networks. Like centered networks, the critical constants do not scale with depth and width. Again, suggesting that a large network output at initialization ∥ f ( x ; θ 0 ) ∥ is related to the opening of the sharpness reduction phase in the early training results shown in Section 2.

2

-

1

2 0

2 1

2 2

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a) centered,

d

= 4

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b) centered,

d

= 8

2

-

1

2 1

2 3

2 5

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(c) centered,

d

= 16

2

-

1

2 0

2 1

2 2

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(d) zero-init,

d

= 4

2

-

1

2 0

2 1

2 2

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(e) zero-init,

d

= 8

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(f) zero-init,

d

= 16

G.3 Insights from uv model trained on ( x, y )

In this section, we gain insights into the effect of setting network output to zero at initialization using uv model trained on an example ( x, y ) . In particular, we show that loss catapults at k loss = 2 and sharpness increases during early training.

Consider the uv model trained on a single training example ( x, y ) with y = 0 9

where ∆ f is the residual. The trace of the Hessian tr( H ) is

The Frobeinus norm can be written in terms of the trace and the network output

The function and residual updates are given by

̸

This simplifies the loss function to

0

5

10

15

20

step

0

.

040

0

.

042

0

.

044

0

.

046

0

.

048

0

.

050

Training loss

s = 0.1

loss:

mse, width:

512, depth: 8

c

0

.

8

1

.

6

2

.

4

3

.

2

4

.

0

(a)

0

5

10

15

20

step

1

2

3

λ H

t

/λ H

0

s = 0.1

loss:

mse, width:

512, depth: 8

c

0

.

8

1

.

6

2

.

4

3

.

2

4

.

0

(b)

0

5

10

15

20

step

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

Training loss

s = 1.0

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(c)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

λ H

t

/λ H

0

s = 1.0

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(d)

0

5

10

15

20

step

0

10

20

30

40

50

Training loss

s = 10.0

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(e)

0

5

10

15

20

step

0

.

0

0

.

5

1

.

0

1

.

5

λ H

t

/λ H

0

s = 10.0

loss:

mse, width:

512, depth: 8

c

1

.

5

3

.

0

4

.

5

6

.

0

7

.

5

(f)

Similarly, we can obtain the trace update equations

Let us analyze them for the networks with zero output at initialization. The loss at the first step increases if

(53)

Setting f 0 = 0 and scaling the learning rate as η = k/ tr( H 0 ) , we see that the loss increases at the first step if k > 2 .

Next, we analyze the change in trace during the first training step. Setting f 0 = 0 , we observe that the trace increases for all learning rates

modulated by the learning rate and width. Finally, we analyze the change in Frobenius norm in the first training step at k = k loss , which implies ∆ f 2 1 = ∆ f 2 0 ,

As tr( H ) increases in the first training step, ∥ H ∥ F also increases in the first training step.

H The effect of output scale on the training dynamics

Given a neural network function f ( x ) with depth d and width w , we define the scaled network as f s ( x ) = αf ( x ) , where α is referred to as the output scale. In this section, we empirically study the impact of the output scale on the early training dynamics. In particular, we show that a large (resp. small) value of ∥ f ( x ; θ 0 ) ∥ relative to the one-hot encodings of the labels causes the sharpness to decrease (resp. increase) during early training. Interestingly, we still observe an increase in ⟨ c loss ⟩ 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.

H.1 The effect of fixed output scale at initialization

In this section, we study the training dynamics of models trained with a fixed output scale at initialization. Given a network output function f ( θ ) , we define the 'scaled network' as

where s is a scalar, fixed throughout training. By construction, the network output norm ∥ f s ( θ 0 ) ∥ equals s . For standard initialization, s = ∥ f ( θ 0 ) ∥ = O ( √ k ) , where k are the number of classes.

Figure 32 shows the training dynamics of FCNs for three different values of the output scale s . The training dynamics of networks with s = 1 . 0 and s = 10 . 0 share qualitative similarities. In contrast, networks initialized with a smaller output scale ( s = 0 . 1 ) exhibit distinctly different dynamics. In particular, we observe that for large output scales ( s ≳ 0 . 5 ) sharpness decreases during early training, while sharpness increases for small output scales 10 . Furthermore, the training dynamics tends to

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

s = 0.1

loss: mse, depth: 4

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(a)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

s = 1.0

loss: mse, depth: 4

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(b)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

s = 10.0

loss: mse, depth: 4

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(c)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

s = 0.1

loss: mse, depth: 8

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(d)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

s = 1.0

loss: mse, depth: 8

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(e)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

s = 10.0

loss: mse, depth: 8

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(f)

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 as 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 c loss increases with d and 1 / w . In comparison, sharpness decreases (increases) for large (small) values of s . These experiments suggest that the output scale primarly influences the increase/decrease in sharpness during early training and does not affect the scaling of c loss 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, 6, 4]. We consider three distinct σ values {-0 . 5 , 0 . 0 , 0 . 5 } , where σ = -0 . 5 represents the lazy regime, σ = 0 . 5 corresponds to feature learning (rich) regime and σ = 0 . 0 correponds to standard (vanilla) initialization.

Figure 34 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 is 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 in 34(e, f)). These results are in parallel to the fixed output scale networks studied in the pervious section.

0

5

10

15

20

step

0

500

1000

1500

2000

Training loss

σ

:

-0.5

loss:

mse, width:

512, depth: 4

c

1

2

3

4

5

(a)

0

5

10

15

20

step

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

λ H

t

/λ H

0

σ

:

-0.5

loss:

mse, width:

512, depth: 4

c

1

2

3

4

5

(b)

0

5

10

15

20

step

0

1

2

3

4

Training loss

σ

:

0.0

loss:

mse, width:

512, depth: 4

c

1

2

3

4

5

(c)

0

5

10

15

20

step

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

λ H

t

/λ H

0

σ

:

0.0

loss:

mse, width:

512, depth: 4

c

1

2

3

4

5

(d)

0

5

10

15

20

step

0

.

040

0

.

045

0

.

050

Training loss

σ

:

0.5

loss:

mse, width:

512, depth: 4

c

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

(e)

0

5

10

15

20

step

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

λ H

t

/λ H

0

σ

:

0.5

loss:

mse, width:

512, depth: 4

c

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

(f)

Figure 35 summarizes the early training dynamics of FCNs with different σ values. We observe similar results as in the previous section. The output scale affects the initial increase/decrease of sharpness but does not affect the scaling trend of c loss with depth and width. Moreover, we observe a systematic pattern of c max scaling with width. In the lazy regime, we observe that c max increases with 1 / w , while c max decreases with 1 / w in the feature learning regime.

I Sharpness curves in the intermediate saturation regime

This section shows additional results for Section 3 for MSE loss. Cross-entropy results are shown in Appendix F. Figures 36 to 40 show the normalized sharpness curves for different depths and widths.

I.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 λ H τ at a time τ in the middle of the intermediate saturation regime. We choose τ so that cτ ≈ 200 , for learning rates c = 2 x , where x ∈ [ -1 . 0 , 4 . 0] in steps of 0 . 1 . The value 200 is chosen such that τ is in the middle of the intermediate saturation regime. Next, we measure sharpness over a range of steps

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

σ

:

-0.5

loss: mse, depth: 4

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(a)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

σ

:

0.0

loss: mse, depth: 4

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(b)

2

-

2

2

-

1

2 0

2 1

2 2

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

σ

:

0.5

loss: mse, depth: 4

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(c)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

σ

:

-0.5

loss: mse, depth: 8

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(d)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

σ

:

0.0

loss: mse, depth: 8

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(e)

2

-

1

2 1

2 3

c

2

-

11

2

-

10

2

-

9

2

-

8

1

/w

σ

:

0.5

loss: mse, depth: 8

〈

c

loss

〉

〈

c

sharp

〉

〈

c

〉

(f)

t ∈ [ τ -5 , τ + 5] and average over t to reduce fluctuations. We repeat this process for various initializations and obtain the average sharpness.

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

256

512

1024

2048

(a)

d

= 4

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

256

512

1024

2048

(b)

d

= 8

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

256

512

1024

2048

(c)

d

= 16

I.2 Estimating the critical constant c crit

This subsection explains how to estimate c crit from sharpness measured at time τ . First, we normalize the sharpness with its initial value, and then average over random initializations. Next, we estimate the critical point c crit 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 c crit . We resolve this issue by estimating the smooth derivatives of the averaged order parameter with the Savitzky-Golay filter [59] using its scipy implementation [63]. The estimated c crit is shown by vertical lines in the sharpness curves in Figures 36 to 40.

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

256

512

1024

2048

(a)

d

= 4

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

256

512

1024

2048

(b)

d

= 8

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

256

512

1024

2048

(c)

d

= 16

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

256

512

1024

2048

(a)

d

= 4

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

256

512

1024

2048

(b)

d

= 8

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

0

.

6

λ H

τ

/λ H

0

width

256

512

1024

2048

(c)

d

= 16

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

128

256

512

(a)

d

= 5

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

128

256

512

(b)

d

= 7

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

128

256

512

(c)

d

= 10

J The effect of batch size on the reported results

J.1 The early transient regime

Figure 41 shows the phase diagrams of early training dynamics of FCNs with d = 4 trained on the CIFAR-10 dataset using two different batch sizes. The phase diagram obtained is consistent with the findings presented in 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

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

32

64

128

(a)

d

= 18

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

λ H

τ

/λ H

0

width

32

64

128

(b)

d

= 34

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

32

64

128

(c)

d

= 50

at relatively smaller values of c . This is reflected in Fig. 41 by ⟨ c sharp ⟩ moving to the left as B is reduced from 512 to 128 . However, this initial increase in sharpness is small compared to the sharpness catapult observed at larger batch sizes. We found that this increase at small batch sizes is due to fluctuations in gradient estimation that can cause sharpness to increase above its initial value by chance.

2

-

1

2 1

2 3

c

2

-

9

2

-

8

2

-

7

2

-

6

d/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

〈

cbarrier

〉

(a)

B

= 128

2

-

1

2 1

2 3

c

2

-

9

2

-

8

2

-

7

2

-

6

d/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

〈

cbarrier

〉

(b)

B

= 512

J.2 The intermediate saturation regime

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

λ H

τ

/λ H

0

width

256

512

1024

2048

B

32

128

512

(a)

2

-

1

2 0

2 1

2 2

2 3

2 4

c

0

.

0

0

.

2

0

.

4

0

.

6

χτ

width

256

512

1024

2048

|

S

|

32.0

128.0

512.0

(b)

2

-

1

2 0

2 1

2 2

2 3

2 4

c

-

0

.

5

0

.

0

0

.

5

1

.

0

χ

′

τ

width

256

512

1024

2048

B

32.0

128.0

512.0

(c)

Figure 42 shows the normalized sharpness, measured at cτ = 200 , and its derivatives for various widths and batch sizes. The results are consistent with those in Section 3, with a lowering in the peak

heights of the derivatives χ and χ ′ 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 c crit 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 a similar phase diagram compared to the no-bias case (compare with Figure 26).

2

-

1

2 1

2 3

c

2

-

9

2

-

8

2

-

7

2

-

6

d/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(a)

d

= 4

2

-

1

2 1

2 3

2 5

c

2

-

8

2

-

7

2

-

6

2

-

5

d/w

〈

closs

〉

〈

csharp

〉

〈

cmax

〉

(b)

d

= 8

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nUnderstanding 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nIt 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\nLIMITATIONS:\nThe limitations of the results presented are discussed well.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- The paper doesn't explain the early-stage sharpness reduction for general neural nets (i.e. beyond the simple u-v model)\n\nQUESTIONS:\n- Have you looked at loss connectivity between the final parameters (at the end of training) and the initialization?\n\nLIMITATIONS:\ndiscussed above", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1.\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\nWEAKNESSES:\n1.\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\nQUESTIONS:\n1.\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\nLIMITATIONS:\nThe limitations are discussed in the paper. This is a theoretical work and therefore does not seem to have negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\n(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\nQUESTIONS:\n**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.\n\nLIMITATIONS:\n**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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}