SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(b)  $\mathrm{Gain} = 0.5$
Layer#	Epoch 1	Epoch 20	Epoch 40
\( R_{med}^{SGD}(t) \)	\( R_{med}^{Adiagrad}(t) \)	\( R_{med}^{SGD}(t) \)	\( R_{med}^{Adiagrad}(t) \)	\( R_{med}^{SGD}(t) \)	\( R_{med}^{Adiagrad}(t) \)
1	69.36	78.60	15.26	7.74	18.22	7.23
2	24.12	24.36	4.05	2.30	3.70	2.04
3	2.83	2.85	3.78	4.98	3.56	4.40
4	5.25	4.74	3.83	5.68	3.11	4.81
5	66.49	67.83	88.75	19.31	63.01	15.64
6	6.54	6.91	3.57	2.08	3.50	1.97
7	3.22	3.73	13.03	3.97	9.55	4.07
8	6.12	5.99	6.73	7.82	5.43	6.98
<sup>11</sup>https://pytorch.org/tutorials/beginner/transformer_tutorial.html
Table 4:  ${R}_{\text{med }}^{\mathrm{{SGD}}}\left( t\right)$  and  ${R}_{\text{med }}^{\text{Adagrad }}\left( t\right)$  under normal initialization with different gains
(a)  $\mathrm{Gain} = 1$
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.76	6.06	8.27	12.28	9.69	11.17
2	9.51	6.61	3.19	1.87	3.21	1.73
3	7.38	7.35	8.61	3.38	9.25	3.94
4	18.02	15.63	6.45	4.86	7.49	4.44
5	12.70	9.35	11.69	11.23	15.07	12.18
6	12.76	11.86	3.84	2.32	3.20	2.09
7	11.79	8.58	17.95	4.32	14.99	4.50
8	17.09	12.73	26.70	5.16	26.91	6.73
(b)  $\mathrm{Gain} = 0.5$
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	9.12	14.46	10.90	8.00	10.19	8.55
2	10.70	15.42	8.52	2.12	8.88	2.04
3	5.73	5.94	10.16	2.80	6.05	2.99
4	16.62	12.94	8.90	3.91	8.12	4.14
5	15.98	16.98	42.57	10.76	18.45	10.16
6	4.84	6.46	7.92	2.66	5.30	2.46
7	6.52	6.55	107.51	3.14	136.38	2.73
8	8.39	8.20	337.34	5.18	315.21	4.48
A.1.2 Experiments on random datasets
Figure 5b shows the training losses on Wikitext-2 but with randomly generated targets and Table 5 shows the corresponding  $R_{\mathrm{med}}^{\mathrm{SGD}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{Adagrad}}(t)$  in different layers.
Table 5:  ${R}_{\text{med }}^{\mathrm{{SGD}}}\left( t\right)$  and  ${R}_{\text{med }}^{\text{Adagrad }}\left( t\right)$  on Wikitext-2 but with randomly generated targets
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	10.88	10.98	9.99	18.66	9.67	22.37
2	9.47	12.15	14.98	4.43	13.01	3.99
3	7.45	8.52	459.71	6.09	451.16	5.11
4	9.84	10.42	135.37	7.22	126.91	6.04
5	7.09	7.88	103.60	353.89	184.61	190.17
6	7.68	8.58	18.38	4.08	18.69	2.73
7	7.81	5.40	294.68	62.72	229.25	29.76
8	13.51	9.16	329.12	20.59	203.70	9.57
A.2 RMSprop and AMSGrad
In this section, we present the results of RMSprop and AMSGrad and compare them with  $\mathrm{SGD + M}$ . The experiments were conducted on the translation task described in Section 4.1. We used learning rate  $2.5\mathrm{e - }5$  for RMSprop, 5e-4 for AMSGrad and 0.03 for  $\mathrm{SGD + M}$ . Both RMSprop and  $\mathrm{SGD + M}$  used momentum parameter 0.9. The two momentum parameters  $(\beta_{1},\beta_{2})$  of AMSGrad were (0.9, 0.98). Figure 6 shows the training losses and Table 6 shows the corresponding  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ .
Figure 6: Training losses of RMSprop, AMSGrad and SGD+M on the translation task described in Section 4.1.
Table 6:  ${R}_{\text{med }}^{\mathrm{{RMSprop}}}\left( t\right) ,{R}_{\text{med }}^{\mathrm{{AMSGrad}}}\left( t\right)$  and  ${R}_{\text{med }}^{\mathrm{{SGDM}}}\left( t\right)$  in some layers,on the translation task described in Section 4.1
Layer#	Epoch 10	Epoch 20	Epoch 40
\( R_{med}^{SGDM}(t) \)	\( R_{med}^{RMSprop}(t) \)	\( R_{med}^{AMSGrad}(t) \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{RMSprop}(t) \)	\( R_{med}^{AMSGrad}(t) \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{RMSprop}(t) \)	\( R_{med}^{AMSGrad}(t) \)
3	3.97	2.69	2.56	2.33	1.89	1.68	2.83	1.62	1.56
5	26.17	21.19	11.36	37.11	17.83	10.85	51.94	10.22	12.31
7	4.10	6.98	6.12	3.94	4.95	2.92	7.58	2.29	2.58
9	29.41	35.72	25.86	37.81	19.89	16.90	30.68	16.24	9.97
12	4.93	6.20	12.67	4.63	6.61	4.64	6.44	5.13	4.06
15	85.06	33.63	19.51	140.99	12.22	6.72	44.07	6.98	5.37
18	8.71	2.99	9.48	3.86	2.44	4.16	3.51	2.10	2.35
21	95.34	11.68	6.62	47.20	6.37	4.74	22.20	4.58	3.58
24	8.70	5.67	6.95	8.13	3.59	5.13	6.46	2.30	2.83
28	4.44	2.42	2.64	4.67	1.85	1.81	2.63	1.46	2.13
A.3 Comparison conditioned on the same loss
In this section, we compare  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  conditioned on the same training loss. More precisely, we make comparison of  $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 SGD+M iterate have the same training loss. The details of the tasks are described in Section 4.1. Table 7 shows the results of  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t')$  in some layers.
Table 7:  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t')$  in some layers. Dataset and task: (a) sentence classification task on BERT-small, (b) translation task on Multi30k. See Section 4.1 for detailed setup.
(a)
Layer#	Loss 0.251	Loss 0.170	Loss 0.133
RGDM(t')	RAdam(t)	RGDM(t')	RAdam(t)	RGDM(t')	RAdam(t)
9	16.77	13.69	14.14	12.71	15.17	9.86
12	16.68	8.29	9.98	8.31	8.90	5.42
15	18.64	7.79	51.39	46.43	80.82	40.97
17	208.29	381.05	464.37	315.58	498.26	313.99
18	14.43	23.56	19.17	19.26	15.76	12.99
22	257.32	88.47	188.55	110.87	197.79	139.48
24	34.22	16.34	16.42	18.08	14.04	15.97
(b)
Layer#	Loss 3.72	Loss 2.78	Loss 1.90
R6GDM (t')	RAdam (t)	R6GDM (t')	RAdam (t)	R6GDM (t')	RAdam (t)
3	4.01	4.45	5.80	3.02	2.44	2.28
5	31.19	27.50	44.29	21.46	57.83	19.52
7	5.80	4.38	7.51	3.71	5.25	2.87
9	21.23	53.65	28.99	20.92	44.26	28.13
13	53.18	17.77	51.17	20.64	35.80	35.49
15	82.30	186.41	34.17	13.76	33.87	5.31
21	100.43	23.66	23.45	5.12	12.96	5.35
26	7.45	3.48	4.69	3.10	3.33	2.83
30	19.14	9.54	10.46	5.48	9.56	5.33
Figure 7: Training losses of SGD, Adam after SGD and Adam for the translation task
Table 8:  ${R}_{\text{med }}^{\mathrm{{OPT}}}\left( t\right)$  of SGD,Adam after SGD and Adam in some layers after roughly 2160 iterations
Layer#	SGD	Adam	Adam after SGD	Adam
13	294.76	150.02	332.96	150.02
14	14.34	5.84	5.33	5.84
15	36.38	16.66	11.86	16.66
16	6.47	7.05	3.76	7.05
17	17.17	6.05	4.76	6.05
26	5.68	3.53	2.30	3.53
27	14.33	15.93	21.76	15.93
28	9.10	1.71	1.71	1.71
29	8.22	3.04	2.82	3.04
30	11.39	5.12	5.29	5.12
A.4 Experiments of switching from SGD to Adam
In this section we describe another learning schedule: the "Adam after SGD" schedule, where we switch from SGD to Adam in the middle to see whether the loss and  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  can catch up with the model trained by Adam from the very beginning. Again, we used the same model as in the translation task in Section 4.1. In this section, we did not add momentum term to SGD in order to get a larger gap between SGD and Adam than the case using momentum. We want to see whether this larger gap can be closed after switching to Adam in the middle.
As is shown in Figure 7 and Table 8, both the loss gap and the gap of  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  were closed after a period of training after switching algorithms, which provides evidence of the connection between convergence speed and uniformity of diagonal of loss Hessian.
A.5 The low rank structure
In this section, we present more results for the experiments in Section 6.
We examined the weights of the model trained for the translation task in Section 4.1. Among roughly 30 layers, we observed that for 12 layers, at least the weight matrices obtained by Adam after training have approximately low rank structures.
Figure 8 shows the examples of layers with or without the low rank structure.
Figure 8: Examples of singular values on (left) layers without low rank structure and (right) layers with approximately low rank structure.
We then studied the uniformity of leading singular vectors of these 12 layers, i.e. computed  $R_{u}$  defined in (B) of Section 6. The observation was that for 10 out of these 12 layers,  $R_{u}$  values of Adam were smaller those of SGD, which implies the uniformity of leading left singular vectors of Adam. Here we may also want to consider the right singular vectors  $v_{1}, v_{2}, \ldots, v_{k}$  and corresponding  $\tilde{v} = [\tilde{v}_{1}, \tilde{v}_{2}, \ldots, \tilde{v}_{d}]^{T} := \sum_{i=1}^{k} \sigma_{i}^{2} v_{i} \odot v_{i}$  and compute  $R_{v} := \frac{\max_{i} \tilde{v}_{i}}{\text{median} \tilde{v}_{i}}$  for Adam and SGD+M. However, on this translation task, among the 12 layers which were approximately low rank, for only 6 of them,  $R_{v}$  of Adam were smaller, i.e. we did not observe uniformity of the leading right singular
vector for Adam. One possible reason is that for a weight matrix, its right singular vectors are closer to the input data than left singular vectors and more easily influenced by the data, therefore may not show uniformity. Figure 9 shows how  $R_{u}$  and  $R_{v}$  changed over time in some layers.
Figure 9:  $R_{u}$  and  $R_{v}$  for Adam and SGD+M in some layers
A.6 How does the (adaptive) gradient align with diagonal of loss Hessian
In this section we present the uniformity of diagonal geometry of adaptive methods from another perspective. Denote  $H_{ii}$  as the  $(i,i)$ -th element of the loss Hessian  $H$  and  $g_{i}$  as the  $i$ -th element of the gradient. It is conjectured that when  $|H_{ii}|$  is large, the corresponding  $|g_{i}|$  is usually large as well. For adaptive methods, we can regard the update per step as the learning rate times the "adaptive gradient". Let's use  $g_{\mathrm{adapt},i}$  to represent the  $i$ -th component of the adaptive gradient. Through experiments on language models, we found that  $|g_{\mathrm{adapt},i}|$  for different  $i$  are quite uniform and do not align with  $|H_{ii}|$  as the true gradient  $|g_{i}|$  does.
In the experiments, we first sorted  $|H_{ii}|$  in the ascent order:  $|H_{i_1,i_1}|\leq |H_{i_2,i_2}|\leq \ldots \leq |H_{i_d,i_d}|$  (suppose  $H\in \mathbb{R}^{d\times d}$ ), and then plotted the corresponding  $|g_{i_k}|$  and  $|g_{\mathrm{adapt},i_k}|$  for  $k\in [d]$ .
A.6.1 SGD vs. Adagrad
Here we compare SGD and Adagrad on the language modeling task on wikitext-2 described in Section A.1. We observed that the figures of all layers are quite similar so we select one layer as an example, as is shown in Figure 10.
A.6.2 SGD with momentum vs. Adam
Here we compare Adam and SGD+M on the tasks described in Section 4.1. Again, we select one layer as an example for each task. Figure 11 shows the results on the sentence classification task and Figure 12 shows the results on the translation task.
Figure 10: How the true gradient  $\left\{\left|g_{i_k}\right|\right\}_{k = 1}^d$  and "adaptive gradient"  $\left\{\left|g_{\mathrm{adapt},i_k}\right|\right\}_{k = 1}^d$  align with diagonal of Hessian  $\left(\left|H_{i_k,i_k}\right|\right)_{k = 1}^d$ . Here coordinates are sorted such that  $|H_{i_1,i_1}| \leq |H_{i_2,i_2}| \leq \ldots \leq |H_{i_d,i_d}|$  (suppose  $H \in \mathbb{R}^{d\times d}$ ). Experiments were conducted on the model described in Section A.1. This figure shows the results on the 12-th layer.
Figure 11: How the true gradient  $\left\{\left|g_{i_k}\right|\right\}_{k = 1}^d$  and "adaptive gradient"  $\left\{\left|g_{\mathrm{adapt},i_k}\right|\right\}_{k = 1}^d$  align with diagonal of Hessian  $\left(\left|H_{i_k,i_k}\right|\right)_{k = 1}^d$ . Here coordinates are sorted such that  $|H_{i_1,i_1}| \leq |H_{i_2,i_2}| \leq \ldots \leq |H_{i_d,i_d}|$  (suppose  $H \in \mathbb{R}^{d\times d}$ ). Experiments were conducted on the sentence classification task described in Section 4.1. This figure shows the results on the 12-th layer.
A.7 Adding regularization and other tricks
In this section, we add weight decay to both Adam and  $\mathrm{SGD + M}$  on the translation task described in Section 4. The momentum parameter  $\beta$  in SGD was set as 0.9. The two momentum parameters  $(\beta_{1},\beta_{2})$  of Adam were set as (0.9, 0.98). For both algorithms, we set the weight decay parameter as 0.001. We trained the model using constant learning rates for 60 epochs (1800 iterations). We tuned and chose the best learning rate 0.03 for  $\mathrm{SGD + M}$ . The learning rate of Adam was set as 0.0001, under which Adam converged faster than  $\mathrm{SGD + M}$  with its best learning rate 0.03. Figure 13a shows the training losses and Table 9 shows the values of  $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)}$  in some randomly selected layers.
Figure 12: How the true gradient  $\left\{\left|g_{i_k}\right|\right\}_{k = 1}^d$  and "adaptive gradient"  $\left\{\left|g_{\mathrm{adapt},i_k}\right|\right\}_{k = 1}^d$  align with diagonal of Hessian  $\left(\left|H_{i_k,i_k}\right|\right)_{k = 1}^d$ . Here coordinates are sorted such that  $|H_{i_1,i_1}| \leq |H_{i_2,i_2}| \leq \ldots \leq |H_{i_d,i_d}|$  (suppose  $H \in \mathbb{R}^{d \times d}$ ). Experiments were conducted on the translation task described in Section 4.1. This figure shows the results on the 5-th layer.
Table 9:  ${R}_{\text{med }}^{\text{Adam }}\left( t\right)$  and  ${R}_{\text{med }}^{\text{SGDM }}\left( t\right)$  (both using weight decay) in some layers for the translation task.
Layer#	Epoch 0	Epoch 30	Epoch 55
\( R_{med}^{\mathrm{SGDM}}(t) \)	\( P_{med}^{\mathrm{Adam}}(t) \)	\( R_{med}^{\mathrm{SGDM}}(t) \)	\( P_{med}^{\mathrm{Adam}}(t) \)	\( \frac{P_{med}^{\mathrm{SGDM}}(t)}{R_{med}^{\mathrm{Adam}}(t)} \)	\( R_{med}^{\mathrm{SGDM}}(t) \)	\( P_{med}^{\mathrm{Adam}}(t) \)	\( \frac{P_{med}^{\mathrm{SGDM}}(t)}{R_{med}^{\mathrm{Adam}}(t)} \)
3	73.09	73.09	17.65	13.11	1.35	13.38	6.28	2.13
5	469.88	469.88	293.48	310.85	0.94	601.68	588.12	1.02
7	80.78	80.78	8.22	39.65	0.21	13.65	4.85	2.81
9	494.27	494.27	150.14	123.79	1.21	301.89	119.53	2.53
15	632.10	632.10	277.18	175.34	1.58	334.48	282.88	1.18
18	55.08	55.08	6.56	4.45	1.47	23.88	4.52	5.29
21	549.62	549.62	257.89	44.78	5.76	515.99	53.79	9.59
24	107.51	107.51	8.54	3.64	2.34	53.79	3.32	16.20
28	13.77	13.77	4.74	2.37	2.00	15.60	2.15	7.24
30	491.62	491.62	6.91	2.66	2.60	9.60	2.02	4.77
A.8 Results on image tasks
Although in this paper we focus on language models where Adam shows significant fast convergence, we also add results in this section on image tasks where  $\mathrm{SGD + M}$  performs better. We trained a ResNet $^{12}$  on CIFAR-10 dataset and compared the convergence speed and  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  of  $\mathrm{SGD + M}$  and Adam. The momentum parameter  $\beta$  in SGD was set as 0.9. The two momentum parameters  $(\beta_{1},\beta_{2})$  of Adam were set as (0.9, 0.98). The model was trained using constant learning rates for 41 epochs (2050 iterations). We tuned and chose the best learning rates for both algorithms: 0.5 for  $\mathrm{SGD + M}$  and 0.005 for Adam. Figure 13b shows the training losses and Table 10 shows the values of  $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)}$ .
As we can see, on this image task, Adam does not converge faster than  $\mathrm{SGD + M}$  and in the meantime,  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  values were no longer smaller than  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  during training. This reveals the connection between the local diagonal geometry and the convergence speed from another perspective. That is, when the diagonal of Hessian of Adam is not more uniform than  $\mathrm{SGD + M}$ , its convergence speed is not better, either.
<sup>12</sup>We borrowed the implementation here https://pytorch-tutorial.readthedocs.io/en/latest/tutorial/chapter03_intermediate/3_2_2_cnn_resnet_cifar10/ and replace the "layers" array [2,2,2] with [1,1,1].
(b)
Figure 13: (a) Training losses of Adam and  $\mathrm{SGD + M}$  for the translation task, both with weight decay. (b) Training losses of Adam and  $\mathrm{SGD + M}$  for a ResNet trained on CIFAR-10.
Table 10:  ${R}_{\text{med }}^{\text{Adam }}\left( t\right)$  and  ${R}_{\text{med }}^{\text{SGDM }}\left( t\right)$  for ResNet on CIFAR-10.
Layer#	Epoch 10	Epoch 20	Epoch 40
\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( \frac{R_{med}^{SGDM}(t)}{R_{med}^{Adam}(t)} \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( \frac{R_{med}^{SGDM}(t)}{R_{med}^{Adam}(t)} \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( \frac{R_{med}^{SGDM}(t)}{R_{media}^{Adam}(t)} \)
1	6.88	25.34	0.27	3.74	39.35	0.09	4.39	15.80	0.28
2	110.19	35.93	3.07	32.97	36.27	0.91	60.69	28.06	2.16
3	40.89	16.92	2.42	13.98	15.92	0.88	11.70	37.01	0.32
4	28.56	23.66	1.21	11.48	13.04	0.88	7.99	14.51	0.55
5	13.47	23.78	0.57	8.64	12.07	0.72	6.52	14.23	0.46
6	18.72	12.49	1.50	12.19	8.80	1.38	8.96	21.69	0.41
7	18.85	39.25	0.48	9.00	12.81	0.70	13.87	11.42	1.22
8	13.79	19.91	0.69	8.87	11.72	0.76	7.48	9.34	0.80
9	12.50	14.85	0.84	9.62	8.06	1.19	11.35	8.08	1.41
10	14.89	14.53	1.02	8.15	5.80	1.41	6.21	8.89	0.70
A.9 Comparison between  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  and singular value-based metrics
In Section 4, through extensive experiments on language models, we demonstrate that when we train the neural network using Adam, the uniformity of diagonal geometry, measured by  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  is smaller than that when we train using SGD+M from the same initialization, for the vast majority of layers. We are aware that people also usually consider Hessian singular values instead of diagonal entries to measure the loss geometry. Hence in this section we make a comparison between our diagonal-based metric and singular value-based metrics.
First, we believe that our metric has a natural connection to the mechanism that underlies adaptive methods. Adaptive methods in practice choose coordinate-wise adaptive learning rates. From a high-level perspective, this procedure can be viewed as adapting to the loss smoothness with respect to each coordinate. The smoothness of certain coordinate is measured by the second derivative with respect to this coordinate and therefore corresponds to the diagonal entries of loss Hessian. Our metric, which measures these diagonal entries, is thus fundamentally intertwined with the mechanism that underlies adaptive methods.
Next, we empirically demonstrate that our metric  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  is a reasonable proxy of singular value-based metrics. Define a singular value-based metric  $S_{\mathrm{med}}^{\mathrm{OPT}}(t) \coloneqq \frac{\max\{\sigma_i(t)\}_{i=1}^d}{\min\{\sigma_i(t)\}_{i=1}^d}$  as an analogy of our diagonal-based metric  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ , where  $\{\sigma_i(t)\}_{i=1}^d$  denotes the singular values of loss Hessian  $H(t) \in \mathbb{R}^{d \times d}$  at the  $t$ th iterate. We compare  $S_{\mathrm{med}}^{\mathrm{OPT}}(t)$  along the trajectories of Adam and SGD+M in the translation task described in Section 4.1. Table 11 suggests that if measured by singular values, Adam is also biased to a region with smaller  $S_{\mathrm{med}}^{\mathrm{OPT}}(t)$  than SGD+M, similar to the observation for  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ . This is expected because in Appendix F, we demonstrate that the loss Hessian approaches diagonal during training. The fact that our diagonal-based metric and singular value-based metric give the same result also reveals the robustness of our observation to the choice of metric, demonstrating that there does exist some geometry bias of Adam towards more uniform regions even when measured by different metrics.
Finally, there is strong reason why our metric is often easier to compute empirically and analyze theoretically than singular value-based metrics such as  $S_{\mathrm{med}}^{\mathrm{OPT}}(t)$ .
1. From the empirical computation perspective, suppose the loss Hessian is  $d \times d$ . Then computing its singular values, in general, requires computing the whole matrix with  $d^2$  elements. However, our metric only requires computing the  $d$  diagonal entries.
2. From the theoretical analysis perspective, in Appendix G, we show that the diagonal of loss Hessian in linear networks can be connected to weight matrices by simple formulas. These straightforward formulas simplify the analysis and allow us to connect our metric to the low-rank structure of weight matrices and the uniformity of their leading singular vectors (see Section 6 for more discussions). However, all these nice connections fail to hold for singular value-based metrics. The formulas of singular values are very complicated even in linear networks, making it almost impossible to theoretically analyze any singular value-based metrics.
Table 11:  $S_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $S_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  in some layers for the translation task.
Layer#	Epoch 0	Epoch 30	Epoch 55
\( S_{med}^{SGDM}(t) \)	\( S_{med}^{Adam}(t) \)	\( S_{med}^{SGDM}(t) \)	\( S_{med}^{Adam}(t) \)	\( S_{med}^{SGDM}(t) \)	\( S_{med}^{Adam}(t) \)	\( S_{med}^{SGDM}(t) \)	\( S_{med}^{SGDM}(t) \)
3	4.53	4.53	6.08	2.87	2.12	5.92	2.58	2.30
5	14.64	14.64	40.01	15.38	2.60	52.28	15.05	3.47
7	6.91	6.91	9.84	5.06	1.94	12.25	4.22	2.90
9	24.12	24.12	42.02	30.89	1.36	33.20	21.54	1.54
12	19.07	19.07	32.41	24.84	1.30	28.83	14.23	2.03
15	47.03	47.03	69.97	11.54	6.06	42.71	7.19	5.94
18	15.96	15.96	26.03	29.73	0.88	18.46	17.94	1.03
21	31.03	31.03	25.84	7.92	3.26	19.71	7.06	2.79
24	35.42	35.42	21.31	18.08	1.18	14.62	10.33	1.41
28	55.38	55.38	6.18	2.77	2.23	4.84	2.01	2.41
A.10 The uniformity of diagonal geometry on in-distribution test data
In this section we compare  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  on the in-distribution test data. The task is the translation task described in Section 4.1. Table 12 validates that on in-distribution test data, Adam is also biased to a region with smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  than  $\mathrm{SGD} + \mathbf{M}$ , similar to what happens on the training data shown in Table 2a. This is expected because of the same distribution. One other thing we want to emphasize is that, in real language tasks, the dataset is typically very large and the model see each training example only once. Hence the training behavior usually implies similar in-distribution test behavior.
Table 12:  $R_{\mathrm{med}}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med}}^{\mathrm{SGDM}}(t)$  in some layers for the translation task on in-distribution test data.
Layer#	Epoch 0	Epoch 30	Epoch 55
\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( R_{med}^{SGDM}(t) \)	\( R_{med}^{Adam}(t) \)	\( R_{med}^{SGDM}(t) \)
3	4.39	4.39	5.80	3.06	1.89	6.79	2.81	2.42
5	7.90	7.90	38.01	11.71	3.24	41.40	10.21	4.06
7	5.77	5.77	6.00	4.61	1.30	5.53	3.20	1.73
9	25.09	25.09	28.81	17.17	1.68	16.67	14.85	1.12
12	10.24	10.24	9.13	8.63	1.06	13.78	9.09	1.52
15	79.71	79.71	77.18	13.56	5.69	37.93	9.91	3.83
18	14.78	14.78	3.94	7.15	0.55	5.42	6.04	0.90
21	83.25	83.25	26.04	5.44	4.79	13.11	5.57	2.36
24	29.91	29.91	6.89	5.42	1.27	6.51	7.16	0.91
28	22.57	22.57	5.39	3.94	1.37	6.13	2.14	2.87
B Supplementary result on the contribution of small  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  to fast optimization
B.1 Supplementary empirical result
As explained in Section 1, Figure 1 demonstrates that Adam biases the trajectory towards a region which benefits fast optimization, and in the meantime, it has more uniform diagonal Hessian (characterized by smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ ). However, since the two trajectories in Figure 1 come from two different algorithms (SGD+M and Adam), it is reasonable to ask whether the benefit of Adam's trajectory is indeed due to smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  or it is due to other properties of Adam and small  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  is just a byproduct. To address this concern, we add the following experiment where we select two trajectories that both come from SGD+M and have different  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ .