SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Phase diagram of early training dynamics in deep networks: effect of the learning rate, depth, and width
Dayal Singh Kalra  $^{* \dagger}$  dayal@umd.edu
Maissam Barkeshli *‡§ maissam@umd.edu
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_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 (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.
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  $\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$  [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  $\lambda^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]?
*Condensed Matter Theory Center, University of Maryland, College Park
†Institute for Physical Science and Technology, University of Maryland, College Park
$^{\ddagger}$ Department of Physics, University of Maryland, College Park
$^{\S}$ Joint Quantum Institute, University of Maryland, College Park
Figure 1: 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.
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 during early training, observing a number of qualitatively distinct phenomena summarized below.
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 \lesssim C$ , where  $C$  depends on architecture class (e.g. for FCNs,  $C \approx 1 / 16$ ) and loss function, but is independent of  $d$ ,  $w$ , and  $\eta$ . Above this ratio, the dynamics become noise-dominated, and separating the underlying deterministic dynamics from random fluctuations becomes challenging, as shown in Appendix E. 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 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 / \eta$  for MSE loss; whereas for cross-entropy loss, sharpness drops after reaching this maximum value while remaining less than  $2 / \eta$  [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  $\eta = c / \lambda_0^H$ , 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_{\text{sharp}} < c < c_{\text{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_{\text{sharp}}$ ,  $c_{\text{max}}$  are random variables that depend on random initialization, SGD batch selection, and architecture. The averages of  $c_{loss}$ ,  $c_{\text{sharp}}$ ,  $c_{\text{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_{\text{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} \approx 2$ , whereas for crossentropy loss,  $4 \gtrsim c_{crit} \gtrsim 2$ . For  $c \ll 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.
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 / \lambda_0^{NTK}$  is the maximum learning rate for convergence as  $w\to \infty$ . 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 / \lambda_0^{NTK}$  and divergent training for larger learning rates. In contrast, we analyze optimization dynamics at training timescales  $t_\star$  that grow with width  $w$ . Specifically, the end of the early time transient period occurs at  $t_\star \sim \log (w)$ .
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 / \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 [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
(f)
Figure 2: 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$ .
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 \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 Fig. 2(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' or 'sharpness cutoff'.
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  $\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
(a)  $d = 8$
(d)  $w = 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 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, see Appendices A and C, respectively. The phase diagram of early training of FCNs with depth for three different widths trained on Fashion-MNIST with MSE loss using SGD.
(b)  $d = 7$
(e)  $w = 1024$
(c)  $d = 18$
(f)  $w = 2048$
such that the loss (sharpness) increases during the initial transient period:
$$
c _ {l o s s} = \min  _ {c} \{c \mid \max  _ {t \in [ 1, T _ {1} ]} \mathcal {L} _ {t} (c) > \mathcal {L} _ {0} (c) \}, \qquad c _ {s h a r p} = \min  _ {c} \{c \mid \max  _ {t \in [ 1, T _ {1} ]} \lambda_ {t} ^ {H} (c) > \lambda_ {0} ^ {H} (c) \},
$$
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]\}$ , 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  $\langle \cdot \rangle$ .
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  $\langle c_{loss}\rangle$ ,  $\langle c_{sharp}\rangle$ , and  $\langle c_{max}\rangle$  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 \approx 2$ . As  $1 / w$  increases,  $\langle c_{loss}\rangle$  increases and eventually converges to  $\langle c_{max}\rangle$  at large  $1 / w$ . By comparison, sharpness starts to catapult at relatively large learning rates at small  $1 / w$ , with  $\langle c_{sharp}\rangle$  continuing to increase with  $1 / w$  while remaining between  $\langle c_{loss}\rangle$  and  $\langle c_{max}\rangle$ . 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  $\langle c_{max}\rangle$  may increase or decrease with  $1 / w$  in different settings while consistently increasing with  $d$ , as shown in Appendices F and H.
<sup>5</sup>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.
(a) FCNs on MNIST
Figure 4: The relationship between critical constants for (a) FCNs, (b) CNNs, and (c) ResNets. Each data point corresponds to a run with varying depth, width, and initialization. The dashed line represents the  $y = x$  line.
(b) CNNs on Fashion-MNIST
(c) ResNets on CIFAR-10
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$  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} \leq c_{sharp} \leq 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  $\langle c_{loss}\rangle \gtrsim 4$  and  $\langle c_{max}\rangle$  consistently decreases with  $1 / w$ , while still satisfying  $c_{loss} \leq c_{sharp} \leq 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.  $(\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_{\text{barrier}}$  as follows.
Definition 3.  $(c_{\text{barrier}})$  Given the initial  $(\theta_0)$  and final parameters  $(\theta_{T_1})$ , we define  $c_{\text{barrier}}$  as the minimum learning rate constant such that there exists a barrier from  $\theta_0$  to  $\theta_{T_1}$ :  $c_{\text{barrier}} = \min_c \{c \mid U(c) > 0\}$ .
Here,  $c_{\text{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. 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_{\text{barrier}} = 5.65$ , while the loss starts to catapult at  $c_{\text{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  $\langle c_{\text{barrier}} \rangle$  and  $1/w$  for various models and datasets. We consistently observe that  $c_{\text{sharp}} \leq c_{\text{barrier}}$ , suggesting that training traverses a barrier only when sharpness starts to catapult during early training. Similar results