SlowGuess/NeurIPS_2023_OCR huggingface dataset:
These two conditions mean that the gap between  $h^*$  and  $h$  will not grow if  $h(s)$  is smaller than  $h^*(s - 1)$ . Then for all  $h(s)$  that meet (i), we have  $h^*(s) - h(s) \leq h^*(s) - h^*(s - 1) \leq h^*(1) - h^*(0) = \eta_Y / M$  from Eqn. 33. And for any  $s \geq 2$ , every time  $h(s)$  transfer from (i) to (ii) exactly at  $s$ , which means that  $h(s - 1)$  meets (i) and thus no smaller than  $h^*(s - 2)$ , we get  $h^*(s) - h(s) \leq h^*(s) - h(s - 1) \leq h^*(s) - h^*(s - 2) \leq h^*(2) - h^*(0) \leq 2\eta_Y / M$ .
Finally from Eqn. 53 in Lemma 9, when  $s \leq \frac{\ln M}{\eta_Y}$ , we get  $h(s) = \Theta(\eta_Y t / M)$  and thus there exist some constant  $c = \Theta(1)$  such that  $h(s + c) \geq h(s) + 2\eta_Y / M \geq h^*(s) \geq h(s)$ .
(b) Assume that there exist  $\epsilon \in (0,2\eta_Y / M]$  such that  $h^{*}(s) - h(s)\geq \epsilon$  for all  $s\in \mathbb{N}$ . Since  $h$  is unbounded, then  $\dot{h} (t)\to 0$  when  $t\to \infty$  from Eqn. 33, so there exist some  $s_0^\prime \in \mathbb{N}$  such that when  $s\geq s_0^{\prime},h(s + 1) - h(s)\leq \epsilon +\ln (1 / 2) / M$ . Also, from Lemma 9 we know that exists  $s_0^{\prime \prime} = \frac{(3 + \delta)\ln(M)}{\eta_Y}$  where  $\delta >0,\delta = \Theta (1)$  such that when  $s\geq s_0^{\prime \prime},\exp (Mh(s)) > 2(M - 1)$ . Since  $s\rightarrow \infty$ , we just consider the case that  $s = \lfloor t\rfloor \geq s_0\coloneqq \max (s_0',s_0'')$ . Then denote  $\Delta_{1}\coloneqq \frac{2(M - 1)}{\exp(Mh(s))} < 1$ , we have:
$$
\begin{array}{l} \dot {h} ^ {\circ} (t) - \dot {h} (t) = \frac {\eta_ {Y}}{M - 1 + \exp (M h ^ {*} (s))} - \frac {\eta_ {Y}}{M - 1 + \exp (M h (t))} \\ \leq \frac {\eta_ {Y}}{M - 1 + \exp (M (h (s) + \epsilon))} - \frac {\eta_ {Y}}{M - 1 + \exp (M h (s + 1))} \\ = - \frac {\eta_ {Y} \exp (M h (s)) \cdot [ \exp (M \epsilon) - \exp (M h (s + 1) - M h (s)) ]}{[ M - 1 + \exp (M (h (s) + \epsilon)) ] \cdot [ M - 1 + \exp (M h (s + 1)) ]} \\ \leq - \frac {\eta_ {Y} \exp (M h (s)) \cdot \exp (M \epsilon)}{2 [ M - 1 + \exp (M (h (s) + \epsilon)) ] \cdot [ M - 1 + \frac {1}{2} \exp (M (h (s) + \epsilon)) ]} \tag {48} \\ \leq - \frac {\eta_ {Y} \exp (M \epsilon)}{(1 + \Delta_ {1}) ^ {2} \exp (M h (s)) \exp (4 \eta_ {Y})}, \quad (s \geq s _ {0} = \max (s _ {0} ^ {\prime}, s _ {0} ^ {\prime \prime})) \\ \leq - \frac {\exp (M \epsilon)}{4 \exp (4 \eta_ {Y}) M} \cdot \frac {1}{t} =: - \frac {C}{t} \\ \end{array}
$$
Here  $C = \frac{\exp(M\epsilon)}{4\exp(4\eta_Y)\bar{M}} > 0$  and for the last inequality, we use the fact that  $t \geq s_0' > \frac{3\ln M}{\eta_Y}$  and thus  $h(s) \leq h(t) = O(\frac{\ln(M\eta_Yt)}{M})$  from Lemma 9. So we get
$$
\left[ h ^ {\circ} (t) - h (t) \right] - \left[ h ^ {\circ} \left(s _ {0}\right) - h \left(s _ {0}\right)\right] \leq - \int_ {t ^ {\prime} = s _ {0}} ^ {\infty} \frac {C d t}{t} \rightarrow - \infty \tag {49}
$$
This contradicts  $h^{\circ}(t) - h(t) \geq 0!$ . So the original assumption doesn't hold, which means that  $h^{*}(s) - h(s) \to 0$  when  $s \to \infty$ .
Remark 2. By some qualitative estimation, we claim that if  $\eta_{Y} = O(1)$ , then there exists some constant  $c = O(\ln M)$  such that  $h(s) \leq h^{*}(s) \leq h(s + c)$  for all  $s > s_1 \coloneqq \frac{2\ln(1 + \omega_1)}{\eta_Y}$  where  $\omega_1 = \Theta (\ln \ln M / \ln M)$  is defined in Lemma 8. Denote  $\delta h(t) \coloneqq h^{\circ}(t) - h(t)$ , when  $\delta h(t) \ll h(t)$ ,
we have  $\dot{h} (t) = \dot{h}^{\circ}(t) - \dot{h} (t)\asymp -\eta_{Y}M\cdot \delta h(t)\cdot \exp (-Mh(t))\asymp -\delta h(t) / t$  by computing the second-order derivative of  $\delta h$ , and thus  $h^\circ (t) - h(t)\asymp 2\eta_Ys_0 / (Mt) = O(\ln M / (Mt))$ . Combining this with the fact that  $h(t) = \Theta (\ln (M\eta_Yt) / M)$  when  $t > s_1$ , we prove our claim. The results of Lemma 7 and Remark 2 are also confirmed by the numerical simulation results as Fig. 8.
So from Lemma 7 and Remark 2, we just assume  $\eta_{Y} < 1$  and replace  $h^{*}$  with  $h$  in the latter parts for convenience. Then we further investigate the properties of Eqn. 45.
Lemma 8. There exists  $\omega_{i},0 < \omega_{i}\ll 1,i = 2,3,$  such that for  $h\in \mathbb{J}_1\coloneqq [\frac{1}{M^{2 - \omega_0}},\frac{(1 + \omega_1)\ln(M)}{M} ],$  we have  $\exp (Mh(t))\leq (M - 1)Mh(t)$  . And for  $h\notin \mathbb{J}_1$  , we have  $\exp (Mh(t)) > (M - 1)Mh(t)$  Here  $\omega_{1} = \Theta (\frac{\ln\ln(M)}{\ln(M)})$  , and if  $M\gg 100$  , we have  $\omega_0\lesssim \left(\frac{1}{M^{0.99}\ln M}\right)\ll 0.01$
Proof. It's obvious that  $\exp (Mh(t)) - (M - 1)Mh(t)$  has two zero points in  $\mathbb{R}^+$ . Let  $h(t) = M^{-(2 - \omega_0)}$ , we get
$$
\omega_ {0} = \frac {1}{\ln M} \left(\ln \left(\frac {M}{M - 1}\right) + \frac {1}{M ^ {1 - \omega_ {0}}}\right) = O \left(\frac {1}{M ^ {0 . 9 9} \ln (M)}\right) \tag {50}
$$
For another zero point, let  $\omega_{1} \in (0, 1)$  to be some constant such that  $h(t) = \frac{(1 + \omega_1)\ln(M)}{M}$  satisfies  $\exp(Mh) = (M - 1)Mh$ , then we get
$$
\begin{array}{l} M ^ {\omega_ {1}} = \left(1 + \omega_ {1}\right) \ln (M) \frac {(M - 1)}{M} = c ^ {\prime} \cdot \ln (M) \frac {(M - 1)}{M} \tag {51} \\ \Rightarrow \omega_ {1} = \Theta (\frac {\ln \ln (M)}{\ln (M)}) \\ \end{array}
$$
where  $c' \in (0.5, 2)$  is some universal constant.
Remark 3. From Lemma 8, if we assume  $M \gg 100$ , then  $\omega_0 \ll 0.01$ , and if we assume  $\eta_Y \gg \frac{1}{M^{1 - \omega_0}} > \frac{1}{M^{0.99}}$ , then  $h(1) \gtrsim \frac{\eta_Y}{M} \gg \frac{1}{M^{2 - \omega_0}}$  and function  $\exp(Mh(t)) - (M - 1)Mh(t)$  has only one zero point  $\frac{(1 + \omega_1) \ln M}{M}$  in  $[1, \infty)$ . For convenience, we just assume  $M \gg 100$  and  $1 > \eta_Y \gg \frac{1}{M^{0.99}}$  and thus focus on the unique zero point  $\frac{(1 + \omega_1) \ln M}{M}$  of  $h$  in the latter parts.
We can then show the properties of speed control coefficient  $\gamma(t) \coloneqq \frac{(M - 1)^2 h(t / K)}{(M - 1) + \exp(Mh(t / K))}$  as below.
Lemma 9. We have two stage for  $h$  and  $\gamma$ :
- When  $t \leq \frac{K \ln(M)}{\eta_Y}$ , we have  $\exp(Mh(t / K)) \leq \min(M - 1, (M - 1)Mh(t / K))$ ,  $h = O(\eta_Y t / (MK))$  and  $\gamma(t) = O(\eta_Y t / K)$ .
- When  $t \geq \frac{2(1 + \omega_1)K\ln(M)}{\eta_Y}$  where  $\omega_{1} = \Theta (\frac{\ln\ln M}{\ln M})$  is defined in Lemma 8, we have  $\exp (Mh(t / K)) \geq \max (M - 1,(M - 1)Mh(t / K))$ ,  $h = O(\frac{1}{M}\ln (M\eta_Yt / K))$  and  $\gamma (t) = O(\frac{K\ln(M\eta_Yt / K)}{\eta_Yt})$ .
Proof. For convenience, we just let  $K = 1$ . And the proof for  $K \neq 1$  is similar. We denote  $\Delta_1(h) := \frac{\exp(Mh)}{M - 1}$  and  $\Delta_2(h) := \frac{\exp(Mh)}{(M - 1)Mh}$ .
Step 1:  $t \leq \frac{\ln(M)}{\eta_Y}$ . If  $h \geq \frac{\ln(M - 1)}{M}$ , from Eqn. 45 we have:
$$
t \geq \frac {M - 2 + (M - 1) \ln (M - 1)}{M \eta_ {Y}} > \frac {\ln (M)}{\eta_ {Y}} \tag {52}
$$
So when  $t \leq \frac{\ln(M)}{\eta_Y}$  we have  $h < \frac{\ln(M - 1)}{M}$ , and thus  $\exp (Mh(t)) \leq \min (M - 1, (M - 1)Mh(t))$ , i.e.,  $\Delta_1, \Delta_2 \leq 1$ . Then from Eqn. 45 we get
$$
h = \frac {M \eta_ {Y} t + 1}{\left(1 + \Delta_ {2}\right) M (M - 1)} = O \left(\frac {1}{M} \eta_ {Y} t\right) \tag {53}
$$
$$
\gamma = \frac {(M - 1) h}{1 + \Delta_ {1}} = \frac {M \eta_ {Y} t + 1}{(1 + \Delta_ {1}) (1 + \Delta_ {2}) M} = O \left(\eta_ {Y} t\right) \tag {54}
$$
Step 2:  $t > \frac{2(1 + \omega_1)\ln(M)}{\eta_Y}$  where  $\omega_{1} = \Theta (\frac{\ln\ln(M)}{\ln(M)})$ . So now  $h > \frac{\ln(M - 1)}{M}$  and thus  $\Delta_1 > 1$  from Eqn. 52. Then if  $\exp (Mh) \leq M(M - 1)h$ , i.e.  $\Delta_2 \leq 1$ , from Lemma 8 we have  $h = \frac{M\eta_Yt + 1}{(1 + \Delta_2)M(M - 1)} \leq \frac{(1 + \omega_1)\ln(M)}{M}$ . Therefore,
$$
t \leq \frac {1}{\eta_ {Y}} \left(\left(1 + \omega_ {1}\right) \left(1 + \Delta_ {2}\right) \frac {M - 1}{M} \ln M - \frac {1}{M}\right) <   \frac {2 \left(1 + \omega_ {1}\right) \ln (M)}{\eta_ {Y}}. \tag {55}
$$
Contradiction! So when  $t \geq \frac{2(1 + \omega_1) \ln(M)}{\eta_Y}$ , we have  $\Delta_2 > 1$ . Then from Eqn. 45 we get:
$$
h = \frac {1}{M} \ln \left(\frac {M \eta_ {Y} t + 1}{1 + \Delta_ {2} ^ {- 1}}\right) = O \left(\frac {1}{M} \ln \left(M \eta_ {Y} t\right)\right) \tag {56}
$$
$$
\gamma = \frac {M - 1}{M} \frac {(M - 1) \ln \left(\frac {M \eta_ {Y} t + 1}{1 + \Delta_ {2} ^ {- 1}}\right)}{\left(1 + \Delta_ {1} ^ {- 1}\right) \left(\frac {M \eta_ {Y} t + 1}{1 + \Delta_ {2} ^ {- 1}}\right)} = O \left(\frac {\ln (M \eta_ {Y} t)}{\eta_ {Y} t}\right) \tag {57}
$$
D.4 The dynamics under multiple uniformly sampled sequence classes
We then generalize our analysis of  $W$  to the case where  $x_{T + 1}$  can be any value in  $[K]$  rather than fixing  $x_{T + 1} = n$  with the key observation that the row vectors of  $W'$  can be independently updated. Before formalizing this result, we first conduct the concentration inequality of the sampling number for each next-token case. Let  $N_{n} \coloneqq \sum_{i = 1}^{N}\mathbb{I}[x_{T + 1} = n]$  to be the number of times the event  $x_{T + 1} = n$  happens, then we have:
Lemma 10. For  $\delta \in (0,1)$ , with probability at least  $1 - \delta$  we have
$$
\left| N _ {n} - \lceil N \mathbb {P} (n) \rceil \right| \leq \sqrt {\frac {N}{2} \ln \left(\frac {2}{\delta}\right)} + 1 <   \sqrt {N \ln \left(\frac {2}{\delta}\right)} \tag {58}
$$
Proof. From Hoeffding's inequality, we have
$$
\mathbb {P} \left(\left| \frac {N _ {n}}{N} - \mathbb {P} (n) \right| > t\right) \leq 2 \exp (- 2 N t ^ {2}) \tag {59}
$$
Let  $t = \sqrt{\frac{1}{2N}\ln\left(\frac{2}{\delta}\right)}$  and we can get the results by direct calculation.
Remark 4. From Lemma 10, if we consider the uniform sampling case where  $\mathbb{P}(n) \equiv \frac{1}{K}$ , then  $N\mathbb{P}(n) = N / K \gg \sqrt{N}$ . So  $N_{n}$  are all concentrated around  $N\mathbb{P}(n)$ . Recall the definition of  $\bar{N} = \lceil N / K \rceil$  and  $\Delta = \lceil \sqrt{N\ln(\frac{1}{\delta})} \rceil$ , with probability at least  $1 - \delta$  we have:
$$
\left| N _ {n} - \bar {N} \right| \lesssim \Delta \ll \bar {N} \tag {60}
$$
We then further investigate the concentration of  $h(N_{n})$ :
Lemma 11. For  $\delta \in (0,1)$ , with probability at least  $1 - \delta$  we have
$$
\left| h \left(N _ {n}\right) - h (\bar {N}) \right| \lesssim h (\bar {N}) \cdot \frac {\Delta}{\bar {N}} \tag {61}
$$
$$
\left| \frac {1}{M - 1 + \exp \left(M h \left(N _ {n}\right)\right)} - \frac {1}{M - 1 + \exp \left(M h (\bar {N})\right)} \right| \tag {62}
$$
$$
\lesssim \frac {1}{M - 1 + \exp (M h (\bar {N}))} \cdot \sigma^ {\prime}
$$
where  $\sigma' > 0$  is some constant such that  $\sigma' \leq \frac{1}{3}\eta_Y\Delta \ll \ln(M)$ . And if  $N \geq \frac{2K(1 + \omega_1)\ln M}{\eta_Y}$  where  $\omega_1$  is defined in Lemma 8, then  $\sigma' \lesssim \frac{\Delta}{N} \ll 1$ .
Proof. First, we note that  $h$  has a decreasing gradient, so  $h(x) \geq \dot{h}(x) \times x$  and  $h(x_1 + x_2) - h(x_1) \leq \dot{h}(x_1) \times x_2$  for any  $x_1, x_2 \geq 0$ . So with probability at least  $1 - \delta$ , we have:
$$
\left| h \left(N _ {n}\right) - h (\bar {N}) \right| \leq h (\bar {N}) - h (\bar {N} - \Delta) \leq \dot {h} (\bar {N} - \Delta) \times \Delta \leq \frac {h (\bar {N}) \Delta}{\bar {N} - \Delta} \asymp h (\bar {N}) \cdot \frac {\Delta}{\bar {N}} \tag {63}
$$
For the second inequality, without loss of generality, we let  $N_{n} > \bar{N}$ . Denote  $g(s) \coloneqq (M - 1 + \exp (Mh(s)))^{-1}$  and note that:
$$
\begin{array}{l} \frac {\mathrm {d} g}{\mathrm {d} s} = \frac {M \exp (M h (s))}{(M - 1 + \exp (M h (s))) ^ {2}} \cdot \frac {\mathrm {d} h}{\mathrm {d} s} \\ = \frac {1}{M - 1 + \exp (M h (s))} \cdot \frac {\eta_ {Y} M \exp (M h (s))}{(M - 1 + \exp (M h (s))) ^ {2}} \tag {64} \\ \leq \frac {1}{M - 1 + \exp (M h (s))} \cdot \frac {M}{(M - 1)} \cdot \frac {\eta_ {Y}}{4} \\ \end{array}
$$
the last equality holds only when  $h(s) = \frac{\ln(M - 1)}{M}$ . So from  $|g(\bar{N} + \Delta) - g(N_n)| \leq \max_{s \in [N_n, N_n + \Delta]} \dot{g}(s) \cdot \Delta$ , we get:
$$
\left| \frac {1}{M - 1 + \exp (M h (\bar {N} + \Delta))} - \frac {1}{M - 1 + \exp (M h (\bar {N}))} \right| \leq \frac {1}{M - 1 + \exp (M h (\bar {N}))} \cdot \frac {1}{3} \eta_ {Y} \Delta \tag {65}
$$
If  $\bar{N} < \frac{2(1 + \omega_1)\ln(M)}{\eta_Y} +\Delta$  with  $\omega_{1} = \Theta (\frac{\ln\ln M}{\ln M})$  defined in Lemma 8, we have  $\sigma^{\prime}\leq \eta_{Y}\Delta /3\ll$ $\eta_Y\bar{N}\lesssim \ln (M)$  . If  $\bar{N}\geq \frac{2(1 + \omega_1)\ln(M)}{\eta_Y} +\Delta$  , we utilize the Eqn.45 and obtain:
$$
\begin{array}{l} | \frac {1}{M - 1 + \exp (M h (\bar {N} + \Delta))} - \frac {1}{M - 1 + \exp (M h (\bar {N}))} | \\ = \frac {1}{M - 1 + \exp (M h (\bar {N}))} \cdot \frac {\left| \exp (M h (\bar {N} + \Delta)) - \exp (M h (\bar {N})) \right|}{M - 1 + \exp (M h (\bar {N} + \Delta))} \\ \leq \frac {1}{M - 1 + \exp (M h (\bar {N}))} \cdot \frac {M \eta_ {Y} \Delta}{M - 1 + \exp (M h (\bar {N} + \Delta))}, \quad (E q n. 4 5) \\ \leq \frac {1}{M - 1 + \exp (M h (\bar {N}))} \cdot \frac {M \eta_ {Y} \Delta}{M + \frac {1}{2} \cdot M \eta_ {Y} (\bar {N} + \Delta)}, \quad (\text {L e m m a} 9, N _ {n} \geq \frac {2 (1 + \omega_ {1}) \ln (M)}{\eta_ {Y}} + \Delta) \\ \lesssim \frac {1}{M - 1 + \exp (M h (\bar {N}))} \cdot \frac {\Delta}{\bar {N}} \\ \end{array}
$$
So  $\sigma' \leq \Delta / \bar{N}$ . When  $N_{n} < \bar{N}$ , with probability at least  $1 - \delta$  we have  $N_{n} \gtrsim \bar{N} - \Delta$ , and similar inequalities also hold for such cases, so we finish the proof.
Recall that  $\zeta_n \in \mathbb{R}^M$  is defined as  $\zeta_n = \frac{M}{M - 1} (\pmb{e}_n - \frac{1}{M}\pmb{1})$ . And we have  $q_1 \coloneqq \zeta_i^\top \zeta_i = 1 + \frac{1}{M - 1}$ ,  $q_0 \coloneqq \zeta_j^\top \zeta_i = -\frac{M}{(M - 1)^2}$  for all  $i, j \in [M]$  where  $i \neq j$ . For convenience, we denote  $\bar{W}'(N) \coloneqq [\bar{\pmb{w}}_1(N), \dots, \bar{\pmb{w}}_K(N), \mathbf{0}, \dots, \mathbf{0}]^\top \in \mathbb{R}^{M \times M}$ , where  $\bar{\pmb{w}}_n(N) \coloneqq (M - 1)h([\bar{N} / K])\zeta_n = (M - 1)h(\bar{N})\zeta_n$ . So using these concentration inequalities, we get:
Lemma 12. Assume the assumptions in Lemma 5 hold but we uniformly sample the training data. Then if the total number of epochs  $N$  satisfies  $N \gg K^2$ , we have  $Y = (F')^{-\top}(I + \Theta')\tilde{W}'(N)$  where  $\Theta' := \mathrm{diag}(\theta_1, \ldots, \theta_K, 0, \ldots, 0) \in \mathbb{R}^{M \times M}$  and with probability at least  $1 - \delta$  we have  $|\theta_i| \lesssim \frac{K}{\sqrt{N}} \sqrt{\ln(\frac{K}{\delta})}, \forall i \in [K]$ .
Proof. From Lemma 5 and the first inequality of Lemma 11, we know that
$$
\begin{array}{l} \boldsymbol {w} _ {n} (N) = (M - 1) h \left(N _ {n}\right) \boldsymbol {\zeta} _ {n} (66) \\ = (M - 1) h (\bar {N}) \zeta_ {n} + (M - 1) \left(h \left(N _ {n}\right) - h (\bar {N})\right) \zeta_ {n} (67) \\ = (1 + \theta_ {n}) \cdot (M - 1) h (\bar {N}) \zeta_ {n} (68) \\ = (1 + \theta_ {n}) \bar {\boldsymbol {w}} _ {n} (N) (69) \\ \end{array}
$$
where for any  $\delta \in (0,1)$ , with probability at least  $1 - \delta$  we have  $|\theta_i| \lesssim \frac{K}{\sqrt{N}} \sqrt{\ln(\frac{K}{\delta})}, \forall n \in [K]$ . Therefore,  $W'(N) = [\pmb{w}_1(N), \dots, \pmb{w}_K(N), \mathbf{0}, \dots, \mathbf{0}]^\top = (I + \Theta') \bar{W}'(N)$ , then from  $W' = (F')^\top Y$ , we finish the proof.
Then, we can give out the exact solution of  $Y$  by pointing out the properties of  $F^{\circ}$  and  $F'$  from the observation that each row of  $Y$  should be the linear combination of vectors in  $\{f_n^\top\}_{n \in [K]}$ :
Theorem 5. If Assumption 2 holds and  $Y(0) = 0$ . Furthermore, we assume the training data is uniformly sampled and the total number of epochs  $N$  satisfies  $N \gg K^2$ . Then the solution of Eqn. 26 will be:
$$
Y = \left(F ^ {\dagger}\right) ^ {\top} (I + \Theta) \bar {W} (N) = F \left(I - E ^ {\prime}\right) (I + \Theta) \bar {W} (N) \tag {70}
$$
Here  $\Theta := \text{diag}(\theta_1, \ldots, \theta_K)$  and for any  $\delta \in (0, 1)$ , with probability at least  $1 - \delta$  we have  $|\theta_i| \lesssim \frac{K}{\sqrt{N}} \sqrt{\ln\left(\frac{K}{\delta}\right)}, \forall i \in [K]$ .
Proof. Let  $\pmb{q}_i, i \in [M]$  be the  $i$ -th row vector of  $(F')^{-1}$ , then we have  $\pmb{q}_j^\top \pmb{f}_i = \mathbb{I}[i = j]$ . From Lemma 12 we get  $Y = (F')^{-\top}(I + \Theta')\bar{W}'(N)$ . And from Eqn. 26, we know all the columns of  $Y$  are the linear combination of  $\pmb{f}_n, n \in [K]$ . Note that  $\bar{W}(N)$  has only top  $K$  rows to be non-zero, so we need to constrain that all the top  $K$  columns of  $(F')^{-\top}$ , i.e.,  $\pmb{q}_i, i \in [K]$ , to be the linear combination of  $\pmb{f}_n, n \in [K]$ , which means that  $\pmb{q}_1, \dots, \pmb{q}_K$  must be the basis of  $\Xi := \operatorname{span}(\pmb{f}_j; j \in [K])$  and thus  $\pmb{q}_{K+1}, \dots, \pmb{q}_M$  are the basis of  $\Xi' := \operatorname{span}(\pmb{f}_j; K \leq j \leq M)$ . Therefore, we get  $\Xi \perp \Xi'$ , and thus  $[\pmb{q}_1, \dots, \pmb{q}_K]$  can only be  $(F^\dagger)^\top$ . So the proof is done.
Actually, we see that the result of Theorem 5 matches the modified gradient update on  $Y$  (Eqn. 26). And we show that using such reparameterization dynamics, we can still approach the critical point of Eqn. 7 in the rate of  $\mathcal{O}\left(\frac{1}{N}\right)$ :
Corollary 1. Assume assumptions in Theorem 5 hold,  $M \gg 100$  and  $\eta_{Y}$  satisfies  $M^{-0.99} \ll \eta_{Y} < 1$ . Then  $\forall n \in [K]$ , we have
$$
\begin{array}{l} \left(\boldsymbol {x} _ {T + 1} - \boldsymbol {\alpha} _ {n}\right) = \frac {M - 1}{(M - 1) + \exp (M h \left(N _ {n}\right))} \boldsymbol {\zeta} _ {n} \\ = \frac {M - 1}{(M - 1) + \exp (M h (\bar {N}))} \cdot (1 + \sigma) \cdot \zeta_ {n} \tag {71} \\ \end{array}
$$
where  $\sigma > -1$  and for any  $\delta \in (0,1)$ , with probability at least  $1 - \delta$  we have  $|\sigma| \lesssim \eta_Y \sqrt{N \ln(\frac{1}{\delta})}$  and when  $N \gg K \left( \sqrt{N \ln(\frac{1}{\delta})} + \frac{2(1 + \omega_1) \ln M}{\eta_Y} \right)$  with  $\omega_1$  defined in Lemma 8,  $|\sigma| \lesssim \frac{K}{\sqrt{N}} \sqrt{\ln(\frac{1}{\delta})}$ . Further, to let  $\| x_{T+1} - \alpha_n \|_2 \leq \epsilon$  with probability at least  $1 - \delta$  for any  $n \in [K]$  and  $\epsilon \ll 1$ , we need the total number of training epochs to be at most  $O\left( \frac{K}{\epsilon \eta_Y} \log \left( \frac{M}{\epsilon} \right) \right)$ .
Proof. Note that  $\pmb{x}_{T + 1} = \pmb{e}_n$ , then we just need to combine Lemma 5 and the second inequality of Lemma 11, to get Eqn. 71. Denote  $S_{n}$  to be the number of training epochs that are needed to let  $\| \pmb{x}_{T + 1} - \pmb{\alpha}_n\| _2\asymp \epsilon$ , then we have
$$
h \left(S _ {n}\right) \asymp \frac {1}{M} \ln \left(\frac {M}{\epsilon}\right) \tag {72}
$$
But note that  $h(t + 1) - h(t) \geq \frac{\eta Y}{M - 1 + \exp(Mh(S_n))} \asymp \frac{\eta Y\epsilon}{M - 1}, \forall t \in [0, S - 1]$  from Eqn. 71, we have
$$
S _ {n} \lesssim \frac {h \left(S _ {n}\right)}{\eta_ {Y} \epsilon / (M - 1)} \asymp \frac {1}{\epsilon \eta_ {Y}} \ln \left(\frac {M}{\epsilon}\right) \tag {73}
$$
Note that  $\epsilon \ll 1$  and we have  $N\gg K^2$ , then we have  $S = \sum_{n}S_{n}\lesssim \frac{K}{\epsilon\eta_{Y}}\ln (\frac{M}{\epsilon})$ .
D.5 Proof of Theorem 1
Finally, we turn to prove Theorem 1. Obviously, all the diagonal elements of  $E$  are zero and all the off-diagonal elements of  $E$  are non-negative since  $c_{l|m,n} \geq 0$ . Note that  $E$  is a real symmetric matrix, then it can be orthogonal diagonalization by  $E = U^{\top}DU$  where  $U := [\pmb{u}_1, \dots, \pmb{u}_K] \in O_{K \times K}$ ,  $D = \mathrm{diag}(\lambda_1, \dots, \lambda_K)$  and  $|\lambda_1| \geq \dots \geq |\lambda_K| \geq 0$ . Then we can get the following properties of  $E$  and  $E'$ :
Lemma 13.  $\max_{i,j\in [K]}(|E_{ij}|)\leq |\lambda_1|$
$$
\left| E _ {i j} \right| = \boldsymbol {u} _ {i} ^ {\top} D \boldsymbol {u} _ {j} \leq \left| \lambda_ {1} \right| \cdot \| \boldsymbol {u} _ {i} \| _ {2} \| \boldsymbol {u} _ {j} \| _ {2}, \quad \forall i, j \in [ K ] \tag {74}
$$
Lemma 14. If  $E \in \mathbb{R}^K$  satisfies  $|\lambda_1| \leq \lambda < 1$ , then  $(I + E)$  is invertible and  $(I + E)^{-1} = I - E'$ , where  $E'$  satisfies  $E' = U^\top D'U$  and  $D' = \text{diag}(\lambda_1', \dots, \lambda_K')$  and  $\lambda_i' = \frac{\lambda_i}{1 + \lambda_i}$ ,  $\forall i \in [K]$ .
Proof. Since  $U$  is orthonormal and  $|\lambda_i| \leq \lambda < 1$ , we have  $E^n = U^\top D^n U \to O$ . Then from the property of the Neumann series, we get  $I + E$  is invertible and
$$
\begin{array}{l} (I + E) ^ {- 1} = I + \sum_ {n = 1} ^ {\infty} (- 1) ^ {n} E ^ {n} (75) \\ = I + U ^ {\top} \left(\sum_ {n = 1} ^ {\infty} (- D ^ {n}) U \right. (76) \\ = I - U ^ {\top} D ^ {\prime} U =: I - E ^ {\prime} (77) \\ \end{array}
$$
Here we define  $D^{\prime} = \mathrm{diag}\big(\lambda_1^{\prime},\dots,\lambda_K^{\prime}\big)$  and use the fact that  $\sum_{n = 1}^{\infty}(-\lambda_i)^n = -\frac{\lambda_i}{1 + \lambda_i}$
Lemma 15. If  $|\lambda_1| \leq \lambda < 1$ , then  $\max_{i \in [K]} |\lambda_i(E')| \leq \frac{1}{1 - \lambda} |\lambda_1| \leq \frac{\lambda}{1 - \lambda}$ .
$$
\max  _ {i \in [ K ]} | \lambda_ {i} \left(E ^ {\prime}\right) | = \max  _ {i \in [ K ]} | - \frac {\lambda_ {i}}{1 + \lambda_ {i}} | \leq \frac {\max  _ {i \in [ K ]} | \lambda_ {i} |}{1 - \max  _ {i \in [ K ]} | \lambda_ {i} |} \leq \frac {1}{1 - \lambda} | \lambda_ {1} | \tag {78}
$$
Lemma 16. Assume that Assumption 2 holds, then all the diagonal elements of  $E'$  are nonpositive, i.e.,  $E_{ii}' \leq 0, \forall i \in [K]$ . Further, if there exist any  $k \neq i \in [K]$  such that  $E_{ki} > 0$ , then  $E_{ii}' < 0$ .
Proof. Note that  $E_{ii} = \sum_{k=1}^{K} \lambda_k u_{ik}^2 = 0$  (here  $u_{ik}$  is the  $k$ -th component of eigenvector  $\mathbf{u}_i$ ) and  $|\lambda_k| < 1$ , we have
$$
E _ {i i} ^ {\prime} = \sum_ {k = 1} ^ {K} \frac {\lambda_ {k}}{1 + \lambda_ {k}} u _ {i k} ^ {2} = \sum_ {k = 1} ^ {K} \lambda_ {k} u _ {i k} ^ {2} - \sum_ {k = 1} ^ {K} \frac {\lambda_ {k} ^ {2}}{1 + \lambda_ {k}} u _ {i k} ^ {2} = - \sum_ {k = 1} ^ {K} \frac {\lambda_ {k} ^ {2}}{1 + \lambda_ {k}} u _ {i k} ^ {2} \leq 0 \tag {79}
$$
When  $E_{ii}^{\prime} = 0$ , then  $\lambda \coloneqq (\lambda_1, \dots, \lambda_K)$  must not have overlapping entries with respect to  $\mathbf{u}_i$ , which results that  $E_{ij} \coloneqq \sum_{k=1}^{K} \lambda_k u_{ik} u_{jk} = 0$  holds for any  $j \in [K]$ . So we prove the results.
Lemma 17. If  $\lambda_1 < 1$ , then  $|E_{nn'}' - E_{nn'}| \leq |\lambda_1|^2 (1 - |\lambda_1|)^{-1}$ .
Proof. From Lemma 14 we have:
$$
\begin{array}{l} \left| E _ {n n ^ {\prime}} ^ {\prime} - E _ {n n ^ {\prime}} \right| = \left| \sum_ {k = 1} ^ {K} \lambda_ {k} u _ {n k} u _ {n ^ {\prime} k} - \sum_ {k = 1} ^ {K} \frac {\lambda_ {k}}{1 + \lambda_ {k}} u _ {n k} u _ {n ^ {\prime} k} \right| \\ = | \sum_ {k = 1} ^ {K} \frac {\lambda_ {k} ^ {2}}{1 + \lambda_ {k}} u _ {n k} u _ {n ^ {\prime} k} | \\ \leq \frac {\left| \lambda_ {1} \right| ^ {2}}{1 - \left| \lambda_ {1} \right|} \sum_ {k = 1} ^ {K} \left| u _ {n k} \right| \left| u _ {n ^ {\prime} k} \right| \tag {80} \\ \leq \frac {\left| \lambda_ {1} \right| ^ {2}}{1 - \left| \lambda_ {1} \right|} \sqrt {\left(\sum_ {k = 1} ^ {K} \left| u _ {n k} \right| ^ {2}\right) \left(\sum_ {k = 1} ^ {K} \left| u _ {n ^ {\prime} k} \right| ^ {2}\right)} = \frac {\left| \lambda_ {1} \right| ^ {2}}{1 - \left| \lambda_ {1} \right|} \\ \end{array}
$$
Finally we can prove our main theorem in Sec. 4.
Proof. Note that if Assumption 2 holds, then  $F^{\dagger} = (I - E')F^{\top}$ . Recall  $q_{1} := 1 + \frac{1}{M - 1} \approx 1$  and  $q_{0} := -\frac{M}{(M - 1)^{2}} \approx 0$ . Then given  $x_{T + 1}[i] = n$ , we get: