SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left(1 + \frac {\left| c ^ {(t)} \right|}{\left\| \boldsymbol {v} ^ {(t)} \right\| _ {2}}\right) \sqrt {\frac {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2}}} \leq \mathcal {O} \left(\frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{\left| X _ {i} \right|}\right), \tag {48} \\ \left(1 + \frac {\left| v _ {k} ^ {(t)} \right|}{\left| c ^ {(t)} \right|}\right) \sqrt {\frac {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2}}} \leq \mathcal {O} \left(\frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{\left| X _ {i} \right|}\right). \\ \end{array}
$$
By the proof in Section D.8, we know that for  $t \in [T_{\mathrm{SGD},1}, T_{\mathrm{SGD},2}]$ ,  $\forall i \in [d] : v_i^{(t)}, c^{(t)}$  are positive and monotonically increasing. On the other hand, the proof in Section D.2 and D.9 tells us that w.h.p.  $\left\| E^{(t)} \right\|_2$  (resp.  $\forall k \in [d]$ ,  $\left| E_k^{(t)} \right|$ ) decreases from  $\Theta(\sqrt{d})$  (resp.  $\Theta(1)$ ) when  $t = T_{\mathrm{SGD},1}$  to  $\mathcal{O}(\sqrt{\epsilon_0 d})$  (resp.  $\mathcal{O}(\sqrt{\epsilon_0})$ ) when  $t = T_{\mathrm{SGD},2}$ . Therefore, the trend of  $\frac{\left\| E^{(t)} \right\|_2}{u_i^2 \left\| v^{(t)} \right\|_2^2}$  and  $\frac{\left| E_k^{(t)} \right|}{(c^{(t)})^2 u_i^2}$  is to decrease over time, and when  $t = T_{\mathrm{SGD},2}$ , we have w.h.p.
$$
\forall k \in [ d ]: \left| E _ {k} ^ {(t)} \right| = \mathcal {O} (\sqrt {\epsilon_ {0}}), \quad \left\| E ^ {(t)} \right\| _ {2} = \mathcal {O} (\sqrt {\epsilon_ {0} d}). \tag {49}
$$
Moreover, when  $t = T_{\mathrm{SGD},2}$ , the inequality in eq. (26) becomes equality, i.e.  $c^2 \| \pmb{u} \|_2^2 = \Theta \left( \sqrt{d} \right)$  and  $\forall j \in [d] : \| \pmb{u} \|_2^2 v_j^2 = \Theta \left( \frac{1}{\sqrt{d}} \right)$ .
Using  $\pmb{u} = X + Y$  and eq. (46), we have
$$
c ^ {2} \| X \| _ {2} ^ {2} = \Theta (\sqrt {d}), \quad \forall j \in [ d ]: \| X \| _ {2} ^ {2} v _ {j} ^ {2} \Theta \left(\frac {1}{\sqrt {d}}\right), \quad \Rightarrow \quad \| X \| _ {2} ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2} = \Theta (\sqrt {d}),
$$
which together with the second inequality in eq. (47) yields
$$
\frac {1}{u _ {i} ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2}} \leq \left(\frac {1}{1 - \delta_ {x y}}\right) ^ {2} \frac {1}{X _ {i} ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2}} = \Theta \left(\frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2} \sqrt {d}}\right),
$$
$$
\frac {1}{c ^ {2} u _ {i} ^ {2}} \leq \left(\frac {1}{1 - \delta_ {x y}}\right) ^ {2} \frac {1}{c ^ {2} X _ {i} ^ {2}} = \Theta \left(\frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2} \sqrt {d}}\right).
$$
Combining with eq. (49), we get that
$$
\frac {\left\| E ^ {(t)} \right\| _ {2}}{u _ {i} ^ {2} \left\| \boldsymbol {v} ^ {(t)} \right\| _ {2} ^ {2}} \leq \mathcal {O} \left(\frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} \cdot \sqrt {\epsilon_ {0}}\right), \quad \frac {\left| E _ {k} ^ {(t)} \right|}{(c ^ {(t)}) ^ {2} u _ {i} ^ {2}} \leq \mathcal {O} \left(\frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} \cdot \sqrt {\frac {\epsilon_ {0}}{d}}\right). \tag {50}
$$
Substituting eq. (48) and (50) into eq. (44) and (45) gives us
$$
\forall 1 \leq i \leq d: \frac {\sqrt {\sum_ {j \neq i} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{| H ^ {(t)} [ i , i ] |} \leq \mathcal {O} \left(\frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{| X _ {i} |}\right) + q _ {1 i} ^ {(t)},
$$
where the trend of  $q_{1i}^{(t)}$  is to decrease over time and  $q_{1i}^{(T_{\mathrm{SGD},2})} \leq \mathcal{O}\left(\frac{\sum_{j=1}^{d} X_j^2}{X_i^2} \cdot \sqrt{\epsilon_0}\right)$ .
$$
\forall 1 \leq i \leq d, 1 \leq k \leq d: \frac {\sqrt {\sum_ {j \neq i d + k} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{\left| H ^ {(t)} [ i d + k , i d + k ] \right|} \leq \mathcal {O} \left(\frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{\left| X _ {i} \right|}\right) + q _ {2 i} ^ {(t)},
$$
where the trend of  $q_{2i}^{(t)}$  is to decrease over time and  $q_{2i}^{(T_{\mathrm{SGD},2})} \leq \mathcal{O}\left(\frac{\sum_{j=1}^{d} X_j^2}{X_i^2} \cdot \sqrt{\frac{\epsilon_0}{d}}\right)$ .
$$
\begin{array}{l} R _ {\mathrm {d i a g}} ^ {\mathrm {S G D M}} (t) = \mathcal {O} \left(\frac {1}{d} \sum_ {i = 1} ^ {d} \frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{| X _ {i} |}\right) + \frac {1}{d ^ {2} + d} \sum_ {i = 1} ^ {d} q _ {1 i} ^ {(t)} + \frac {d}{d ^ {2} + d} \sum_ {i = 1} ^ {d} q _ {2 i} ^ {(t)} \\ := \mathcal {O} \left(\frac {1}{d} \sum_ {i = 1} ^ {d} \frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{| X _ {i} |}\right) + q ^ {(t)}, \\ \end{array}
$$
where the trend of  $q^{(t)}$  is to decrease over time and
$$
\begin{array}{l} q ^ {(T _ {\mathrm {s g d}, 2})} \leq \frac {1}{d ^ {2} + d} \sum_ {i = 1} ^ {d} \mathcal {O} \left(\frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} \cdot \sqrt {\epsilon_ {0}}\right) + \frac {d}{d ^ {2} + d} \sum_ {i = 1} ^ {d} \mathcal {O} \left(\frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} \cdot \sqrt {\frac {\epsilon_ {0}}{d}}\right) \\ \leq \mathcal {O} \left(\frac {1}{d ^ {2} + d} \sum_ {i = 1} ^ {d} \frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} \cdot \sqrt {\epsilon_ {0} d}\right) = \mathcal {O} \left(\frac {1}{d} \sum_ {i = 1} ^ {d} \frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} \cdot \sqrt {\epsilon_ {0}} \right). \\ \end{array}
$$
Denote  $\sigma^2$  as the variance of  $X_{i}$  for  $i\in [d]$ . By concentration of chi-squared distribution, we know that with probability at least  $1 - \delta$  for  $\delta >0$
$$
\sum_ {i = 1} ^ {d} X _ {i} ^ {2} \leq \sigma^ {2} d + \sigma^ {2} \mathcal {O} \left(\sqrt {d \log \frac {1}{\delta}}\right).
$$
By Lemma 36 in Appendix H, we know that with constant probability  $\frac{1}{d}\sum_{i=1}^{d}\frac{1}{|X_i|} = \mathcal{O}\left(\frac{1}{\sigma}\log d\right)$ . Then with constant probability,  $\frac{1}{d}\sum_{i=1}^{d}\frac{1}{X_i^2} \leq \frac{1}{d}\left(\sum_{i=1}^{d}\frac{1}{|X_i|}\right)^2 = \mathcal{O}\left(\frac{d}{\sigma^2}\log^2 d\right)$ . Hence
$$
\frac {1}{d} \sum_ {i = 1} ^ {d} \frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{| X _ {i} |} = \tilde {\mathcal {O}} (\sqrt {d}), \quad \frac {1}{d} \sum_ {i = 1} ^ {d} \frac {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}{X _ {i} ^ {2}} = \tilde {\mathcal {O}} \left(d ^ {2}\right).
$$
Therefore with constant probability,
$$
R _ {\text {d i a g}} ^ {\text {S G D M}} (t) = \tilde {\mathcal {O}} (\sqrt {d}) + q ^ {(t)},
$$
where the trend of  $q^{(t)}$  is to decrease over time and  $q^{(T_{\mathrm{SGD},2})} \leq \tilde{\mathcal{O}}\left(d\sqrt{\epsilon_0d}\right)$ . For any  $p > 0$ , by picking the same hyperparameters as in Theorem 1, we have  $\epsilon_0d \leq \tilde{\mathcal{O}}\left(\frac{1}{d^p}\right)$  and hence  $q^{(T_{\mathrm{SGD},2})} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{p/2-1}}\right) = o(d)$ .
F.3.3 Proof of eq. (43)
By the analyses in Section E.1, we know that for  $t \in [T_{\mathrm{Adam},1}, T_{\mathrm{Adam},2}]$ , the weights obtained by Adam satisfy
$$
W _ {1} ^ {(t)} = \boldsymbol {u} \boldsymbol {v} ^ {(t) T} + R _ {1} ^ {(t)},
$$
$$
W _ {2} ^ {(t)} = c ^ {(t)} \boldsymbol {u} ^ {T} + R _ {2} ^ {(t) T},
$$
where  $\forall i\in [d]:u_i = \mathrm{sign}(w_{2i}^{(0)})\in \{\pm 1\}$  and
$$
\forall 1 \leq i, j \leq d: \quad \frac {\left| R _ {1} ^ {(t)} [ i , j ] \right|}{\left| u _ {i} v _ {j} ^ {(t)} \right|} \leq \delta := \tilde {\mathcal {O}} \left(\eta^ {\frac {1}{4}} + \frac {1}{d ^ {\frac {\alpha}{2} - \frac {1}{4}}}\right), \quad \frac {\left| R _ {2 i} ^ {(t)} \right|}{\left| c ^ {(t)} u _ {i} \right|} \leq \delta .
$$
Hence by Lemma 29, when  $t \in [T_{\mathrm{Adam},1}, T_{\mathrm{Adam},2}]$ , we have for  $1 \leq i \leq d$
$$
\begin{array}{l} \frac {\sqrt {\sum_ {j \neq i} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{\left| H ^ {(t)} [ i , i ] \right|} \leq \frac {1 + \delta}{1 - \delta} \left(1 + \frac {\left| c ^ {(t)} \right|}{\left\| \boldsymbol {v} ^ {(t)} \right\| _ {2}}\right) \sqrt {\frac {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2}}} + \frac {\left\| E ^ {(t)} \right\| _ {2}}{\left. (1 - \delta) ^ {2} u _ {i} ^ {2} \left\| \boldsymbol {v} ^ {(t)} \right\| _ {2} ^ {2} \right.} \tag {51} \\ = \mathcal {O} \left(1 + \frac {\left| c ^ {(t)} \right|}{\left\| \boldsymbol {v} ^ {(t)} \right\| _ {2}}\right) \sqrt {d} + \mathcal {O} \left(\frac {\left\| E ^ {(t)} \right\| _ {2}}{\left\| \boldsymbol {v} ^ {(t)} \right\| _ {2} ^ {2}}\right), \\ \end{array}
$$
$$
\begin{array}{l} \frac {\sqrt {\sum_ {j \neq i d + k} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{\left| H ^ {(t)} [ i d + k , i d + k ] \right|} \leq \frac {1 + \delta}{1 - \delta} \left(1 + \frac {\left| v _ {k} ^ {(t)} \right|}{\left| c ^ {(t)} \right|}\right) \sqrt {\frac {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2}}} + \frac {\left| E _ {k} ^ {(t)} \right|}{\left(1 - \delta\right) ^ {2} \left(c ^ {(t)}\right) ^ {2} u _ {i} ^ {2}} \tag {52} \\ = \mathcal {O} \left(1 + \frac {\left| v _ {k} ^ {(t)} \right|}{\left| c ^ {(t)} \right|}\right) \sqrt {d} + \mathcal {O} \left(\frac {\left| E _ {k} ^ {(t)} \right|}{\left(c ^ {(t)}\right) ^ {2}}\right). \\ \end{array}
$$
Recall the following facts of Adam.
(A) By Lemma 15, we know that for  $t \in [T_{\mathrm{Adam},1}, T_1]$  (where  $T_1$  is defined in Definition 4), w.h.p.  $\forall k \in [d] : v_k^{(t)} = c^{(t)} = \eta(t - t_{\mathrm{inc}})$ . Specially, when  $t = T_{\mathrm{Adam},1}$ ,  $\forall k \in [d] : v_k^{(t)} = c^{(t)} = \frac{1}{d^{\frac{\alpha}{2}}}$ . Lemma 25 and 27 tell us that for  $t \in [T_1, T_{\mathrm{Adam},2}]$  w.h.p.  $\forall i, j \in [d] : |W_1[i, j]| = \tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right), |w_{2i}| = \tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right)$ , which gives us  $\forall k \in [d] : |v_k^{(t)}| = \tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right)$  and  $|c^{(t)}| = \tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right)$ . That means when  $t \in [T_{\mathrm{Adam},1}, T_{\mathrm{Adam},2}]$ ,  $\forall k \in [d] : |v_k^{(t)}|$  and  $|c^{(t)}|$  increase from  $\frac{1}{d^{\frac{\alpha}{2}}}$  to  $\tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right)$  and  $\frac{|v_k^{(t)}}{|c^{(t)}} = \tilde{\Theta}(1)$ ,  $\frac{|c^{(t)}}{\|v^{(t)}\|_2} = \tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right)$ .
(B) Lemma 15 and 26 tell us that w.h.p.  $\left\| E^{(t)}\right\| _2$  (resp.  $\forall k\in [d],\left|E_k^{(t)}\right|$ ) decreases from  $\Theta (d)$  (resp.  $\Theta (1)$ ) when  $t = T_{\mathrm{Adam},1}$  to  $\tilde{\mathcal{O}}\left(d^2\sqrt{\eta}\right)$  (resp.  $\tilde{\mathcal{O}}\left(d\sqrt{\eta d}\right)$ ) when  $t = T_{\mathrm{Adam},2}$ .
Combining (A) and (B), we get that the trend of  $\frac{\left\|E^{(t)}\right\|_2}{\left\|\boldsymbol{v}^{(t)}\right\|_2^2}$  and  $\frac{\left|E_k^{(t)}\right|}{\left(c^{(t)}\right)^2}$  is to decrease over time, and when  $t = T_{\mathrm{Adam},2}$ , we have w.h.p.
$$
\frac {\left\| E ^ {(t)} \right\| _ {2}}{\left\| \boldsymbol {v} ^ {(t)} \right\| _ {2} ^ {2}} \leq \tilde {\mathcal {O}} \left(d ^ {2} \sqrt {\eta}\right), \quad \frac {\left| E _ {k} ^ {(t)} \right|}{\left(c ^ {(t)}\right) ^ {2}} \leq \tilde {\mathcal {O}} \left(d ^ {2} \sqrt {\eta d}\right). \tag {53}
$$
Substituting (A) and eq. (53) into eq. (51) and (52) gives us w.h.p.,
$$
\forall 1 \leq i \leq d: \frac {\sqrt {\sum_ {j \neq i} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{\left| H ^ {(t)} [ i , i ] \right|} \leq \mathcal {O} \left(\sqrt {d}\right) + r _ {1 i} ^ {(t)},
$$
where the trend of  $r_{1i}^{(t)}$  is to decrease over time and  $r_{1i}^{(T_{\mathrm{Adam},2})}\leq \tilde{\mathcal{O}}\left(d^2\sqrt{\eta}\right)$
$$
\forall 1 \leq i \leq d, 1 \leq k \leq d: \frac {\sqrt {\sum_ {j \neq i d + k} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{\left| H ^ {(t)} [ i d + k , i d + k ] \right|} \leq \tilde {\mathcal {O}} \left(\sqrt {d}\right) + r _ {2 i} ^ {(t)},
$$
where the trend of  $r_{2i}^{(t)}$  is to decrease over time and  $r_{2i}^{(T_{\mathrm{Adam},2})}\leq \tilde{\mathcal{O}}\left(d^2\sqrt{\eta d}\right)$
Hence  $R_{\mathrm{diag}}^{\mathrm{Adam}}(t) = \tilde{\mathcal{O}}\left(\sqrt{d}\right) + \frac{1}{d^2 + d}\sum_{i=1}^{d}r_{1i}^{(t)} + \frac{d}{d^2 + d}\sum_{i=1}^{d}r_{2i}^{(t)} := \tilde{\mathcal{O}}\left(\sqrt{d}\right) + r^{(t)}$  where the trend of  $r^{(t)}$  is to decrease over time and
$$
r ^ {\left(T _ {\text {A d a m}, 2}\right)} \leq \frac {1}{d ^ {2} + d} \sum_ {i = 1} ^ {d} \tilde {\mathcal {O}} \left(d ^ {2} \sqrt {\eta}\right) + \frac {d}{d ^ {2} + d} \sum_ {i = 1} ^ {d} \tilde {\mathcal {O}} \left(d ^ {2} \sqrt {\eta d}\right) \leq \tilde {\mathcal {O}} \left(d ^ {2} \sqrt {\eta d}\right).
$$
For any  $p > 0$ , by picking the same hyperparameters as in Theorem 1, we have  $\eta d^4 \leq \tilde{\mathcal{O}}\left(\frac{1}{d^p}\right)$  and hence  $r^{(T_{\mathrm{Adam},2})} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{p - 1}{2}}}\right) = o\left(\sqrt{d}\right)$ .
F.4 Proof of Lemma 29
By the assumed weight structure, we get that
$$
\begin{array}{l} \forall i \in [ d ]: (1 - \delta) ^ {2} \left(c u _ {i}\right) ^ {2} \leq \left(w _ {2 i}\right) ^ {2} \leq (1 + \delta) ^ {2} \left(c u _ {i}\right) ^ {2}, \\ (1 - \delta) ^ {2} (u _ {i}) ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2} \leq \| W _ {1} [ i,: ] \| _ {2} ^ {2} \leq (1 + \delta) ^ {2} (u _ {i}) ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2}. \\ \end{array}
$$
For the  $i$ -th row where  $1 \leq i \leq d$ , i.e. the  $i$ -th row of the submatrix  $[H_{22} \quad H_{21}^T]$ , by triangle inequality, we have
$$
\begin{array}{l} \sqrt {\sum_ {j \neq i} H ^ {2} [ i , j ]} \leq \sqrt {\sum_ {j \neq i} H _ {2 2} ^ {2} [ i , j ]} + \sqrt {\sum_ {j = 1} ^ {d ^ {2}} H _ {2 1} ^ {2} [ j , i ]} \\ \leq \sqrt {\sum_ {j \neq i} \langle W _ {1} [ i , : ] , W _ {1} [ j , : ] \rangle^ {2}} + \sqrt {\sum_ {j = 1} ^ {d} w _ {2 j} ^ {2} \sum_ {k = 1} ^ {d} W _ {1} ^ {2} [ i , k ]} + \| E \| _ {2} \\ \leq \| W _ {1} [ i,: ] \| _ {2} \left(\sqrt {\sum_ {j \neq i} \| W _ {1} [ j , : ] \| _ {2} ^ {2}} + \sqrt {\sum_ {j = 1} ^ {d} w _ {2 j} ^ {2}}\right) + \| E \| _ {2}. \\ \end{array}
$$
$$
\begin{array}{l} \frac {\sqrt {\sum_ {j \neq i} H ^ {2} [ i , j ]}}{| H [ i , i ] |} \leq \frac {\| W _ {1} [ i , : ] \| _ {2} \sqrt {\sum_ {j \neq i} \| W _ {1} [ j , : ] \| _ {2} ^ {2}} + \sqrt {\sum_ {j = 1} ^ {d} w _ {2 j} ^ {2}}}{\| W _ {1} [ i , : ] \| _ {2} ^ {2}} + \frac {\| E \| _ {2}}{\| W _ {1} [ i , : ] \| _ {2} ^ {2}} \\ = \sqrt {\frac {\sum_ {j \neq i} \| W _ {1} [ j , : ] \| _ {2} ^ {2}}{\| W _ {1} [ i , : ] \| _ {2} ^ {2}}} + \sqrt {\frac {\sum_ {j = 1} ^ {d} w _ {2 j} ^ {2}}{\| W _ {1} [ i , : ] \| _ {2} ^ {2}}} + \frac {\| E \| _ {2}}{\| W _ {1} [ i , : ] \| _ {2} ^ {2}} \\ \leq \sqrt {\frac {(1 + \delta) ^ {2}}{(1 - \delta) ^ {2}} \cdot \frac {\sum_ {j \neq i} u _ {j} ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2}}{u _ {i} ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2}}} + \sqrt {\frac {(1 + \delta) ^ {2}}{(1 - \delta) ^ {2}} \cdot \frac {c ^ {2} \sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2} \| \boldsymbol {v} \| _ {2} ^ {2}}} + \frac {\| E \| _ {2}}{(1 - \delta) ^ {2} u _ {i} ^ {2} \| v \| _ {2} ^ {2}} \\ \leq \frac {1 + \delta}{1 - \delta} \left(1 + \frac {| c |}{\| \boldsymbol {v} \| _ {2}}\right) \sqrt {\frac {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2}}} + \frac {\| E \| _ {2}}{(1 - \delta) ^ {2} u _ {i} ^ {2} \| v \| _ {2} ^ {2}}. \\ \end{array}
$$
For the  $(id + k)$ -th row where  $1 \leq i \leq d, 1 \leq k \leq d$ , i.e. the  $((i - 1)d + k)$ -th row of the submatrix  $[H_{21} H_{11}]$ , by triangle inequality again, we have
$$
\begin{array}{l} \sqrt {\sum_ {j \neq i d + k} H ^ {2} [ i , j ]} \leq \sqrt {\sum_ {j \neq (i - 1) d + k} H _ {1 1} ^ {2} [ (i - 1) d + k , j ]} + \sqrt {\sum_ {j = 1} ^ {d} H _ {2 1} ^ {2} [ (i - 1) d + k , j ]} \\ \leq \sqrt {\sum_ {j \neq i} w _ {2 i} ^ {2} w _ {2 j} ^ {2}} + \sqrt {\sum_ {j = 1} ^ {d} w _ {2 i} ^ {2} W _ {1} ^ {2} [ j , k ]} + | E _ {k} | \\ = | w _ {2 i} | \left(\sqrt {\sum_ {j \neq i} w _ {2 j} ^ {2}} + \sqrt {\sum_ {j = 1} ^ {d} W _ {1} ^ {2} [ j , k ]}\right) + | E _ {k} |. \\ \end{array}
$$
$$
\begin{array}{l} \frac {\sqrt {\sum_ {j \neq i d + k} H ^ {2} [ i , j ]}}{| H [ i d + k , i d + k ] |} \leq \frac {| w _ {2 i} | \sqrt {\sum_ {j \neq i} w _ {2 j} ^ {2}} + \sqrt {\sum_ {j = 1} ^ {d} W _ {1} ^ {2} [ j , k ]}}{w _ {2 i} ^ {2}} + \frac {| E _ {k} |}{w _ {2 i} ^ {2}} \\ = \sqrt {\frac {\sum_ {j \neq i} w _ {2 j} ^ {2}}{w _ {2 i} ^ {2}}} + \sqrt {\frac {\sum_ {j = 1} ^ {d} W _ {1} ^ {2} [ j , k ]}{w _ {2 i} ^ {2}}} + \frac {| E _ {k} |}{w _ {2 i} ^ {2}} \\ \leq \sqrt {\frac {(1 + \delta) ^ {2}}{(1 - \delta) ^ {2}} \cdot \frac {\sum_ {j \neq i} c ^ {2} u _ {j} ^ {2}}{c ^ {2} u _ {i} ^ {2}}} + \sqrt {\frac {(1 + \delta) ^ {2}}{(1 - \delta) ^ {2}} \cdot \frac {v _ {k} ^ {2} \sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{c ^ {2} u _ {i} ^ {2}}} + \frac {| E _ {k} |}{(1 - \delta) ^ {2} c ^ {2} u _ {i} ^ {2}} \\ \leq \frac {1 + \delta}{1 - \delta} \left(1 + \frac {| v _ {k} |}{| c |}\right) \sqrt {\frac {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}{u _ {i} ^ {2}}} + \frac {| E _ {k} |}{(1 - \delta) ^ {2} c ^ {2} u _ {i} ^ {2}}. \\ \end{array}
$$
G Connection between diagonal of loss Hessian and weights
The partial derivative at  $W_{i}$  of the cost function for each  $i$  is given by:
$$
\nabla_ {W _ {i}} L (W) = W _ {i + 1} ^ {T} \dots W _ {H + 1} ^ {T} \left(W _ {H + 1} W _ {H} \dots W _ {1} - A\right) W _ {1} ^ {T} \dots W _ {i - 1} ^ {T} \tag {54}
$$
In our experiments, we were interested in the diagonal elements of the hessian. These are given by:
$$
\nabla_ {(W _ {i}) _ {a, b}} (\nabla_ {W _ {i}} L (W)) _ {a, b} = \nabla_ {(W _ {i}) _ {a, b}} (W _ {i + 1} ^ {T} \dots W _ {H + 1} ^ {T} (W _ {H + 1} W _ {H} \dots W _ {1} - A) W _ {1} ^ {T} \dots W _ {i - 1} ^ {T}) _ {a, b}
$$
for each possible  $i, a, b$ . For ease in notation, define for each  $i$ , the quantities  $M_i \coloneqq W_{i+1}^T \ldots W_{H+1}^T$  and  $N_i \coloneqq W_1^T \ldots W_{i-1}^T$ . Then we have the following lemma.
Lemma 30. The diagonal elements of the hessian of the cost function are given by:
$$
\nabla_ {(W _ {i}) _ {a, b}} (\nabla_ {W _ {i}} L (W)) _ {a, b} = \left(M _ {i} M _ {i} ^ {T}\right) _ {a, a} \left(N _ {i} ^ {T} N _ {i}\right) _ {b, b}
$$
for each possible  $i, a, b$ .
$$
\begin{array}{l} \nabla_ {W _ {i}} L (W) = W _ {i + 1} ^ {T} \dots W _ {H + 1} ^ {T} (W _ {H + 1} W _ {H} \dots W _ {1} - A) W _ {1} ^ {T} \dots W _ {i - 1} ^ {T} \\ = M _ {i} \left(W _ {H + 1} W _ {H} \dots W _ {1} - A\right) N _ {i} \\ = M _ {i} W _ {H + 1} W _ {H} \dots W _ {1} N _ {i} - M _ {i} Y N _ {i}. \\ \end{array}
$$
$$
\begin{array}{l} \nabla_ {\left(W _ {i}\right) _ {a, b}} \left(\nabla_ {W _ {i}} L (W)\right) _ {a, b} = \nabla_ {\left(W _ {i}\right) _ {a, b}} \left(M _ {i} W _ {H + 1} W _ {H} \dots W _ {1} N _ {i} - M _ {i} Y N _ {i}\right) _ {a, b} \\ = \nabla_ {(W _ {i}) _ {a, b}} \left(M _ {i} W _ {H + 1} W _ {H} \dots W _ {1} N _ {i}\right) _ {a, b}, \\ \end{array}
$$
where the last step follows since  $M_{i}$  and  $N_{i}$  are not functions of  $W_{i}$ .
Since  $M_{i} \coloneqq W_{i + 1}^{T}\ldots W_{H + 1}^{T}$ ,  $N_{i} \coloneqq W_{1}^{T}\ldots W_{i - 1}^{T}$ , by defining  $C_i \coloneqq M_iW_{H + 1}W_H\ldots W_{i + 1} = M_iM_i^T$  and  $D_{i} \coloneqq W_{i - 1}\ldots W_{2}W_{1}N_{i} = N_{i}^{T}N_{i}$  we have that:
$$
\nabla_ {(W _ {i}) _ {a, b}} (\nabla_ {W _ {i}} L (W)) _ {a, b} = \nabla_ {(W _ {i}) _ {a, b}} (C _ {i} W _ {i} D _ {i}) _ {a, b},
$$
where  $C_i$  and  $D_i$  are not functions of  $W_i$ . Now, Equation 74 in the Matrix Cookbook<sup>13</sup> shows us that for any matrices  $A$  and  $X$  we have:
$$
\nabla_ {X _ {m n}} (X A) _ {i j} = \delta_ {i m} A _ {n j}.
$$
Note that  $W_{i}\in \mathbb{R}^{d_{i}\times d_{i - 1}}$ , then we can apply this to obtain that:
$$
\begin{array}{l} \nabla_ {(W _ {i}) _ {a, b}} (\nabla_ {W _ {i}} L (W)) _ {a, b} = \nabla_ {(W _ {i}) _ {a, b}} (C _ {i} W _ {i} D _ {i}) _ {a, b} \\ = \nabla_ {\left(W _ {i}\right) _ {a, b}} \left[ \sum_ {k = 1} ^ {d _ {i}} \left(C _ {i}\right) _ {a, k} \left(W _ {i} D _ {i}\right) _ {k, b} \right] \\ = \sum_ {k = 1} ^ {d _ {i}} \left(C _ {i}\right) _ {a, k} \nabla_ {\left(W _ {i}\right) _ {a, b}} \left(W _ {i} D _ {i}\right) _ {k, b} \\ = \sum_ {k = 1} ^ {d _ {i}} \left(C _ {i}\right) _ {a, k} \delta_ {a k} \left(D _ {i}\right) _ {b, b} \\ = (C _ {i}) _ {a, a} (D _ {i}) _ {b, b} \\ = \left(M _ {i} M _ {i} ^ {T}\right) _ {a, a} \left(N _ {i} ^ {T} N _ {i}\right) _ {b, b}. \\ \end{array}
$$
<sup>13</sup>https://www.math.uwaterloo.ca/~hwolkowi/matrixcookbook.pdf
For ease of notation, let's now drop the superscript OPT and  $(t)$  and write  $R_{\mathrm{med},1}^{\mathrm{OPT}}(t)$  as  $R_{\mathrm{med},1}$  and  $R_{\mathrm{med},2}^{\mathrm{OPT}}(t)$  as  $R_{\mathrm{med},2}$ . For a 2-layer linear network,  $H = 1$ . Consider the Hessian w.r.t  $W_{1}$ , we have  $M_1M_1^T = W_2^T W_2$  and  $N_1^T N_1$  is an identity matrix. Under Assumption 1, we know that  $W_{2}$  is a row vector, which can be denoted as  $W_{2} = [w_{21},w_{22},\dots,w_{2d_{1}}]$ . Then we have
$$
(M _ {1} M _ {1} ^ {T}) _ {a, a} = w _ {2 a} ^ {2}, (N _ {1} ^ {T} N _ {1}) _ {b, b} = 1, \quad \Rightarrow \quad R _ {\mathrm {m e d , 1}} = \frac {\max _ {i} (w _ {2 i}) ^ {2}}{\operatorname* {m e d i a n} (w _ {2 i}) ^ {2}}.
$$
Similarly, consider the Hessian w.r.t.  $W_{2}$ , we have that  $M_1M_1^T$  is an identity matrix and  $N_1^T N_1 = W_1W_1^T$ . Therefore,
$$
(M _ {1} M _ {1} ^ {T}) _ {a, a} = 1, (N _ {1} ^ {T} N _ {1}) _ {b, b} = \| W _ {1} [ b,: ] \| _ {2} ^ {2}, \quad \Rightarrow \quad R _ {\mathrm {m e d}, 2} = \frac {\max _ {i} \| W _ {1} [ i , : ] \| _ {2} ^ {2}}{\operatorname* {m e d i a n} \| W _ {1} [ i , : ] \| _ {2} ^ {2}}.
$$
Hence we have related the uniformity of diagonal Hessian to that of weight matrices. In the detailed analysis, for both GD and Adam, we can prove that  $W_{1}$  converges to an approximately rank 1 matrix. The following lemma allows us to use this rank 1 structure to compute  $R_{\mathrm{med},1}$  and  $R_{\mathrm{med},2}$ .
Lemma 31. Suppose  $W_{1} \in \mathbb{R}^{d \times d}$  and  $W_{2} \in \mathbb{R}^{1 \times d}$  have the following structure:
$$
W _ {2} = c \boldsymbol {u} ^ {T} + R _ {2},
$$
where  $\pmb{u} \in \mathbb{R}^d$ ,  $\pmb{v} \in \mathbb{R}^d$ ,  $R_1 \in \mathbb{R}^{d \times d}$ ,  $R_2 \in \mathbb{R}^{1 \times d}$  and that
$$
\forall 1 \leq i, j \leq d: \quad \frac {| R _ {1} [ i , j ] |}{| u _ {i} v _ {j} |} \leq \delta , \quad \frac {| R _ {2 i} |}{| c u _ {i} |} \leq \delta , \quad \delta \in (0, 1).
$$
$$
R _ {m e d, 1}, R _ {m e d, 2} \in \left[ \frac {(1 - \delta) ^ {2}}{(1 + \delta) ^ {2}} \cdot \frac {\operatorname* {m a x} _ {i} u _ {i} ^ {2}}{\operatorname* {m e d i a n} u _ {i} ^ {2}}, \frac {(1 + \delta) ^ {2}}{(1 - \delta) ^ {2}} \cdot \frac {\operatorname* {m a x} _ {i} u _ {i} ^ {2}}{\operatorname* {m e d i a n} u _ {i} ^ {2}} \right].
$$
Proof. Let's first consider  $R_{\mathrm{med},1}$ . We have
$$
\forall i \in [ d ]: (1 - \delta) ^ {2} \left(c u _ {i}\right) ^ {2} \leq w _ {2 i} ^ {2} \leq (1 + \delta) ^ {2} \left(c u _ {i}\right) ^ {2}
$$
$$
\Rightarrow \quad (1 - \delta) ^ {2} \max  _ {i} \left(c u _ {i}\right) ^ {2} \leq \max  _ {i} w _ {2 i} ^ {2} \leq (1 + \delta) ^ {2} \max  _ {i} \left(c u _ {i}\right) ^ {2}
$$
$$
(1 - \delta) ^ {2} \text {m e d i a n} (c u _ {i}) ^ {2} \leq \text {m e d i a n} w _ {2 i} ^ {2} \leq (1 + \delta) ^ {2} \text {m e d i a n} (c u _ {i}) ^ {2},
$$
which yields
$$
\frac {(1 - \delta) ^ {2}}{(1 + \delta) ^ {2}} \cdot \frac {\max _ {i} u _ {i} ^ {2}}{\operatorname {m e d i a n} u _ {i} ^ {2}} \leq R _ {\text {m e d}, 1} = \frac {\max _ {i} w _ {2 i} ^ {2}}{\operatorname {m e d i a n} w _ {2 i} ^ {2}} \leq \frac {(1 + \delta) ^ {2}}{(1 - \delta) ^ {2}} \cdot \frac {\max _ {i} u _ {i} ^ {2}}{\operatorname {m e d i a n} u _ {i} ^ {2}}.
$$
Similarly, for  $R_{\mathrm{med},2}$ . We have that