SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(b)
Figure 5: (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 I.2.
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_{\text{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  $\lambda_{\tau}^{H}$  at a time  $\tau$  in the middle of the intermediate saturation regime. We choose  $\tau$  so that  $c\tau \approx 200$ . For further details on sharpness measurement, see Appendix I.1. Fig. 5(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  $d$  and  $1/w$  increase ( $c < 2$  in Fig. 5(a)). (2) A crossover regime where  $\lambda_{\tau}^{H}$  is dropping significantly ( $2 < c < 2^3$  in Fig. 5(a)). (3) A saturation stage where  $\lambda_{\tau}^{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,  $\lambda_{\tau}^{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  $\lambda_{\tau}^{H}$ . To do this, we consider the averaged normalized sharpness over initializations and denote it by  $\langle \lambda_{\tau}^{H} / \lambda_{0}^{H} \rangle$ . The first two derivatives of the averaged normalized sharpness,  $\chi_{\tau} = -\frac{\partial}{\partial c} \langle \lambda_{\tau}^{H} / \lambda_{0}^{H} \rangle$  and  $\chi_{\tau}' = -\frac{\partial^2}{\partial c^2} \langle \lambda_{\tau}^{H} / \lambda_{0}^{H} \rangle$ , characterize the change in sharpness with learning rate. The extrema of  $\chi_{\tau}'$  quantitatively define the boundaries between the three stages described above. In particular, using the maximum of  $\chi_{\tau}'$ , we define  $\langle c_{crit} \rangle$ , which marks the beginning of the sharp decrease in  $\lambda_{\tau}^{H}$  with the learning rate.
${}^{6}$  time-step  $\tau  = {}^{200}/c$  is in the middle of regime (ii) for the models studied. Normalizing by  $c$  allows proper comparison for different learning rates.
(a) vanilla
(b) centered
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.
(c) last layer zero at init
Definition 4.  $(\langle c_{crit}\rangle)$  Given the averaged normalized sharpness  $\langle \lambda_{\tau}^{H} / \lambda_{0}^{H}\rangle$  measured at  $\tau$ , we define  $c_{crit}$  to be the learning rate constant that minimizes its second derivative:  $\langle c_{crit}\rangle = \arg \max_c\chi_\tau^\prime$
Here, we use  $\langle .\rangle$  to denote that the critical constant is obtained from the averaged normalized sharpness. Fig. 5(b, c) show  $\chi_{\tau}$  and  $\chi_{\tau}^{\prime}$  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  $\langle c_{crit}\rangle$  has slight fluctuations as  $d$  and  $1 / w$  are changed but generally stay in the vicinity of  $c = 2$ . The peak in  $\chi_{\tau}^{\prime}$  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  $\langle c_{crit}\rangle$ ,  $\langle c_{loss}\rangle$  increase with  $d$  and  $1 / w$ , implying the opening of the sharpness reduction phase  $\langle c_{crit}\rangle < c < \langle c_{loss}\rangle$  as  $d$  and  $1 / w$  increase. In Appendix F, we show that cross-entropy loss shows qualitatively similar results, but with  $2\lesssim \langle c_{crit}\rangle \lesssim 4$ .
4 Effect of network output at initialization on early training
Here we discuss the effect of network output  $f(x; \theta_t)$  at initialization on the early training dynamics.  $x$  is the input and  $\theta_t$  denotes the set of parameters at time  $t$ . We consider setting the network output to zero at initialization,  $f(x; \theta_0) = 0$ , by either (1) considering the "centered" network:  $f_c(x; \theta) = f(x; \theta) - f(x; \theta_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  $\langle c_{loss} \rangle$ , 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; \theta)$ , we define the scaled network as  $f_{s}(x; \theta) = \alpha f(x; \theta)$ , where  $\alpha$  is a scalar, fixed throughout training. In Appendix H, we show that a large (resp. small) value of  $\| f(x; \theta_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  $\langle c_{loss} \rangle$  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) = \frac{1}{\sqrt{w}} v^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) = 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;\theta_0)\|$  relative to the output labels. We minimize the loss using gradient descent (GD) with learning rate  $\eta$ . 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.
(c)
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$  ( $k = 2$ ) for reference. Each data point is an average over 500 random initializations.
The update equations of the  $uv$  model in function space can be written in terms of the trace of the Hessian  $\operatorname{tr}(H)$
$$
f _ {t + 1} = f _ {t} \left(1 - \eta \operatorname {t r} \left(H _ {t}\right) + \frac {\eta^ {2} f _ {t} ^ {2}}{w}\right), \quad \operatorname {t r} \left(H _ {t + 1}\right) = \operatorname {t r} \left(H _ {t}\right) + \frac {\eta f _ {t} ^ {2}}{w} (\eta \operatorname {t r} \left(H _ {t}\right) - 4). \tag {1}
$$
From the above equations, it is natural to scale the learning rate as  $\eta = k / \mathrm{tr}(H_0)$ . Note that  $c = \eta \lambda_0^H = k\lambda_0^H /\mathrm{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  $\langle .\rangle$  to denote an average over initialization. Figure 7(b) shows the phase diagram of early training, with  $\mathrm{tr}(H_t)$  replaced with  $\lambda_t^H$  as the measure of sharpness and with the learning rate scaled as  $\eta = k / \mathrm{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  $\mathrm{tr}(H)$  does not increase during training (see Equation 1). We also observe  $\langle k_{max}\rangle = 4$ , independent of width.
In Appendix B.3, we show that 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$  increases with  $1/w$ . Combined with  $\langle k_{crit} \rangle \approx 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  $\mathrm{tr}(H)$  does not increase for  $0 < k < 4$ . As  $\lambda_t^H$  is difficult to analyze, we consider the Frobenius norm  $\| H \|_F = \sqrt{\operatorname{tr}(H^T H)}$  as a proxy for sharpness. 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. In Appendix B.4, we show that the critical value of  $k$  for which  $\langle \| H_1 \|_F^2 - \| H_0 \|_F^2 \rangle > 0$  increases from  $\langle k_{loss} \rangle$  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  $\mathrm{tr}(H)$  and weight correlation  $\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  $\mathrm{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  $\langle k_{crit}\rangle < k < \langle k_{loss}\rangle$ , 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_{\text{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 \neq 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
Figure 8: Training trajectories of the  $uv$  model trained on  $(x, y) = (1, 0)$ , 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.
$y \neq 0$  in the dynamical equations results in loss catapult at  $k = 2$ , implying  $\langle k_{loss} \rangle \approx \langle k_{crit} \rangle \approx 2$ , irrespective of  $w$ .
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  $\eta = c / \lambda_0^H$ . 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  $\langle c_{loss}\rangle < c < \langle c_{barrier}\rangle$ , 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 \approx 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  $\eta_{loss} = 4 / \lambda_0^{NTK}$ . 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.
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.
[1] Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep networks. In International Conference on Learning Representations, 2019.
[2] Atish Agarwala, Fabian Pedregosa, and Jeffrey Pennington. Second-order regression models exhibit progressive sharpening to the edge of stability. *ArXiv*, abs/2210.04860, 2022.
[3] Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. In International Conference on Machine Learning, 2022.
[4] Alexander Atanasov, Blake Bordelon, Sabarish Sainathan, and Cengiz Pehlevan. The onset of variance-limited behavior for networks in the lazy and rich regimes. In *The Eleventh International Conference on Learning Representations*, 2023.
[5] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[6] Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks. ArXiv, abs/2304.03408, 2023.
[7] 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.
[8] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021.
[9] Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022.
[10] Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141-142, 2012.
[11] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[12] Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Gal Elidan, Kristian Kersting, and Alexander Ihler, editors, Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017. AUAI Press, 2017.
[13] 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.
[14] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
[68] Sho Yaida. Meta-principled family of hyperparameter scaling strategies. ArXiv, abs/2210.04909, 2022.
[69] Greg Yang and Edward J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11727-11737. PMLR, 18-24 Jul 2021.
[70] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub W. Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. In Neural Information Processing Systems, 2022.
[71] Hanxu Zhou, Qixu Zhou, Zhenyuan Jin, Tao Luo, Yaoyu Zhang, and Zhi-Qin John Xu. Empirical phase diagram for three-layer neural networks with infinite width. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[72] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for understanding neural network dynamics. ArXiv, abs/2205.11787, 2022.
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.
1. 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
$$
h _ {i} ^ {l + 1} = \gamma^ {l} \sum_ {j} ^ {w} W _ {i j} ^ {l} \phi \left(h _ {j} ^ {l}\right), \tag {2}
$$
where  $\phi (.)$  is the ReLU activation and  $\gamma^l$  is a constant. For NTP,  $\gamma^l = 2 / \sqrt{w}$  and the weights  $W^{l}$  are initialized using normal distribution, i.e.,  $W_{ij}^{l}\sim \mathcal{N}(0,1)$ . For SP,  $\gamma^l = 1$  and the weights  $W^{l}$  are initialized as  $W_{ij}^{l}\sim \mathcal{N}(0,2 / w)$ . For the last layer, we have  $\gamma^L = 1 / \sqrt{w}$  for NTP and  $W_{ij}^{L}\sim \mathcal{N}(0,1 / w)$  for SP.
For  $d / w \gtrsim 1 / 16$ , the dynamics is noisier, and it becomes challenging to separate the underlying deterministic dynamics from random fluctuations (see Appendix E).
2. CNNs: We considered Myrtle family ReLU CNNs [61] without any bias or normalization in Standard Parameterization (SP), initialized using He initialization [29]. The above model uses a fixed number of channels in each layer, which we refer to as the width of the network. In this case, the forward pass equations for the pre-activations from layer  $l$  to layer  $l + 1$  are given by
$$
h _ {i} ^ {l + 1} (\alpha) = \sum_ {j} ^ {w} \sum_ {\beta \in k e r} W _ {i j} ^ {l + 1} (\beta) \phi \left(h _ {i} ^ {l} (\alpha + \beta)\right), \tag {3}
$$
where  $\alpha, \beta$  label the spacial location. The weights are initialized as  $W_{ij}^{l}(\beta) \sim \mathcal{N}(0, 2 / k^{2}w)$ , where  $k$  is the filter size.
3. ResNets: We considered version 1 ResNet [29] implementations from Flax examples without Batch Norm or regularization. For ResNets, width corresponds to the number of channels in the first block. For example, the standard ResNet-18 has four blocks with widths  $[w, 2w, 4w, 8w]$ , with  $w = 64$ . We refer to  $w$  as the width or the widening factor. We considered ResNet-18 and ResNet-34.