SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Note that the above inequality holds for any  $j_0 \in S^c$ , which means  $\forall j \in S^c : \left|E_j^{(\tilde{T})}\right| \leq \tilde{\mathcal{O}}\left(d\sqrt{\eta d}\right)$ . Combining with the fact that  $\forall j \in S : \left|E_j^{(\tilde{T})}\right| \leq \tilde{\mathcal{O}}\left(d\sqrt{\eta d}\right)$  completes the proof.
E.16 Proof of Lemma 28
Consider certain  $j \in [d]$ , when  $t < \min \left\{\tilde{T}, T_{f,j}\right\}$ , we have that  $E_j^{(t)} < -\sqrt{\eta d}$ . Therefore we can use the same argument as in Section E.13 to prove that  $\left|W_1^{(t)}[i,j]\right|$  keeps increasing, and sign  $\left(W_1^{(t)}[i,j]\right) = \operatorname{sign}\left(w_{2i}^{(0)}\right)$  for all  $i \in [d]$ .
At time the "flip time"  $t = T_{f,j}$ , by definition,  $E_j^{(t)} \geq -\sqrt{\eta d}$ . After that  $E_j^{(t)}$  may oscillate. Now we prove that once  $E_j^{(t)} \geq \sqrt{\eta d}$  (or  $E_j^{(t)} \leq -\sqrt{\eta d}$ ), after a short period  $E_j^{(t)}$  will decrease (or increase) until  $E_j^{(t)} \leq \tilde{\mathcal{O}}(\sqrt{\eta d})$  (or  $E_j^{(t)} \geq -\tilde{\mathcal{O}}(\sqrt{\eta d})$ ). Moreover, during this period,  $E_j^{(t)}$  won't change too much.
We first recall that when  $T_1 \leq t < \tilde{T}$ , Lemma 27 gives us for all  $i \in [d]$ ,  $\left|w_{2i}^{(t)}\right| = \tilde{\Theta}\left(\frac{1}{\sqrt{d}}\right)$  and  $\left|W_1^{(t)}[i,j]\right| = \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right)$ . Then eq.(35) we obtained in the first phase analysis still holds, which tells us that the change of  $E_j^{(t)}$  per step satisfies  $\left|E_j^{(t + 1)} - E_j^{(t)}\right| = \tilde{\mathcal{O}}\left(\eta \sqrt{d}\right)$  for all  $T_1 \leq t < \tilde{T}$ .
We divide the analysis into two cases, based on whether these  $E_{j}^{(t)} \geq \sqrt{\eta d}$  or  $E_{j}^{(t)} \leq -\sqrt{\eta d}$ . By Lemma 25, we know that when  $T_{1} \leq t < \tilde{T}, \forall i \in [d]$ ,  $\left|W_{1}^{(t)}[i,j]\right| = \tilde{\Omega}\left(\frac{1}{\sqrt{d}}\right)$ . Since the update per step  $\left|\Delta W_{1}^{(t)}[i,j]\right| \leq \tilde{\mathcal{O}}(\eta)$ , we know that sign  $\left(W_{1}^{(t)}[i,j]\right)$  remains unchanged during this period and sign  $\left(W_{1}^{(t)}[i,j]\right) = \operatorname{sign}\left(W_{1}^{(T_{1})}[i,j]\right) = \operatorname{sign}\left(w_{2i}^{(0)}\right)$ .
By the analysis of  $w_{2i}^{(t)}$  in Lemma 24, we have for all  $i \in [d]$ ,  $w_{2i}^{(t+1)} = w_{2i}^{(t)} + \mathrm{sign}\left(w_{2i}^{(0)}\right)\Delta_{2i}^{(t)}$ , where  $\Delta_{2i}^{(t)} = \eta\left(1 \pm \tilde{\mathcal{O}}(\sqrt{\eta})\right)$ .
Case 1. Consider some time point  $t$  such that  $E_{j}^{(t)} \leq -\sqrt{\eta d}$ . Note that for all  $i \in [d]$ ,  $\left|g_1^{(t)}[i,j]\right| = \left|w_{2i}^{(t)}E_j^{(t)}\right| = \tilde{\Omega}\left(\sqrt{\eta}\right)$  and that  $\mathrm{sign}\left(g_1^{(t)}[i,j]\right) = -\mathrm{sign}\left(w_{2i}^{(t)}\right) = -\mathrm{sign}\left(w_{2i}^{(0)}\right)$ . By Lemma 24, for all  $i \in [d]$  we have  $W_1^{(t + 1)}[i,j] = W_1^{(t)}[i,j] + \mathrm{sign}\left(w_{2i}^{(0)}\right)\Delta_1^{(t)}[i,j]$  with  $\Delta_1^{(t)}[i,j] = \eta \left(1 \pm \tilde{\mathcal{O}}\left(\sqrt{\eta}\right)\right)$ . That gives us
$$
\begin{array}{l} E _ {j} ^ {(t + 1)} = \sum_ {i = 1} ^ {d} w _ {2 i} ^ {(t + 1)} W _ {1} ^ {(t + 1)} [ i, j ] - A _ {j} \\ = \sum_ {i = 1} ^ {d} \left(w _ {2 i} ^ {(t)} + \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \Delta_ {2 i} ^ {(t)}\right) \left(W _ {1} ^ {(t)} [ i, j ] + \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \Delta_ {1} ^ {(t)} [ i, j ]\right) - A _ {j} \\ = \sum_ {i = 1} ^ {d} \left(w _ {2 i} ^ {(t)} W _ {1} ^ {(t)} [ i, j ] + \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \left(w _ {2 i} ^ {(t)} \Delta_ {1} ^ {(t)} [ i, j ] + \Delta_ {2 i} ^ {(t)} W _ {1} ^ {(t)} [ i, j ]\right) + \Delta_ {2 i} ^ {(t)} \Delta_ {1} ^ {(t)} [ i, j ]\right) - A _ {j} \\ \stackrel {(i)} {=} E _ {j} ^ {(t)} + \sum_ {i = 1} ^ {d} \left(\left| w _ {2 i} ^ {(t)} \right| \Delta_ {1} ^ {(t)} [ i, j ] + \Delta_ {2 i} ^ {(t)} \left| W _ {1} ^ {(t)} [ i, j ] \right| + \Delta_ {2 i} ^ {(t)} \Delta_ {1} ^ {(t)} [ i, j ]\right), \\ \Rightarrow E _ {j} ^ {(t + 1)} > E _ {j} ^ {(t)}, \\ \end{array}
$$
where  $(i)$  is because  $\mathrm{sign}\left(w_{2i}^{(t)}\right) = \mathrm{sign}\left(W_1^{(t)}[i,j]\right) = \mathrm{sign}\left(w_{2i}^{(0)}\right)$ . Therefore we have proved that  $E_j^{(t)}$  will increase in the next step. After that for  $\tau \geq t + 1$ , as long as  $E_j^{(\tau)} \leq -\sqrt{\eta d}$ , the above analysis will hold and  $E_j^{(\tau)}$  will keep increasing until  $E_j^{(\tau)} > -\sqrt{\eta d}$  or we reach  $\tilde{T}$ .
Case 2. Consider some time point  $t$  such that  $E_{j}^{(t)} \geq \sqrt{\eta d}$ . We will prove that  $E_{j}^{(t)}$  will decrease after a short period, and during this period, the change of it is at most  $\tilde{\mathcal{O}}\left(\sqrt{\eta d}\right)$ .
By similar arguments as in Case 1, we can get that  $W_1^{(t + 1)}[i,j] = W_1^{(t)}[i,j] - \mathrm{sign}\left(w_{2i}^{(0)}\right)\Delta_1^{(t)}[i,j]$ , where  $\Delta_1^{(t)}[i,j] = \eta \left(1 \pm \tilde{\mathcal{O}}(\sqrt{\eta})\right)$ , Then
$$
\begin{array}{l} E _ {j} ^ {(t + 1)} = \sum_ {i = 1} ^ {d} w _ {2 i} ^ {(t + 1)} W _ {1} ^ {(t + 1)} [ i, j ] - A _ {j} \\ = \sum_ {i = 1} ^ {d} \left(w _ {2 i} ^ {(t)} + \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \Delta_ {2 i} ^ {(t)}\right) \left(W _ {1} ^ {(t)} [ i, j ] - \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \Delta_ {1} ^ {(t)} [ i, j ]\right) - A _ {j} \\ = \sum_ {i = 1} ^ {d} \left(w _ {2 i} ^ {(t)} W _ {1} ^ {(t)} [ i, j ] - \operatorname {s i g n} \left(w _ {2 i} ^ {(0)}\right) \left(w _ {2 i} ^ {(t)} \Delta_ {1} ^ {(t)} [ i, j ] - \Delta_ {2 i} ^ {(t)} W _ {1} ^ {(t)} [ i, j ]\right) - \Delta_ {2 i} ^ {(t)} \Delta_ {1} ^ {(t)} [ i, j ]\right) - A _ {j} \\ \stackrel {(i)} {=} E _ {j} ^ {(t)} - \sum_ {i = 1} ^ {d} \left(\left| w _ {2 i} ^ {(t)} \right| \Delta_ {1} ^ {(t)} [ i, j ] - \Delta_ {2 i} ^ {(t)} \left| W _ {1} ^ {(t)} [ i, j ] \right| + \Delta_ {2 i} ^ {(t)} \Delta_ {1} ^ {(t)} [ i, j ]\right), \\ \end{array}
$$
where  $(i)$  is because  $\mathrm{sign}\left(w_{2i}^{(t)}\right) = \mathrm{sign}\left(W_1^{(t)}[i,j]\right) = \mathrm{sign}\left(w_{2i}^{(0)}\right)$ .  $E_j^{(t + 1)}$  may not be smaller than  $E_j^{(t)}$ , but we will show that after at most  $t_s$  steps for some  $t_s$ , we will have  $E_j^{(t + t_s + 1)} < E_j^{(t + t_s)}$ .
To see this, first note that by the bounds of  $\Delta_1^{(t)}[i,j]$  and  $\Delta_{2i}^{(t)}$ , we get  $\Delta_1^{(t)}[i,j] \geq \Delta_{2i}^{(t)} - \eta \tilde{\mathcal{O}} (\sqrt{\eta})$ . Since  $\left|w_{2i}^{(t)}\right|$  increases by  $\tilde{\Theta} (\eta)$  per step, and  $\left|W_1^{(t)}[i,j]\right|$  keeps decreasing, then we have either i) after  $t_s$  steps for some  $t_s, \forall i \in [d]: \left|w_{2i}^{(t + t_s)}\right| \geq \left|W_1^{(t + t_s)}[i,j]\right| + \sqrt{\eta}$  or ii) we reach  $\tilde{T}$ .
For i), if  $E_{j}^{(t + t_{s})} < \sqrt{\eta d}$ , then it's already what we want. Otherwise we will have  $\Delta_1^{(t + t_s)}[i,j] = \eta \left(1\pm \tilde{\mathcal{O}} (\sqrt{\eta})\right)$ . Hence
$$
\begin{array}{l} E _ {j} ^ {(t + t _ {s})} - E _ {j} ^ {(t + t _ {s} + 1)} \\ = \sum_ {i = 1} ^ {d} \left(\left| w _ {2 i} ^ {(t + t _ {s})} \right| \Delta_ {1} ^ {(t + t _ {s})} [ i, j ] - \Delta_ {2 i} ^ {(t + t _ {s})} \left| W _ {1} ^ {(t + t _ {s})} [ i, j ] \right| + \Delta_ {2 i} ^ {(t + t _ {s})} \Delta_ {1} ^ {(t + t _ {s})} [ i, j ]\right) \\ \geq \sum_ {i = 1} ^ {d} \left(\left| W _ {1} ^ {(t + t _ {s})} [ i, j ] \right| \left(\Delta_ {1} ^ {(t + t _ {s})} [ i, j ] - \Delta_ {2 i} ^ {(t + t _ {s})}\right) + \sqrt {\eta} \Delta_ {1} ^ {(t + t _ {s})} [ i, j ] + \Delta_ {2 i} ^ {(t + t _ {s})} \Delta_ {1} ^ {(t + t _ {s})} [ i, j ]\right) \\ \geq \sum_ {i = 1} ^ {d} \left(- \eta \tilde {\mathcal {O}} (\sqrt {\eta}) \left| W _ {1} ^ {(t + t _ {s})} [ i, j ] \right| + \eta \sqrt {\eta} + \Delta_ {2 i} ^ {(t + t _ {s})} \Delta_ {1} ^ {(t + t _ {s})} [ i, j ]\right) > 0, \\ \end{array}
$$
where the last inequality uses  $\forall i,j\in [d]:\left|W_1^{(t + t_s)}[i,j]\right| = \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{d}}\right)$ . Therefore  $E_{j}^{(t + t_{s} + 1)} < E_{j}^{(t + t_{s})}$ . After that for  $\tau \geq t + t_s + 1$ , as long as  $E_{j}^{(\tau)}\geq \sqrt{\eta d}$ , the above analysis will hold and  $E_{j}^{(\tau)}$  will keep decreasing until  $E_{j}^{(\tau)} < \sqrt{\eta d}$  or we reach  $\tilde{T}$ .
Now we prove that during these  $t_s$  steps, the change of  $E_j$  is  $\tilde{\mathcal{O}}\left(\sqrt{\eta d}\right)$ . Since at each step the difference  $|w_{2i}| - |W_1[i,j]|$  will be enlarged by  $\tilde{\Omega} (\eta)$ , then we know that  $t_s = \sqrt{\eta} /\tilde{\Omega} (\eta) = \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{\eta}}\right)$ .
Combining with the fact that for all  $T_{1} \leq \tau \leq \tilde{T}$ ,  $\left|E_{j}^{(\tau + 1)} - E_{j}^{(\tau)}\right| = \tilde{\mathcal{O}}\left(\eta \sqrt{d}\right)$  gives us
$$
E _ {j} ^ {(t + t _ {s})} - E _ {j} ^ {(t)} \leq \tilde {\mathcal {O}} \left(\eta t _ {s} \sqrt {d}\right) = \tilde {\mathcal {O}} \left(\sqrt {\eta d}\right).
$$
For ii), we reach  $\tilde{T}$  before  $\forall i\in [d]:\left|w_{2i}^{(t + t_s)}\right|\geq \left|W_1^{(t + t_s)}[i,j]\right| + \sqrt{\eta}$ . Then we have  $\tilde{T} -t\leq \sqrt{\eta} /\tilde{\Omega} (\eta) = \tilde{\mathcal{O}}\left(\frac{1}{\sqrt{\eta}}\right)$ , which yields  $E_j^{(\tilde{T})} - E_j^{(t)}\leq \tilde{\mathcal{O}}\left(\eta (\tilde{T} -t)\sqrt{d}\right)\leq \tilde{\mathcal{O}}\left(\sqrt{\eta d}\right)$ .
Combining the above two cases, we find that if for some  $t$ ,  $E_j^{(t)} \geq \sqrt{\eta d}$ , then after at most  $t_s$  steps  $E_j$  will decrease and keeps decreasing until  $E_j < \sqrt{\eta d}$  or we reach  $\tilde{T}$ . During these steps,  $E_j$  can increase at most  $\tilde{\mathcal{O}}(\sqrt{\eta d})$ . If for some  $t$ ,  $E_j^{(t)} \leq -\sqrt{\eta d}$ , then after one step it will increase and keeps increasing until  $E_j > \sqrt{\eta d}$  or we reach  $\tilde{T}$ . That means once for some coordinate  $j$ ,  $E_j$  overshoots, it will zigzag in a small region around zero, which is  $\left[-\tilde{\mathcal{O}}(\sqrt{\eta d}), \tilde{\mathcal{O}}(\sqrt{\eta d})\right]$ .
F Hessian tends to become more and more diagonal during training
In this section, we empirically demonstrate that the trend of loss Hessian in practice is to become more and more diagonal during training. We also give a rigorous theoretical analysis on a two-layer network under Assumption 1 and 2.
F.1 Empirical Results
Let's first define the diagonal domination of the  $i$ -th coordinate at time  $t$ .
$$
r _ {\mathrm {d i a g}, i} ^ {\mathrm {O P T}} (t) := \frac {\sqrt {\sum_ {j \neq i} \left(H ^ {(t)} [ i , j ]\right) ^ {2}}}{\left| H ^ {(t)} [ i , i ] \right|}.
$$
To measure the diagonal domination of the whole Hessian, we need to consider the distribution of  $r_{\mathrm{diag},i}^{\mathrm{OPT}}(t)$  for different  $i$ . Figure 15 shows the mean and median of  $r_{\mathrm{diag},i}^{\mathrm{SGDM}}(t)$  and  $r_{\mathrm{diag},i}^{\mathrm{Adam}}(t)$  on the sentence classification task (see Section 4.1). Here we chose 4 layers (Layer #6, 12, 17 and 22) and computed the Hessians across these 4 layers. Since the number of parameters is very large, we did the computation by random sampling. As we can see, for both  $r_{\mathrm{diag},i}^{\mathrm{SGDM}}(t)$  and  $r_{\mathrm{diag},i}^{\mathrm{Adam}}(t)$ , the trend of their mean or median is to decrease over time, although there might be some oscillation.
(a) Mean
(b) Median
Figure 15: Mean and median of  $r_{\mathrm{diag},i}^{\mathrm{SGDM}}(t)$  and  $r_{\mathrm{diag},i}^{\mathrm{Adam}}(t)$  for the full hessian across the four layers (#6,12,17,22)
F.2 Theoretical Analysis
To simplify the theoretical analysis, we consider the mean of  $r_{\mathrm{diag},i}^{\mathrm{OPT}}(t)$  over all coordinate and define
$$
R _ {\text {d i a g}} ^ {\mathrm {O P T}} (t) := \operatorname {m e a n} \left(r _ {\text {d i a g}, i} ^ {\mathrm {O P T}} (t)\right). \tag {40}
$$
We consider a 2-layer network under Assumption 1 and 2, and have two goals in our proof:
1. To show that  $R_{\mathrm{diag}}^{\mathrm{OPT}}(t)$  after training is smaller than that before training  $(t = 0)$ .
2. Note that in our setting (see in Assumption 1), the Hessian is a  $(d^{2} + d)\times (d^{2} + d)$  matrix. For a completely "uniform" matrix with the same size, we have that  $R_{\mathrm{diag}}^{\mathrm{OPT}}(t) = \Theta (\sqrt{d^2 + d}) = \Theta (d)$ . Hence our second goal is to show that the  $R_{\mathrm{diag}}^{\mathrm{OPT}}(t)$  after training is on lower order than  $\Theta (d)$ .
Theorem 2. Consider the ratio  $R_{\text{diag}}^{\text{OPT}}(t)$  defined in eq. (40). Under Assumption 1 and 2, we have that before training  $(t = 0)$ , with high probability,
$$
R _ {\text {d i a g}} ^ {O P T} (0) \geq \tilde {\Omega} \left(d ^ {4 \alpha - \frac {3}{2}}\right). \tag {41}
$$
For  $SGD + M$  defined in eq. (3). For any  $p > 0$ , by picking the same hyperparameters as in Theorem 1, for  $T_{SGD,1}, T_{SGD,2}$  mentioned in Theorem 1, we have with constant probability, for any  $t \in [T_{SGD,1}, T_{SGD,2}]$ ,
$$
R _ {\text {d i a g}} ^ {S G D M} (t) \leq \tilde {\mathcal {O}} (\sqrt {d}) + q ^ {(t)}, \tag {42}
$$
where the trend of  $q^{(t)}$  is to decrease over time and  $q^{(T_{SGD,2})} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{p/2-1}}\right) = o(d)$ .
For Adam defined in eq. (3). For any  $p > 0$ , by picking the same hyperparameters as in Theorem 1, for  $T_{Adam,1}, T_{Adam,2}$  mentioned in Theorem 1, we have with high probability, for any  $t \in [T_{Adam,1}, T_{Adam,2}]$ ,
$$
R _ {\text {d i a g}} ^ {\text {A d a m}} (t) \leq \tilde {\mathcal {O}} (\sqrt {d}) + r ^ {(t)}, \tag {43}
$$
where the trend of  $r^{(t)}$  is to decrease over time and  $r^{(T_{Adam,2})} \leq \tilde{\mathcal{O}}\left(\frac{1}{d^{\frac{p - 1}{2}}}\right) = o\left(\sqrt{d}\right)$ .
F.3 Proof of Theorem 2
Lemma 4.3 of [Kaw16] gives us the following forms of Hessian.
For any  $k\in \{1,2,\dots,H + 1\}$ , we know that  $\nabla_{vec(W_k)}(\nabla_{vec(W_k)}L(W))$  equals
$$
\left(\left(W _ {H + 1} \dots W _ {k + 1}\right) ^ {T} \left(W _ {H + 1} \dots W _ {k + 1}\right) \otimes \left(W _ {k - 1} \dots W _ {1}\right) \left(W _ {k - 1} \dots W _ {1}\right) ^ {T}, \right.
$$
and for  $k\in \{2,3,\ldots ,H + 1\}$
$$
\begin{array}{l} \nabla_ {v e c \left(W _ {k}\right)} \left(\nabla_ {v e c \left(W _ {1}\right)} L (W)\right) \\ = \left(C ^ {T} \left(W _ {H + 1} \dots W _ {k + 1}\right) \otimes \left(W _ {k - 1} \dots W _ {1}\right) ^ {T}\right) \\ + \left[ \left(W _ {k - 1} \dots W _ {2}\right) ^ {T} \otimes I \right] \left[ I _ {d _ {k - 1}} \otimes \left(r \left(W _ {H + 1} \dots W _ {k + 1}\right)\right), _ {1} \dots I _ {d _ {k - 1}} \otimes \left(r \left(W _ {H + 1} \dots W _ {k + 1}\right)\right), _ {, d _ {k}} \right], \\ \end{array}
$$
where  $r = (W_{H + 1}\dots W_1 - A)^T, C = W_{H + 1}W_H\cdot \cdot \cdot W_2$
For the 2-layer linear network, write the Hessian as
$$
H := \left[ \begin{array}{c c} H _ {2 2} & H _ {2 1} ^ {T} \\ H _ {2 1} & H _ {1 1} \end{array} \right],
$$
then we have that
$$
H _ {1 1} = \left(W _ {2} ^ {T} W _ {2}\right) \otimes I _ {d} \in \mathbb {R} ^ {d ^ {2} \times d ^ {2}},
$$
$$
H _ {2 2} = W _ {1} W _ {1} ^ {T} \in \mathbb {R} ^ {d \times d},
$$
$$
H _ {2 1} = W _ {2} ^ {T} \otimes W _ {1} ^ {T} + I _ {d} \otimes \left(W _ {2} W _ {1} - A\right) ^ {T} \in \mathbb {R} ^ {d ^ {2} \times d}.
$$
Intuitively, before training the elements of  $W_{1}$  and  $W_{2}$  are very close to zero, and  $W_{2}W_{1} - A \approx -A$ . Since the elements of  $A$  are  $\Theta(1)$ , we know that the magnitudes of elements of  $H_{21}$  are much bigger than those of  $H_{11}$  and  $H_{22}$ .
After training, for both  $\mathrm{SGD + M}$  and Adam,  $W_{2}W_{1} - A\approx 0$ . Then  $H_{21}\approx (W_2)^T\otimes (W_1)^T$  and the magnitudes of its elements are no longer much larger than those of  $H_{11}$  and  $H_{22}$ . From the formula of  $H_{11}$ , we know that all the diagonal entries are nonzero, and among the  $d^4 -d^2$  off-diagonal entries, there are only  $d^{3} - d^{2}$  nonzero entries, which helps us to bound  $R_{\mathrm{diag}}^{\mathrm{OPT}}(t)$ .
F.3.1 Proof of eq. (41)
Let's first analyze the weights and Hessian before training  $(t = 0)$ . For ease of notation, we omit the superscript  $(t)$ .
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}]$ , we have
$$
\begin{array}{l} \sum_ {j \neq i} H ^ {2} [ i, j ] = \sum_ {j \neq i} H _ {2 2} ^ {2} [ i, j ] + \sum_ {j = 1} ^ {d ^ {2}} H _ {2 1} ^ {2} [ j, i ] \\ \geq \sum_ {j = (i - 1) d} ^ {i d} H _ {2 1} ^ {2} [ j, i ] = \sum_ {j = 1} ^ {d} \left(w _ {2 i} W _ {1} [ i, j ] + \left(W _ {2} W _ {1} - A\right) _ {j}\right) ^ {2} = \Theta (d). \\ \end{array}
$$
On the other hand, for the diagonal elements, we have w.h.p.
$$
| H [ i, i ] | = | H _ {2 2} [ i, i ] | = \| W _ {1} [ i,: ] \| _ {2} ^ {2} = \sum_ {j = 1} ^ {d} W _ {1} ^ {2} [ i, j ] \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {4 \alpha - 1}}\right).
$$
Then we have that for  $1 \leq i \leq d$ ,
$$
\frac {\sqrt {\sum_ {j \neq i} H ^ {2} [ i , j ]}}{| H [ i , i ] |} \geq \frac {\sqrt {\Omega (d)}}{\tilde {\mathcal {O}} \left(\frac {1}{d ^ {4 \alpha - 1}}\right)} = \tilde {\Omega} \left(d ^ {4 \alpha - \frac {1}{2}}\right).
$$
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}]$ , we have
$$
\begin{array}{l} \sum_ {j \neq i d + k} H ^ {2} [ i, j ] = \sum_ {j \neq (i - 1) d + k} H _ {1 1} ^ {2} [ (i - 1) d + k, j ] + \sum_ {j = 1} ^ {d} H _ {2 1} ^ {2} [ (i - 1) d + k, j ] \\ \geq H _ {2 1} ^ {2} \left[ (i - 1) d + k, i \right] = \left(w _ {2 i} W _ {1} [ i, k ] + \left(W _ {2} W _ {1} - A\right) _ {k}\right) ^ {2} = \Theta (1). \\ \end{array}
$$
$$
\left| H [ i d + k, i d + k ] \right| = \left| H _ {1 1} [ (i - 1) d + k, (i - 1) d + k ] \right| = w _ {2 i} ^ {2} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {2 \alpha}}\right).
$$
Then we have that for  $1 \leq i \leq d, 1 \leq k \leq d$ ,
$$
\frac {\sqrt {\sum_ {j \neq i d + k} H ^ {2} [ i , j ]}}{| H [ i d + k , i d + k ] |} \geq \frac {\sqrt {\Omega (1)}}{\tilde {\mathcal {O}} \left(\frac {1}{d ^ {2 \alpha}}\right)} = \tilde {\Omega} \left(d ^ {2 \alpha}\right).
$$
Taking the average, we obtain that before training, i.e. when  $t = 0$
$$
R _ {\text {d i a g}} ^ {\mathrm {O P T}} (0) \geq \frac {d \tilde {\Omega} \left(d ^ {4 \alpha - \frac {1}{2}}\right) + d ^ {2} \tilde {\Omega} \left(d ^ {2 \alpha}\right)}{d ^ {2} + d} = \tilde {\Omega} \left(d ^ {4 \alpha - \frac {3}{2}}\right).
$$
F.3.2 Proof of eq. (42)
The proof is based on the lemma below.
Lemma 29. Suppose the weight matrices have the following structure:
$$
W _ {1} = \boldsymbol {u} \boldsymbol {v} ^ {T} + R _ {1},
$$
$$
W _ {2} = c \boldsymbol {u} ^ {T} + R _ {2} ^ {T},
$$
where  $\forall 1\leq i,j\leq d:\frac{|R_1[i,j]|}{|u_iv_j|}\leq \delta ,\frac{|R_{2i}|}{|cu_i|}\leq \delta ,\delta \in (0,1).$
Then we have for  $1 \leq i \leq d$ ,
$$
\frac {\sqrt {\sum_ {j \neq i} H ^ {2} [ i , j ]}}{| H [ i , i ] |} \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}},
$$
and for  $1\leq i\leq d,1\leq k\leq d$
$$
\frac {\sqrt {\sum_ {j \neq i d + k} H ^ {2} [ i , j ]}}{| H [ i d + k , i d + k ] |} \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}}.
$$
Now we are ready to prove eq. (42).
By the analyses in Section D.1, we know that for  $t \in [T_{\mathrm{SGD},1}, T_{\mathrm{SGD},2}]$ , the weights obtained by GD with momentum satisfy
$$
W _ {2} ^ {(t)} = c ^ {(t)} \pmb {u} ^ {(T _ {1}) T} + R _ {2} ^ {(t) T},
$$
where  $T_{\mathrm{SGD},1} = T_1$  and
$$
\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}).
$$
Here  $\epsilon_0$  is defined in Definition 2. Since  $\pmb{u}^{(T_1)}$  doesn't depend on time  $t$  in the period  $(T_{\mathrm{SGD},1}, T_{\mathrm{SGD},2}]$ , we write  $\pmb{u}^{(T_1)}$  as  $\pmb{u}$  for ease of notation.
Hence by Lemma 29, when  $t \in [T_{\mathrm{SGD},1}, T_{\mathrm{SGD},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 + \tilde {\mathcal {O}} (\epsilon_ {0})}{1 - \tilde {\mathcal {O}} (\epsilon_ {0})} \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 - \tilde {\mathcal {O}} (\epsilon_ {0})\right) ^ {2} u _ {i} ^ {2} \left\| \boldsymbol {v} ^ {(t)} \right\| _ {2} ^ {2}} \\ = \mathcal {O} \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}}} + \mathcal {O} \left(\frac {\left\| E ^ {(t)} \right\| _ {2}}{u _ {i} ^ {2} \left\| \boldsymbol {v} ^ {(t)} \right\| _ {2} ^ {2}}\right), \tag {44} \\ \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 + \tilde {\mathcal {O}} (\epsilon_ {0})}{1 - \tilde {\mathcal {O}} (\epsilon_ {0})} \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 - \tilde {\mathcal {O}} (\epsilon_ {0})\right) ^ {2} \left(c ^ {(t)}\right) ^ {2} u _ {i} ^ {2}} \\ = \mathcal {O} \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}}} + \mathcal {O} \left(\frac {\left| E _ {k} ^ {(t)} \right|}{\left(c ^ {(t)}\right) ^ {2} u _ {i} ^ {2}}\right). \tag {45} \\ \end{array}
$$
By Lemma 4, we have  $\pmb{u} = X + Y$  where  $X_{i}, i \in [d]$  are i.i.d Gaussian random variables and w.h.p.,
$$
\forall i \in [ d ]: \frac {\left| Y _ {i} \right|}{\left| X _ {i} \right|} \leq \tilde {\mathcal {O}} \left(\frac {1}{d ^ {\frac {1}{4} \alpha - \frac {1}{2}}}\right) := \delta_ {x y}, \tag {46}
$$
$$
\forall i \in [ d ]: \frac {\sqrt {\sum_ {j = 1} ^ {d} u _ {j} ^ {2}}}{| u _ {i} |} \leq \left(\frac {1 + \delta_ {x y}}{1 - \delta_ {x y}}\right) \frac {\sqrt {\sum_ {j = 1} ^ {d} X _ {j} ^ {2}}}{| X _ {i} |}, \quad \frac {1}{u _ {i} ^ {2}} \leq \left(\frac {1}{1 - \delta_ {x y}}\right) ^ {2} \frac {1}{X _ {i} ^ {2}}. \tag {47}
$$
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. The induction in Section D.9 further gives us that for  $t \in [T_{\mathrm{SGD},1}, T_{\mathrm{SGD},2}]$ , w.h.p.  $\forall k \in [d] : \frac{v_k^{(t)}}{c^{(t)}} = \Theta \left(\frac{1}{\sqrt{d}}\right)$ , which yields  $\frac{c^{(t)}}{\|\pmb{v}^{(t)}\|_2} = \Theta(1)$ . Combining with eq. (47), we obtain