Phase diagram of early training dynamics in deep neural networks: effect
  of the learning rate, depth, and width
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.
Dayal Singh Kalra, Maissam Barkeshli
# Phase diagram of training dynamics in deep neural networks:

###### Abstract

We systematically analyze optimization dynamics in deep neural networks (DNNs) trained with stochastic gradient descent (SGD) over long time scales and study the effect of learning rate, depth, and width of the neural network. By analyzing the maximum eigenvalue \(\lambda_{t}^{H}\) 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 finally (iv) a late time "edge of stability" regime. The early and intermediate regimes (i) and (ii) exhibit a rich phase diagram depending on learning rate \(\eta\equiv c/\lambda_{0}^{H}\), depth \(d\), and width \(w\). We identify several critical values of \(c\) which separate qualitatively distinct phenomena in the early time dynamics of training loss and sharpness, and extract their dependence on \(d/w\). Our results have implications for how to scale the learning rate with DNN depth and width in order to remain in the same phase of learning.

Machine Learning, Deep Neural Networks, Deep Neural Networks, Neural Networks

## 1 Introduction

The optimization dynamics of deep neural networks (DNNs) is a rich problem that is of fundamental and practical interest. Seemingly small changes to the hyperparameters can dramatically alter the effectiveness of training and even lead to training failures. 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 oft-made suggestion that in stochastic gradient descent (SGD), the learning rate \(\eta\) should satisfy \(\eta<2/\lambda^{H}\), where \(\lambda^{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 \(\eta>2/\lambda^{H}\)(Wu et al., 2018; Lewkowycz et al., 2020; Zhu et al., 2022). 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 \(\lambda^{H}\) where training stabilizes. Indeed, during the early training phase, the local curvature of the loss landscape changes rapidly (Keskar et al., 2016; Achille et al., 2017; Jastrzebski et al., 2018; Fort and Ganguli, 2019), and the learning rate plays a crucial role in determining the convergence basin (Jastrzebski et al., 2018). Flatter basins are believed to be preferable because they potentially lead to lower generalization error (Hochreiter and Schmidhuber, 1994; 1997; Keskar et al., 2016; Dziugaite and Roy, 2017; Jiang et al., 2019; Foret et al., 2020) 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 (Kaplan et al., 2020). 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 architecture size is taken to infinity (Jacot et al., 2018; Lee et al., 2018; 2019; Dyer and Gur-Ari, 2020; Yang and Hu, 2021; Roberts et al., 2022; Yang et al., 2022; Yaida, 2022)?

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 (\(\lambda^{H}\)) trajectories out to late times, observing a number of qualitatively distinct phenomena summarized below.
### 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 parameterize FCNs in the Neural Tangent Parameterization (NTP) (Jacot et al., 2018) and CNNs and ResNets in Standard Parameterization (SP) (Sohl-Dickstein et al., 2020); other parameterizations, such as muP (Yang and Hu, 2021), can give different results. For further details, see Appendix A. We study both mean-squared error (MSE) and cross-entropy loss functions, and the datasets CIFAR-10, MNIST, Fashion-MNIST. The basic phenomena we report are reliably reproduced across these architectures and datasets. We use sharpness to refer to \(\lambda_{t}^{H}\), the maximum eigenvalue of \(H\) at time-step \(t\), and flatness refers to \(1/\lambda_{t}^{H}\).

By monitoring the sharpness, we find four clearly separated, qualitatively distinct regimes throughout the entire span of the training trajectory. Fig. 1 shows an example from a CNN architecture. The four observed regimes are: (i) an early time transient regime, (ii) an intermediate saturation regime where the sharpness has lowered and remains relatively constant, (iii) a progressive sharpening regime, and finally (iv) a late time "edge of stability" regime where the sharpness saturates around \(2/\eta\)(Cohen et al., 2021).

Next, we analyze the early transient and intermediate saturation regimes. For MSE loss, a clear picture emerges, and leads to a rich phase diagram, sketched in Figure 2, as learning rate, \(d\), and \(w\) are tuned. Given the learning rate scaled as \(\eta=c/\lambda_{0}^{H}\), we characterize four distinct behaviors in the training dynamics in the early transient regime (i):

1. **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 \(2\approx 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. The regime \(c_{crit}<c<c_{loss}\) opens up significantly with increasing \(d/w\), which is a new result of this work.
2. **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.
3. **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.
4. **Divergent phase (\(c>c_{max}\)):** The learning rate is too

Figure 1: Full training trajectories of the (a) training loss, (b) sharpness, and (c) training accuracy of CNNs (\(d=5\) and \(w=512\)) trained on CIFAR-10 with MSE loss using vanilla SGD with learning rates \(\eta=c/\lambda_{0}^{H}\) and batch size \(B=512\). Vertical dashed lines approximately separate the different training regimes. Horizontal dashed lines in (b) denote the \(2/\eta\) threshold for each learning rate. In Appendix G.1 we show the same trajectories with the time axis rescaled as \(ct\).

Figure 2: Sketched phase diagram of early training
large for training to converge 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}\) are shown in the phase diagrams and show strong systematic dependence on \(d/w\).

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. 6. By identifying an appropriate order parameter we can extract a sharp peak corresponding to \(c_{crit}\approx 2/\lambda_{0}^{H}\). For \(c\ll c_{crit}\), the network is effectively in a lazy training regime, with increasing fluctuations as \(d/w\) is increased.

Finally, we show that a single hidden layer linear network - the \(uv\) model - displays the same phenomena discussed above for the early regimes (i) and (ii), and we analyze the phase diagram in this minimal model.

### Related works

A significant amount of research has identified various training regimes using diverse criteria, e.g., (Erhan et al., 2010; Achille et al., 2017; Fort et al., 2020; Jastrzebski et al., 2018; Frankle et al., 2020; Leclerc and Madry, 2020; Jastrzebski et al., 2020; Cohen et al., 2021; Iyer et al., 2022). Here we focus on studies that characterize training regimes with sharpness and learning rates.

Several studies have analyzed sharpness at different training times (Jastrzebski et al., 2018; Fort and Ganguli, 2019; Jastrzebski et al., 2020; Cohen et al., 2021; Iyer et al., 2022). (Cohen et al., 2021) studied sharpness at late training times and showed how _large-batch_ gradient descent shows progressive sharpening followed by the edge of stability: For MSE loss, the sharpness increases until it reaches a maximum value just above \(2/\eta\) and oscillates near it. This has motivated various theoretical studies (Damian et al., 2022; Agarwala et al., 2022; Arora et al., 2022). We add to this line of work by demonstrating a wide variety of architectures trained with SGD with relatively small batch sizes \(B=512\), \(\approx 1\%\) of the full training set, that clearly show progressive sharpening and edge of stability phenomena.

(Jastrzebski et al., 2018) 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/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 (Iyer et al., 2022).

(Jastrzebski et al., 2020) examined the effect of hyperparameters on sharpness at late training times. (Gilmer et al., 2021) studied the optimization dynamics of SGD with momentum using sharpness. (Leclerc and Madry, 2020) classify training into 2 different regimes using training loss, providing a significantly coarser description of training dynamics than provided here. (Iyer et al., 2022) studied the scaling of the maximum learning rate with \(d\) and \(w\) during early training in FCNs and its relationship with sharpness at initialization.

(Lewkowycz et al., 2020) analyzed the curvature during early training using the top eigenvalue of the neural tangent kernel (NTK) and demonstrated the existence of a new early training phase, which they dubbed the "catapult" phase, \(2/\lambda_{0}^{NTK}<\eta<\eta_{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 (Zhu et al., 2022; Meltzer and Liu, 2023). Our work extends the above analysis by studying the combined effect of learning rate, depth, and width, demonstrating the opening 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/w\) is increased, analyzing linear mode connectivity in the catapult phase, and also examining different qualitative behaviors in the intermediate saturation regime (ii) mentioned above.

## 2 The four regimes of training

Figure 1 illustrates the training trajectories of a 5-layer deep (\(d=5\)) CNN from the Myrtle family (Shankar et al., 2020), trained on the CIFAR-10 dataset using MSE loss. Each network layer has \(w=128\) channels. The networks were trained using SGD with various learning rates and a batch size of \(B=512\). These results indicate that the training dynamics can be divided into four distinct regimes based on training loss and sharpness:

1. **Early time transient: (\(0<t<T_{1}\)).** Within a few (\(\sim 10\)) training steps, training selects the convergence basin, which depends on the optimization hyperparameters. The loss and sharpness may drastically change, and training eventually settles down in a basin by the end of this short period.
* **Intermediate saturation: (\(T_{1}<t<T_{2}\)).** The learning rate is small compared to sharpness, resulting in a slow transient period. The sharpness changes slowly until it reaches a tipping point \(T_{2}\), which marks this regime's end.
* **Progressive sharpening: (\(T_{2}<t<T_{3}\)).** The sharpness steadily rises until it reaches a maximum value \(\approx 2/\eta\)(Jastrzebski et al., 2020; Cohen et al., 2021).
* **Late time: (\(T_{3}<t\)).** Training dynamics show different phenomena depending on the loss and optimization parameters. For networks trained with MSE loss and large enough batch sizes, training enters a stable cycle in which sharpness oscillates around this maximum value, and the training loss fluctuates while decreasing overall. In comparison, sharpness may drop after reaching this maximum value for cross-entropy loss while remaining less than \(2/\eta\).

We confirm the existence of the four training regimes across various architectures, datasets, and loss functions as demonstrated in Appendix G. Our results show that FCNs, CNNs, and ResNets, when trained on CIFAR-10, MNIST, and Fashion-MNIST datasets, using MSE and cross-entropy loss, exhibit all four training regimes. An exception to this observation is that FCNs did not exhibit the fourth regime (iv), even after training for \(1000\) epochs.

Note the log scale in Fig. 1 highlights the early regimes (i) and (ii); in absolute terms these are much shorter in time than regimes (iii) and (iv).

## 3 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\)(Lewkowycz et al., 2020). Figure 3(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\geq 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 Figure 3 (b)). However, this initial spike in sharpness occurs at relatively higher learning rates (\(c\geq 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\). Figure 3(c, d) 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. For cross-entropy loss, we found

Figure 3: Early training dynamics of (a, b, c) a shallow (\(d=5\), \(w=512\)) and (d, e, f) a deep CNN (\(d=10,w=128\)) trained on CIFAR-10 with MSE loss for \(t=10\) steps using SGD for various learning rates \(\eta=c/\lambda_{0}^{H}\) and batch size \(B=512\). (a, d) training loss, (b, e) sharpness, and (c, f) interpolated loss between the initial and final parameters after \(10\) steps for the respective models. For the shallow CNN, \(c_{loss}=2.82,c_{sharp}=5.65,c_{max}=17.14\) and for the deep CNN, \(c_{loss}=36.75,c_{sharp}=39.39,c_{max}=48.50\).
that the early time dynamics is more chaotic and it is unclear to what extent it can be separated into different phases.

### 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 3.1**.: (\(c_{loss},c_{sharp},c_{max}\)) For learning rate \(\eta=c/\lambda_{0}^{H}\), let the training loss and sharpness at step \(t\) be denoted by \(\mathcal{L}_{t}(c)\) and \(\lambda_{t}^{H}(c)\). We define \(c_{loss}(c_{sharp})\) as minimum learning rates constants such that the loss (sharpness) increases during the initial transient period

\[c_{loss}=\min_{c}\{c\mid\max_{t\in[1,T_{1}]}\mathcal{L}_{t}(c)> \mathcal{L}_{0}(c)\} \tag{1}\] \[c_{sharp}=\min_{c}\{c\mid\max_{t\in[1,T_{1}]}\lambda_{t}^{H}(c)> \lambda_{0}^{H}(c)\}, \tag{2}\]

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\mid\mathcal{L}_{t}(c)<K,\forall t\in[1,T_{1}]\}, \tag{3}\]

