SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} W _ {1} ^ {(t)} = W _ {1} ^ {(t - 1)} + \frac {\eta}{1 - \beta} W _ {2} ^ {(t - 1) T} A + \frac {\eta}{1 - \beta} r _ {1} ^ {(t - 1)} \\ = W _ {1} ^ {(0)} + \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} W _ {2} ^ {(\tau) T} A + \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} \\ = W _ {1} ^ {(0)} + \frac {\eta}{1 - \beta} \left(C _ {1} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {1} ^ {\tau} + C _ {2} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} + \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} r _ {5} ^ {(\tau)}\right) ^ {T} A + \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} \\ = W _ {1} ^ {(0)} + \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} + \frac {\eta}{1 - \beta} \left(C _ {1} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {1} ^ {\tau} + \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} r _ {5} ^ {(\tau)}\right) ^ {T} A + \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} C _ {2} ^ {T} A \\ := W _ {1} ^ {(0)} + \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} + \left(C _ {1} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {1} ^ {\tau} + \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} r _ {5} ^ {(\tau)}\right) ^ {T} \boldsymbol {v} ^ {(t) T} + \boldsymbol {u} ^ {(t)} \boldsymbol {v} ^ {(t) T}, \\ \end{array}
$$
where  $\pmb{v}^{(t)T} = \frac{\eta}{1 - \beta}\pmb{A}$  and
$$
\boldsymbol {u} ^ {(t)} = \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} C _ {2} ^ {T}. \tag {8}
$$
Moreover, we have that
$$
W _ {2} ^ {(t)} = C _ {1} \lambda_ {1} ^ {t} + \left(C _ {2} + r _ {5} ^ {(t)}\right) \lambda_ {2} ^ {t} := C _ {1} \lambda_ {1} ^ {t} + r _ {5} ^ {(t)} \lambda_ {2} ^ {t} + c ^ {(t)} \pmb {u} ^ {(t) T}, \quad c ^ {(t)} = \frac {\lambda_ {2} ^ {t}}{\sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau}}.
$$
For  $t \leq T_1$ , by eq. (7), we get that w.h.p.,
$$
\begin{array}{l} \forall 1 \leq i, j \leq d: \quad \frac {\left| \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} r _ {5 i} ^ {(\tau)} v _ {j} ^ {(t)} \right|}{\left| u _ {i} ^ {(t)} v _ {j} ^ {(t)} \right|} \leq \frac {\tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 1}}\right) \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau}}{\tilde {\Omega} \left(\frac {1}{d ^ {\frac {5}{4} \alpha}}\right) \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau}} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - 1}}\right), \\ \frac {\left| \lambda_ {2} ^ {t} r _ {5 i} ^ {(t)} \right|}{\left| c ^ {(t)} u _ {i} ^ {(t)} \right|} = \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). \\ \end{array}
$$
For  $t \leq T_1$ , by Lemma 6,  $\forall 1 \leq i,j \leq d: \left|r_1^{(t)}[i,j]\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3}{2}\alpha - 1}}\right)$ . Then we have that w.h.p.
$$
\begin{array}{l} \frac {\left| \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} [ i , j ] \right|}{\left| u _ {i} ^ {(t)} v _ {j} ^ {(t)} \right|} = \frac {\left| \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} [ i , j ] \right|}{\left| \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau} C _ {2 i} A _ {j} \right|} \leq \frac {\left| \sum_ {\tau = 0} ^ {t - 1} r _ {1} ^ {(\tau)} [ i , j ] \right|}{\left| \sum_ {\tau = 0} ^ {t - 1} C _ {2 i} A _ {j} \right|} \leq \frac {\sum_ {\tau = 0} ^ {t - 1} \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 1}}\right)}{\sum_ {\tau = 0} ^ {t - 1} \tilde {\Omega} \left(\frac {1}{d ^ {\frac {5}{4} \alpha}}\right)} \\ = \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - 1}}\right). \\ \end{array}
$$
Here we used  $\forall i\in [d]:A_{i} = \Theta (1)$  by Assumption 1.
Since  $\lambda_1 = 1 - \frac{\eta}{1 - \beta}\| A\| _2$  , we have that  $|C_{1i}\lambda_1^t |\leq |C_{1i}|\leq \tilde{\mathcal{O}}\left(\frac{1}{d^\alpha}\right)$  and that
$$
\left| C _ {1 i} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {1} ^ {\tau} v _ {j} ^ {(t)} \right| = \frac {\eta A _ {j}}{1 - \beta} \left| C _ {1 i} \sum_ {\tau = 0} ^ {t - 1} \lambda_ {1} ^ {\tau} \right| \leq \frac {\eta A _ {j} | C _ {1 i} |}{(1 - \beta) (1 - \lambda_ {1})} \leq \frac {A _ {j} | C _ {1 i} |}{\| A \| _ {2}} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\alpha + \frac {1}{2}}}\right).
$$
Using the Gaussian tail bound and union bound, we have w.h.p.  $\forall 1\leq i,j\leq d:\left|W_1^{(0)}[i,j]\right| = \tilde{\mathcal{O}}\left(\frac{1}{d^{2\alpha}}\right)$ . Combining the above bounds together yields that for  $t\leq T_{1}$  and  $\forall i,j\in [d]$ ,
$$
\begin{array}{l} W _ {1} ^ {(t)} [ i, j ] = R _ {1 1} ^ {(t)} [ i, j ] + u _ {i} ^ {(t)} v _ {j} ^ {(t)} \left(1 + e _ {1} ^ {(t)} [ i, j ]\right), \tag {9} \\ w _ {2 i} ^ {(t)} = R _ {2 1. i} ^ {(t)} + c ^ {(t)} u _ {i} ^ {(t)} (1 + e _ {2 i} ^ {(t)}). \\ \end{array}
$$
where for  $\forall i,j\in [d]$ $\left|R_{11}^{(t)}[i,j]\right|\leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\alpha + \frac{1}{2}}}\right),\left|R_{21,i}^{(t)}\right|\leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\alpha}}\right)$  and  $\left|e_1^{(t)}[i,j]\right|, \left|e_{2i}^{(t)}\right|\leq$ $\tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{1}{4}\alpha - 1}}\right)$ .
Further we notice that for  $t \leq T_1$ , we have  $\forall j \in [d]$
$$
\frac {\left| v _ {j} ^ {(t)} \right|}{\left| c ^ {(t)} \right|} = \frac {\eta A _ {j}}{1 - \beta} \cdot \frac {\sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {\tau}}{\lambda_ {2} ^ {t}} = \frac {\eta A _ {j}}{1 - \beta} \frac {\lambda_ {2} ^ {t} - 1}{\lambda_ {2} ^ {t} (\lambda_ {2} - 1)} = \frac {A _ {j} \left(\lambda_ {2} ^ {t} - 1\right)}{\lambda_ {2} ^ {t} \| A \| _ {2}} \leq \frac {A _ {j}}{\| A \| _ {2}} = \mathcal {O} \left(\frac {1}{\sqrt {d}}\right).
$$
which yields that  $\left|u_{i}^{(t)}v_{j}^{(t)}\right| \leq \mathcal{O}\left(\frac{1}{\sqrt{d}}\right)\left|c^{(t)}u_{i}^{(t)}\right|$ . Together with eq. (9) gives us that  $\left|w_{2i}^{(t)}\right|$  reaches  $\frac{1}{d^{\alpha / 2}}$  for some  $i \in [d]$  before  $\left|W_1^{(t)}[k,j]\right|$  for  $\forall k,j \in [d]$ , i.e.  $T_{1} = \inf \left\{t \geq 0 : \exists i \in [d] : \left|w_{2i}^{(t)}\right| \geq \frac{1}{d^{\frac{\alpha}{2}}} \right\}$ .
Further, we know that at time  $T_{1}$ ,  $\left|c^{(T_1)}u_{i_0}^{(T_1)}\right| = |C_{2i_0}|\lambda_2^{T_1} = \Theta \left(\frac{1}{d^{\alpha / 2}}\right)$  for some  $i_0 \in [d]$ , which means w.h.p.
$$
\begin{array}{l} \frac {\Theta \left(\frac {1}{d ^ {\frac {\alpha}{2}}}\right)}{\tilde {\mathcal {O}} \left(\frac {1}{d ^ {\alpha}}\right)} \leq \lambda_ {2} ^ {T _ {1}} \leq \frac {\Theta \left(\frac {1}{d ^ {\frac {\alpha}{2}}}\right)}{\tilde {\Omega} \left(\frac {1}{d ^ {\frac {5}{4} \alpha}}\right)}, \quad \Rightarrow \quad \tilde {\Omega} \left(d ^ {\frac {\alpha}{2}}\right) \leq \lambda_ {2} ^ {T _ {1}} = \left(1 + \frac {\eta}{1 - \beta} \| A \| _ {2}\right) ^ {T _ {1}} \leq \tilde {\mathcal {O}} \left(d ^ {\frac {3}{4} \alpha}\right), \\ \Rightarrow \quad T _ {1} = \Theta \left(\frac {\log d}{\eta \| A \| _ {2}}\right). \tag {10} \\ \end{array}
$$
This is the length of the first phase. As for  $c^{(T_1)}u_i^{(T_1)}$  and  $u_{i}^{(T_{1})}v_{j}^{(T_{1})}$  for other coordinates, we have that w.h.p.  $\forall 1\leq i,j\leq d$
$$
\begin{array}{l} \left| u _ {i} ^ {(T _ {1})} v _ {j} ^ {(T _ {1})} \right| = \frac {\eta}{1 - \beta} \sum_ {\tau = 0} ^ {T _ {1} - 1} \lambda_ {2} ^ {\tau} | C _ {2 i} A _ {j} | = \frac {\eta}{1 - \beta} \cdot \frac {\lambda_ {2} ^ {T _ {1}} - 1}{\lambda_ {2} - 1} | C _ {2 i} A _ {j} | \stackrel {(i)} {=} \frac {\lambda_ {2} ^ {T _ {1}} - 1}{\| A \| _ {2}} | C _ {2 i} A _ {j} | \\ \geq \frac {\tilde {\Omega} \left(d ^ {\alpha / 2}\right)}{\Theta \left(\sqrt {d}\right)} \tilde {\Omega} \left(\frac {1}{d ^ {\frac {5}{4} \alpha}}\right) = \tilde {\Omega} \left(\frac {1}{d ^ {\frac {3}{4} \alpha + \frac {1}{2}}}\right), \\ \end{array}
$$
$$
\left| c ^ {(T _ {1})} u _ {i} ^ {(T _ {1})} \right| = | C _ {2 i} | \lambda_ {2} ^ {T _ {1}} \geq \tilde {\Omega} \left(d ^ {\alpha / 2}\right) \tilde {\Omega} \left(\frac {1}{d ^ {\frac {5}{4} \alpha}}\right) = \tilde {\Omega} \left(\frac {1}{d ^ {\frac {3}{4} \alpha}}\right).
$$
Here in (i) we used  $\lambda_{2} = 1 + \frac{\eta}{1 - \beta}\| A\|_{2}$ . Then we have at time  $T_{1}, \forall i, j \in [d]$ ,  $\left|\frac{R_{11}^{(T_1)}[i,j]}{u_i^{(T_1)}v_j^{(T_1)}}\right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{1}{4}\alpha}}\right)$  and that  $\frac{\left|R_{21,i}^{(T_1)}\right|}{\left|c^{(T_1)}u_i^{(T_1)}\right|} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{1}{4}\alpha}}\right)$ . Together with eq. (9), we have the following weight structure:
where 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| c ^ {(T _ {1})} u _ {i} ^ {(T _ {1})} \right|} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - 1}}\right).
$$
Finally, we consider the loss. Since  $\forall j\in [d]:\left(W_2^{(T_1)}W_1^{(T_1)}\right)_j - A_j = -\Theta (1)$ , we know that  $\bar{L}\left(W^{(T_1)}\right) = \Theta (d)$ .
D.3 Proof of Lemma 4
Eq. (8) tells us that  $\pmb{u}^{(T_1)} = \sum_{\tau=0}^{T_1-1} \lambda_2^\tau C_2^T$ . Lemma 8 tells us that  $C_2$  can be written as  $C_2 := \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)$ . Combining these two facts together finishes the proof.
D.4 Proof of Lemma 5
Replacing  $t$  by  $t - 1$  in eq. (6), we get
$$
W _ {2} ^ {(t)} = W _ {2} ^ {(t - 1)} + \frac {\eta}{1 - \beta} A W _ {1} ^ {(t - 1) T} + \frac {\eta}{1 - \beta} r _ {2} ^ {(t - 1)}. \tag {11}
$$
Eq. (6)-(11) and substituting eq. (5) yield
$$
\begin{array}{l} W _ {2} ^ {(t + 1)} - W _ {2} ^ {(t)} = W _ {2} ^ {(t)} - W _ {2} ^ {(t - 1)} + \frac {\eta^ {2}}{(1 - \beta) ^ {2}} \| A \| _ {2} ^ {2} W _ {2} ^ {(t - 1)} + \frac {\eta^ {2}}{(1 - \beta) ^ {2}} A r _ {1} ^ {(t - 1) T} \\ + \frac {\eta}{1 - \beta} \left(r _ {2} ^ {(t)} - r _ {2} ^ {(t - 1)}\right), \\ \Rightarrow W _ {2} ^ {(t + 1)} = 2 W _ {2} ^ {(t)} - \left(1 - \frac {\eta^ {2}}{(1 - \beta) ^ {2}} \| A \| _ {2} ^ {2}\right) W _ {2} ^ {(t - 1)} + r _ {3} ^ {(t)}, \\ \end{array}
$$
where  $r_3(t) \coloneqq \frac{\eta^2}{(1 - \beta)^2} Ar_1^{(t - 1)T} + \frac{\eta}{1 - \beta}\left(r_2^{(t)} - r_2^{(t - 1)}\right)$ .
For the equation  $x^{2} - 2x + 1 - \frac{\eta^{2}}{(1 - \beta)^{2}} \| A \|_{2}^{2} = 0$ , the roots are  $\lambda_{1} = 1 - \frac{\eta}{1 - \beta} \| A \|_{2}$  and  $\lambda_{2} = 1 + \frac{\eta}{1 - \beta} \| A \|_{2}$ . We have that
$$
\begin{array}{l} W _ {2} ^ {(t + 1)} - \lambda_ {2} W _ {2} ^ {(t)} = \lambda_ {1} \left(W _ {2} ^ {(t)} - \lambda_ {2} W _ {2} ^ {(t - 1)}\right) + r _ {3} ^ {(t)} \\ \Rightarrow W _ {2} ^ {(t)} - \lambda_ {2} W _ {2} ^ {(t - 1)} = \lambda_ {1} ^ {t - 1} \left(W _ {2} ^ {(1)} - \lambda_ {2} W _ {2} ^ {(0)}\right) + \sum_ {\tau = 1} ^ {t - 1} \lambda_ {1} ^ {t - 1 - \tau} r _ {3} ^ {(\tau)} \\ := \lambda_ {1} ^ {t - 1} \left(W _ {2} ^ {(1)} - \lambda_ {2} W _ {2} ^ {(0)}\right) + r _ {4} ^ {(t)}. \\ \end{array}
$$
$$
\begin{array}{l} W _ {2} ^ {(t)} = \lambda_ {2} ^ {t} W _ {2} ^ {(0)} + \sum_ {\tau = 0} ^ {t - 1} \lambda_ {2} ^ {t - 1 - \tau} \lambda_ {1} ^ {\tau} \left(W _ {2} ^ {(1)} - \lambda_ {2} W _ {2} ^ {(0)}\right) + \sum_ {\tau = 1} ^ {t} \lambda_ {2} ^ {t - \tau} r _ {4} ^ {(\tau)} \\ = \lambda_ {2} ^ {t} W _ {2} ^ {(0)} + \frac {\lambda_ {2} ^ {t} - \lambda_ {1} ^ {t}}{\lambda_ {2} - \lambda_ {1}} \left(W _ {2} ^ {(1)} - \lambda_ {2} W _ {2} ^ {(0)}\right) + \sum_ {\tau = 1} ^ {t} \lambda_ {2} ^ {t - \tau} r _ {4} ^ {(\tau)} \\ = C _ {1} \lambda_ {1} ^ {t} + C _ {2} \lambda_ {2} ^ {t} + \sum_ {\tau = 1} ^ {t} \lambda_ {2} ^ {t - \tau} r _ {4} ^ {(\tau)} \\ = C _ {1} \lambda_ {1} ^ {t} + \left(C _ {2} + r _ {5} ^ {(t)}\right) \lambda_ {2} ^ {t}, \\ \end{array}
$$
where  $r_5^{(t)} = \sum_{\tau = 1}^t\lambda_2^{-\tau}r_4^{(\tau)},C_1 = -\frac{W_2^{(1)} - \lambda_2W_2^{(0)}}{\lambda_2 - \lambda_1}$  and  $C_2 = \frac{W_2^{(1)} - \lambda_1W_2^{(0)}}{\lambda_2 - \lambda_1}$ .
D.5 Proof of Lemma 6
Write  $r_1^{(t)} = -\beta^{t + 1}W_2^{(t)T}A + q_{12}^{(t)} + q_{13}^{(t)} + q_{14}^{(t)}$  where  $q_{12}^{(t)} = (1 - \beta)\sum_{\tau = 0}^{t}\beta^{t - \tau}\left(W_2^{(\tau)T} - W_2^{(t)T}\right)A,$ $q_{13}^{(t)} = -(1 - \beta)\sum_{\tau = 0}^{t}\beta^{t - \tau}W_{2}^{(\tau)T}W_{2}^{(\tau)}W_{1}^{(\tau)}$  and  $q_{14}^{(t)} = -(1 - \beta)\sum_{\tau = 0}^{t}\beta^{t - \tau}Dg_{1}^{(\tau)}.$  And write  $r_2^{(t)} = -\beta^{t + 1}AW_1^{(t)T} + q_{22}^{(t)} + q_{23}^{(t)} + q_{24}^{(t)}$  , where  $q_{22}^{(t)} = (1 - \beta)\sum_{\tau = 0}^{t}\beta^{t - \tau}A\left(W_1^{(\tau)T} - W_1^{(t)T}\right),q_{23}^{(t)} = -(1 - \beta)\sum_{\tau = 0}^{t}\beta^{t - \tau}W_2^{(\tau)}W_1^{(\tau)}W_1^{(\tau)T}$
and  $q_{24}^{(t)} = -(1 - \beta)\sum_{\tau = 0}^{t}\beta^{t - \tau}Dg_2^{(\tau)}$
Let's first try to bound  $\left|q_{12}^{(t)}[i,j]\right|$  and  $\left|q_{22,i}^{(t)}\right|$ . For any  $\tau \leq T_1$ , we have that
$$
\forall i \in [ d ]: \quad \left| \left(W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)}\right) _ {i} \right| = \left| \sum_ {j = 1} ^ {d} w _ {2 j} ^ {(\tau)} W _ {1} ^ {(\tau)} [ j, i ] \right| \leq \sum_ {j = 1} ^ {d} \left| w _ {2 j} ^ {(\tau)} \right| \left| W _ {1} ^ {(\tau)} [ j, i ] \right| \leq \sum_ {i = 1} ^ {d} \frac {1}{d ^ {\alpha}} = \frac {1}{d ^ {\alpha - 1}},
$$
and thus  $\forall i\in [d]:\left|E_i^{(\tau)}\right| = \mathcal{O}(1)$ . Then we have for all  $i,j\in [d]$
$$
\begin{array}{l} \left| W _ {1} ^ {(\tau + 1)} [ i, j ] - W _ {1} ^ {(\tau)} [ i, j ] \right| \leq \eta \sum_ {k = 0} ^ {\tau} \beta^ {\tau - k} \left| w _ {2 i} ^ {(k)} E _ {j} ^ {(k)} \right| \leq \eta \sum_ {k = 0} ^ {\tau} \beta^ {\tau - k} \mathcal {O} \left(\frac {1}{d ^ {\alpha / 2}}\right) \\ = \eta \mathcal {O} \left(\frac {1}{d ^ {\alpha / 2}}\right), \\ \end{array}
$$
$$
\begin{array}{l} \left| w _ {2 i} ^ {(\tau + 1)} - w _ {2 i} ^ {(\tau)} \right| \leq \eta \sum_ {k = 0} ^ {\tau} \beta^ {\tau - k} \sum_ {j = 1} ^ {d} \left| E _ {j} ^ {(k)} W _ {1} ^ {(k)} [ i, j ] \right| \leq \eta \sum_ {k = 0} ^ {\tau} \beta^ {\tau - k} \mathcal {O} \left(\frac {1}{d ^ {\alpha / 2 - 1}}\right) \\ = \eta \mathcal {O} \left(\frac {1}{d ^ {\alpha / 2 - 1}}\right). \\ \end{array}
$$
That gives us  $\forall i,j\in [d]$
$$
\begin{array}{l} \left| q _ {1 2} ^ {(t)} [ i, j ] \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left| \left(w _ {2 i} ^ {(\tau)} - w _ {2 i} ^ {(t)}\right) A _ {j} \right| \leq \eta (1 - \beta) \sum_ {\tau = 0} ^ {t} \mathcal {O} \left(\frac {\beta^ {t - \tau} (t - \tau)}{d ^ {\alpha / 2 - 1}}\right) \\ = \mathcal {O} \left(\frac {\eta}{d ^ {\alpha / 2 - 1}}\right), \\ \end{array}
$$
$$
\begin{array}{l} \left| q _ {2 2, i} ^ {(t)} \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \sum_ {j = 1} ^ {d} \left| A _ {j} \left(W _ {1} ^ {(\tau)} [ i, j ] - W _ {1} ^ {(t)} [ i, j ]\right) \right| \\ \leq \eta (1 - \beta) \sum_ {\tau = 0} ^ {t} \mathcal {O} \left(\frac {\beta^ {t - \tau} (t - \tau)}{d ^ {\alpha / 2 - 1}}\right) = \mathcal {O} \left(\frac {\eta}{d ^ {\alpha / 2 - 1}}\right). \\ \end{array}
$$
Then we bound  $\left|q_{13}^{(t)}[i,j]\right|$  and  $\left|q_{23,i}^{(t)}\right|$ . We have for  $\forall i, j \in [d]$ ,
$$
\begin{array}{l} \left| q _ {1 3} ^ {(t)} [ i, j ] \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left| w _ {2 i} ^ {(\tau)} \left(W _ {2} ^ {(\tau)} W _ {1} ^ {(\tau)}\right) _ {j} \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \frac {1}{d ^ {\frac {\alpha}{2}}} \cdot \frac {1}{d ^ {\alpha - 1}} \\ = \mathcal {O} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 1}}\right), \\ \end{array}
$$
$$
\begin{array}{l} \left| q _ {2 3, i} ^ {(t)} \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \sum_ {j = 1} ^ {d} \left| \left(W _ {2} ^ {(t)} W _ {1} ^ {(t)}\right) _ {j} W _ {1} ^ {(t)} [ i, j ] \right| \\ \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \sum_ {i = 1} ^ {d} \frac {1}{d ^ {\alpha - 1 + \frac {\alpha}{2}}} = \mathcal {O} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 2}}\right). \\ \end{array}
$$
Finally we use Lemma 32 to bound  $\left|q_{14}^{(t)}[i,j]\right|$  and  $\left|q_{24,i}^{(t)}\right|$ . For  $t \leq T_1$ , the  $M_1^{(t)}, M_2^{(t)}$  in Lemma 32 are upper bounded by  $\frac{1}{d^{\frac{\alpha}{2}}}$ . In the theorem we consider the training period before  $T_{\mathrm{SGD},2}$  so the time  $T$  in Lemma 32 is set as  $T_{\mathrm{SGD},2}$ . In the following sections, we will prove that  $T_{\mathrm{SGD},2} \leq \mathcal{O}\left(\frac{d^{\alpha}\log(\sqrt{d / \epsilon})}{\eta}\right)$ . Then by Lemma 32, we have with probability at least  $1 - \frac{1}{d}$ , for  $\forall t \leq T_1$  and
$$
\forall i, j \in [ d ],
$$
$$
\begin{array}{l} \left| D g _ {1} ^ {(t)} [ i, j ] \right| = \left| \tilde {g} _ {1} ^ {(t)} [ i, j ] - g _ {1} ^ {(t)} [ i, j ] \right| \leq \mathcal {O} \left(\frac {1}{d ^ {\frac {3 \alpha}{2} - 3}} \sigma \sqrt {\frac {d ^ {\alpha + 1}}{\eta} \log \frac {d}{\epsilon}}\right) + \mathcal {O} \left(\frac {1}{d ^ {\frac {\alpha}{2}}} \sigma \sqrt {\frac {d ^ {\alpha + 2}}{\eta} \log \frac {d}{\epsilon}}\right) \\ \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {\alpha}{2}}} \sigma \sqrt {\frac {d ^ {\alpha + 2}}{\eta}}\right), \\ \end{array}
$$
$$
\begin{array}{l} \left| D g _ {2 i} ^ {(t)} \right| = \left| \tilde {g} _ {2 i} ^ {(t)} - g _ {2 i} ^ {(t)} \right| \leq \mathcal {O} \left(\frac {1}{d ^ {\frac {3 \alpha}{2} - 4}} \sigma \sqrt {\frac {d ^ {\alpha + 1}}{\eta} \log \frac {d}{\epsilon}}\right) + \mathcal {O} \left(\frac {1}{d ^ {\frac {\alpha}{2} - 1}} \sigma \sqrt {\frac {d ^ {\alpha + 2}}{\eta} \log \frac {d}{\epsilon}}\right) \\ \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {\alpha}{2} - 1}} \sigma \sqrt {\frac {d ^ {\alpha + 2}}{\eta}}\right). \\ \end{array}
$$
By picking  $\sigma \leq \frac{\eta^{3/2}}{d^{\alpha/2+1}}$ , we have w.h.p. for  $\forall t \leq T_1$  and  $\forall i, j \in [d]$ ,  $|Dg_1^{(t)}[i,j]| \leq \eta \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{\alpha}{2}}}\right)$  and  $|Dg_{2i}^{(t)}| \leq \eta \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{\alpha}{2}-1}}\right)$ , which yields
$$
\begin{array}{l} \left| q _ {1 4} ^ {(t)} [ i, j ] \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left| D g _ {1} ^ {(\tau)} [ i, j ] \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \eta \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {\alpha}{2}}}\right) = \eta \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {\alpha}{2}}}\right), \\ \left| q _ {2 4, i} ^ {(t)} \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \left| D g _ {2 i} ^ {(\tau)} \right| \leq (1 - \beta) \sum_ {\tau = 0} ^ {t} \beta^ {t - \tau} \eta \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {\alpha}{2} - 1}}\right) = \eta \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {\alpha}{2} - 1}}\right). \\ \end{array}
$$
Combining all the above bounds and substituting  $\eta \leq \mathcal{O}\left(\frac{1}{d^{\alpha}}\right)$  gives us for  $\forall t \leq T_{1}$  and  $\forall i, j \in [d]$ ,
$$
\left| r _ {1} ^ {(t)} [ i, j ] \right| \leq \beta^ {t + 1} \left| w _ {2 i} ^ {(t)} A _ {j} \right| + \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 1}}\right), \quad \left| r _ {2 i} ^ {(t)} \right| \leq \beta^ {t + 1} \left| \sum_ {j = 1} ^ {d} A _ {j} W _ {1} ^ {(t)} [ i, j ] \right| + \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 2}}\right). \tag {12}
$$
For  $t \leq T_1$ , we have  $\forall i, j \in [d], \left|w_{2i}^{(t)}A_j\right| \leq \mathcal{O}\left(\frac{1}{d^{\alpha/2}}\right)$  and  $\left|\sum_{j=1}^{d}A_jW_1^{(t)}[i,j]\right| \leq \mathcal{O}\left(\frac{1}{d^{\alpha/2-1}}\right)$ , which gives us  $\left|r_1^{(t)}[i,j]\right| \leq \mathcal{O}\left(\frac{1}{d^{\alpha/2}}\right)$  and  $\left|r_{2i}^{(t)}\right| \leq \mathcal{O}\left(\frac{1}{d^{\alpha/2-1}}\right)$ . Substituting into eq. (5) and eq. (6) yields that for  $t \leq T_1$  and  $\forall i, j \in [d]$ ,
$$
\left| W _ {1} ^ {(t + 1)} [ i, j ] - W _ {1} ^ {(t)} [ i, j ] \right| \leq \mathcal {O} \left(\frac {\eta}{d ^ {\alpha / 2}}\right), \quad \left| w _ {2 i} ^ {(t + 1)} - w _ {2 i} ^ {(t)} \right| \leq \mathcal {O} \left(\frac {\eta}{d ^ {\alpha / 2 - 1}}\right).
$$
Hence for  $t \leq \min \left\{\frac{\alpha \log d}{\log(1 / \beta)}, T_1\right\}$ , we have  $\forall i, j \in [d]$
$$
\begin{array}{l} \left| W _ {1} ^ {(t)} [ i, j ] \right| \leq \left| W _ {1} ^ {(0)} [ i, j ] \right| + \frac {\alpha \log d}{\log (1 / \beta)} \mathcal {O} \left(\frac {\eta}{d ^ {\alpha / 2}}\right) \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3 \alpha}{2}}}\right), \\ \left| w _ {2 i} ^ {(t)} \right| \leq \left| w _ {2 i} ^ {(0)} \right| + \frac {\alpha \log d}{\log (1 / \beta)} \mathcal {O} \left(\frac {\eta}{d ^ {\alpha / 2 - 1}}\right) \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3 \alpha}{2} - 1}}\right). \\ \end{array}
$$
Then we know that  $T_{1} > \frac{\alpha\log d}{\log(1 / \beta)}$  and also get tighter bounds of  $\left|W_1^{(t)}[i,j]\right|, \left|w_{2i}^{(t)}\right|$  for  $t \leq \frac{\alpha\log d}{\log(1 / \beta)}$ . Now we use these new bounds to analyze  $\left|r_1^{(t)}[i,j]\right|$  and  $\left|r_{2i}^{(t)}\right|$  again.
When  $t \leq \frac{\alpha \log d}{\log(1 / \beta)}$ , we have for all  $i, j \in [d]$ ,  $\beta^{t + 1} \left| w_{2i}^{(t)} A_j \right| \leq \left| w_{2i}^{(t)} A_j \right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3\alpha}{2} - 1}}\right)$  and  $\beta^{t + 1} \left| \sum_{j = 1}^{d} A_j W_1^{(t)}[i,j] \right| \leq \left| \sum_{j = 1}^{d} A_j W_1^{(t)}[i,j] \right| \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3\alpha}{2} - 1}}\right)$ . When  $\frac{\alpha \log d}{\log(1 / \beta)} < t \leq T_1$ , we have  $\beta^{t + 1} \leq \frac{1}{d^{\alpha}}$ , suggesting that  $\forall i, j \in [d]$ ,  $\beta^{t + 1} \left| w_{2i}^{(t)} A_j \right| \leq \frac{1}{d^{\alpha}} \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{\alpha}{2}}}\right) \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3\alpha}{2}}}\right)$  and
$\beta^{t + 1}\left|\sum_{j = 1}^{d}A_{j}W_{1}^{(t)}[i,j]\right|\leq \frac{1}{d^{\alpha}}\tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{\alpha}{2} - 1}}\right)\leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{3\alpha}{2} - 1}}\right).$  Substituting into (12) completes the proof.
D.6 Proof of Lemma 7
Based on the bound in Lemma 6, we have
$$
\begin{array}{l} \left| r _ {3 i} ^ {(t)} \right| = \left| \frac {\eta^ {2}}{(1 - \beta) ^ {2}} \sum_ {j = 1} ^ {d} A _ {j} r _ {1} ^ {(t - 1)} [ i, j ] + \frac {\eta}{1 - \beta} \left(r _ {2 i} ^ {(t)} - r _ {2 i} ^ {(t - 1)}\right) \right| \\ \leq \frac {\eta^ {2}}{(1 - \beta) ^ {2}} \sum_ {j = 1} ^ {d} \left| A _ {j} r _ {1} ^ {(t - 1)} [ i, j ] \right| + \frac {\eta}{1 - \beta} \left| r _ {2 i} ^ {(t)} \right| + \frac {\eta}{1 - \beta} \left| r _ {2 i} ^ {(t - 1)} \right| \\ \leq \eta^ {2} \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 2}}\right) + 2 \eta \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 2}}\right) = \eta \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 2}}\right). \\ \end{array}
$$
Since  $\lambda_1 = 1 - \frac{\eta}{1 - \beta}\| A\| _2, \lambda_2 = 1 + \frac{\eta}{1 - \beta}\| A\| _2$ , and note that  $\| A\| _2 = \Theta \left(\sqrt{d}\right)$ , we have that
$$
\left| r _ {4 i} ^ {(t)} \right| = \left| \sum_ {\tau = 1} ^ {t - 1} \lambda_ {1} ^ {t - 1 - \tau} r _ {3 i} ^ {(\tau)} \right| \leq \sum_ {\tau = 1} ^ {t - 1} \lambda_ {1} ^ {t - 1 - \tau} \tilde {\mathcal {O}} \left(\frac {\eta}{d ^ {\frac {3}{2} \alpha - 2}}\right) \leq \frac {\eta}{1 - \lambda_ {1}} \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} \alpha - 2}}\right) = \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {3}{2} (\alpha - 1)}}\right),
$$