SlowGuess/NeurIPS_2023_OCR huggingface dataset:
By Lemma 15, with probability at least  $1 - \exp \left(-\Omega \left(\left(N - 1\right)^{-\hat{c}} d_{s}^{-1}\right)\right)$ , one has  $\lambda_0 \geq \Theta \left(\eta_V / d_s\right) = \Theta \left(d^{-1} d_s^{-1}\right)$ . Thus, when  $d_m \geq \tilde{\Omega}(N^2)$ , all of the above conditions hold. As a result, the conditions in Eqs. (45) and (46) are satisfied and the convergence of training Transformer is guaranteed as in Eq. (4). Note that one can achieve the same width requirement and probability for He initialization, and the proof bears resemblance to the LeCun initialization.
C.4 Proof of Theorem 2 (NTK analysis)
Proof. Below, we present the proof for NTK initialization. We select  $C_Q = C_K = C_V = C_O = 1$ , then by Lemma 5, with probability at least  $1 - 8e^{-d_m / 2}$ , we have:
$$
\bar {\lambda} _ {V} = \mathcal {O} (\sqrt {d _ {m}} + \sqrt {d}), \quad \bar {\lambda} _ {O} = \mathcal {O} (\sqrt {d _ {m}}), \tag {56}
$$
$$
\bar {\lambda} _ {Q} = \mathcal {O} (\sqrt {d _ {m}} + \sqrt {d}), \quad \bar {\lambda} _ {K} = \mathcal {O} (\sqrt {d _ {m}} + \sqrt {d}).
$$
When  $d_m \geq d$ , plugging Eq. (56) into Eqs. (2) and (3), it suffices to prove the following equations.
$$
\alpha^ {2} \geq \mathcal {O} \left(\sqrt {N} d _ {s} ^ {5 / 2} C _ {x} ^ {3}\right) \sqrt {2 \ell \left(\boldsymbol {\theta} ^ {0}\right)}, \tag {57}
$$
$$
\alpha^ {3} \geq \mathcal {O} \left(N d _ {s} ^ {3} C _ {x} ^ {2} \left(1 + 4 C _ {x} ^ {4} d _ {s} ^ {2}\right)\right) \sqrt {2 \ell \left(\boldsymbol {\theta} ^ {0}\right)}. \tag {58}
$$
$$
\boldsymbol {\Phi} = \left[ \boldsymbol {X} _ {1} ^ {\top} \beta_ {1, 1}, \dots , \boldsymbol {X} _ {N} ^ {\top} \beta_ {1, N} \right] ^ {\top}, \tag {59}
$$
By Eq. (48), w.p at least  $1 - \delta$ ,  $\lambda_{\min}(\hat{F}_{\mathrm{pre}}\hat{F}_{\mathrm{pre}}^{\top}) \geq d_m\lambda_0 / 4$ , as long as it holds  $d_m \geq \tilde{\Omega}(N / \lambda_0)$  where  $\lambda_0 = \lambda_{\min}\left(\mathbb{E}_{\boldsymbol{w} \sim \mathcal{N}(0,\mathbb{I}_d)}[\sigma_r(\Phi \boldsymbol{w})\sigma_r(\Phi \boldsymbol{w})^\top]\right)$ , and  $\tilde{\Omega}$  hides logarithmic factors depending on  $\delta_1$ . Lastly, w.p. at least  $1 - \delta_3$ , one has  $\sqrt{2\ell(\pmb{\theta}^0)} \leq \tilde{\mathcal{O}}(\sqrt{N})$ . Plugging back Eqs. (2) and (3), it suffices to prove the following inequality.
$$
d _ {m} \lambda_ {0} / 4 \geq \tilde {\mathcal {O}} \left(N d _ {s} ^ {5 / 2} C _ {x} ^ {3}\right), \tag {60}
$$
$$
\left(d _ {m} \lambda_ {0} / 4\right) ^ {3 / 2} \geq \tilde {\mathcal {O}} \left(N ^ {3 / 2} d _ {s} ^ {3} C _ {x} ^ {2} \left(1 + 4 C _ {x} ^ {4} d _ {s} ^ {2}\right)\right). \tag {61}
$$
By Lemma 15, with probability at least  $1 - \exp \left(-\Omega ((N - 1)^{-\hat{c}}d_s^{-1})\right)$ , we have  $\lambda_0\geq \Omega (1)$ . Thus, all of the above conditions hold when  $d_{m} = \tilde{\Omega} (N)$ . As a result, the conditions in Eqs. (57) and (58) are satisfied and the convergence of training Transformer is guaranteed.
C.5 Discussion for different initialization schemes
Recall that the convergence result in Theorem 2 shows:
$$
\ell (\boldsymbol {\theta} ^ {\top}) \leq \left(1 - \gamma \frac {\alpha^ {2}}{2}\right) ^ {t} \ell (\boldsymbol {\theta} ^ {0}).
$$
Thus, to discuss the convergence speed for different initialization, we need to check the lower bound for  $\alpha^2$ . From the proofs for different initialization schemes above, we have the following lower bound for  $\alpha^2$ , i.e.,
$$
\alpha^ {2} = \lambda_ {\mathrm {m i n}} (\hat {\boldsymbol {F}} _ {\mathrm {p r e}} \hat {\boldsymbol {F}} _ {\mathrm {p r e}} ^ {\top}) \geq \tau_ {1} ^ {2} d _ {m} \lambda_ {0} / 4 \geq \tau_ {1} ^ {2} \eta_ {V} d _ {m} \Omega (N / d),
$$
with high probability. Plugging the value of  $\tau_{1}$  and  $\eta_V$ , we observe that for LeCun initialization and He initialization:  $\alpha^{2} \geq \Omega(d_{m}N / d)$  while for NTK initialization:  $\alpha^{2} \geq \Omega(N / d)$ . Thus, the convergence speed of LeCun initialization and He initialization is faster than NTK initialization. As a result, faster step-size is required for NTK initialization.
C.6 Discussion for  $\tau_0 = d_m^{-1}$  and  $\tau_0 = d_m^{-1/2}$
Eq. (4) indicates that the convergence speed is affected by  $\alpha$ , there we compare the lower bound for  $\alpha$  for these two scaling. In Appendix C.3, we have proved that under the LeCun initialization, one has  $\alpha^2 \geq d_m \lambda_0 / 4 \geq d_m \eta_V \mu (\sigma_r)^2 \Theta (\| \beta_{1,k} \|_2^2)$ . Note that this bound holds for these two different scaling, which is inside  $\beta$ . Thus in the next, we need to see the difference between the lower bound of  $\| \beta_{1,k} \|_2^2$  in the case of these two scalings. Specifically, for  $\tau_0 = d_m^{-1/2}$  scaling, we have proved that  $\| \beta_{1,k} \|_2^2 \geq 1 / d_s$  by Lemma 8. However, for the case of  $\tau_0 = d_m^{-1}$  scaling, when the width is large enough, the value inside the softmax tends to zero, as a result,  $\| \beta_{1,k} \|_2^2 \approx 1 / d_s$ . Thus, we can see that as the width increases, the convergence speed of  $\tau_0 = d_m^{-1/2}$  could be faster. Lastly, we remark on the difference in the step size for these two scales, which can be seen from the definition of  $C$  and its corresponding  $c_1, c_2, c_3$  in Appendix C.2.
C.7 Discussion for extension to deep Transformer and residual Transformer
Extension from our shallow Transformer to deep Transformer is not technically difficult as they share the same analysis framework. Nevertheless, the extension requires several tedious derivations and calculations involving the query, key, and value matrices along different layers. Here we point out the proof roadmap for this extension. The first step is following Proposition 1 to provide the sufficient condition for the convergence guarantee, e.g., Eqs. (2) and (3). The second step is similar to Theorem 2, where we need to verify the aforementioned assumptions for different initialization. In the second part, the main task is to prove the lower bound of  $\alpha \coloneqq \sigma_{\mathrm{min}}\left(\boldsymbol{F}_{\mathrm{pre}}^{0}\right)$ , where  $\boldsymbol{F}_{\mathrm{pre}}$  is the output of the last hidden layer. One can apply concentration inequality to bound the difference between  $\sigma_{\mathrm{min}}\left(\boldsymbol{F}_{\mathrm{pre}}^{0}\right)$  and  $\sigma_{\mathrm{min}}\left(\boldsymbol{F}_{\mathrm{pre}}^{\star 0}\right)$ , where the latter is the corresponding limit in infinite width. Lastly, one needs to plug the lower bound into the assumptions in order to obtain the width requirement.
Our proof framework is general and can handle the following residual Transformer. Here we give a proof sketch to show how to achieve this. Specifically, we consider the residual block in the self-attention layer:
$$
\mathbf {A _ {1}} = \operatorname {S e l f - a t t e n t i o n} (\mathbf {X}) \triangleq \sigma_ {\mathbf {s}} \left(\tau_ {\mathbf {0}} \left(\mathbf {X W _ {Q} ^ {\top}}\right) \left(\mathbf {X W _ {K} ^ {\top}}\right) ^ {\top}\right) \left(\mathbf {X W _ {V} ^ {\top}}\right) + \mathbf {X}.
$$
As a result, the output becomes
$$
f (\boldsymbol {X}) = \tau_ {1} \boldsymbol {w} _ {O} ^ {\top} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} + \left(\boldsymbol {X} ^ {(i:)}\right) ^ {\top}\right).
$$
To prove the convergence of the above residual Transformer, the first part that will be modified is the proof for Proposition 1. The formula for  $f_{\mathrm{pre}}$  becomes as follows:
$$
\boldsymbol {f} _ {p r e} = \tau_ {1} \boldsymbol {w} _ {O} ^ {\top} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} + \left(\boldsymbol {X} ^ {(i:)}\right) ^ {\top}\right).
$$
Lemmas 7 and 8 remain unchanged. In Lemma 9, only the first step in the proof changes while the remaining part does not change because the term  $X^{(i:)}^{\top}$  in two adjacent time steps cancels out. In Lemma 10, we have
$$
| | \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\pmb {W} _ {V} ^ {t ^ {\prime}} \pmb {X} _ {1} ^ {\top} \pmb {\beta} _ {i, 1} ^ {t ^ {\prime}} (\pmb {X} ^ {(i)})\right) | | _ {2} \leq \tau_ {1} d _ {s} \left(| | \pmb {W} _ {V} ^ {t ^ {\prime}} | | _ {2} d _ {s} ^ {1 / 2} C _ {x} + C _ {x}\right).
$$
Similarly, in the remaining lemmas, we need to add the additional term  $C_x$  for the upper bound of  $||X^{(i)}||_2$ .
The second part is the proof for Theorem 1 regarding the lower bound for  $\alpha_0$ . By Weyl's inequality:
$$
\alpha_ {0} = \sigma_ {\min } \left(\boldsymbol {F} _ {\text {p r e}}\right) \geq \sigma_ {\min } \left(\boldsymbol {F} _ {\text {p r e}} ^ {*}\right) - \left\| \boldsymbol {F} _ {\text {p r e}} - \boldsymbol {F} _ {\text {p r e}} ^ {*} \right\| _ {2},
$$
where we denote by  $\mathbf{F}_{\mathrm{pre}} = [f_{\mathrm{pre}}(\mathbf{X}_1), \dots, f_{\mathrm{pre}}(\mathbf{X}_N)]$ , and  $\mathbf{F}_{\mathrm{pre}}^*$  is the corresponding one without the residual connection. Then we can upper bound the second term can be bounded as follows:
$$
\begin{array}{l} \left\| \boldsymbol {F} _ {\text {p r e}} - \boldsymbol {F} _ {\text {p r e}} ^ {*} \right\| _ {2} \leq \left\| \boldsymbol {F} _ {\text {p r e}} - \boldsymbol {F} _ {\text {p r e}} ^ {*} \right\| _ {F} \\ \leq \sqrt {N} \left\| \tau_ {1} \boldsymbol {w} _ {O} ^ {\top} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} + \left(\boldsymbol {X} ^ {(i:)}\right) ^ {\top}\right) - \tau_ {1} \boldsymbol {w} _ {O} ^ {\top} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right) \right\| _ {2} \tag {62} \\ \leq \sqrt {N} \tau_ {1} d _ {s} | | \boldsymbol {w} _ {O} | | _ {2} | | \boldsymbol {W} _ {V} | | _ {2} C _ {x}. \\ \end{array}
$$
The remains step follows the same as previous analysis.
C.8 Linear over-parametrization and attention module behaving as a pooling layer
In this section, we discuss the link between the linear over-parametrization and attention module behaving as a pooling layer under the  $d_m^{-1}$  scaling. First, due to the  $d_m^{-1}$  scaling, the attention module degenerates to a pooling layer according to the law of large numbers. In this case, the nonlinearity on  $\pmb{X}$  disappears and thus the minimum eigenvalue of  $\Theta \Theta^{\top}$  can be estimated via  $\pmb{X} \pmb{X}^{\top}$ . Accordingly, this leads to the minimum eigenvalue in the order of  $\Omega(N / d)$ , and thus linear over-parameterization is enough.
D Proof for NTK
In this section, we elaborate the proof for Lemma 1 in Appendix D.1, the proof for Theorem 3 in Appendix D.3, respectively.
Proof. We will compute the inner product of the Jacobian of each weight separately. Firstly, we analyze  $\pmb{w}_O$ . Let us denote by  $\beta_{i} := \sigma_{s}\left(\tau_{0}\pmb{X}^{(i,:)}\pmb{W}_{Q}^{\top}\pmb{W}_{K}\pmb{X}^{\top}\right)^{\top} \in \mathbb{R}^{d_{s}}$ . Then:
$$
\frac {\partial f (\boldsymbol {X})}{\partial \boldsymbol {w} _ {O}} = \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right).
$$
The inner product of the gradient is:
$$
\begin{array}{l} \lim  _ {d _ {m} \rightarrow \infty} \left\langle \frac {\partial f (\boldsymbol {X})}{\partial \boldsymbol {w} _ {O}}, \frac {\partial f (\boldsymbol {X} ^ {\prime})}{\partial \boldsymbol {w} _ {O}} \right\rangle \\ = \tau_ {1} ^ {2} \sum_ {i = 1, j = 1} ^ {d _ {s}} \lim  _ {d _ {m} \rightarrow \infty} \left(\sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\prime \top} \beta_ {j} ^ {\prime}\right)\right) ^ {\top} \left(\sigma_ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \beta_ {i}\right)\right) \tag {63} \\ = d _ {s} ^ {2} \mathbb {E} _ {\boldsymbol {w} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I})} \left(\sigma_ {r} \left(\boldsymbol {w} ^ {\top} \boldsymbol {X} ^ {\prime \top} \boldsymbol {1} _ {d _ {s}}\right)\right) \left(\sigma_ {r} \left(\boldsymbol {w} ^ {\top} \boldsymbol {X} ^ {\top} \boldsymbol {1} _ {d _ {s}}\right)\right), \\ \end{array}
$$
where the second equality uses the law of large numbers. Secondly, we analyze  $\mathbf{W}_Q$ :
$$
\frac {\partial f (\boldsymbol {X})}{\partial W _ {Q} ^ {(p , q)}} = \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) ^ {\top} \boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {X} \boldsymbol {W} _ {K} ^ {\top} \boldsymbol {e} _ {p} \boldsymbol {e} _ {q} ^ {\top} \boldsymbol {X} ^ {(i,:) \top}.
$$
$$
\frac {\partial f (\boldsymbol {X})}{\partial \boldsymbol {W} _ {Q}} = \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {W} _ {k} \boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} \left(\beta_ {i}\right) - \beta_ {i} \beta_ {i} ^ {\top}\right) \boldsymbol {X} \boldsymbol {W} _ {V} ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \beta_ {i}\right)\right) \boldsymbol {X} ^ {(i,:)}.
$$
$$
\begin{array}{l} \lim  _ {d _ {m} \rightarrow \infty} \left\langle \frac {\partial f (\boldsymbol {X})}{\partial \boldsymbol {W} _ {Q}}, \frac {\partial f (\boldsymbol {X} ^ {\prime})}{\partial \boldsymbol {W} _ {Q}} \right\rangle \\ = \lim  _ {d _ {m} \rightarrow \infty} \tau_ {0} ^ {2} \tau_ {1} ^ {2} \sum_ {i = 1, j = 1} ^ {d _ {s}} \operatorname {T r a c e} \left(\boldsymbol {W} _ {k} \boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {X} \boldsymbol {W} _ {V} ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) \boldsymbol {X} ^ {(i,:)} \right. \\ \end{array}
$$
$$
\begin{array}{l} \boldsymbol {X} ^ {(j,:) \top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\prime \top} \beta_ {j} ^ {\prime}\right)\right) ^ {\top} \boldsymbol {W} _ {V} \boldsymbol {X} ^ {\prime \top} \left(\operatorname {d i a g} \left(\beta_ {j} ^ {\prime}\right) - \beta_ {j} ^ {\prime} \beta_ {j} ^ {\prime \top}\right) \boldsymbol {X} ^ {\prime} \boldsymbol {W} _ {k} ^ {\top}) \\ = \lim  _ {d _ {m} \rightarrow \infty} \sum_ {i = 1, j = 1} ^ {d _ {s}} \boldsymbol {X} ^ {(i,:)} \boldsymbol {X} ^ {(j,:) \top} \left\langle \tau_ {0} \tau_ {1} \left(\boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} (\beta_ {i}) - \beta_ {i} \beta_ {i} ^ {\top}\right) \boldsymbol {X} \boldsymbol {W} _ {V} ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \beta_ {i}\right)\right), \right.\right. \\ \end{array}
$$
$$
\tau_ {0} \tau_ {1} \boldsymbol {W} _ {k} ^ {\top} \boldsymbol {W} _ {k} \boldsymbol {X} ^ {\prime \top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {j} ^ {\prime}\right) - \beta_ {j} ^ {\prime} \boldsymbol {\beta} _ {j} ^ {\prime \top}\right) \boldsymbol {X} ^ {\prime} \boldsymbol {W} _ {V} ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime}\right)\right)) \rangle .
$$
For the first term of the inner product, we have:
$$
\begin{array}{l} \lim  _ {d _ {m} \rightarrow \infty} \tau_ {0} \tau_ {1} \boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} \left(\beta_ {i}\right) - \beta_ {i} \beta_ {i} ^ {\top}\right) \boldsymbol {X} \boldsymbol {W} _ {V} ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) \\ = \lim  _ {d _ {m} \rightarrow \infty} \boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} (\beta_ {i}) - \beta_ {i} \beta_ {i} ^ {\top}\right) \boldsymbol {X} \left[\begin{array}{c}\tau_ {0} \tau_ {1} \sum_ {k = 1} ^ {d _ {m}} W _ {V} ^ {(k, 1)} w _ {O} ^ {(k)} \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \beta_ {i}\right) ^ {(k)}\\\vdots\\\tau_ {0} \tau_ {1} \sum_ {k = 1} ^ {d _ {m}} W _ {V} ^ {(k, d)} w _ {O} ^ {(k)} \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \beta_ {i}\right) ^ {(k)}\end{array}\right] \\ = \boldsymbol {X} ^ {\top} \left(\operatorname {d i a g} \left(\mathbf {1} _ {d _ {s}}\right) - \mathbf {1} _ {d _ {s}} \mathbf {1} _ {d _ {s}} ^ {\top}\right) \boldsymbol {X} \left[ \begin{array}{c} 0 \\ \vdots \\ 0 \end{array} \right] = 0  , \\ \end{array}
$$
where the second equality is by the law of large numbers and  $\mathbb{E}w = 0$  for a random variable  $w\sim \mathcal{N}(0,1)$ . Thus:  $\lim_{d_m\to \infty}\left\langle \frac{\partial f(\pmb{X})}{\partial\pmb{W}_Q},\frac{\partial f(\pmb{X}')}{\partial\pmb{W}_Q}\right\rangle = 0$ . Similarly,  $\lim_{d_m\to \infty}\left\langle \frac{\partial f(\pmb{X})}{\partial\pmb{W}_K},\frac{\partial f(\pmb{X}')}{\partial\pmb{W}_K}\right\rangle = 0$ . Lastly, we analyze  $\pmb{W}_{V}$ :
$$
\frac {\partial f (\boldsymbol {X})}{\partial W _ {V} ^ {(p , q)}} = \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) ^ {\top} \boldsymbol {e} _ {p} \boldsymbol {e} _ {q} ^ {\top} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}.
$$
$$
\frac {\partial f (\boldsymbol {X})}{\partial \boldsymbol {W} _ {V}} = \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {X}.
$$
$$
\begin{array}{l} \lim  _ {d _ {m} \rightarrow \infty} \left\langle \frac {\partial f (\boldsymbol {X})}{\partial \boldsymbol {W} _ {V}}, \frac {\partial f (\boldsymbol {X} ^ {\prime})}{\partial \boldsymbol {W} _ {V}} \right\rangle \\ = \lim  _ {d _ {m} \rightarrow \infty} \tau_ {1} ^ {2} \sum_ {i = 1, j = 1} ^ {d _ {s}} \operatorname {T r a c e} \left(\left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {X} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime}\right)\right) ^ {\top}\right) \\ = \lim  _ {d _ {m} \rightarrow \infty} \tau_ {1} ^ {2} \sum_ {i = 1, j = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime}\right)\right) ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right)\right) \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {X} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime} \\ = \lim  _ {d _ {m} \rightarrow \infty} \tau_ {1} ^ {2} \sum_ {i = 1, j = 1} ^ {d _ {s}} \sum_ {k = 1} ^ {d _ {m}} (w _ {O} ^ {(k)}) ^ {2} \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {(k,:)} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime}\right) \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {(k,:)} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i}\right) \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {X} \boldsymbol {X} ^ {\prime \top} \boldsymbol {\beta} _ {j} ^ {\prime} \\ = d _ {s} ^ {2} \mathbb {E} _ {\boldsymbol {w} \sim \mathcal {N} (\boldsymbol {0}, I)} \left(\dot {\sigma} _ {r} \left(\boldsymbol {w} ^ {\top} \boldsymbol {X} ^ {\prime \top} \mathbf {1} _ {d _ {s}}\right)\right) \left(\dot {\sigma} _ {r} \left(\boldsymbol {w} ^ {\top} \boldsymbol {X} ^ {\top} \mathbf {1} _ {d _ {s}}\right)\right) \left(\mathbf {1} _ {d _ {s}} ^ {\top} \boldsymbol {X} \boldsymbol {X} ^ {\prime \top} \mathbf {1} _ {d _ {s}}\right), \tag {64} \\ \end{array}
$$
where the last equality uses the law of large numbers.
D.2 NTK minimum eigenvalue
Lemma 16. Given  $\Phi$  defined in Eq. (53) and  $\Phi^{\star}$  defined in Lemma 1, denote  $\lambda_{*} = \lambda_{\min}(\Phi^{\star}\Phi^{\star^{\top}})$ , and suppose the width satisfies  $d_m = \Omega \left( \frac{N^2 \sqrt{\log(2d^2 N^2 / \delta)}}{\lambda_*^2} \right)$ , then with probability at least  $1 - \delta$ , one has  $\left\| \Phi \Phi^{\top} - \Phi^{\star} \Phi^{\star^{\top}} \right\|_{\mathrm{F}} \leq \frac{\lambda_*}{4}$  and  $\lambda_{\min}(\Phi \Phi^{\top}) \geq \frac{3\lambda_*}{4}$ .
Proof. In the following content, the variable with  $\star$  indicates the corresponding variable with infinite width  $d_m$ . According to the definition of  $\Phi$  in Eq. (53) and  $\Phi^{\star}$  in Lemma 1, for each entry in  $\Phi \Phi^{\top}$  and  $\Phi^{\star} \Phi^{\star^{\top}}$ , we have
$$
\begin{array}{l} \left| \left(\Phi \Phi^ {\top} - \Phi^ {\star} \Phi^ {\star \top}\right) ^ {(n, r)} \right| = \left| \beta_ {1, n} ^ {\top} \mathbf {X} _ {n} \mathbf {X} _ {r} ^ {\top} \beta_ {1, r} - \beta_ {1, n} ^ {\star \top} \mathbf {X} _ {n} \mathbf {X} _ {r} ^ {\top} \beta_ {1, r} ^ {\star} \right| \\ \leq \left\| \boldsymbol {X} _ {n} \boldsymbol {X} _ {r} ^ {\top} \left(\beta_ {1, n} - \beta_ {1, n} ^ {\star}\right) \right\| _ {2} + \left\| \boldsymbol {X} _ {n} \boldsymbol {X} _ {r} ^ {\top} \left(\beta_ {1, r} - \beta_ {1, r} ^ {\star}\right) \right\| _ {2} \leq C _ {x} ^ {2} d _ {s} \left(\left\| \beta_ {1, n} - \beta_ {1, n} ^ {\star} \right\| _ {2} + \left\| \beta_ {1, r} - \beta_ {1, r} ^ {\star} \right\| _ {2}\right). \tag {65} \\ \end{array}
$$
Next, we will bound  $\left\| \beta_{1,n} - \beta_{1,n}^{\star}\right\|_{2}$
$$
\begin{array}{l} \left\| \boldsymbol {\beta} _ {1, n} - \boldsymbol {\beta} _ {1, n} ^ {\star} \right\| _ {2} = \left\| \sigma_ {s} \left(\tau_ {0} \boldsymbol {X} _ {n} ^ {(1,:)} \boldsymbol {W} _ {Q} ^ {\top} \boldsymbol {W} _ {K} \boldsymbol {X} _ {n} ^ {\top}\right) ^ {\top} - \sigma_ {s} \left(\tau_ {0} \boldsymbol {X} _ {n} ^ {(1,:)} \boldsymbol {W} _ {Q} ^ {\star \top} \boldsymbol {W} _ {K} ^ {\star} \boldsymbol {X} _ {n} ^ {\top}\right) ^ {\top} \right\| _ {2} \tag {66} \\ \leq 2 C _ {x} ^ {2} d _ {s} \left\| \tau_ {0} \boldsymbol {W} _ {Q} ^ {\top} \boldsymbol {W} _ {K} - \tau_ {0} \boldsymbol {W} _ {Q} ^ {\star \top} \boldsymbol {W} _ {K} ^ {\star} \right\| _ {\mathrm {F}} \quad (\text {B y L e m m a 7}). \\ \end{array}
$$
Let first consider the absolute value of each element of  $\tau_0\pmb{W}_Q^\top \pmb{W}_K - \tau_0\pmb{W}_Q^{\star \top}\pmb{W}_K^{\star}$ , i.e.,
$$
| \tau_ {0} \pmb {W} _ {Q} ^ {\top} \pmb {W} _ {K} - \tau_ {0} \pmb {W} _ {Q} ^ {\star \top} \pmb {W} _ {K} ^ {\star} | ^ {(i, j)} = | \tau_ {0} \sum_ {q = 1} ^ {d _ {m}} W _ {Q} ^ {(q, i)} W _ {k} ^ {(q, j)} - \tau_ {0} \sum_ {q = 1} ^ {d _ {m}} W _ {Q} ^ {\star (q, i)} W _ {k} ^ {\star (q, j)} |.
$$
Since  $W_{Q}^{(q,i)}W_{k}^{(q,j)}\sim SE(\sqrt{2}\eta_{QK},\sqrt{2}\eta_{QK})$  is sub-exponential random variable, by Lemma 2 and Lemma 3,  $\tau_0\sum_{q = 1}^{d_m}W_Q^{(q,i)}W_k^{(q,j)}\sim SE(\tau_0\eta_{QK}\sqrt{2d_m},\sqrt{2}\tau_0\eta_{QK})$ , by Bernstein's inequality, when  $d_{m}\geq 2\log (2 / \delta)$ , the following inequality holds with probability at least  $1 - \delta$ :
$$
\left| \tau_ {0} \sum_ {q = 1} ^ {d _ {m}} W _ {Q} ^ {(q, i)} W _ {k} ^ {(q, j)} - \tau_ {0} \sum_ {q = 1} ^ {d _ {m}} W _ {Q} ^ {\star (q, i)} W _ {k} ^ {\star (q, j)} \right| \leq 2 \tau_ {0} \eta_ {Q K} \sqrt {d _ {m} \log (2 / \delta)}. \tag {67}
$$
Substituting  $\delta = \frac{\delta'}{d^2}$  and applying union bound, we can obtain that when  $d_m \geq 2\log \frac{2d^2}{\delta'}$ , with probability at least  $1 - \delta'$ :
$$
\left\| \boldsymbol {\beta} _ {1, n} - \boldsymbol {\beta} _ {1, n} ^ {\star} \right\| _ {2} \leq 4 \tau_ {0} \eta_ {Q K} C _ {x} ^ {2} d _ {s} d \sqrt {d _ {m} \log \left(2 d ^ {2} / \delta^ {\prime}\right)}. \tag {68}
$$