where \(K\) is a fixed large constant 1.

Footnote 1: We use \(K=10^{5}\) to estimate \(c_{max}\). In all our experiments, \(\mathcal{L}_{0}=\mathcal{O}(1)\) (see Appendix A), which justifies the use of a fixed value.

Note that the definition of \(c_{max}\) allows for more flexibility than previous studies (Iyer et al., 2022) 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 \(\langle\cdot\rangle\).

Figure 4 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 \(\langle c_{loss}\rangle\), \(\langle c_{sharp}\rangle\), and \(\langle c_{max}\rangle\) are affected by network depth and width when plotted in the subspace represented by the depth-to-width ratio (\(d/w\)). This choice is motivated by various studies that show how in FCNs, the ratio \(d/w\) controls the fluctuations (Hanin and Nica, 2019; Dyer and Gur-Ari, 2020; Roberts et al., 2022). The results show that the averaged values of all the critical constants increase significantly with \(d/w\) (note the log scale). At small values of \(d/w\), the loss starts to catapult at \(c\approx 2\). As \(d/w\) increases, \(\langle c_{loss}\rangle\) increases and eventually converges to \(\langle c_{max}\rangle\) at large \(d/w\). By comparison, sharpness starts to

Figure 4: The phase diagrams of the early training of three different types of neural networks trained with MSE loss function using SGD. (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 dataset (without batch normalization). Each data point in the figure represents an average of ten distinct initialization, and the solid lines represent a smooth curve fitted to the raw data points. The vertical dotted line shows \(c=2\) for comparison, and various colors are filled in between the various curves for better visualization. For experimental details and additional results for different depths, see Appendices A and C, respectively.

Figure 5: The relationship between critical constants for ResNets with \(d\in\{18,34\}\), \(w\in\{32,64,128\}\), trained on the CIFAR-10 dataset. Each data point corresponds to a randomly initialized model. The dashed line represents the \(y=x\) line. For experimental details and additional results, see Appendices A and D.
catapult at relatively large learning rates at small \(d/w\), with \(\langle c_{sharp}\rangle\) continuing to increase with \(d/w\) while remaining between \(\langle c_{loss}\rangle\) and \(\langle c_{max}\rangle\). Moreover, although \(\langle c_{max}\rangle\) increases with \(d/w\), the bare learning rate \(\eta_{max}\equiv c_{max}/\lambda_{0}^{H}\) decreases on average with \(d/w\).

While we plotted the averaged quantities \(\langle c_{loss}\rangle\), \(\langle c_{sharp}\rangle\), \(\langle c_{max}\rangle\), we have observed that their variance also increases significantly with \(d/w\); in Appendix C we show standard deviations about the averages for different random initializations. Nevertheless, we have found that the following inequality typically holds, for any given initialization and batch sequences:

\[c_{loss}\leq c_{sharp}\leq c_{max}, \tag{4}\]

except for some outliers due to high fluctuations when the averaged critical curves start merging at large \(d/w\). Figure 5 shows evidence of this for ResNets trained on the CIFAR-10 dataset. Appendix D presents extensive additional results across various architectures.

### 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. A naive intuition would predict a barrier from initialization to the final point after the 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 3.2**.: \((\mathcal{L}_{int}(s,c),U(c))\) Let \(\theta_{0}\) represent the initial set of parameters, and let \(\theta_{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 \(\mathcal{L}_{int}(s,c)=\mathcal{L}[(1-s)\ \theta_{0}+s\ \theta_{T_{1}}]\), where \(s\in[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\in[0,1]}\mathcal{L}_{int}(s)-\mathcal{L}(\theta_{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.3**.: (\(c_{barrier}\)) Given the initial (\(\theta_{0}\)) and final parameters (\(\theta_{T_{1}}\)), we define \(c_{barrier}\) as the minimum learning rate constant such that there exists a barrier from \(\theta_{0}\) to \(\theta_{T_{1}}\)

\[c_{barrier}=\min_{c}\{c\ |\ U(c)>0\} \tag{5}\]

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 \(\langle.\rangle\) as before. Figure 3(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 3. For the network with smaller \(d/w\), 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/w\). Figure 4 shows the relationship between \(\langle c_{barrier}\rangle\) and \(d/w\) for various models and datasets. We consistently observe that \(c_{sharp}\leq c_{barrier}\), suggesting that training traverses a barrier only when sharpness starts to catapult during early training. We chose not to characterize

Figure 6: (a) Normalized sharpness measured at \(c\tau=200\) against the learning rate constant for \(7\)-layer CNNs trained on the CIFAR-10 dataset, with varying widths. Each data point is an average over 5 initializations, where the shaded region depicts the standard deviation around the mean trend. (b, c) Smooth estimations of the first two derivatives, \(\chi_{\tau}\) and \(\chi_{\tau}^{\prime}\), of the averaged normalized sharpness wrt the learning rate constant. The vertical lines denote \(c_{crit}\) estimated using the maximum of \(\chi_{\tau}^{\prime}\). For smoothening details, see Appendix E.2.
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}\). Instead, 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.

