SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We used the same model and momentum parameters as in the translation task described in Section 4.1 but generated random integers as targets. Similar to the setting on real targets, the model was trained using constant learning rates (0.015 for  $\mathrm{SGD + M}$  and 5e-5 for Adam) for 60 epochs (1800 iterations), and we repeated the experiments for 8 times starting from the same initialization. Figure 3(right) shows the training losses for one among them. Table 2b shows the averaged  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$ ,  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  and  $\frac{R_{\mathrm{med}}^{\mathrm{SGDM}}(t)}{R_{\mathrm{med}}^{\mathrm{Adam}}(t)}$  (with standard deviation in the brackets) of the same 10 layers as in Table 2a.
4.3 Summarization of the empirical results and discussion
Overall, through extensive experiments on language models, we demonstrate that starting from the same initialization, for the vast majority of layers, the  $R_{\mathbf{med}}^{\mathrm{OPT}}(t)$  values found by Adam are smaller than those found by SGD+M. This suggests that Adam biases the trajectory towards a region with more uniform Hessian diagonal than SGD+M. In Appendix A.10 we also validate this observation on the in-distribution test data.
Contribution of uniform Hessian diagonal to fast convergence. We observe that on dataset with random targets,  $\mathrm{SGD + M}$  plateaus after about 400 steps and thus converges much slower when compared to Adam than on real dataset (see Figure 3). On the other hand, the gaps of  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  and
To prevent  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  from getting too large due to tiny median, we added an additional term  $0.001\max \{|H_{ii}^{(t)}|\}_{i = 1}^{d}$  to the denominator of eq. (1) when computing.
$R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  are more significant on random data than on real data (see Table 2a and Table 2b) as well. In Appendix A.4, we conduct another experiment where we switch from SGD to Adam in the middle and compare it with the model trained by Adam from the beginning. The observation is that both the loss gap and the gap of  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  are gradually closed after switching (see Figure 7 and Table 8). Hence we find a positive correlation between fast convergence and uniform Hessian diagonal (small  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ ). In Appendix A we study other adaptive algorithms (Adagrad, RMSprop and AMSGrad) and get similar observation: all these adaptive methods converge faster than SGD or  $\mathrm{SGD + M}$  and also bias the trajectory to regions with smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ , suggesting the universality of our observation.
Despite the above positive correlation, it is reasonable to ask whether small  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  indeed contributes to fast convergence or is just a byproduct of adaptive methods. To address this concern, in Appendix B.1, we add a supplementary experiment similar to that in Figure 1. We select two iterates  $x_{1}$  and  $x_{2}$  from two trajectories that both come from SGD+M (instead of one from Adam and one from SGD+M in Figure 1), such that the loss  $f(x_{1}) = f(x_{2})$  but  $x_{2}$  has a smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  than  $x_{1}$ . We then run SGD+M with the same configuration twice, once from  $x_{1}$  and once from  $x_{2}$ . Under this setting, we get similar observation as before: running SGD+M from  $x_{2}$  (with smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ ) achieves faster convergence than from  $x_{1}$ . This suggests that the uniformity of the diagonal of loss Hessian (measured by  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ ) reveals some intrinsic trajectory property beyond the algorithm choice and is indeed a contributing factor to fast optimization. In Appendix B.2 we theoretically prove the contribution of small  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  to fast optimization in a simplified setting.
More discussions on the trajectory difference. Considering the fact that our comparison between  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  is conditioned on the same iteration when  $\mathrm{SGD} + \mathbf{M}$  has larger training loss than Adam, there is a potential alternative explanation of the Hessian diagonal uniformity. That is, the global minimum has uniform Hessian, and Adam simply converges faster to it than  $\mathrm{SGD} + \mathbf{M}$ . To rule out this possibility, in Appendix A.3 we add a comparison of our measurements  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t')$ , where  $t, t'$  are picked such that  $t$ th Adam iterate and  $t'$ th  $\mathrm{SGD} + \mathbf{M}$  iterate have the same training loss. The results (in Table 7) show that  $R_{\mathrm{med}}^{\mathrm{Adam}}(t) < R_{\mathrm{med}}^{\mathrm{SGDM}}(t')$  for most layers, thus demonstrating that the trajectories of Adam and  $\mathrm{SGD} + \mathbf{M}$  are truly different and that the difference is because Adam biases the local geometry (as opposed to faster convergence).
Adding regularization. People in practice usually add weight decay (equivalent to  $l_{2}$  regularization) to encourage better generalization ability. In Appendix A.7 we compare SGD+M and Adam when both using small weight decay values (0.001). The results in Figure 13a and Table 9 suggest that in this case, our observation still holds: Adam converges faster than SGD+M and in the vast majority of layers,  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  values are smaller than  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$ . This reveals the robustness of our observation under weak regularization. However, under large weight decay parameters, we observed cases where Adam still converged faster but  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  values were larger rather than smaller. In the case of strong regularization, the adaptivity of Adam requires further exploration and we hope to find new mechanisms in the future.
In Section 4, we empirically demonstrate the uniformity of diagonal geometry. In this section, we theoretically analyze this property for large batch Adam and  $\mathrm{SGD + M}$  on a two-layer linear network with 1-dimensional output. Although simple, the choice of 2-layer linear networks to understand learning dynamics is common in prior works (e.g. [TCG21]). Moreover, in language transformer models when the weights are small, the softmax in the key-value-query structure is near the linear regime. Then this structure might be approximated by the product of 3 linear operators, similar to a three-layer linear network. Hence the theoretical analysis of this phenomenon on linear networks would be a good starting point for further understanding of more complicated language models.
5.1 Problem setup
Notation Let  $[d] = \{1,2,\dots,d\}$ . We use  $\| \cdot \|_2$  to denote the  $l_2$  norm of a vector, and  $\| \cdot \|_F$  to denote the Frobenius norm of a matrix. Let  $\langle \cdot, \cdot \rangle$  be the Euclidean inner product between vectors or matrices. Let  $\mathcal{N}(\mu, \sigma^2)$  be the one-dimensional Gaussian distribution with mean  $\mu$  and variance  $\sigma^2$ . For a scalar (vector, matrix)  $A$  which evolves over time, we use  $A^{(t)}$  to denote its value at time  $t$ .
Let there be  $m$  data points. The data matrix is  $X \in \mathbb{R}^{d_x \times m}$  and the label matrix is  $Y \in \mathbb{R}^{d_y \times m}$ . We assume that the input dataset is whitened, i.e.  $\Lambda_{xx} \coloneqq \frac{1}{m} XX^T \in \mathbb{R}^{d_x \times d_x}$  is an identity matrix.
The parameters of a 2-layer linear network are given by  $W \coloneqq (W_2, W_1)$ . Assume  $W_i \in \mathbb{R}^{d_i \times d_{i-1}}$  for  $i = 1, 2$ . We have  $d_2 = d_y$ ,  $d_0 = d_x$ . We consider the square loss  $L(W) \coloneqq \frac{1}{2m} \| W_2 W_1 X - Y \|_F^2$ .
Denote  $A \coloneqq \frac{1}{m} YX^T \in \mathbb{R}^{d_y \times d_x}$ . [AGCH19] show that with whitened dataset,
$$
L (W) := \frac {1}{2 m} \| W _ {2} W _ {1} X - Y \| _ {F} ^ {2} = \bar {L} (W) + c, \quad \bar {L} (W) := \frac {1}{2} \| W _ {2} W _ {1} - A \| _ {F} ^ {2}. \tag {2}
$$
where  $c$  does not depend on  $W$ . We consider the following model with small Gaussian initialization.
Assumption 1 (Setup). The input covariance  $\Lambda_{xx} \coloneqq \frac{1}{m} XX^T \in \mathbb{R}^{d_x \times d_x}$  is an identity matrix. The input and hidden layers are both of dimension  $d$ , i.e.  $d_1 = d_0 = d$ . Without loss of generality, we can assume that  $A$  is a row vector (i.e.  $d_2 = 1$ ) whose coordinates are positive and  $\Theta(1)$  in terms of  $d$ .
Assumption 2 (Gaussian Initialization).  $\forall i,j:w_{2i}^{(0)}\sim \mathcal{N}(0,\frac{1}{d^{2\alpha}}),W_1^{(0)}[i,j]\sim \mathcal{N}(0,\frac{1}{d^{4\alpha}})$  are independently initialized with sufficiently large  $\alpha >0$
Denote  $\tilde{A}$  and  $\tilde{\Lambda}_{xx}$  as the batch versions of  $A$  and  $\Lambda_{xx}$ . We make the following large-batch assumption. We emphasize that large batches are commonly used in NLP tasks (e.g. [BMR $^{+}$ 20]).
Assumption 3 (Large Batch). For the randomly selected batches, assume  $\mathbb{E}[\tilde{A}] = A$ ,  $\mathbb{E}[\tilde{\Lambda}_{xx}] = \Lambda_{xx}$ .  $\forall i,j\in [d]:\mathbb{E}\left[(\tilde{A}_i - A_i)^2\right]\leq \sigma^2$ $\mathbb{E}\left[(\tilde{\Lambda}_{xx}[i,j] - \Lambda_{xx}[i,j])^2\right]\leq \sigma^2$ , and  $\sigma^2 = \mathcal{O}\left(\frac{1}{poly(d)}\right)$ .
Denote  $\tilde{g}^{(t)}$  as the batch gradient at time  $t$ . The update rules of  $\mathrm{SGD + M}$  and Adam are given by
$$
\begin{array}{l} \mathbf {S G D} + \mathbf {M}: u ^ {(t + 1)} = \beta u ^ {(t)} + \tilde {g} ^ {(t)}, W ^ {(t + 1)} = W ^ {(t)} - \eta u ^ {(t)}, \\ \text {A d a m :} \quad \eta_ {t} = \eta \cdot \sqrt {1 - \beta_ {2} ^ {t + 1}} / \left(1 - \beta_ {1} ^ {t + 1}\right), \quad m ^ {(t + 1)} = \beta_ {1} m ^ {(t)} + \left(1 - \beta_ {1}\right) \tilde {g} ^ {(t)}, \tag {3} \\ v ^ {(t + 1)} = \beta_ {2} v ^ {(t)} + (1 - \beta_ {2}) \tilde {g} ^ {(t)} \odot \tilde {g} ^ {(t)}, \quad W ^ {(t + 1)} = W ^ {(t)} - \eta_ {t} \frac {m ^ {(t)}}{\sqrt {v ^ {(t)}} + \xi}, \\ \end{array}
$$
where  $\eta$  is the learning rate,  $\beta, \beta_1, \beta_2$  are momentum parameters, and  $\xi$  is for numerical stability. All operations on vectors are element-wise. Here and throughout, the notation  $f(x) = \mathcal{O}(g(x))$  (resp.  $f(x) = \Omega(g(x)), f(x) = \Theta(g(x)))$  means that  $\exists C_1, C_2 > 0$  such that  $f(x) \leq C_2 g(x)$  (resp.  $f(x) \geq C_1 g(x), C_1 g(x) \leq f(x) \leq C_2 g(x)$ ). Here  $C_1, C_2$  may depend on  $\beta, \beta_1, \beta_2$ . We will also use the notation with  $\sim$ , i.e.  $\tilde{\mathcal{O}}(\cdot), \tilde{\Omega}(\cdot), \tilde{\Theta}(\cdot)$  to hide factors of order  $\log d$ . In our theoretical analysis, "with high probability", or "w.h.p." for short, means that with probability at least  $1 - \frac{1}{\mathrm{poly}(d)}$ .
Since the weights and Hessians in different layers may have different magnitudes, we compute the  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  layer by layer. Denote  $R_{\mathrm{med},k}^{\mathrm{SGDM}}(t)$  (resp.  $R_{\mathrm{med},k}^{\mathrm{Adam}}(t)$ ) as the  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  found by  $\mathrm{SGD + M}$  (resp. Adam) w.r.t.  $W_{k}$  at time  $t$  where  $k = 1,2$ .
5.2 Main results
Theorem 1. Under Assumption 1, 2 and 3, consider the weights  $\left\{W_{SGD}^{(t)}\right\}_{t\geq 0}$  (resp.  $\left\{W_{Adam}^{(t)}\right\}_{t\geq 0}$ ) obtained by  $SGD + M$  (resp. Adam) defined in (3).
1. For any  $p > 0$ , pick  $0 < \epsilon < \frac{1}{d^p}$ ,  $\eta \leq \mathcal{O}\left(\frac{\epsilon}{d^{7\alpha / 4 + 4}}\right)$  and  $\alpha \geq 4(p + 2)$ . Suppose  $\sigma \leq \frac{\eta^{3 / 2}}{d^{\alpha / 2 + 1}}$ , then there exists  $T_{SGD,1}, T_{SGD,2}$  such that w.h.p.,  $\bar{L}\left(W_{SGD}^{(T_{SGD,1})}\right) = \Theta (d)$ ,  $\bar{L}\left(W_{SGD}^{(T_{SGD,2})}\right) \leq \tilde{\mathcal{O}}\left(\frac{1}{d^p}\right)$ , and
$$
\forall t \in \left[ T _ {S G D, 1}, T _ {S G D, 2} \right]: \quad R _ {\text {m e d}, k} ^ {S G D M} (t) = \Omega (\log d), \quad k = 1, 2.
$$
<sup>9</sup>In Assumption 2 we assume Gaussian initialization. Due to the rotational invariance of Gaussian distribution, we can assume that all coordinates of  $A$  are positive without loss of generality.
2. For any  $p > 0$ , pick  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right)$ ,  $\xi \leq \sqrt{\frac{\eta}{d^{3\alpha - 1}}}$ ,  $\alpha \geq \frac{p + 4}{3}$  and  $\beta_{2} = \beta_{1}^{2}$ . Suppose  $\sigma \leq \frac{\eta^{3/2}\xi^{2}}{d^{13/4}}$ . Then  $\exists T_{Adam,1}, T_{Adam,2}$  such that w.h.p.,  $\bar{L}\left(W_{Adam}^{(T_{Adam,1})}\right) = \Theta(d)$ ,  $\bar{L}\left(W_{Adam}^{(T_{Adam,2})}\right) \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{p}}\right)$ , and
$$
\forall t \in [ T _ {A d a m, 1}, T _ {A d a m, 2} ]: R _ {m e d, k} ^ {A d a m} (t) = 1 + \tilde {\mathcal {O}} \left(\eta^ {\frac {1}{4}} + \frac {1}{d ^ {\frac {\alpha}{2} - \frac {1}{4}}}\right), \quad k = 1, 2.
$$
Remark Theorem 1 holds for all values of hyperparameters (such as  $\alpha, \sigma$ ) in certain ranges instead of just particular values. The ranges of  $\mathrm{SGD + M}$  and Adam overlap with each other. That means we can choose the same hyperparameters for  $\mathrm{SGD + M}$  and Adam in the overlapped region to make a fair comparison, for example, same  $\alpha, \sigma$  such that  $\alpha \geq 4(p + 2)$  and  $\sigma \leq \min \left\{\frac{\eta^{3/2}}{d^{\alpha/2 + 1}}, \frac{\eta^{3/2} \xi^2}{d^{13/4}}\right\}$ .
An immediate corollary of this theorem below gives the difference between iterates of Adam and  $\mathrm{SGD + M}$  that have the same loss.
Corollary 1. Under the setup in Theorem 1, w.h.p., for any  $t \in [T_{SGD,1}, T_{SGD,2}]$  and  $t' \in [T_{Adam,1}, T_{Adam,2}]$  such that  $\bar{L}\left(W_{SGD}^{(t)}\right) = \bar{L}\left(W_{Adam}^{(t')}\right) \in \left[\tilde{\Omega}\left(\frac{1}{d^p}\right), \Theta(d)\right]$ , we have
$$
R _ {m e d, k} ^ {S G D M} (t) = \Omega (\log d), \quad R _ {m e d, k} ^ {A d a m} (t ^ {\prime}) = 1 + \tilde {\mathcal {O}} \left(\eta^ {\frac {1}{4}} + \frac {1}{d ^ {\frac {\alpha}{2} - \frac {1}{4}}}\right), \quad k = 1, 2.
$$
Theorem 1 and Corollary 1 tell us that during a long training period when the loss decreases from  $\Theta(d)$  to  $\tilde{\mathcal{O}}\left(\frac{1}{d^p}\right)$ , the diagonal of loss Hessian for Adam keeps nice uniformity:  $R_{\mathrm{med},k}^{\mathrm{Adam}}(t) = 1 \pm o(1), k = 1,2$ . On the other hand, the diagonal of loss Hessian for  $\mathrm{SGD + M}$  is less uniform. Appendix C gives a proof sketch of Theorem 1. The detailed proof can be found in Appendix D and E.
6 Low-rank weight matrices and uniformity of leading singular vectors
The proof sketch in Appendix C highlights one crucial intuition of Theorem 1: After  $T_{\mathrm{SGD},1}$  (resp.  $T_{\mathrm{Adam},1}$ ) steps,  $W_{1}$  of  $\mathrm{SGD + M}$  (resp. Adam) becomes an approximately rank-1 matrix. Consider the left singular vector  $\pmb{u} \coloneqq [u_1, u_2, \dots, u_d]^T$  which corresponds to the leading singular value  $\sigma_{1}$ . We can show that the distribution of  $u_{1}^{2}, u_{2}^{2}, \dots, u_{d}^{2}$  for Adam is more uniform than that of  $\mathrm{SGD + M}$ . This property, we call the uniformity of the leading singular vector, is related to the uniformity of the diagonal of loss Hessian, see Appendix G for more details.
Similar low rank bias after training has been studied in prior works (e.g.  $\mathrm{[GWB^{+}17}$ , LMZ18, CGMR20]). For more complicated models, we want to check whether the weight matrices also have low rank structures and if so, whether we can still observe the uniformity of leading singular vectors. More formally, consider the weight matrix in some layer  $W\in \mathbb{R}^{m\times n}$ , we want to check
(A) Whether  $W\in \mathbb{R}^{m\times n}$  is approximately a rank  $k$  matrix with  $k\ll \min \{m,n\}$ , and if true,
(B) Consider the top  $k$  singular values  $\sigma_{1},\ldots,\sigma_{k}$  and left singular vectors  $\boldsymbol{u}_1,\dots,\boldsymbol{u}_k$ . Define  $\tilde{\boldsymbol{u}} := \sum_{i=1}^{k}\sigma_i^2\boldsymbol{u}_i \odot \boldsymbol{u}_i := [\tilde{u}_1,\dots,\tilde{u}_d]^T$  and compute  $R_u := \frac{\max_i\tilde{u}_i}{\text{median}\tilde{u}_i}$ , a generalized version of  $\frac{\max_i u_i^2}{\text{median} u_i^2}$  in the rank-1 case. We want to see whether  $R_u$  obtained by Adam is smaller than that of  $\mathrm{SGD} + \mathbf{M}$ .
After reviewing the weight matrices we got in different settings, we observed that (A) and (B) hold for many layers in those models. For example, on the translation task mentioned in Section 4.1, we found 12 layers which had approximately low rank structures and for 10 of them,  $R_{u}$  values (defined in (B)) obtained by Adam were smaller than those found by SGD+M. Figure 4 shows the result on one typical layer. Results of more layers can be found in Appendix A.5.
Remark The definition of  $R_{u}$  is based on the connection between diagonal of loss Hessian and weight matrices. Appendix G shows that for a 2-layer linear network,  $R_{\mathrm{med},2}^{\mathrm{OPT}}(t) = \frac{\max_i\|W_1^{(t)}[i,:]\|_2^2}{\mathrm{median}\|W_1^{(t)}[i,:]\|_2^2}$ . When  $W_{1}\in \mathbb{R}^{m\times n}$  is approximately rank  $k$ , i.e.  $W_{1}\approx \sum_{i = 1}^{k}\sigma_{i}\pmb{u}_{i}\pmb{v}_{i}^{T}$ , denote  $\pmb {u}_i = [u_{i1},u_{i2},\dots,u_{im}]^T$  and  $\pmb {v}_i = [v_{i1},v_{i2},\dots,v_{in}]^T$ , we have that for the  $j$ -th row,
<sup>10</sup>The assumption  $\beta_{2} = \beta_{1}^{2}$  is only for convenience to make the proof easier but is not essential. Our result can also generalize to cases with other  $\beta_{1}$  and  $\beta_{2}$ . See Appendix C for more discussions.
$\| W_{1}[j,:]\|_{2}^{2}\approx \left\| \sum_{i = 1}^{k}\sigma_{i}u_{ij}\pmb{v}_{i}^{T}\right\|_{2}^{2} = \sum_{i = 1}^{k}\sigma_{i}^{2}u_{ij}^{2}$ . By defining  $\tilde{\pmb{u}} = [\tilde{u}_1,\tilde{u}_2,\dots,\tilde{u}_d]^T := \sum_{i = 1}^{k}\sigma_{i}^{2}\pmb{u}_{i}\odot \pmb{u}_{i}$ , we have that  $\| W_{1}[j,:]\|_{2}^{2}\approx \tilde{u}_{j}$ . Although in multi-layer nonlinear neural networks, the connection between diagonal of loss Hessian and the weight matrices is more complicated and  $R_{\mathrm{med},2}^{\mathrm{OPT}}(t)$  may depend on the product of many weight matrices rather than one single matrix, we still believe that this definition of  $R_{u}$  is a reasonable ratio to consider.
7 Conclusion and future work
We demonstrate that adaptive optimization methods bias the training trajectory towards a region where the diagonal of loss Hessian is more uniform, through extensive experiments on language models and theoretical analysis in a simplified setting of two-layer linear networks. Although our findings may not directly lead to an improved algorithm for practical use, they provide a new way of thinking when designing new algorithms: in contrast with the traditional view which tries to design a method that performs better in the bad loss geometry, our findings suggest that we can design algorithms which implicitly avoid regions with bad geometry. There are a lot of future directions along this line. For example, our theoretical results on the two-layer linear networks may be able to generalize to multi-layer networks. As is discussed in Section 5, the key-value-query structure in the transformer models might be approximated by a three-layer linear network and the analysis of multi-layer networks might provide more connection to real deep models and could be an interesting and challenging future direction. Moreover, it is also possible to relax our large-batch assumption (Assumption 3) and prove similar results in the general stochastic setting.
This project was supported by NSF CCF 2145703.
[ACH18] Sanjeev Arora, Nadav Cohen, and Elad Hazan. On the optimization of deep networks: Implicit acceleration by overparameterization. In 35th International Conference on Machine Learning, ICML 2018, pages 372-389. International Machine Learning Society (IMLS), 2018.
[ADH+19] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019.
[AGCH19] Sanjeev Arora, Noah Golowich, Nadav Cohen, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
[AZLL19] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[AZLS19] Zeyuan Allen-Zhu, Yuzhhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[BDR21] Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristics, 2021.
$\left[\mathrm{BMR}^{+20}\right]$  Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner,
A More experiments of the uniformity of diagonal geometry
A.1 Vanilla SGD vs. Adagrad
In this section, we present the  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  values defined in eq. (1) obtained by vanilla SGD and Adagrad on a language modeling task<sup>11</sup>. The task is to assign a probability for the likelihood of a given word (or a sequence of words) to follow a sequence of words. We trained a transformer model to solve this problem on Wikitext-2 [MXBS17](CC BY-SA 3.0) with real targets and randomly generated targets. This model has roughly 8 layers (not counting normalization and dropout layers)
The setup is the same as in Section 3.2. We used the same learning rate schedule (constant or decreasing) for SGD and Adagrad. We tuned and chose the best (initial) learning rate of SGD. The (initial) learning rate of Adagrad was set as a value under which Adagrad converged faster than SGD with its best (initial) learning rate. We used large batch sizes to make the training procedure more stable. When computing Hessian, we also used large batch sizes. Due to the extremely large dimension, we did the computation on some uniformly selected coordinates, more precisely, 200 coordinates per layer.
We tried different initialization (normal and uniform) by using different gains of the Pytorch initialization schedule.
A.1.1 Experiments on real datasets
Figure 5a shows the training losses on Wikitext-2. Table 3 (resp. Table 4) shows the  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  for Adagrad and SGD under uniform (resp. normal) initialization with different gains.
(b)
Figure 5: Training losses of Adagrad and SGD. (a) on Wikitext-2 and (b) on Wikitext-2 but with randomly generated targets.
Table 3:  ${R}_{\text{med }}^{\mathrm{{SGD}}}\left( t\right)$  and  ${R}_{\text{med }}^{\text{Adagrad }}\left( t\right)$  under uniform initialization with different gains
(a)  $\mathrm{Gain} = 2$
Layer#	Epoch 1	Epoch 20	Epoch 40
\( R_{med}^{SGD}(t) \)	\( R_{med}^{Adagrad}(t) \)	\( R_{med}^{SGD}(t) \)	\( R_{med}^{Adagrad}(t) \)	\( R_{med}^{SGD}(t) \)	\( R_{med}^{Adagrad}(t) \)
1	6.07	6.77	5.91	9.77	5.16	10.37
2	4.60	6.26	3.43	1.66	3.44	1.88
3	5.15	6.84	4.35	4.34	4.84	3.60
4	9.47	10.78	9.76	3.54	8.67	3.14
5	12.54	13.96	10.31	6.59	9.79	6.98
6	4.92	5.25	7.21	2.33	7.94	2.28
7	5.73	5.45	40.56	4.57	21.24	4.76
8	9.39	8.87	37.95	4.50	46.03	3.19