SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
I _ {r} = \int_ {0} ^ {\infty} d R R ^ {2 w - 1} e ^ {- R ^ {2} / 2} \int_ {0} ^ {\pi / 2} d \theta \cos^ {w + 1} \theta \sin^ {w + 1} \theta \tag {31}
$$
$$
I _ {r} = \frac {\sqrt {\pi}}{2 ^ {3}} \frac {\Gamma (w) \Gamma \left(\frac {w + 2}{2}\right)}{\Gamma \left(\frac {w + 3}{2}\right)}, \tag {32}
$$
where  $\Gamma (.)$  denotes the Gamma function. The angular integral  $I_{\varphi}$  is
$$
I _ {\varphi} = \int d \varphi_ {1} d \varphi_ {2} \dots d \varphi_ {w - 1} \sin^ {w - 2} \varphi_ {1} \cos^ {2} \varphi_ {1} \sin^ {w - 3} \varphi_ {2} \dots \sin \varphi_ {w - 2} \tag {33}
$$
$$
I _ {\varphi} = \int d \varphi_ {1} d \varphi_ {2} \dots d \varphi_ {w - 1} \sin^ {w - 2} \varphi_ {1} \sin^ {w - 3} \varphi_ {2} \dots \sin \varphi_ {w - 2} \frac {\int_ {0} ^ {\pi} d \varphi_ {1} \sin^ {w - 2} \varphi_ {1} \cos^ {2} \varphi_ {1}}{\int_ {0} ^ {\pi} d \varphi_ {1} \sin^ {w - 2} \varphi_ {1}} \tag {34}
$$
$$
I _ {\varphi} = \frac {\pi^ {w / 2}}{\Gamma \left(\frac {w + 2}{2}\right)}. \tag {35}
$$
Plugging in Equations 32 and 35 into Equation 29, we obtain a very simple expression
$$
\left\langle \frac {f _ {0} ^ {2}}{T r (H _ {0}) ^ {2}} \right\rangle = \frac {w ^ {2}}{2 ^ {w + 3}} \frac {\sqrt {\pi} \Gamma (w)}{\Gamma (\frac {w + 2}{2}) \Gamma (\frac {w + 3}{2})} = \frac {w}{4 (w + 1)}. \tag {36}
$$
The other integral  $\left\langle \frac{f_0^4}{Tr(H_0)^4}\right\rangle$  can be obtained by generalizing the above approach as described below
$$
\left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle = w ^ {2} \int_ {- \infty} ^ {\infty} \prod_ {i = 1} ^ {w} \left(\frac {d v _ {i} d u _ {i}}{2 \pi}\right) \exp \left(- \frac {\| u \| ^ {2} + \| v \| ^ {2}}{2}\right) \frac {\sum_ {j , k , l , m = 1} ^ {w} u _ {j} v _ {j} u _ {k} v _ {k} u _ {l} v _ {l} u _ {m} v _ {m}}{\left(\| u \| ^ {2} + \| v \| ^ {2}\right) ^ {4}}. \tag {37}
$$
The integral is zero if either  $j = k$  and  $l = m$  or  $j = k = l = m$ , which we consider separately. Without loss of generality, we find the following integrals
$$
\left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle_ {2 2} = w ^ {2} \int_ {- \infty} ^ {\infty} \prod_ {i = 1} ^ {w} \left(\frac {d v _ {i} d u _ {i}}{2 \pi}\right) \exp \left(- \frac {\| u \| ^ {2} + \| v \| ^ {2}}{2}\right) \frac {u _ {1} ^ {2} u _ {2} ^ {2} v _ {1} ^ {2} v _ {2} ^ {2}}{\left(\| u \| ^ {2} + \| v \| ^ {2}\right) ^ {4}} \tag {38}
$$
$$
\left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle_ {4} = w ^ {2} \int_ {- \infty} ^ {\infty} \prod_ {i = 1} ^ {w} \left(\frac {d v _ {i} d u _ {i}}{2 \pi}\right) \exp \left(- \frac {\| u \| ^ {2} + \| v \| ^ {2}}{2}\right) \frac {u _ {1} ^ {4} v _ {1} ^ {4}}{\left(\| u \| ^ {2} + \| v \| ^ {2}\right) ^ {4}}, \tag {39}
$$
which have the following expressions
$$
\left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle_ {2 2} = \frac {w ^ {2}}{1 6} \frac {\Gamma (w)}{\Gamma (w + 4)} \tag {40}
$$
$$
\left\langle \frac {f _ {0} ^ {4}}{T r \left(H _ {0}\right) ^ {4}} \right\rangle_ {4} = \frac {9 w ^ {2}}{1 6} \frac {\Gamma (w)}{\Gamma (w + 4)}, \tag {41}
$$
(a)  $d = 4$
(b)  $d = 8$
Figure 14: Phase diagrams of FCNs in NTP with varying depths trained on the MNIST dataset using MSE loss.
(c)  $d = 16$
where  $\Gamma(.)$  denotes the gamma function. On combining the expressions with their multiplicities, we obtain the final result
$$
\left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle = 3 w (w - 1) \left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle_ {2 2} + w \left\langle \frac {f _ {0} ^ {4}}{T r (H _ {0}) ^ {4}} \right\rangle_ {4} \tag {42}
$$
$$
\left\langle \frac {f _ {0} ^ {4}}{T r \left(H _ {0}\right) ^ {4}} \right\rangle = \frac {3 (w + 2) w ^ {3}}{1 6} \frac {\Gamma (w)}{\Gamma (w + 4)} \tag {43}
$$
B.9 Insights into the catapult effect in crossentropy loss using  $uv$  model
In this section, we summarize the main intuition behind the discrepancy in the values of  $c_{loss}$  for cross-entropy loss at large widths. We consider the  $uv$  model trained on a classification task using logistic loss, as presented in [50].
Consider the  $uv$  model trained on a binary classification task using the logistic loss on two training examples  $(x_{1},y_{1}) = (1,1)$  and  $(x_{2},y_{2}) = (1, - 1)$ . Then, the total loss is  $\mathcal{L}(f) = \frac{1}{2}\log (2 + 2\cosh (f))$ . Hence, the loss grows monotonically as the output function  $|f|$  increases. The update equation of the function is given by:
$$
f _ {t + 1} = f _ {t} \left(1 - \frac {\eta \operatorname {t r} \left(H _ {t}\right) \mathcal {L} ^ {\prime} (f)}{f _ {t}} + \frac {\eta^ {2} \mathcal {L} ^ {\prime} \left(f _ {t}\right) ^ {2}}{n}\right), \tag {44}
$$
where  $\eta$  is the learning rate and  $\mathcal{L}'(.)$  is the derivative of the loss. At large width, if the condition  $|1 - \eta \operatorname{tr}(H)\mathcal{L}'(f) / f| < 1$  holds, then output function continues to decrease. Given that  $\mathcal{L}'(f) / f \leq 1/2$  in the above case, this decrease persists for  $\eta \operatorname{tr}(H) < 4$ . This result provides some intuition behind the discrepancy.
C Phase diagrams of early training
This section describes experimental details and shows additional phase diagrams of early training. The results include (1) FCNs in NTP trained on MNIST, Fashion-MNIST, and CIFAR-10 datasets, (2) CNNs in SP trained on Fashion-MNIST and CIFAR-10, and (3) ResNets in SP trained on CIFAR-10 datasets using MSE loss. Figures 14 to 19 show these results. The depths and widths are the same as specified in Appendix A. Each data point is an average over 10 initializations. The horizontal bars around the average data point indicate the region between  $25\%$  and  $75\%$  quantile. Phase diagrams for cross-entropy results are shown in Appendix F.
Additional experimental details : We train each model for  $t = 10$  steps using SGD with learning rates  $\eta = c / \lambda_0^H$  and batch size of 512, where  $c = 2^x$  with  $x \in \{0.0, \dots, x_{max}\}$  in steps of 0.1. Here,  $x_{max}$  is related to the maximum trainable learning rate constant as  $c_{max} = 2^{x_{max}}$ . We have considered 10 random initializations for each model. As mentioned in Appendix A, we do not
(b)  $d = 8$
Figure 15: Phase diagrams of FCNs in NTP with varying depths trained on the Fashion-MNIST dataset.
(a)  $d = 4$
(a)  $d = 5$
(b)  $d = 7$
Figure 16: Phase diagrams of FCNs in NTP with varying depths trained on the CIFAR-10 dataset.
(c)  $d = 10$
Figure 17: Phase diagrams of Convolutional Neural Networks (CNNs) in SP with varying depths trained on the Fashion-MNIST dataset.
(a)  $d = 5$
Figure 18: Phase diagrams of Convolutional Neural Networks (CNNs) in SP with varying depths trained on the CIFAR-10 dataset.
(c)  $d = 10$
(a)  $d = 18$
(b)  $d = 34$
Figure 19: Phase diagrams of Resnets in SP with different depths trained on the CIFAR-10 dataset.
(a)  $w = 512$
Figure 20: Phase diagrams of FCNs in NTP with varying widths trained on the CIFAR-10 dataset.
(b)  $w = 1024$
(c)  $w = 2048$
consider averages over SGD runs as the randomness from initialization outweighs it. Hence, we obtain 10 values for each of the critical values in the following results. For each initialization, we compute the critical constants using Definitions 1 and 3. To avoid a random increase in loss and sharpness due to fluctuations, we round off the values of  $\lambda_t^H / \lambda_0^H$  and  $\mathcal{L}_t / \mathcal{L}_0$  to their second decimal places before comparing with 1. We denote the average values using data points and variation using horizontal bars around the average data points, which indicate the region between  $25\%$  and  $75\%$  quantile. The smooth curves are obtained by fitting a two-degree polynomial  $y = a + bx + cx^2$  with  $x = 1/w$  and  $y$  can take on one of three values:  $c_{loss}$ ,  $c_{sharp}$  and  $c_{max}$ .
Phase diagrams with depth Figure 20: shows the phase diagrams with depth for FCNs in NTP trained on the CIFAR-10 dataset. The phase diagrams look qualitatively similar compared to the  $\frac{1}{w}$  phase diagrams.
D Relationship between various critical constants
Figure 21 illustrates the relationship between the early training critical constants for models and datasets. The experimental setup is the same as in Appendix C. Typically, we find that  $c_{loss} \leq c_{sharp} \leq c_{sharp}$  holds true. However, there are some exceptions, which are observed at high values of  $d / w$  (see 21 (d, e)), where the trends of the critical constants converge, and large fluctuations can cause deviations from the inequality.
E The effect of  $d / w$  on the noise in dynamics
In this section, we demonstrate that for FCNs with  $d / w \gtrsim 1 / 16$ , the dynamics becomes noise-dominated. This aspect makes it challenging to distinguish the underlying deterministic dynamics from random fluctuations. To demonstrate this, we consider FCNs trained on CIFAR-10 using MSE and cross-entropy loss and use 4096 training examples for estimating sharpness.
(b) FCNs on Fashion-MNIST
(a) FCNs on MNIST
(d) CNNs on Fashion-MNIST
(e) CNNs on CIFAR-10
(c) FCNs on CIFAR-10
(f) ResNets on CIFAR-10
Figure 21: The relationship between various critical constants for various models and datasets. Each data point corresponds to a model with random initialization. The dashed line denotes the values where  $x = y$ .
Figure 22: Training loss trajectories of ReLU FCNs with  $d = 16$  trained on the CIFAR-10 dataset with MSE loss using SGD with learning rate  $\eta = c / \lambda_0^H$  and batch size  $B = 512$ .
Figures 22 and 23 show the training loss and sharpness of FCNs with  $d = 16$  and varying widths, trained on CIFAR-10 using MSE loss. We observe that the sharpness dynamics becomes noisier for  $w \lesssim 64$ .
Figures 24 and 25 shows the training dynamics with loss switched to cross-entropy, while keeping the initialization and SGD batch sequence the same as in the MSE loss case. In comparison to MSE
Figure 23: Sharpness trajectories of ReLU FCNs with  $d = 16$  trained on the CIFAR-10 dataset with MSE loss using SGD with learning rate  $\eta = c / \lambda_0^H$  and batch size  $B = 512$ .
Figure 24: Training loss trajectories of ReLU FCNs with  $d = 16$  trained on the CIFAR-10 dataset with cross-entropy loss using SGD with learning rate  $\eta = c / \lambda_0^H$  and batch size  $B = 512$ .
loss, the training loss and sharpness dynamics show a higher level of noise, especially for  $w \lesssim 256$ . As a result, it becomes difficult to characterize the training dynamics for  $d / w \gtrsim 1 / 16$ .
F Crossentropy
In this section, we provide additional results for models trained with cross-entropy (xent) loss and compare them with MSE results. Broadly speaking, models trained with cross-entropy loss show similar characteristics to those trained with MSE loss, such as, (i) sharpness reduction during early training, (ii) an increase in critical constants  $c_{loss}$ ,  $c_{sharp}$  with  $d$  and  $1/w$ , (iii)  $c_{loss} \leq c_{sharp} \leq c_{max}$ . However, the dynamics of models trained with cross-entropy loss is noisier compared to MSE as shown in the previous section, and characterizing these dynamics can be more complex. In the
Figure 25: Sharpness trajectories of ReLU FCNs with  $d = 16$  trained on the CIFAR-10 dataset with cross-entropy loss using SGD with learning rate  $\eta = c / \lambda_0^H$  and batch size  $B = 512$ .
(a) MSE,  $d = 4$
(d) xent,  $d = 4$
Figure 26: The phase diagrams of early training of FCNs trained on the CIFAR-10 dataset using (a, b, c) MSE and (d, e, f) cross-entropy loss. Each data point is an average over 10 initializations, and solid lines represent a smooth curve fitted to raw data points. The horizontal bars around the averaged data point indicates the region between  $25\%$  and  $75\%$  quantile. For cross-entropy phase diagrams, the  $c = 2$  line is shown for reference only and does not relate to  $c_{crit}$ .
(e) xent,  $d = 8$
(f) xent,  $d = 16$
following experiments, we consider models trained on the CIFAR-10 dataset and used 4096 training examples to estimate sharpness.
F.1 Phase diagrams
Figure 26 compares the phase diagrams of FCNs in SP trained on the CIFAR-10 dataset, using both MSE and cross-entropy loss. The estimated critical constants for cross-entropy loss are generally
(b)
Figure 27: Comparison of the relationship between critical constants for FCNs in SP trained on CIFAR-10 using MSE and cross-entropy loss. Each data point corresponds to a randomly initialized model with depths and widths mentioned in Appendix A.
more noisy, as quantified by the confidence intervals. In comparison to phase diagrams of models trained with MSE loss, we observe a few notable differences. First, the loss starts to catapult at a value appreciably larger than  $c = 2$  at large widths. Primarily,  $4 \lesssim c_{loss} \lesssim 8$ . Additionally,  $c_{max}$  generally decreases with  $1/w$ . This decreasing trend becomes less sharp at large depths.
Despite these differences, the phase diagrams for both loss functions share various similarities. First, we observe sharpness reduces during early training for  $c < c_{\text{sharp}}$  (see the first row of Figure 25). Next, we observe that the inequality  $c_{\text{loss}} \leq c_{\text{sharp}} \leq c_{\text{max}}$  generally holds for both loss functions as demonstrated in Figure 27, barring some exceptions.
Figure 28 shows the phase diagrams for CNNs and ResNets trained on the CIFAR-10 dataset using cross-entropy loss. The observed critical constants are much noisier as quantified by the confidence intervals. Nevertheless, the phase diagram shows similar trends as mentioned above. For large  $\frac{1}{w}$  models, we found that progressive sharpening begins after 5 - 10 training steps. For these cases, we only use the first 5 steps to measure sharpness to avoid progressive sharpening. For CNNs, we observed that the dynamics becomes difficult to characterize for  $w \lesssim 32$  and  $d \gtrsim 10$ , due to large fluctuations. Consequently, we've opted not to include these particular results.
(a) CNN,  $d = 5$
Figure 28: Phase diagrams of (a, b) CNNs and (c) ResNets trained on the CIFAR-10 dataset with cross-entropy loss using SGD with  $\eta = c / \lambda_0^H$  and  $B = 512$ .
(b) CNN,  $d = 7$
(c) ResNet,  $d = 18$
F.2 Intermediate saturation regime
Figure 29 shows the normalized sharpness measured at  $c\tau = 100$  for FCNs trained on CIFAR-10 using cross-entropy loss. Similar to MSE loss, we observe an abrupt drop in sharpness at large learning rates. However, this abrupt drop occurs at  $2 \lesssim c_{crit} \lesssim 4$ . The estimated sharpness is noisier (compare with Figure 38), which hinders a reliable estimation of  $c_{crit}$ . We speculate that we require a
The time step  $\tau = 100 / c$  is in the middle of the intermediate saturation regime for most of the models. For further details on estimating sharpness, see Appendix I.1.
(b)  $d = 8$
Figure 29: Sharpness measured at  $c\tau = 100$  against the learning rate constant for FCNs trained on the CIFAR-10 dataset using cross-entropy loss, 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 dashed line shows  $c = 2$  for reference.
large number of averages for a reliable estimation of  $c_{crit}$  for cross-entropy loss. We leave the precise characterization of  $c_{crit}$  for cross-entropy loss for future work.
G The effect of setting model output to zero at initialization
In this section, we demonstrate the effect of network output  $f(x; \theta_t)$  at initialization on the early training dynamics. In particular, we set the network output to zero at initialization,  $f(x; \theta_0) = 0$ , by (1) 'centering' the network by its initial value  $f_c(x; \theta_t) = f(x; \theta) - f(x; \theta_0)$  or (2) setting the last layer weights to zero at initialization. We show that both (1) and (2) remove the opening of the sharpness reduction phase with  $1/w$ . Resultantly, the average onset of loss catapult occurs at  $c_{loss} \approx 2$ , independent of depth and width.
Throughout this section, we use 'vanilla' networks to refer to networks initialized in the standard way. For simplicity, we train FCNs using full batch gradient descent with MSE loss using a subset consisting of 4096 examples of the CIFAR-10 dataset.
G.1 The effect of centering networks
Given a network function  $f(x; \theta_t)$ , we define the centered network  $f_c(x; \theta_t)$  as
$$
f _ {c} (x; \theta_ {t}) = f (x; \theta_ {t}) - f (x; \theta_ {0}), \tag {45}
$$
where  $f(x; \theta_0)$  is the network output at initialization. By construction, the network output is zero at initialization. It is noteworthy that centering a network is an unusual way of training deep networks as it doubles the cost of training because of two forward passes.
Figure 30 compares the training loss and sharpness dynamics of vanilla networks and centered networks. Unlike vanilla networks, we do not observe a decrease in sharpness for  $c < c_{loss}$  during early training. Rather, we observe a slight increase in sharpness. To distinguish this slight increase from sharpness catapult, we introduce a threshold  $\epsilon$ , comparing normalized sharpness  $\lambda_t^H / \lambda_0^H$  with  $1 + \epsilon$ , to define a sharpness catapult. As demonstrated in Appendix G.3, the  $uv$  model trained on a single training example  $(x, y)$  with  $y \neq 0$  sheds lights on this initial increase in sharpness.
Interestingly, irrespective of depth and width, we observe that loss catapults at  $c_{loss} \approx 2$ , as demonstrated in the phase diagrams in Figure 31(a, b, c). These findings suggest a strong correlation between a large network output at initialization  $\| f(x; \theta_0) \|$  and the opening of the sharpness reduction phase discussed in Section 2.
G.2 The effect of setting the last layer to zero
An alternative way to train networks with  $f(x; \theta_0) = 0$  is by setting the last layer to zero at initialization. The principle of criticality at initialization [55, 58, 68] does not put any constraints on
In experiments, we set  $\epsilon = 0.05$ . We use the same threshold for zero-init networks.
(f)
Figure 30: Comparison of the early training dynamics of (a, b) vanilla, (c, d) centered, and (e, f) zero-initialized FCNs (with depth  $= 8$  and width  $= 512$ ), trained on the CIFAR-10 dataset with MSE loss using gradient descent for 20 steps.
the last layer weights. Hence, setting the last layer to zero does not affect signal/gradient propagation at initialization. Yet, setting the last layer to zero results in initialization in a flat curvature region at initialization, resulting in access to larger learning rates. We refer to these networks as 'zero-init' networks.
Figure 30 compares the training dynamics of zero-init networks with vanilla and centered networks. We observe that the dynamics is quite similar to the centered networks: (i) sharpness does not reduce for small learning rates and (ii) loss catapults  $c_{loss} \approx 2$ , irrespective of depth and width. Figure 31(d, e, f) show the phase diagrams of networks with zero-initialized networks. Like centered networks, the critical constants do not scale with depth and width. Again, suggesting that a large network output at initialization  $\| f(x; \theta_0) \|$  is related to the opening of the sharpness reduction phase in the early training results shown in Section 2.
(a) centered,  $d = 4$
(d) zero-init,  $d = 4$
Figure 31: The phase diagrams of early training dynamics of (a, b, c) centered and (d, e, f) zero-init networks trained on CIFAR-10 using MSE 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.
(b) centered,  $d = 8$
(e) zero-init,  $d = 8$
(c) centered,  $d = 16$
(f) zero-init,  $d = 16$
G.3 Insights from  $uv$  model trained on  $(x, y)$
In this section, we gain insights into the effect of setting network output to zero at initialization using  $uv$  model trained on an example  $(x,y)$ . In particular, we show that loss catapults at  $k_{loss} = 2$  and sharpness increases during early training.
Consider the  $uv$  model trained on a single training example  $(x,y)$  with  $y\neq 0$
$$
f (x) = \frac {1}{\sqrt {w}} \sum_ {i} ^ {w} u _ {i} v _ {i} x.
$$
This simplifies the loss function to
$$
\mathcal {L} = \frac {1}{2} (f (x) - y) ^ {2} = \frac {1}{2} \Delta f ^ {2}, \tag {46}
$$
where  $\Delta f$  is the residual. The trace of the Hessian  $\operatorname{tr}(H)$  is
$$
\operatorname {t r} (H) = \frac {x ^ {2}}{w} \left(\| v \| ^ {2} + \| u \| ^ {2}\right). \tag {47}
$$
The Frobenius norm can be written in terms of the trace and the network output
$$
\| H \| _ {F} ^ {2} = \lambda^ {2} + 2 x ^ {2} \Delta f ^ {2} \left(1 + \frac {2 f}{w \Delta f}\right). \tag {48}
$$
The function and residual updates are given by
Note that for  $y = 0$ , the network is already at a minimum for  $f_0 = 0$ .
(f)
Figure 32: The early training dynamics of FCNs with a fixed output scale trained on the CIFAR-10 dataset with MSE loss using gradient descent.
$$
f _ {t + 1} = f _ {t} - \eta \operatorname {t r} \left(H _ {t}\right) + \frac {\eta^ {2} x ^ {2}}{w} f _ {t} \Delta f _ {t} ^ {2} \tag {49}
$$
$$
\Delta f _ {t + 1} = \Delta f _ {t} \left(1 - \eta \operatorname {t r} \left(H _ {t}\right) + \frac {\eta^ {2} x ^ {2}}{w} f _ {t} \Delta f _ {t}\right). \tag {50}
$$
Similarly, we can obtain the trace update equations
$$
\operatorname {t r} \left(H _ {t + 1}\right) = \operatorname {t r} \left(H _ {t}\right) + \frac {\eta \Delta f _ {t} ^ {2} x ^ {2}}{w} \left(\eta \operatorname {t r} \left(H _ {t}\right) - 4 \frac {f _ {t}}{\Delta f _ {t}}\right). \tag {51}
$$
Let us analyze them for the networks with zero output at initialization. The loss at the first step increases if
$$
\left\langle \frac {L _ {1}}{L _ {0}} \right\rangle = \left\langle \left(1 - \eta \operatorname {t r} \left(H _ {0}\right) + \frac {\eta^ {2} x ^ {2}}{n} f _ {0} \Delta f _ {0}\right) ^ {2} \right\rangle > 1 \tag {53}
$$
Setting  $f_0 = 0$  and scaling the learning rate as  $\eta = k / \operatorname{tr}(H_0)$ , we see that the loss increases at the first step if  $k > 2$ .
$$
\left\langle \frac {L _ {1}}{L _ {0}} \right\rangle = \left\langle (1 - k) ^ {2} \right\rangle > 1 \tag {54}
$$
Next, we analyze the change in trace during the first training step. Setting  $f_0 = 0$ , we observe that the trace increases for all learning rates
$$
\operatorname {t r} \left(H _ {1}\right) = \operatorname {t r} \left(H _ {0}\right) + \frac {\eta^ {2} x ^ {2}}{w} \Delta f _ {0} ^ {2} \operatorname {t r} \left(H _ {0}\right), \tag {55}
$$
modulated by the learning rate and width. Finally, we analyze the change in Frobenius norm in the first training step at  $k = k_{loss}$ , which implies  $\Delta f_1^2 = \Delta f_0^2$
$$
\left\langle \Delta \| H _ {1} \| ^ {2} \right\rangle = \left\langle \operatorname {t r} \left(H _ {1}\right) ^ {2} - \operatorname {t r} \left(H _ {0}\right) ^ {2} + 2 x ^ {2} \left(\Delta f _ {1} ^ {2} - \Delta f _ {0} ^ {2}\right) \right\rangle . \tag {56}
$$
As  $\operatorname{tr}(H)$  increases in the first training step,  $\| H\| _F$  also increases in the first training step.
H The effect of output scale on the training dynamics
Given a neural network function  $f(x)$  with depth  $d$  and width  $w$ , we define the scaled network as  $f_{s}(x) = \alpha f(x)$ , where  $\alpha$  is referred to as the output scale. In this section, we empirically study the impact of the output scale on the early training dynamics. In particular, we show that a large (resp. small) value of  $\| f(x; \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. For simplicity, we train FCNs using gradient descent with MSE loss using a subset consisting of 4096 examples of the CIFAR-10 dataset, as in the previous section.