SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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(\theta)$ , we define the 'scaled network' as
$$
f _ {s} (\theta) = \frac {s f (\theta)}{\| f \left(\theta_ {0}\right) \|}, \tag {57}
$$
where  $s$  is a scalar, fixed throughout training. By construction, the network output norm  $\| f_s(\theta_0)\|$  equals  $s$ . For standard initialization,  $s = \| f(\theta_0)\| = \mathcal{O}(\sqrt{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 \gtrsim 0.5$ ) sharpness decreases during early training, while sharpness increases for small output scales  $^{10}$ . Furthermore, the training dynamics tends to
<sup>10</sup>We empirically observed that sharpness reduces for output scales as small as  $s \sim 0.5$ , which is relatively small compared to  $\sqrt{k}$ .
(f)
Figure 33: The phase diagrams of early training dynamics for ReLU FCNs with fixed output scale trained on a subset of the CIFAR-10 dataset using MSE loss using gradient descent. Each data point is an average over 10 initializations. The horizontal bars around the average data point indicate the region between  $25\%$  and  $75\%$  quantile.
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 primarily 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 discussed 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  $\alpha = w^{-\sigma}$ , which is commonly used in the literature [19, 6, 4]. We consider three distinct  $\sigma$  values  $\{-0.5, 0.0, 0.5\}$ , where  $\sigma = -0.5$  represents the lazy regime,  $\sigma = 0.5$  corresponds to feature learning (rich) regime and  $\sigma = 0.0$  corresponds to standard (vanilla) initialization.
Figure 34 shows the training loss and sharpness trajectories of FCNs trained on for different  $\sigma$  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(\mathrm{e},\mathrm{f}))$ . These results are in parallel to the fixed output scale networks studied in the pervious section.
(f)
Figure 34: The early training dynamics of FCNs with output scale  $\alpha = w^{-\sigma}$  trained on the CIFAR-10 dataset with MSE loss using gradient descent.
Figure 35 summarizes the early training dynamics of FCNs with different  $\sigma$  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  $\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 learning rates  $c = 2^x$ , where  $x \in [-1.0, 4.0]$  in steps of 0.1. The value 200 is chosen such that  $\tau$  is in the middle of the intermediate saturation regime. Next, we measure sharpness over a range of steps
(a)
(d)
Figure 35: The phase diagrams of early training dynamics for ReLU FCNs with varying depths and output scale.
$t \in [\tau - 5, \tau + 5]$  and average over  $t$  to reduce fluctuations. We repeat this process for various initializations and obtain the average sharpness.
(a)  $d = 4$
Figure 36: Sharpness measured at  $c\tau = 200$  against the learning rate constant for FCNs trained on the MNIST dataset, with varying depths and widths. Each curve is an average over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote  $c_{crit}$  estimated using the maximum of  $\chi_{\tau}^{\prime}$ .
I.2 Estimating the critical constant  $c_{crit}$
This subsection explains how to estimate  $c_{crit}$  from sharpness measured at time  $\tau$ . 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.
Figure 37: Sharpness measured at  $c\tau = 200$  against the learning rate constant for FCNs trained on the Fashion-MNIST dataset, with varying depths and widths. Each curve is an average over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote  $c_{crit}$  estimated using the maximum of  $\chi_{\tau}^{\prime}$ .
(a)  $d = 4$
(a)  $d = 5$
Figure 39: Sharpness measured at  $c\tau = 200$  against the learning rate constant for Myrtle-CNNs trained on the CIFAR-10 dataset, with varying depths and widths. Each curve is an average of over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote  $c_{crit}$  estimated using the maximum of  $\chi_{\tau}^{\prime}$ .
Figure 38: Sharpness measured at  $c\tau = 200$  against the learning rate constant for FCNs trained on the CIFAR-10 dataset, with varying depths and widths. Each curve is an average over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote  $c_{crit}$  estimated using the maximum of  $\chi_{\tau}^{\prime}$ .
(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
(b)  $d = 34$
Figure 40: Sharpness measured at  $c\tau = 200$  against the learning rate constant for ResNets trained on the CIFAR-10 dataset, with varying depths and widths. Each curve is an average of over ten initializations, where the shaded region depicts the standard deviation around the mean trend. The vertical lines denote  $c_{crit}$  estimated using the maximum of  $\chi_{\tau}^{\prime}$ .
(c)  $d = 50$
at relatively smaller values of  $c$ . This is reflected in Fig. 41 by  $\langle c_{\text{sharp}} \rangle$  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.
(a)  $B = 128$
Figure 41: The phase diagram of early training for FCNs with  $d = 4$  trained on the CIFAR-10 dataset with MSE loss using SGD with different batch sizes.
(b)  $B = 512$
J.2 The intermediate saturation regime
(a)
Figure 42: (a) Normalized sharpness measured at  $c\tau = 200$  against the learning rate constant for FCNs with  $d = 4$  trained on the CIFAR-10 dataset, with varying widths. Each data point is an average over 10 initializations, where the shaded region depicts the standard deviation around the mean trend. (b, c) Smooth estimations of the first two derivatives,  $\chi_{\tau}$  and  $\chi_{\tau}^{\prime}$ , of the averaged normalized sharpness wrt the learning rate constant.
Figure 42 shows the normalized sharpness, measured at  $c\tau = 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  $\chi$  and  $\chi^{\prime}$  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_{\mathrm{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).
(b)  $d = 8$
Figure 43: The phase diagram of early training for FCNs with bias trained on the CIFAR-10 dataset with MSE loss using SGD with different depths.