## 4 Intermediate saturation regimes

In the intermediate saturation regime, the learning rate is small compared to the local sharpness (\(\eta<2/\lambda_{t}\)). As a result, 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.

We measure the sharpness \(\lambda_{\tau}^{H}\) at a time \(\tau\) in the middle of the intermediate saturation regime. We choose \(\tau\) so that \(cr\approx 200\).2 For further details on sharpness measurement, see Appendix E.1. Figure 6(a) illustrates the relationship between \(\lambda_{\tau}^{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 \(\lambda_{\tau}^{H}\) on learning rate can be grouped into three distinct stages. (1) At small learning rates, \(\lambda_{\tau}^{H}\) remains relatively constant, with fluctuations increasing as the ratio of \(d/w\) increases (\(c<2\) in Fig. 6(a)). (2) A crossover regime where \(\lambda_{\tau}^{H}\) is dropping significantly (\(2<c<2^{3}\) in Fig. 6(a)). (3) A saturation stage where \(\lambda_{\tau}^{H}\) stays small and constant with learning rate (\(c>2^{3}\)) in Fig. 6(a)). In Appendix E, 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, \(\lambda_{\tau}^{H}\) decreases with increasing \(d/w\). In other words, for small \(c\) and in the intermediate saturation regime, the loss is locally flatter as \(d/w\) increases.

Footnote 2: In practice time-step \(\tau=200/c\) is in the middle of regime (ii) for the models we studied. Normalizing by \(c\) allows proper comparison for different learning rates.

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 \(\lambda_{\tau}^{H}\). To do this, we consider the averaged normalized sharpness over initializations and denote it by \(\langle\frac{\lambda_{\tau}^{H}}{\lambda_{\theta}^{H}}\rangle\). The first two derivatives of the averaged normalized sharpness, \(\chi_{\tau}=-\frac{\partial}{\partial c}\langle\frac{\lambda_{\tau}^{H}}{ \lambda_{\theta}^{H}}\rangle\) and \(\chi_{\tau}^{\prime}=-\frac{\partial^{2}}{\partial c^{2}}\langle\frac{ \lambda_{\tau}^{H}}{\lambda_{\tau}^{H}}\rangle\), characterize the change in sharpness with learning rate. The extrema of \(\chi_{\tau}^{\prime}\) quantitatively define the boundaries between the three stages described above. In particular, using the maximum of \(\chi_{\tau}^{\prime}\), we define \(\langle c_{crit}\rangle\), which marks the beginning of the sharp decrease in \(\lambda_{\tau}^{H}\) with the learning rate.

**Definition 4.1**.: (\(\langle c_{crit}\rangle\)) Given the averaged normalized sharpness \(\langle\frac{\lambda_{\tau}^{H}}{\lambda_{\theta}^{H}}\rangle\) estimated using sharpness measured at \(\tau\), we define \(c_{crit}\) to be the learning rate constant that minimizes the second derivative of \(\langle\frac{\lambda_{\tau}^{H}}{\lambda_{\theta}^{H}}\rangle\)

\[\langle c_{crit}\rangle=\operatorname*{arg\,min}_{c}\frac{\partial^{2}}{ \partial c^{2}}\left\langle\frac{\lambda_{\tau}^{H}}{\lambda_{0}^{H}}\right\rangle =\operatorname*{arg\,max}_{c}\chi_{\tau}^{\prime}. \tag{6}\]

Here, we use \(\langle.\rangle\) to denote that the critical constant is obtained from the averaged normalized sharpness. Figures 6(b, c) show \(\chi_{\tau}\) and \(\chi_{\tau}^{\prime}\) obtained from the results in Figure 6(a). We observe similar results across various architectures and datasets, as shown in Appendix E. Our results show

Figure 7: The phase diagram of the \(uv\) model trained with MSE loss using gradient descent with (a) the top eigenvalue of Hessian \(\lambda_{t}^{H}\), (b) the trace of Hessian \(\operatorname{Tr}(H_{t})\) and (c) the square of the Frobenius norm \(\operatorname{Tr}(H_{t}^{T}H_{t})\) used as a measure of sharpness. In (a), the learning rate is scaled as \(\eta=c/\lambda_{0}^{H}\), while in (b) and (c), the learning rate is scaled as \(\eta=k/\operatorname{Tr}(H_{0})\). The vertical dashed line shows \(c=2\) for reference. Each data point is an average over \(500\) random initializations.
that \(\langle c_{crit}\rangle\) has slight fluctuations as \(d/w\) is changed, but generally stays in the vicinity of \(c=2\). The peak in \(\chi_{\tau}^{\prime}\) becomes wider as \(d/w\) increases, 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 \(\langle c_{crit}\rangle\), \(\langle c_{loss}\rangle\) increases with \(d/w\), implying the opening of the sharpness reduction phase \(\langle c_{crit}\rangle<c<\langle c_{loss}\rangle\) as \(d/w\) increases.

## 5 Insights from a simple model

Here we analyze a two-layer linear network (Rennie and Srebro, 2005; Saxe et al., 2014; Lewkowycz et al., 2020), the \(uv\) model, which shows much of the phenomena presented above. Define \(f(x)=\frac{1}{\sqrt{w}}w^{T}ux\), with \(x,f(x)\in\mathbb{R}\). Here, \(u,v\in\mathbb{R}^{w}\) are the trainable parameters, initialized using the normal distribution, \(u_{i},v_{i}\sim\mathcal{N}(0,1)\) for \(i\in\{1,\dots,w\}\). The model is trained with MSE loss on a single training example \((x,y)=(1,0)\), which simplifies the loss to \(\mathcal{L}(u,v)=\frac{1}{2}f^{2}\). We minimize the loss using gradient descent (GD) with learning rate \(\eta=c/\lambda_{0}^{H}\).

This model exhibits the (i) early time transient and (ii) intermediate saturation regimes, but not regimes (iii) and (iv).3 The early time phase diagram also shows similar features to those described in preceding sections (compare Fig. 7(a) and Fig. 4). Below we develop an understanding of this early time phase diagram in the \(uv\) model.

Footnote 3: For other choices where \(y\neq 0\), the model exhibits different phenomenology than discussed here.

The update equations of the \(uv\) model in function space can be written in terms of the trace of the Hessian \(\operatorname{Tr}(H)\).

\[\begin{split}& f_{t+1}=f_{t}\left(1-\eta\operatorname{Tr}(H_{t})+ \frac{\eta^{2}f_{t}^{2}}{w}\right)\\ &\operatorname{Tr}(H_{t+1})=\operatorname{Tr}(H_{t})+\frac{\eta f _{t}^{2}}{w}\left(\eta\operatorname{Tr}(H_{t})-4\right).\end{split} \tag{7}\]

In this model, \(\operatorname{Tr}(H)\) is equal to the NTK. From the above equations, it is natural to scale the learning rate as \(\eta=k/\operatorname{Tr}(H_{0})\). Note that \(c=\eta\lambda_{0}^{H}=k\frac{\lambda_{0}^{H}}{\operatorname{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 3.1 and 4.1 on replacing sharpness with trace. Figure 7(b) shows the phase diagram of early training, with \(\operatorname{Tr}(H_{t})\) replaced with \(\lambda_{t}^{H}\) as the measure of sharpness and with the learning rate scaled as \(\eta=k/\operatorname{Tr}(H_{0})\). Similar to Figure 7(a), we observe a new phase \(\langle k_{crit}\rangle<k<\langle k_{loss}\rangle\) opening up at small width. However, we do not observe the loss-sharpness catapult phase as \(\operatorname{Tr}(H)\) does not increase during training (see Equation 7). We also observe \(\langle k_{max}\rangle=4\), independent of width.

In Appendix B.2, we show that

\[\begin{split}\left\langle\frac{\mathcal{L}_{1}}{\mathcal{L}_{0} }\right\rangle=&(1-k)^{2}+\frac{k^{2}(1-k)}{2(w+1)}+\frac{3k^{4}}{ 16(w+1)(w+3)}\end{split} \tag{8}\]

where the average is over initializations. As explained in Appendix B.2, the critical value of \(k\) for which \(\langle\mathcal{L}_{1}/\mathcal{L}_{0}\rangle>1\) increases with \(1/w\), which explains why \(\langle k_{loss}\rangle\) and \(\langle c_{loss}\rangle\) increase with \(1/w\). Combined with \(k_{crit}\lesssim 2\), this implies the opening up of the sharpness reduction phase as \(w\) is decreased.

Figure 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 \(\operatorname{Tr}(H)\) and weight correlation \(\frac{\langle v,u\rangle}{||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 \(\operatorname{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. This means that the loss catapult requires a larger learning rate at a smaller width. 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. We speculate that similar reasoning holds for real-world models for the opening up of the sharpness reduction phase with increasing \(d/w\).

To understand the loss-sharpness catapult phase, we require

Figure 8: Training trajectories of the uv model, with (a, b) large and (c, d) small width, in a two-dimensional slice of the parameters defined by the trace of Hessian \(\operatorname{Tr}(H)\) and weight correlation, trained with (a, c) small (\(c=0.5\)) and (b, d) large (\(c=2.5\)) learning rates. The colors correspond to the training loss \(\mathcal{L}\), with darker colors representing a smaller loss.
some other measure as \(\operatorname{Tr}(H)\) does not increase for \(0<k<4\). \(\lambda_{t}^{H}=||H||_{\infty}\) corresponds to the \(p=\infty\) Schatten norm, which is difficult to analyze. Instead, we analyze the \(p=2\) Schatten norm, which is the Frobenius norm \(\|H\|_{F}=\sqrt{\operatorname{Tr}(H^{T}H)}\). For the \(uv\) model,

\[||H_{t}||_{F}^{2}=\operatorname{Tr}(H_{t}^{T}H_{t})=\operatorname{Tr}(H_{t})^{ 2}+4\left(1+\frac{2}{w}\right)\mathcal{L}_{t} \tag{9}\]

We define \(k_{frob}\) as the minimum learning rate such that \(||H_{t}||_{F}^{2}\) increases during early training. Figure 7(c) shows the phase diagram of the \(uv\) model, with \(||H_{t}||_{F}^{2}\) as the measure of sharpness, while the learning rate is scaled as \(\eta=k/\operatorname{Tr}(H_{0})\). We observe the loss-sharpness catapult phase at small widths, which we understand as follows. In Appendix B.3, we show \(\Delta||H_{1}||_{F}^{2}\equiv||H_{1}||_{F}^{2}-||H_{0}||_{F}^{2}\) averaged over initializations is

\[\left\langle\Delta||H_{1}||_{F}^{2}\right\rangle= I(k,w)+4\left(1+\frac{2}{w}\right)\left\langle\mathcal{L}_{1}- \mathcal{L}_{0}\right\rangle, \tag{10}\]

\(I(k,w)=\frac{k(k-4)}{w}\left[\frac{3k(k-4)}{4(w+3)}+2\right]\). As demonstrated in Appendix B.3, the critical value for which \(\left\langle\Delta||H_{1}||_{F}^{2}\right\rangle>0\) increases from \(k_{loss}\) as \(1/w\) increases. This explains the existence of the loss-sharpness catapult phase, and the opening up of the loss catapult phase at small \(w\) in Fig. 7 (c).

For \(k>4\), both the loss and trace increase initially, causing training to diverge irrespective of width. This yields \(c_{max}=4\lambda_{0}^{H}/\operatorname{Tr}(H)\). Figure 10(c) in Appendix B.2 shows that the ratio \(\lambda_{0}^{H}/\operatorname{Tr}(H_{0})\) increases with \(1/w\), which explains the increase in \(\left\langle c_{max}\right\rangle\) with \(1/w\).

The loss landscape of the \(uv\) model shown in Figure 8 reveals interesting insights into the loss landscape connectivity results in Section 3.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.

## 6 Discussion

We have studied the effect of learning rate, depth, and width on the training dynamics in DNNs trained using SGD with learning rate scaled as \(\eta=c/\lambda_{0}^{H}\). For various models and datasets, we demonstrated the existence of four training regimes. 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. We also find that the critical \(c\)'s increase significantly with \(d/w\) on average. 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 and increase of \(c_{max}\) with \(d/w\).

We performed a preliminary analysis on the effect of batch size on the presented results in Appendix F. The sharpness trajectories of models trained with a smaller batch size (\(B=32\) vs. \(B=512\)) show all four training regimes. The main effect of lowering batch size is that in the late time regime (iv), sharpness settles down to a value significantly smaller than \(2/\eta\), which is consistent with (Jastrzebski et al., 2020) and (Cohen et al., 2021). 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}\).

Throughout this work, we have scaled the learning rate using the top eigenvalue of the Hessian as a measure of sharpness. In practice, various other choices can be made. For example, (Lewkowycz et al., 2020) scaled the learning rate with the top eigenvalue of the NTK, which also shows similar results but a few differences. At early training times, we observed that the NTK may not catapult during early training, which results in the absence of the sharpness catapult phase. Moreover, the deviation of various critical constants, \(\left\langle c_{loss}\right\rangle\) and \(\left\langle c_{max}\right\rangle\), with \(d/w\), is less pronounced. These differences can be observed in the \(uv\) model analysis presented in Section 5, as \(\operatorname{Tr}(H)=\) NTK in that case.

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 (Poole et al., 2016) as it allows for proper gradient flow at initialization (Hanin, 2018; Roberts et al., 2022), and studied vanilla SGD for simplicity. However, other initializations (LeCun et al., 2012), parameterizations (Yang and Hu, 2021; Yang et al., 2022), and optimization procedures (Goyal et al., 2017) may show dissimilarities with the reported phase diagram of early training.

## Acknowledgements

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. This work is supported by an NSF CAREER grant (DMR-1753240) and the Laboratory for Physical Sciences through the Condensed Matter Theory Center.