SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We ran  $\mathrm{SGD + M}$  twice on the translation task described in Section 4.1 with learning rates 0.01 and 0.03 and got two trajectories. Then we selected an iterate  $x_{1}$  on the trajectory with learning rate 0.01 and  $x_{2}$  on the trajectory with learning rate 0.03 such that  $x_{1}$  and  $x_{2}$  have the same training loss and different  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ . We then ran  $\mathrm{SGD + M}$  with the same configuration twice, once starting from  $x_{1}$  and once from  $x_{2}$ . The observation is that  $x_{2}$  has a more uniform diagonal Hessian (characterized by smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ , see Table 13) than  $x_{1}$ , and in the meantime,  $\mathrm{SGD + M}$  starting from  $x_{2}$  also converges faster (see Figure 14). Hence we demonstrate that a small value of  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  is not just a byproduct of adaptive algorithms (since  $\mathrm{SGD + M}$  with different learning rates can also yield trajectories with different  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$ ), and it is indeed a contributing factor of fast optimization.
Figure 14: Training losses of  $\mathrm{SGD + M}$  starting from  $x_{1}$  and  $x_{2}$
Table 13:  ${R}_{\text{med }}^{\mathrm{{OPT}}}\left( t\right)$  values of  ${x}_{1}$  and  ${x}_{2}$  in some layers
Layer#	ROPT1med (on the iterate x1)	ROPT2med (on the iterate x2)	ROPT1med/ROPT2med
3	2.89	2.73	1.06
5	69.73	51.44	1.36
7	4.39	4.81	0.91
9	42.93	35.89	1.20
12	7.09	6.18	1.15
15	97.25	68.92	1.41
18	3.50	3.74	0.94
21	71.69	45.07	1.59
24	7.49	5.42	1.38
28	4.85	4.00	1.21
B.2 Theoretical analysis on a simple setting
In this section, we add the following lemma on a simplified setting where smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}(t)$  provably contributes to faster optimization.
Lemma 1. Consider a diagonal quadratic loss  $f(x) = \sum_{i=1}^{d} \frac{\lambda_i}{2} x_i^2$  where  $\lambda_1 \geq \lambda_2 \geq \ldots \geq \lambda_d \geq 0$ . Run  $GD$  with learning rate  $\eta = \frac{1}{\lambda_1}$  and produce  $\{x_i^{(t)}\}_{t \geq 0}, 1 \leq i \leq d$ . We initialize the parameters as  $x_i^{(0)} = \pm 1, \forall i \in [d]$ . Without loss of generality, suppose  $d$  is even. Then we have that for any  $p > 0$ , after  $T = \frac{p}{2} \cdot R_{med}^{OPT} \cdot \log d = \frac{p}{2} \cdot \frac{\lambda_1}{\text{median} \{x_i\}_{i=1}^d} \log d$  steps, the loss  $f(x)$  will shrink by a half. More precisely,
$$
f \left(x ^ {(T)}\right) \leq \left(\frac {1}{2} + d ^ {- p}\right) f \left(x ^ {(0)}\right).
$$
Lemma 1 demonstrates that under this simplified setting, even when we just use the vanilla GD with no adaptive methods, a more uniform diagonal geometry (characterized by smaller  $R_{\mathrm{med}}^{\mathrm{OPT}}$ ) can still lead to faster optimization.
Proof. By the update of GD with learning rate  $\eta = \frac{1}{\lambda_1}$ , we have that
$$
\begin{array}{l} \forall i \in [ d ]: x _ {i} ^ {(t + 1)} = \left(1 - \frac {\lambda_ {i}}{\lambda_ {1}}\right) x _ {i} ^ {(t)}, \\ \Rightarrow \left(x _ {i} ^ {(t)}\right) ^ {2} = \left(1 - \frac {\lambda_ {i}}{\lambda_ {1}}\right) ^ {2 t} \left(x _ {i} ^ {(0)}\right) ^ {2} \leq \exp \left(- \frac {2 \lambda_ {i} t}{\lambda_ {1}}\right) \left(x _ {i} ^ {(0)}\right) ^ {2}. \\ \end{array}
$$
Then after  $T = \frac{p}{2} \cdot \frac{\lambda_1}{\operatorname*{median}\{\lambda_i\}_{i=1}^d} \log d$  steps, we have that
$$
\forall i \leq \frac {d}{2}: \quad \left(x _ {i} ^ {(T)}\right) ^ {2} \leq d ^ {- p} \left(x _ {i} ^ {(0)}\right) ^ {2}, \Rightarrow \sum_ {i = 1} ^ {d / 2} \lambda_ {i} \left(x _ {i} ^ {(T)}\right) ^ {2} \leq d ^ {- p} \sum_ {i = 1} ^ {d / 2} \lambda_ {i} \left(x _ {i} ^ {(0)}\right) ^ {2} \leq d ^ {- p} f \left(x ^ {(0)}\right).
$$
$$
\sum_ {i = d / 2 + 1} ^ {d} \lambda_ {i} \left(x _ {i} ^ {(T)}\right) ^ {2} \leq \sum_ {i = d / 2 + 1} ^ {d} \lambda_ {i} \left(x _ {i} ^ {(0)}\right) ^ {2} \stackrel {(i)} {\leq} \sum_ {i = 1} ^ {d / 2} \lambda_ {i} \left(x _ {i} ^ {(0)}\right) ^ {2},
$$
where  $(i)$  is by definition of median and the initialization  $\forall i\in [d]:x_i^{(0)} = \pm 1$ . Combining with the fact  $\sum_{i=d/2+1}^{d}\lambda_i(x_i^{(0)})^2 + \sum_{i=1}^{d/2}\lambda_i(x_i^{(0)})^2 = f(x^{(0)})$  yields that  $\sum_{i=d/2+1}^{d}\lambda_i(x_i^{(T)})^2 \leq \frac{1}{2}f(x^{(0)})$ . Hence we have
$$
f \left(x ^ {(T)}\right) = \sum_ {i = 1} ^ {d / 2} \lambda_ {i} \left(x _ {i} ^ {(T)}\right) ^ {2} + \sum_ {i = d / 2 + 1} ^ {d} \lambda_ {i} \left(x _ {i} ^ {(T)}\right) ^ {2} \leq \left(\frac {1}{2} + d ^ {- p}\right) f \left(x ^ {(0)}\right).
$$
C Proof sketch of Theorem 1
Now we give a proof sketch of Theorem 1, which contains three major steps. The detailed proof can be found in Appendix G, D and E.
First we relate the diagonal of Hessian to weight matrices  $W_{1}, W_{2}$ . Under Assumption 1, denote  $W_{1}[i,:]$  as the  $i$ -th row of  $W_{1}$  and  $W_{2} \coloneqq [w_{2i}, w_{22}, \dots, w_{2d}]$ . Since the input dataset is whitened, we can show that
$$
R _ {\mathrm {m e d}, 1} ^ {\mathrm {O P T}} (t) = \frac {\max _ {i} \left(w _ {2 i} ^ {(t)}\right) ^ {2}}{\operatorname {m e d i a n} \left(w _ {2 i} ^ {(t)}\right) ^ {2}}, \quad R _ {\mathrm {m e d}, 2} ^ {\mathrm {O P T}} (t) = \frac {\max _ {i} \left\| W _ {1} ^ {(t)} [ i , : ] \right\| _ {2} ^ {2}}{\operatorname {m e d i a n} \left\| W _ {1} ^ {(t)} [ i , : ] \right\| _ {2} ^ {2}}.
$$
Next, due to the one-dimensional output, we can prove that  $W_{1}$  converges to an approximately rank-1 matrix. More precisely, we have
$$
W _ {1} ^ {(t)} = \boldsymbol {u} ^ {(t)} \boldsymbol {v} ^ {(t) T} + R _ {1} ^ {(t)},
$$
$$
W _ {2} ^ {(t)} = c ^ {(t)} \boldsymbol {u} ^ {(t) T} + R _ {2} ^ {(t) T}.
$$
where  $c^{(t)}$  is a scalar,  $\pmb{u}^{(t)},\pmb{v}^{(t)},R_2^{(t)}\in \mathbb{R}^d$  and  $R_{1}^{(t)}\in \mathbb{R}^{d\times d}$ . Denote the  $i$ -th coordinate of  $\pmb{u}^{(t)},\pmb{v}^{(t)},R_2^{(t)}$  as  $u_i^{(t)},v_i^{(t)},R_{2i}^{(t)}$ , respectively. Denote the  $(i,j)$ -th element of  $R_{1}^{(t)}$  as  $R_{1}^{(t)}[i,j]$ . We have that  $\forall i,j\in [d]:\quad \left|R_{2i}^{(t)}\right|\ll c^{(t)}\left|u_i^{(t)}\right|$  and  $\left|R_1^{(t)}[i,j]\right|\ll \left|u_i^{(t)}v_i^{(t)}\right|$ .
Using the rank 1 structure, we can further simplify  $R_{\mathrm{med},1}^{\mathrm{OPT}}(t)$  and  $R_{\mathrm{med},2}^{\mathrm{OPT}}(t)$  by
$$
R _ {\text {m e d}, k} ^ {\mathrm {O P T}} (t) \approx \frac {\max  _ {i} \left(u _ {i} ^ {(t)}\right) ^ {2}}{\operatorname {m e d i a n} \left(u _ {i} ^ {(t)}\right) ^ {2}}, k = 1, 2. \tag {4}
$$
The final step is the detailed analysis of  $\pmb{u}^{(t)}$ .
For  $\mathrm{SGD + M}$ , we can prove that  $\pmb{u}^{(t)}\approx C(t)[X_1,X_2,\dots,X_d]^T$  where  $C(t)\in \mathbb{R}$  and  $X_{i},i\in [d]$  are i.i.d. Gaussian variables. Then we have with high probability,  $\frac{\max_i(u_i^{(t)})^2}{\operatorname*{median}(u_i^{(t)})^2} = \Omega (\log d)$ .
For Adam, we can try to convert it to signed descent (e.g. in eq. (29) and Lemma 24). The main intuition is that as long as the learning rate is small enough, the movement of the gradient will be slow and Adam will be similar to signed descent. Since the update of signed descent per step is  $\pm \eta$ , we can prove that  $\forall i \in [d] : u_i^{(t)} \in \{\pm 1\}$ , which gives us  $\frac{\max_i(u_i^{(t)})^2}{\operatorname{median}(u_i^{(t)})^2} = 1$ . Substituting into eq. (4) completes the proof.
Finally, we want to emphasize that the assumption  $\beta_{2} = \beta_{1}^{2}$  is not essential. As discussed above, as long as the learning rate is sufficiently small, Adam can be approximated to signed descent regardless of the ratio between  $\beta_{1},\beta_{2}$ . Hence our result can generalize to cases with other  $\beta_{1},\beta_{2}$ , for example,  $\beta_{2} > \beta_{1}^{2}$  where Zhang et al.  $[\mathrm{ZCS}^{+}22]$  proved the convergence of Adam to stationary points.
D Analysis of  $\mathbf{SGD} + \mathbf{M}$
Note that  $A = \frac{1}{m} YX^T$ ,  $\Lambda_{xx} \coloneqq \frac{1}{m} XX^T$ . Denote  $g_k^{(t)} \coloneqq \nabla_{W_k} L(W^{(t)})$ ,  $k = 1,2$ . We have that
$$
g _ {1} ^ {(t)} = W _ {2} ^ {(t) T} \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)} - A\right), \quad g _ {2} ^ {(t)} = \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)} - A\right) W _ {1} ^ {(t) T}.
$$
Let  $\tilde{A}^{(t)},\tilde{\Lambda}_{xx}^{(t)}$  and  $\tilde{g}_k^{(t)},k = 1,2$  be the corresponding batch versions at time  $t$ . Let  $E^{(t)}\coloneqq W_2^{(t)}W_1^{(t)} - A$ , and use  $E_{i}^{(t)},A_{i}$  and  $\left(W_2^{(t)}W_1^{(t)}\right)_i$  to represent the  $i$ -th coordinates of  $E^{(t)},A$  and  $W_{2}^{(t)}W_{1}^{(t)}$ , respectively. By eq. (2), the update rules of  $W_{1}$  and  $W_{2}$  for  $\mathrm{SGD + M}$  are given by:
$$
W _ {1} ^ {(t + 1)} = W _ {1} ^ {(t)} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau) T} \left(W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - A\right) - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {1} ^ {(\tau)},
$$
$$
W _ {2} ^ {(t + 1)} = W _ {2} ^ {(t)} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left(W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - A\right) W _ {1} ^ {(\tau) T} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {2} ^ {(\tau)},
$$
$$
D g _ {1} ^ {(t)} := \tilde {g} _ {1} ^ {(t)} - g _ {1} ^ {(t)} = W _ {2} ^ {(t) T} \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)} \left(\tilde {\Lambda} _ {x x} ^ {(t)} - \Lambda_ {x x}\right) - \left(\tilde {A} ^ {(t)} - A\right)\right),
$$
$$
D g _ {2} ^ {(t)} := \tilde {g} _ {2} ^ {(t)} - g _ {2} ^ {(t)} = \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)} \left(\tilde {\Lambda} _ {x x} ^ {(t)} - \Lambda_ {x x}\right) - \left(\tilde {A} ^ {(t)} - A\right)\right) W _ {1} ^ {(t) T}.
$$
Based on the magnitude of  $W_{2}$  and  $W_{1}$ , we can intuitively divide the training procedure into 2 phases.
1. First phase: the first several iterations when  $W_{1}$  and  $W_{2}$  are "small" so that  $W_{2}W_{1} - A \approx -A$ .
2. Second phase: later iterations when  $W_{2}W_{1}$  cannot be ignored.
More formally, the boundary between the first and second phase is defined below.
Definition 1 (End of the first phase). The end of the first phase (denoted as  $T_{1}$ ) is defined as  $T_{1} := \inf \left\{ t \geq 0 : \exists i, j \in [d] : \left| w_{2i}^{(t)} \right| \geq \frac{1}{d^{\frac{\alpha}{2}}} \text{ or } \left| W_{1}^{(t)}[i,j] \right| \geq \frac{1}{d^{\frac{\alpha}{2}}} \right\}$ .
By Assumption 2 and the assumption that  $\forall j\in [d]:A_j > 0,A_j = \Theta (1)$ , at the beginning, w.h.p.,  $\forall j\in [d]:(W_2W_1)_j - A_j < 0$ . During the training, each  $(W_{2}W_{1})_{j}$  increases and approaches  $A_{j}$ . We hope that by choosing a small learning rate, when  $(W_{2}W_{1})_{j}$  overshoots for some coordinate  $j$ , i.e.  $(W_{2}W_{1})_{j} > A_{j}$ , it will be close to convergence. To analyze this overshooting issue more carefully, let's first define the following "almost overshooting time".
Definition 2 (Almost overshooting time). For  $\epsilon > 0$ , denote  $\epsilon_0 := \frac{1}{d^{\frac{1}{4}\alpha - 1}} + \epsilon \log \sqrt{\frac{d}{\epsilon}}$ . Define  $T_2 := \inf \left\{t \geq 0 : \exists j \in [d] : \left(W_2^{(t)}W_1^{(t)}\right)_j - A_j \geq -\sqrt{\epsilon_0} \right\}$ .
Definition 3 (Convergence time). For  $\epsilon >0$ , we define the "convergence time"
$$
T _ {3} := \inf  \left\{t \geq 0: \left\| E ^ {(t)} \right\| _ {2} ^ {2} \leq \epsilon \right\}.
$$
We can first show that after the first phase, i.e. when  $t = T_{1}$ ,  $W_{1}$  will become an approximately rank-1 matrix, as described in the following lemma.
Lemma 2. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}}{d^{\alpha/2+1}}$ . By picking  $\eta \leq \mathcal{O}\left(\frac{1}{d^{\alpha}}\right)$ , we have that when  $t = T_1$ ,  $\bar{L}\left(W^{(T_1)}\right) = \Theta(d)$ , and that
$$
W _ {1} ^ {(T _ {1})} = R _ {1} ^ {(T _ {1})} + \boldsymbol {u} ^ {(T _ {1})} \boldsymbol {v} ^ {(T _ {1}) T},
$$
$$
W _ {2} ^ {(T _ {1})} = R _ {2} ^ {(T _ {1}) T} + c ^ {(T _ {1})} \boldsymbol {u} ^ {(T _ {1}) T},
$$
where  $c^{(T_1)} \in \mathbb{R}$ ,  $\pmb{u}^{(T_1)}, \pmb{v}^{(T_1)}, R_2^{(T_1)} \in \mathbb{R}^d$  and  $R_1^{(T_1)} \in \mathbb{R}^{d \times d}$ . Denote the  $i$ -th coordinate of  $\pmb{u}^{(T_1)}, \pmb{v}^{(T_1)}, R_2^{(T_1)}$  as  $u_i^{(T_1)}, v_i^{(T_1)}, R_{2i}^{(T_1)}$ , respectively, and the  $(i,j)$ -th element of  $R_1^{(T_1)}$  as  $R_1^{(T_1)}[i,j]$ . Then w.h.p.,
$$
\forall 1 \leq i, j \leq d: \quad \frac {\left| R _ {1} ^ {(T _ {1})} [ i , j ] \right|}{\left| u _ {i} ^ {(T _ {1})} v _ {j} ^ {(T _ {1})} \right|} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - 1}}\right), \quad \frac {\left| R _ {2 i} ^ {(T _ {1})} \right|}{\left| e ^ {(T _ {1})} u _ {i} ^ {(T _ {1})} \right|} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - 1}}\right).
$$
The following lemma tells us that this approximate rank-1 structure is preserved when  $T_{1} \leq t \leq \min \{T_{2}, T_{3}\}$ .
Lemma 3. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}}{d^{\alpha/2+1}}$ . By picking  $\eta \leq \mathcal{O}\left(\frac{\epsilon}{d^{\frac{7\alpha}{4}+4}}\right)$ , we have that w.h.p. for  $T_1 \leq t \leq \min\{T_2, T_3\}$ ,
$$
W _ {1} ^ {(t)} = \boldsymbol {u} ^ {(T _ {1})} \boldsymbol {v} ^ {(t) T} + R _ {1} ^ {(t)},
$$
$$
W _ {2} ^ {(t)} = c ^ {(t)} \boldsymbol {u} ^ {(T _ {1}) T} + R _ {2} ^ {(t) T}.
$$
$$
\forall 1 \leq i, j \leq d: \quad \frac {\left| R _ {1} ^ {(t)} [ i , j ] \right|}{\left| u _ {i} ^ {(T _ {1})} v _ {j} ^ {(t)} \right|} \leq \tilde {\mathcal {O}} (\epsilon_ {0}), \quad \frac {\left| R _ {2 i} ^ {(t)} \right|}{\left| c ^ {(t)} u _ {i} ^ {(T _ {1})} \right|} \leq \tilde {\mathcal {O}} (\epsilon_ {0}),
$$
and  $\epsilon_0$  is defined in Definition 2. Moreover, when  $t = \min \{T_2,T_3\}$ ,  $\bar{L}\left(W^{(t)}\right) = \mathcal{O}(\epsilon_0d)$ .
The following lemma gives us a more detailed description of  $\pmb{u}^{(T_1)}$ .
Lemma 4. The  $\mathbf{u}^{(T_1)}$  in Lemma 2 and 3 can be written as  $\mathbf{u}^{(T_1)} = X + Y$  where  $X_{i}, i \in [d]$  are i.i.d Gaussian random variables and that w.h.p.  $\forall i \in [d]: \frac{|Y_i|}{|X_i|} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{1}{4}\alpha - \frac{1}{2}}}\right)$ .
Now we are ready to prove the  $\mathrm{SGD + M}$  part of Theorem 1.
D.1 Proof of the  $\mathbf{SGD} + \mathbf{M}$  part of Theorem 1
Define  $T_{\mathrm{SGD},1} = T_1, T_{\mathrm{SGD},2} = \min \{T_2, T_3\}$ . By picking  $\eta \leq \mathcal{O}\left(\frac{\epsilon}{d^{\frac{7\alpha}{4} + 4}}\right)$ , we can apply Lemma 2 and 3 to conclude that  $\bar{L}\left(W^{(T_{\mathrm{SGD},1})}\right) = \Theta(d)$  and  $\bar{L}\left(W^{(T_{\mathrm{SGD},2})}\right) = \mathcal{O}(\epsilon_0 d)$ . For any  $p > 0$ , by picking  $0 < \epsilon < \frac{1}{d^p}$  and  $\alpha \geq 4(p + 2)$ , we have  $\bar{L}\left(W^{(T_{\mathrm{SGD},2})}\right) = \mathcal{O}(\epsilon_0 d) \leq \tilde{\mathcal{O}}\left(\frac{1}{d^p}\right)$ .
Moreover, when  $t \in [T_{\mathrm{SGD},1}, T_{\mathrm{SGD},2}]$ , the conditions in Lemma 31 are satisfied with  $\delta = \tilde{\mathcal{O}}(\epsilon_0)$ . Then we can apply Lemma 31 and get that
$$
R _ {\text {m e d}, 1} ^ {\text {S G D M}} (t), R _ {\text {m e d}, 2} ^ {\text {S G D M}} (t) \geq \left(\frac {1 - \tilde {\mathcal {O}} (\epsilon_ {0})}{1 + \tilde {\mathcal {O}} (\epsilon_ {0})}\right) ^ {2} \cdot \frac {\max  _ {i} \left(u _ {i} ^ {(T _ {1})}\right) ^ {2}}{\operatorname {m e d i a n} \left(u _ {i} ^ {(T _ {1})}\right) ^ {2}}.
$$
By Lemma 4,  $\pmb{u}^{(T_1)} = X + Y$  where w.h.p.  $\forall i\in [d]:\frac{|Y_i|}{|X_i|}\leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{1}{4}\alpha - \frac{1}{2}}}\right)$ . This fact yields
$$
\forall i \in [ d ]: \frac {\operatorname* {m a x} _ {i} \left(u _ {i} ^ {(T _ {1})}\right) ^ {2}}{\operatorname* {m e d i a n} \left(u _ {i} ^ {(T _ {1})}\right) ^ {2}} \geq \left(\frac {1 - \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - \frac {1}{2}}}\right)}{1 + \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - \frac {1}{2}}}\right)}\right) ^ {2} \frac {\operatorname* {m a x} _ {i} X _ {i} ^ {2}}{\operatorname* {m e d i a n} X _ {i} ^ {2}}.
$$
Here  $X_{i}, i \in [d]$  are i.i.d Gaussian random variables by Lemma 4. To prove the concentration of median  $X_{i}^{2}$ , we borrow the Proposition 12 in Chapter 2.3 of [Ler19]. By setting  $K = N = d$  in this proposition, we have
$$
\mathbb {P} \left(\left| \operatorname {m e d i a n} X _ {i} ^ {2} - \mathbb {E} \left[ X _ {1} ^ {2} \right] \right| > 2 \sqrt {\operatorname {V a r} \left(X _ {1} ^ {2}\right)}\right) \leq e ^ {- \frac {d}{8}}.
$$
Denote  $\sigma^2$  as the variance of  $X_{i}$ ,  $i\in [d]$ . Then  $\mathbb{E}[X_i^2 ] = \sigma^2$  and  $\operatorname {Var}(X_i^2) = 2\sigma^4$ . Hence
$$
\mathbb {P} \left(\left| \text {m e d i a n} X _ {i} ^ {2} - \sigma^ {2} \right| > 2 \sqrt {2} \sigma^ {2}\right) \leq e ^ {- \frac {d}{8}}.
$$
That means with high probability, median  $X_{i}^{2} \leq C\sigma^{2}$  for some  $C > 0$ . By Lemma 35 in Appendix H, we know that w.h.p.
$$
\max _ {1 \leq i \leq d} X _ {i} ^ {2} = \sigma^ {2} \Omega (\log d),
$$
which gives us w.h.p.
$$
\frac {\max  _ {1 \leq i \leq d} X _ {i} ^ {2}}{\text {m e d i a n} X _ {i} ^ {2}} = \Omega (\log d).
$$
Hence we have proved that  $R_{\mathrm{med},1}^{\mathrm{SGDM}}(t), R_{\mathrm{med},2}^{\mathrm{SGDM}}(t) \geq \Omega (\log d)$ .
D.2 Proof of Lemma 2
In the first phase,  $W_{2}W_{1}$  is "small", and we write the update equations in the following way
$$
\begin{array}{l} W _ {1} ^ {(t + 1)} = W _ {1} ^ {(t)} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau) T} \left(W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - A\right) - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {1} ^ {(\tau)} \\ = W _ {1} ^ {(t)} + \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau) T} A - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau) T} W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {1} ^ {(\tau)} \\ = W _ {1} ^ {(t)} + \eta W _ {2} ^ {(t) T} A \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} + \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left(W _ {2} ^ {(\tau) T} - W _ {2} ^ {(t) T}\right) A \tag {5} \\ - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau) T} W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {1} ^ {(\tau)} \\ = W _ {1} ^ {(t)} + \frac {\eta}{1 - \beta} W _ {2} ^ {(t) T} A + \frac {\eta}{1 - \beta} r _ {1} ^ {(t)}, \\ \end{array}
$$
$$
\begin{array}{l} r _ {1} ^ {(t)} = - \beta^ {t + 1} W _ {2} ^ {(t) T} A + (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left(W _ {2} ^ {(\tau) T} - W _ {2} ^ {(t) T}\right) A \\ - (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau) T} W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {1} ^ {(\tau)}. \\ \end{array}
$$
$$
W _ {2} ^ {(t + 1)} = W _ {2} ^ {(t)} - \eta \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left(W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} - A\right) W _ {1} ^ {(\tau) T} = W _ {2} ^ {(t)} + \frac {\eta}{1 - \beta} A W _ {1} ^ {(t) T} + \frac {\eta}{1 - \beta} r _ {2} ^ {(t)}, \tag {6}
$$
$$
\begin{array}{l} r _ {2} ^ {(t)} = - \beta^ {t + 1} A W _ {1} ^ {(t) T} + (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} A \left(W _ {1} ^ {(\tau) T} - W _ {1} ^ {(t) T}\right) \\ - (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)} W _ {1} ^ {(\tau) T} - (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} D g _ {2} ^ {(\tau)}. \\ \end{array}
$$
The following lemma gives us an explicit formula of  $W_2^{(t)}$ .
Lemma 5. Let  $\lambda_1 < \lambda_2$  be the two roots of the quadratic equation  $x^2 - 2x + 1 - \frac{\eta^2}{(1 - \beta)^2} \| A \|_2^2 = 0$ . Pick  $\eta < \frac{1 - \beta}{\|A\|_2}$ , then we have that
$$
W _ {2} ^ {(t)} = C _ {1} \lambda_ {1} ^ {t} + \left(C _ {2} + r _ {5} ^ {(t)}\right) \lambda_ {2} ^ {t},
$$
where  $C_1 = -\frac{W_2^{(1)} - \lambda_2W_2^{(0)}}{\lambda_2 - \lambda_1}$ ,  $C_2 = \frac{W_2^{(1)} - \lambda_1W_2^{(0)}}{\lambda_2 - \lambda_1}$ .  $r_5^{(t)}$  will be specified in the proof.
We can prove that in the first phase,  $r_5^{(t)}$  is "small". More specifically, denote its  $i$ -th coordinate as  $r_{5i}^{(t)}$ , and the  $i$ -th coordinate of  $C_2$  as  $C_{2i}$ . Then the following lemmas tell us that  $\forall i \in [d], \left|r_{5i}^{(t)}\right| \leq \mathcal{O}\left(\frac{1}{d^{p(\alpha)}}\right)$ , where w.h.p.  $\mathcal{O}\left(\frac{1}{d^{p(\alpha)}}\right) \ll \min_{i \in [d]} |C_{2i}|$ .
We first have the following bounds of  $\left|r_{1i}^{(t)}\right|, \left|r_{2i}^{(t)}\right|$  and  $\left|r_{5i}^{(t)}\right|$  for  $i \in [d]$ .
Lemma 6. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}}{d^{\alpha/2+1}}$  and pick  $\eta \leq \mathcal{O}\left(\frac{1}{d^\alpha}\right)$ . We have w.h.p. for all  $t \leq T_1$ ,  $\forall i \in [d]: \left|r_1^{(t)}[i,j]\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3}{2}\alpha-1}}\right), \left|r_{2i}^{(t)}\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3}{2}\alpha-2}}\right)$ .
Lemma 7. Under conditions of Lemma 6, we have that w.h.p. for all  $t \leq T_1$ ,  $\forall i \in [d]: \left|r_{5i}^{(t)}\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3}{2}\alpha - 1}}\right)$ .
Next we prove upper and lower bounds of  $|C_{1i}|$  and  $|C_{2i}|$  for  $i \in [d]$ .
Lemma 8. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}}{d^{\alpha/2+1}}$ . Pick  $\eta < \frac{1-\beta}{\|A\|_2}$ , we have that
i) w.h.p.,  $\forall i\in [d]:|C_{1i}|\leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\alpha}}\right),|C_{2i}|\leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\alpha}}\right)$ ;
ii)  $C_2$  can be written as  $C_2 \coloneqq \frac{1}{2} (C_3 + C_4)$  where  $C_{3i}, i \in [d]$  are i.i.d Gaussian random variables and that w.h.p.  $\forall i \in [d]: \frac{|C_{4i}|}{|C_{3i}|} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{1}{4}\alpha - \frac{1}{2}}}\right)$ ;
iii) w.h.p.,  $\forall i\in [d],|C_{1i}|\geq \tilde{\Omega}\left(\frac{1}{d^{\frac{5}{4}\alpha}}\right),|C_{2i}|\geq \tilde{\Omega}\left(\frac{1}{d^{\frac{5}{4}\alpha}}\right).$
Now we are ready to prove Lemma 2. Lemma 5 tells us that
where  $\lambda_1 = 1 - \frac{\eta}{1 - \beta}\| A\| _2$  and  $\lambda_{2} = 1 + \frac{\eta}{1 - \beta}\| A\|_{2}$
Under the conditions of Theorem 1 and pick  $\eta \leq \mathcal{O}\left(\frac{1}{d^{\alpha}}\right)$ , by Lemma 7 and 8, we know that w.h.p.  $\forall t \leq T_1, \forall 1 \leq i \leq d$ ,
$$
\left| r _ {5 i} ^ {(t)} \right| \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 1}}\right), \quad | C _ {2 i} | \geq \tilde {\Omega} \left(\frac {1}{d ^ {\frac {5}{4} \alpha}}\right), \quad | C _ {2 i} | \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\alpha}}\right), \quad \frac {\left| r _ {5 i} ^ {(t)} \right|}{\left| C _ {2 i} \right|} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - 1}}\right). \tag {7}
$$
We first prove that  $\left|w_{2i}^{(t)}\right|$  reaches  $\frac{1}{d^{\alpha / 2}}$  for some coordinate  $i$  before  $\left|W_1^{(t)}[k,j]\right|$  for  $\forall k,j\in [d]$ . To see this, first note that