SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 denote  $E_{j}^{(t)}$  as the  $j$ -th component of  $E^{(t)}$ . We also denote  $\Delta w_{2i}^{(t)}\coloneqq w_{2i}^{(t + 1)} - w_{2i}^{(t)}$ ,  $\Delta W_1^{(t)}[i,j]\coloneqq W_1^{(t + 1)}[i,j] - W_1^{(t)}[i,j]$ . By eq. (2), the update equations of Adam are given by
$$
\eta_ {t} = \eta \cdot \frac {\sqrt {1 - \beta_ {2} ^ {t + 1}}}{1 - \beta_ {1} ^ {t + 1}}, \quad g _ {1} ^ {(t)} [ i, j ] = w _ {2 i} ^ {(t)} E _ {j} ^ {(t)}, \quad g _ {2 i} ^ {(t)} = \left<   E ^ {(t)}, W _ {1} ^ {(t)} [ i,: ] \right > ,
$$
$$
\begin{array}{l} W _ {1} ^ {(t + 1)} [ i, j ] - W _ {1} ^ {(t)} [ i, j ] = - \eta_ {t} \frac {m _ {1} ^ {(t)} [ i , j ]}{\sqrt {v _ {1} ^ {(t)} [ i , j ]}} = - \eta_ {t} \frac {\left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {t} \beta_ {1} ^ {t - \tau} \tilde {g} _ {1} ^ {(\tau)} [ i , j ]}{\sqrt {\left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {t} \beta_ {2} ^ {t - \tau} \left(\tilde {g} _ {1} ^ {(\tau)} [ i , j ]\right) ^ {2}} + \xi}, \\ = - \eta_ {t} \frac {\left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {t} \beta_ {1} ^ {t - \tau} g _ {1} ^ {(\tau)} [ i , j ] + r _ {1 n} ^ {(t)} [ i , j ]}{\sqrt {\left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {t} \beta_ {2} ^ {t - \tau} \left(g _ {1} ^ {(\tau)} [ i , j ]\right) ^ {2} + r _ {1 d} ^ {(t)} [ i , j ]} + \xi}, \\ w _ {2 i} ^ {(t + 1)} - w _ {2 i} ^ {(t)} = - \eta_ {t} \frac {m _ {2 i} ^ {(t)}}{\sqrt {v _ {2 i} ^ {(t)}}} = - \eta_ {t} \frac {\left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {t} \beta_ {1} ^ {t - \tau} \tilde {g} _ {2 i} ^ {(\tau)}}{\sqrt {\left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {t} \beta_ {2} ^ {t - \tau} \left(\tilde {g} _ {2 i} ^ {(\tau)}\right) ^ {2}} + \xi}, \\ = - \eta_ {t} \frac {\left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {t} \beta_ {1} ^ {t - \tau} g _ {2 i} ^ {(\tau)} + r _ {2 n , i} ^ {(t)}}{\sqrt {\left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {t} \beta_ {2} ^ {t - \tau} \left(g _ {2 i} ^ {(\tau)}\right) ^ {2} + + r _ {2 d , i} ^ {(t)}} + \xi}. \tag {27} \\ \end{array}
$$
where  $Dg_{1}^{(t)}\coloneqq \tilde{g}_{1}^{(t)} - g_{1}^{(t)}$  and  $Dg_{2}^{(t)}\coloneqq \tilde{g}_{2}^{(t)} - g_{2}^{(t)}$  ,and
$$
r _ {1 n} ^ {(t)} [ i, j ] := (1 - \beta_ {1}) \sum_ {\tau = 0} ^ {t} \beta_ {1} ^ {t - \tau} D g _ {1} ^ {(\tau)} [ i, j ],
$$
$$
r _ {1 d} ^ {(t)} [ i, j ] = \left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {t} \beta_ {2} ^ {t - \tau} \left(2 g _ {1} ^ {(\tau)} [ i, j ] D g _ {1} ^ {(\tau)} [ i, j ] + \left(D g _ {1} ^ {(\tau)} [ i, j ]\right) ^ {2}\right), \tag {28}
$$
$$
r _ {2 n, i} ^ {(t)} := \left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {t} \beta_ {1} ^ {t - \tau} D g _ {2 i} ^ {(\tau)},
$$
$$
r _ {2 d, i} ^ {(t)} = (1 - \beta_ {2}) \sum_ {\tau = 0} ^ {t} \beta_ {2} ^ {t - \tau} \left(2 g _ {2 i} ^ {(\tau)} D g _ {2 i} ^ {(\tau)} + \left(D g _ {2 i} ^ {(\tau)}\right) ^ {2}\right).
$$
Denote the  $i$ -th coordinate of  $W_{2}W_{1}$  and  $A$  as  $(W_{2}W_{1})_{i}$  and  $A_{i}$ , respectively. By Assumption 2 and the assumption that  $\forall i \in [d]: A_{i} > 0, A_{i} = \Omega(1)$ , at the beginning, w.h.p.,  $\forall i \in [d]: (W_{2}W_{1})_{i} - A_{i} < 0$ . Based on this, we divide the training procedure into two phases (note that these two phases are different from those of  $\mathrm{SGD + M}$ ).
1. First phase: when the error  $(W_{2}W_{1})_{i} - A_{i}$  is negative and its absolute value is big for all  $i\in [d]$ .
2. Second phase: when  $(W_{2}W_{1})_{i} - A_{i}$  is close to zero for some coordinate  $i\in [d]$ .
More formally, we define the boundary between the two phases below.
Definition 4 (End of the first phase). The end of the first phase (denoted as  $T_{1}$ ) is defined as  $T_{1} = \inf \left\{t > 0:\exists i\in [d]:E_{i}^{(t)}\geq -\sqrt{\eta d}\right\}$
In the second phase, we define some time points.
Definition 5. Define  $T_{g} \coloneqq \inf \left\{t > T_{1} : \exists i \in [d] : \left| g_{2i}^{(t)} \right| \leq d\sqrt{\eta} \right\}$ .
For  $t < T_1$ , we have  $\forall i \in [d] : E_i^{(t)} < 0$  by Definition 4. For  $t > T_1$ , some  $E_i^{(t)}$  may flip the sign and become positive. For certain coordinate  $i$ , we define the following "flip time".
Definition 6. Define  $T_{f,i} \coloneqq \inf \left\{t > T_1 : E_i^{(t)} \geq -\sqrt{\eta d}\right\}$ . Define  $T_f \coloneqq \max_i T_{f,i}$  as the largest "flip time" over all  $i \in [d]$ , i.e. the "flip time" of the last  $E_i$  which flips the sign. Moreover, denote  $\tilde{T} \coloneqq \min \{T_g, T_f\}$ .
We can first show that after a few steps in the first phase,  $W_{1}$  will become an approximately rank-1 matrix, as described in the following lemma.
Lemma 15. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}\xi^2}{d^{13/4}}$ . By picking  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right), \xi \leq \sqrt{\frac{\eta}{d^{3\alpha-1}}}$ , and  $\beta_2 = \beta_1^2$ , there exists  $t_{inc} > 0$  such that w.h.p. for  $t_{inc} \leq t < T_1$ ,
$$
\forall i, j \in [ d ]: \quad w _ {2 i} ^ {(t)} = \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \eta \left(t - t _ {i n c}\right) + R _ {2 i} ^ {(t)},
$$
$$
W _ {1} ^ {(t)} [ i, j ] = \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \eta (t - t _ {i n c}) + R _ {1} ^ {(t)} [ i, j ],
$$
where  $\frac{\left|R_1^{(t)}[i,j]\right|}{\eta(t - t_{inc})} = \tilde{\mathcal{O}}\left(\sqrt{\eta} +\frac{1}{\eta(t - t_{inc})d^{\alpha}}\right),\quad \frac{\left|R_{2i}^{(t)}\right|}{\eta(t - t_{inc})} = \tilde{\mathcal{O}}\left(\sqrt{\eta} +\frac{1}{\eta(t - t_{inc})d^{\alpha}}\right).$
Specially, when  $t = T_1$ , we have that
$$
\forall i, j \in [ d ]: \quad w _ {2 i} ^ {(T _ {1})} = \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \eta \left(T _ {1} - t _ {i n c}\right) + R _ {2 i} ^ {(T _ {1})},
$$
$$
W _ {1} ^ {(T _ {1})} [ i, j ] = \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \eta (T _ {1} - t _ {i n c}) + R _ {1} ^ {(T _ {1})} [ i, j ],
$$
where  $\eta (T_1 - t_{inc}) = \Theta \left(\frac{1}{\sqrt{d}}\right)$  and
$$
\frac {\left| R _ {1} ^ {(T _ {1})} [ i , j ] \right|}{\eta (T _ {1} - t _ {i n c})} = \tilde {\mathcal {O}} \left(\sqrt {\eta} + \frac {1}{d ^ {\alpha - \frac {1}{2}}}\right), \quad \frac {\left| R _ {2 i} ^ {(T _ {1})} \right|}{\eta (T _ {1} - t _ {i n c})} = \tilde {\mathcal {O}} \left(\sqrt {\eta} + \frac {1}{d ^ {\alpha - \frac {1}{2}}}\right).
$$
The following lemma tells us that this approximate rank-1 structure is preserved when  $T_{1} \leq t \leq \tilde{T}$ .
Lemma 16. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}\xi^2}{d^{13/4}}$ . By picking  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right), \xi \leq \sqrt{\frac{\eta}{d^{3\alpha-1}}}$ , and  $\beta_2 = \beta_1^2$ , we have w.h.p. for  $T_1 \leq t < \tilde{T}$
$$
\begin{array}{l} \forall i, j \in [ d ]: \quad w _ {2 i} ^ {(t)} = \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) c ^ {(t)} + R _ {2 i} ^ {(t)}, \\ W _ {1} ^ {(t)} [ i, j ] = \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) V _ {j} ^ {(t)} + R _ {1} ^ {(t)} [ i, j ], \\ \end{array}
$$
$$
\text {w h e r e} \quad \frac {\left| R _ {2 i} ^ {(t)} \right|}{\left| c ^ {(t)} \right|} = \tilde {\mathcal {O}} \left(\sqrt {\eta} + \frac {1}{d ^ {\alpha - 1 / 2}}\right), \quad \frac {\left| R _ {1} ^ {(t)} [ i , j ] \right|}{\left| V _ {j} ^ {(t)} \right|} \leq \tilde {\mathcal {O}} \left(\eta^ {\frac {1}{4}} + \frac {1}{d ^ {\frac {\alpha}{2} - \frac {1}{4}}}\right),
$$
and that  $L\left(W^{(\tilde{T})}\right)\leq \tilde{\mathcal{O}}\left(\eta d^4\right)$
Now we are ready to prove the Adam part of Theorem 1.
E.1 Proof of the Adam part of Theorem 1
Define  $T_{\mathrm{Adam},1} = t_{\mathrm{inc}} + \frac{1}{\eta d^{\frac{\alpha}{2}}}$ . Note that this choice of  $T_{\mathrm{Adam},1}$  gives  $\eta (T_{\mathrm{Adam},1} - t_{\mathrm{inc}}) = \frac{1}{d^{\frac{\alpha}{2}}}$ . By picking  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right), \xi \leq \sqrt{\frac{\eta}{d^{3\alpha - 1}}}$  and  $\beta_{2} = \beta_{1}^{2}$ , we can apply Lemma 15 to get that  $\forall i, j \in [d] : \left|w_{2i}^{(T_{\mathrm{Adam},1})}\right| = \Theta \left(\frac{1}{d^{\frac{\alpha}{2}}}\right), \left|W_{1}^{(T_{\mathrm{Adam},1})}[i,j]\right| = \Theta \left(\frac{1}{d^{\frac{\alpha}{2}}}\right)$ , and therefore  $\forall i \in [d]: E_i^{(T_{\mathrm{Adam},1})} = -\Theta(1)$  and  $L(W^{(T_{\mathrm{Adam},1})}) = \Theta(d)$ . Define  $T_{\mathrm{Adam},2} = \tilde{T}$ . By Lemma 16, we have  $L(W^{(T_{\mathrm{Adam},2})}) = \tilde{\mathcal{O}}(\eta d^4)$ . For any  $p > 0$ , by picking  $\alpha \geq \frac{p + 4}{3}$ , we have  $L(W^{(T_{\mathrm{Adam},2})}) = \tilde{\mathcal{O}}(\eta d^4) \leq \tilde{\mathcal{O}}(\frac{1}{d^p})$ .
Moreover, combining Lemma 15 and 16, we get that when  $t \in [T_{\mathrm{Adam},1}, T_{\mathrm{Adam},2}]$ , the conditions in Lemma 31 are satisfied with  $\delta = \tilde{\mathcal{O}}\left(\eta^{\frac{1}{4}} + \frac{1}{d^{\frac{\alpha}{2} - \frac{1}{4}}}\right)$ . The  $i$ -th component of the  $\boldsymbol{u}$  vector (denoted as  $u_i$ ) is sign  $\left(w_{2i}^{(0)}\right)$ . That means  $\forall i \in [d] : u_i^2 = 1$  and  $\frac{\max_i(u_i)^2}{\operatorname{median}(u_i)^2} = 1$ . Then we can apply Lemma 31 and get that
$$
\begin{array}{l} R _ {\text {m e d}, 1} ^ {\text {A d a m}} (t), R _ {\text {m e d}, 2} ^ {\text {A d a m}} (t) \in \left[ \left(\frac {1 - \delta}{1 + \delta}\right) ^ {2} \frac {\max  _ {i} \left(u _ {i}\right) ^ {2}}{\operatorname {m e d i a n} \left(u _ {i}\right) ^ {2}}, \left(\frac {1 + \delta}{1 - \delta}\right) ^ {2} \frac {\max  _ {i} \left(u _ {i}\right) ^ {2}}{\operatorname {m e d i a n} \left(u _ {i}\right) ^ {2}} \right] \\ = \left[ \left(\frac {1 - \delta}{1 + \delta}\right) ^ {2}, \left(\frac {1 + \delta}{1 - \delta}\right) ^ {2} \right]. \\ \end{array}
$$
Note that by definition,  $R_{\mathrm{med},1}^{\mathrm{Adam}}(t)$  and  $R_{\mathrm{med},2}^{\mathrm{Adam}}(t)$  are always larger than or equal to 1, then we have
$$
R _ {\mathrm {m e d}, 1} ^ {\mathrm {A d a m}} (t), R _ {\mathrm {m e d}, 2} ^ {\mathrm {A d a m}} (t) = 1 + \mathcal {O} (\delta) = 1 + \tilde {\mathcal {O}} \left(\eta^ {\frac {1}{4}} + \frac {1}{d ^ {\frac {\alpha}{2} - \frac {1}{4}}}\right).
$$
For some time  $t$ , we introduce two conditions.
$$
\forall \tau \in [ H ]: s i g n \left(g _ {1} ^ {(t - \tau)} [ i, j ]\right) = s _ {1} ^ {(t)} [ i, j ], \quad (1 - \beta_ {1}) \left| \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {(\tau)} g _ {1} ^ {(t - \tau)} [ i, j ] \right| \geq \Omega (\xi).
$$
$$
\forall \tau \in [ H ]: \operatorname {s i g n} \left(g _ {2 i} ^ {(t - \tau)}\right) = s _ {2 i} ^ {(t)}, \quad (1 - \beta_ {1}) \left| \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {(\tau)} g _ {2 i} ^ {(t - \tau)} \right| \geq \Omega (\xi).
$$
Next prove that, under Assumption 1 and 2, by picking  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right), \xi \leq \sqrt{\frac{\eta}{d^{3\alpha - 1}}}$ , and  $\beta_{2} = \beta_{1}^{2}$  there exists  $t_{\mathrm{inc}} > 0$  such that for  $t_{\mathrm{inc}} \leq t < T_{1}$ , the update of Adam can be approximated as that of signed descent.
$$
W _ {1} ^ {(t + 1)} [ i, j ] = W _ {1} ^ {(t)} [ i, j ] - \eta \left(\operatorname {s i g n} \left(g _ {1} ^ {(t)} [ i, j ]\right) + e _ {1} ^ {(t)} [ i, j ]\right), \tag {29}
$$
$$
w _ {2 i} ^ {(t + 1)} = w _ {2 i} ^ {(t)} - \eta \left(\operatorname {s i g n} \left(g _ {2 i} ^ {(t)}\right) + e _ {2 i} ^ {(t)}\right),
$$
where  $\left|e_1^{(t)}[i,j]\right| = \tilde{\mathcal{O}}\left(\sqrt{\eta}\right)$ ,  $\left|e_{2i}^{(t)}\right| = \tilde{\mathcal{O}}\left(\sqrt{\eta}\right)$ .
Before we dive into the proof, let's introduce some useful lemmas.
The following lemma reflects our key idea: converting the exponential average in Adam to a finite-step average, and trying to bound the stochastic error terms in eq. (28).
Lemma 17. Under Assumption 1, 2 and 3 and pick  $\beta_{2} = \beta_{1}^{2}$ . Let  $M_1^{(t)} := \max_{i,j\in [d],\tau \leq t}\left|W_1^{(\tau)}[i,j]\right|$ ,  $M_2^{(t)} := \max_{i,j\in [d],\tau \leq t}\left|w_{2i}^{(\tau)}\right|$ ,  $G_1^{(t)} := \max_{i,j\in [d],\tau \leq t}\left|g_1^{(\tau)}[i,j]\right|$  and  $G_2^{(t)} := \max_{i,j\in [d],\tau \leq t}\left|g_{2i}^{(\tau)}\right|$ . We have that w.h.p., for all  $t \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d\eta}}\right)$  and  $\forall i,j \in [d]$ ,
$$
\Delta W _ {1} ^ {(t)} [ i, j ] = - \eta_ {t} \frac {(1 - \beta_ {1}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} g _ {1} ^ {(t - \tau)} [ i , j ] + \epsilon_ {1 n} ^ {(t)} [ i , j ]}{\sqrt {(1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} (g _ {1} ^ {(t - \tau)} [ i , j ]) ^ {2} + \epsilon_ {1 d} ^ {(t)} [ i , j ]} + \xi},
$$
$$
\Delta w _ {2 i} ^ {(t)} = - \eta_ {t} \frac {(1 - \beta_ {1}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} g _ {2 i} ^ {(t - \tau)} + \epsilon_ {2 n , i} ^ {(t)}}{\sqrt {(1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} (g _ {2 i} ^ {(t - \tau)}) ^ {2} + \epsilon_ {2 d , i} ^ {(t)}} + \xi},
$$
where  $H\geq \frac{1}{1 - \beta_1}\log \frac{\max\left\{G_1^{(t)},G_2^{(t)},\left(G_1^{(t)}\right)^2,\left(G_2^{(t)}\right)^2\right\}}{\eta\xi^2}$  and
$$
\left| \epsilon_ {1 n} ^ {(t)} [ i, j ] \right| \leq \mathcal {O} (\eta \xi^ {2}) + \mathcal {O} \left(D _ {1} ^ {(t)}\right), \quad \left| \epsilon_ {1 d} ^ {(t)} [ i, j ] \right| \leq \mathcal {O} (\eta \xi^ {2}) + \mathcal {O} \left(D _ {1} ^ {(t)} G _ {1} ^ {(t)} + \left(D _ {1} ^ {(t)}\right) ^ {2}\right),
$$
$$
\left| \epsilon_ {2 n, i} ^ {(t)} \right| \leq \mathcal {O} (\eta \xi^ {2}) + \mathcal {O} \left(D _ {2} ^ {(t)}\right), \quad \left| \epsilon_ {2 d, i} ^ {(t)} \right| \leq \mathcal {O} (\eta \xi^ {2}) + \mathcal {O} \left(D _ {2} ^ {(t)} G _ {2} ^ {(t)} + \left(D _ {2} ^ {(t)}\right) ^ {2}\right),
$$
$$
D _ {1} ^ {(t)} \leq \tilde {\mathcal {O}} \left(d ^ {3} M _ {1} ^ {(t)} \left(M _ {2} ^ {(t)}\right) ^ {2} \sigma \sqrt {\frac {d ^ {1 / 2}}{\eta}}\right) + \tilde {\mathcal {O}} \left(M _ {2} ^ {(t)} \sigma \sqrt {\frac {d ^ {3 / 2}}{\eta}}\right),
$$
$$
D _ {2} ^ {(t)} \leq \tilde {\mathcal {O}} \left(d ^ {4} \left(M _ {1} ^ {(t)}\right) ^ {2} M _ {2} ^ {(t)} \sigma \sqrt {\frac {d ^ {1 / 2}}{\eta}}\right) + \tilde {\mathcal {O}} \left(d M _ {1} ^ {(t)} \sigma \sqrt {\frac {d ^ {3 / 2}}{\eta}}\right).
$$
Corollary 2. Under the conditions of Lemma 17 and suppose  $\sigma \leq \frac{\eta^{3/2}\xi^2}{d^{13/4}}$ . Consider any  $t \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}\eta}\right)$ . If  $M_1^{(t)}, M_2^{(t)} \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right)$ ,  $G_1^{(t)} \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right)$ ,  $G_2^{(t)} \leq \tilde{\mathcal{O}}\left(\sqrt{d}\right)$ , then  $H$  in Lemma 17 can be picked as  $\frac{1}{1 - \beta_1} \log \frac{d}{\eta\xi^2}$  and we can get that  $\forall i, j \in [d]$ ,  $\left|\epsilon_{1n}^{(t)}[i,j]\right|, \left|\epsilon_{1d}^{(t)}[i,j]\right|, \left|\epsilon_{2n,i}^{(t)}\right|, \left|\epsilon_{2d,i}^{(t)}\right| \leq \tilde{\mathcal{O}}(\eta\xi^2)$ .
The following lemma analyzes the magnitude of weights during a short period at the beginning.
Lemma 18. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}\xi^2}{d^{13/4}}$ . Pick  $\xi \leq \frac{1}{d^{\frac{3}{2}\alpha}}$ , then there exists some time point  $t_{inc} \in (H, T_1)$ , such that w.h.p., for  $t \leq t_{inc}$ , for every  $i, j \in [d]$ ,
$$
\left| \Delta W _ {1} ^ {(t)} [ i, j ] \right| \leq \tilde {\mathcal {O}} (\eta), \left| \Delta w _ {2 i} ^ {(t)} \right| \leq \tilde {\mathcal {O}} (\eta),
$$
$$
\left| W _ {1} ^ {(t)} [ i, j ] \right| \leq \mathcal {O} \left(\frac {1}{d ^ {\frac {3}{2} \alpha + 1}}\right), \Omega \left(\frac {1}{d ^ {\frac {3}{2} \alpha}}\right) \leq \left| w _ {2 i} ^ {(t)} \right| \leq \mathcal {O} \left(\frac {1}{d ^ {\alpha}}\right).
$$
Specifically, when  $t = t_{inc}$ , we have  $\text{sign}\left(w_{2i}^{(t_{inc})}\right) = \text{sign}\left(W_1^{(t_{inc})}[i,j]\right) = \text{sign}\left(w_{2i}^{(0)}\right)$ ,  $\left|W_1^{(t_{inc})}[i,j]\right| = \Theta \left(\frac{1}{d^{\frac{3}{2}\alpha + 1}}\right)$  and  $\left|g_1^{(t_{inc})}[i,j]\right| \geq \Omega (\xi), \left|g_{2i}^{(t_{inc})}\right| \geq \Omega (\xi)$ . Moreover, Condition 1 and 2 are satisfied for  $t = t_{inc}$ . The  $s_1^{(t)}[i,j]$  and  $s_{2i}^{(t)}$  in the conditions are both -sign  $\left(w_{2i}^{(0)}\right)$ .
The following lemma gives us lower bounds of  $\left|g_1^{(t)}[i,j]\right|$  and  $\left|g_{2i}^{(t)}\right|$ .
Lemma 19. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}\xi^2}{d^{13/4}}$ . Pick  $\xi \leq \sqrt{\frac{\eta}{d^{3\alpha-1}}}$ ,  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right)$ . Consider  $t_{inc}$  in Lemma 18. We have w.h.p. for any  $t \in [t_{inc}, T_1)$ , and for  $\forall i, j \in [d]$ ,  $\text{sign}\left(\Delta W_1^{(t)}[i, j]\right) = \text{sign}\left(\Delta w_{2i}^{(t)}\right) = \text{sign}\left(w_{2i}^{(0)}\right)$  and that  $\forall i, j \in [d] : \left|g_1^{(t)}[i, j]\right| \geq \tilde{\Omega}(\sqrt{\eta}), \left|g_{2i}^{(t)}\right| \geq \tilde{\Omega}(\sqrt{\eta}d)$ . Moreover, we have  $\forall \tau \leq t$ ,  $\forall i, j \in [d] : \left|W_1^{(\tau)}[i, j]\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right), \left|w_{2i}^{(\tau)}\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right)$  and  $\left|g_1^{(\tau)}[i, j]\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right), \left|g_{2i}^{(\tau)}\right| \leq \tilde{\mathcal{O}}\left(\sqrt{d}\right)$ .
The following lemma shows that when  $t_{\mathrm{inc}} \leq t < T_1$ , we have  $\forall i, j \in [d] : \left| g_{2i}^{(t)} \right| \gg \left| g_{2i}^{(t)} - g_{2i}^{(t-1)} \right|$  and that  $\left| g_1^{(t)}[i,j] \right| \gg \left| g_1^{(t)}[i,j] - g_1^{(t-1)}[i,j] \right|$ .
Lemma 20. Under Assumption 1, 2 and 3, suppose  $\sigma \leq \frac{\eta^{3/2}\xi^2}{d^{13/4}}$ . Pick  $\xi \leq \sqrt{\frac{\eta}{d^{3\alpha-1}}}$ ,  $\eta \leq \mathcal{O}\left(\frac{1}{d^{3\alpha}}\right)$ . For  $t_{inc}$  in Lemma 18, we have that w.h.p. for  $t_{inc} \leq t < T_1$  and  $\tau \leq t$ ,  $\forall i, j \in [d]$ ,
$$
\frac {\left| g _ {1} ^ {(t)} [ i , j ] - g _ {1} ^ {(t - \tau)} [ i , j ] \right|}{\left| g _ {1} ^ {(t)} [ i , j ] \right|} = \tilde {\mathcal {O}} (\sqrt {\eta} \tau), \quad \frac {\left| g _ {2 i} ^ {(t)} - g _ {2 i} ^ {(t - \tau)} \right|}{\left| g _ {2 i} ^ {(t)} \right|} = \tilde {\mathcal {O}} (\sqrt {\eta} \tau), \tag {30}
$$
$$
\frac {\left| \left(g _ {1} ^ {(t)} [ i , j ]\right) ^ {2} - \left(g _ {1} ^ {(t - \tau)} [ i , j ]\right) ^ {2} \right|}{\left(g _ {1} ^ {(t)} [ i , j ]\right) ^ {2}} = \tilde {\mathcal {O}} (\sqrt {\eta} \tau) + \tilde {\mathcal {O}} (\eta \tau^ {2}), \tag {31}
$$
$$
\frac {\left| \left(g _ {2 i} ^ {(t)}\right) ^ {2} - \left(g _ {2 i} ^ {(t - \tau)}\right) ^ {2} \right|}{\left(g _ {2 i} ^ {(t)}\right) ^ {2}} = \tilde {\mathcal {O}} \left(\sqrt {\eta} \tau\right) + \tilde {\mathcal {O}} \left(\eta \tau^ {2}\right).
$$
Equipped with these lemmas, now let's prove eq. (29).
For any  $t \in [t_{\mathrm{inc}}, T_1)$ , by Lemma 19, we know that  $M_1^{(t)}, M_2^{(t)} \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right)$ , and that  $G_1^{(t)} \leq \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right), G_2^{(t)} \leq \tilde{\mathcal{O}}\left(\sqrt{d}\right)$ . At the end of the proof for this lemma, we will show that  $T_1 = \Theta\left(\frac{1}{\sqrt{d}\eta}\right)$ . Then we can pick  $H := \frac{1}{1 - \beta_1} \log \frac{d}{\eta\xi^2}$  and apply Lemma 17 and Corollary 2 to get that, w.h.p., for all  $t \in [t_{\mathrm{inc}}, T_1)$  and  $\forall i, j \in [d]$ , eq. (27) can be written as
$$
\Delta W _ {1} ^ {(t)} [ i, j ] = - \eta_ {t} \frac {\left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} g _ {1} ^ {(t - \tau)} [ i , j ] + \epsilon_ {1 n} ^ {(t)} [ i , j ]}{\sqrt {\left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} \left(g _ {1} ^ {(t - \tau)} [ i , j ]\right) ^ {2} + \epsilon_ {1 d} ^ {(t)} [ i , j ]} + \xi}, \tag {32}
$$
$$
\Delta w _ {2 i} ^ {(t)} = - \eta_ {t} \frac {\left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} g _ {2 i} ^ {(t - \tau)} + \epsilon_ {2 n , i} ^ {(t)}}{\sqrt {\left(1 - \beta_ {2}\right) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} \left(g _ {2 i} ^ {(t - \tau)}\right) ^ {2} + \epsilon_ {2 d , i} ^ {(t)}} + \xi},
$$
where  $\forall i,j\in [d],\left|\epsilon_{1n}^{(t)}[i,j]\right|,\left|\epsilon_{1d}^{(t)}[i,j]\right|,\left|\epsilon_{2n,i}^{(t)}\right|,\left|\epsilon_{2d,i}^{(t)}\right|\leq \tilde{\mathcal{O}} (\eta \xi^2).$
Let's first look at the update of  $W_1^{(t)}[i,j]$ . For  $t$  in the first phase, we write the RHS of eq. (32) as
$$
\begin{array}{l} \frac {(1 - \beta_ {1}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} g _ {1} ^ {(t - \tau)} [ i , j ] + \epsilon_ {1 n} ^ {(t)} [ i , j ]}{\sqrt {(1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} \left(g _ {1} ^ {(t - \tau)} [ i , j ]\right) ^ {2} + \epsilon_ {1 d} ^ {(t)} [ i , j ]} + \xi} \\ = \frac {(1 - \beta_ {1}) g _ {1} ^ {(t)} [ i , j ] \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} + (1 - \beta_ {1}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} \left(g _ {1} ^ {(t - \tau)} [ i , j ] - g _ {1} ^ {(t)} [ i , j ]\right) + \epsilon_ {1 n} ^ {(t)} [ i , j ]}{\sqrt {(1 - \beta_ {2}) \left(g _ {1} ^ {(t)} [ i , j ]\right) ^ {2} \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} + (1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} \left(\left(g _ {1} ^ {(t - \tau)} [ i , j ]\right) ^ {2} - \left(g _ {1} ^ {(t)} [ i , j ]\right) ^ {2}\right) + \epsilon_ {1 d} ^ {(t)} [ i , j ]} + \xi}, \\ := \frac {g _ {1} ^ {(t)} [ i , j ] (1 - \beta_ {1} ^ {H + 1}) + e _ {1 n} ^ {(t)} [ i , j ] + \epsilon_ {1 n} ^ {(t)} [ i , j ]}{\sqrt {\left(g _ {1} ^ {(t)} [ i , j ]\right) ^ {2} (1 - \beta_ {2} ^ {H + 1}) + e _ {1 d} ^ {(t)} [ i , j ] + \epsilon_ {1 d} ^ {(t)} [ i , j ]} + \xi}, \\ \end{array}
$$
$$
e _ {1 n} ^ {(t)} [ i, j ] := \left(1 - \beta_ {1}\right) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} \left(g _ {1} ^ {(t - \tau)} [ i, j ] - g _ {1} ^ {(t)} [ i, j ]\right),
$$
$$
e _ {1 d} ^ {(t)} [ i, j ] := (1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {2} ^ {\tau} \left(\left(g _ {1} ^ {(t - \tau)} [ i, j ]\right) ^ {2} - \left(g _ {1} ^ {(t)} [ i, j ]\right) ^ {2}\right).
$$
We have already shown that  $\left|\epsilon_{1n}^{(t)}[i,j]\right|, \left|\epsilon_{1d}^{(t)}[i,j]\right| \leq \tilde{\mathcal{O}}(\eta\xi^2)$ . By Lemma 20, we have that  $\forall i, j \in [d]$ ,
$$
\begin{array}{l} \left| e _ {1 n} ^ {(t)} [ i, j ] \right| \leq (1 - \beta_ {1}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} \left| g _ {1} ^ {(t - \tau)} [ i, j ] - g _ {1} ^ {(t)} [ i, j ] \right| \\ \leq \left| g _ {1} ^ {(t)} [ i, j ] \right| \tilde {\mathcal {O}} (\sqrt {\eta}) (1 - \beta_ {1}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} \tau = \left| g _ {1} ^ {(t)} [ i, j ] \right| \tilde {\mathcal {O}} (\sqrt {\eta}). \\ \end{array}
$$
Similarly, we have  $\forall i,j\in [d]$
$$
\begin{array}{l} \left| e _ {1 d} ^ {(t)} [ i, j ] \right| \leq \left(g _ {1} ^ {(t)} [ i, j ]\right) ^ {2} \tilde {\mathcal {O}} (\sqrt {\eta}) (1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} \tau + \left(g _ {1} ^ {(t)} [ i, j ]\right) ^ {2} \tilde {\mathcal {O}} (\sqrt {\eta}) (1 - \beta_ {2}) \sum_ {\tau = 0} ^ {H} \beta_ {1} ^ {\tau} \tau^ {2} \\ = \left(g _ {1} ^ {(t)} [ i, j ]\right) ^ {2} \tilde {\mathcal {O}} (\sqrt {\eta}). \\ \end{array}
$$
By Lemma 19, we know that  $\left|g_1^{(t)}[i,j]\right| = \Omega\left(\sqrt{\eta}\right)$ . Then we have that
$$
\forall i, j \in [ d ]: \quad \left| \epsilon_ {1 n} ^ {(t)} [ i, j ] \right| \leq \tilde {\mathcal {O}} (\eta \xi^ {2}) \leq \tilde {\mathcal {O}} (\sqrt {\eta}) \left| g _ {1} ^ {(t)} [ i, j ] \right|, \quad \left| \epsilon_ {1 d} ^ {(t)} [ i, j ] \right| \leq \tilde {\mathcal {O}} (\sqrt {\eta}) \xi^ {2}.
$$
Therefore by Lemma 34 in Appendix H, we have