SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Substituting back to Eq. (65), the following inequality holds with the same width requirement and probability.  $|(\Phi \Phi^{\top} - \Phi^{\star}\Phi^{\star \top})^{(n,r)}| \leq 8\tau_{0}\eta_{QK}C_{x}^{4}d_{s}d\sqrt{d_{m}\log(2d^{2} / \delta^{\prime})}$  Applying the union bound over the index  $(n,r)$  for  $n \in [N]$  and  $r \in [N]$ , and substituting  $\delta' = \frac{\delta''}{N^2}$ , we obtain that when the width  $d_m \geq 2\log \frac{2d^2N^2}{\delta''}$ , the following inequality holds with probability at least  $1 - \delta''$ :
$$
\left\| \Phi \Phi^ {\top} - \Phi^ {\star} \Phi^ {\star \top} \right\| _ {\mathrm {F}} = \sqrt {\sum_ {n = 1} ^ {N} \sum_ {r = 1} ^ {N} | (\Phi \Phi^ {\top} - \Phi^ {\star} \Phi^ {\star \top}) ^ {(n , r)} | ^ {2}} \leq 8 N \tau_ {0} \eta_ {Q K} C _ {x} ^ {4} d _ {s} d \sqrt {d _ {m} \log (2 d ^ {2} N ^ {2} / \delta^ {\prime \prime})}.
$$
In the case of LeCun initialization when  $d_m = \Omega \left(\frac{N^2\sqrt{\log(2d^2N^2 / \delta^{\prime\prime})}}{\lambda_0^2}\right)$ ,  $\left\| \Phi \Phi^\top - \Phi^\star \Phi^{\star \top} \right\|_{\mathrm{F}} \leq \frac{\lambda_*}{4}$ . Lastly, one has:
$$
\lambda_ {\mathrm {m i n}} (\Phi \Phi^ {\top}) \geq \lambda_ {*} - \left\| \Phi \Phi^ {\top} - \Phi^ {\star} \Phi^ {\star^ {\top}} \right\| _ {2} \geq \lambda_ {*} - \left\| \Phi \Phi^ {\top} - \Phi^ {\star} \Phi^ {\star^ {\top}} \right\| _ {\mathrm {F}} \geq \frac {3 \lambda_ {*}}{4}.
$$
where the first inequality is by Weyl's inequality. One can easily check that the same result also holds for He initialization and NTK initialization.
Lemma 17. Given the  $\Phi^{\star}$  defined in Lemma 1, then when  $N\geq \Omega (d^{4})$  , with probability at least  $1 - e^{-d}$  one has:  $\lambda_{\mathrm{min}}(\Phi^{\star}\Phi^{\star \top})\geq \Theta (N / d)$
Proof. Firstly, let us define  $Z^{\star} = \Phi^{\star}\Sigma^{-1 / 2}$ , then:
$$
\begin{array}{l} \lambda_ {\min } \left(\Phi^ {\star} \Phi^ {\star^ {\top}}\right) = \lambda_ {\min } \left(\Phi^ {\star^ {\top}} \Phi^ {\star}\right) = \lambda_ {\min } \left(\Sigma^ {1 / 2} Z ^ {\star^ {\top}} Z ^ {\star} \Sigma^ {1 / 2 ^ {\top}}\right) \\ \geq \lambda_ {\min } \left(\boldsymbol {Z} ^ {\star \top} \boldsymbol {Z} ^ {\star}\right) \lambda_ {\min } (\boldsymbol {\Sigma}) \geq \lambda_ {\min } \left(\boldsymbol {Z} ^ {\star \top} \boldsymbol {Z} ^ {\star}\right) C _ {\Sigma}, \\ \end{array}
$$
where the last inequality is by Assumption 4). Note that we can reformulate  $Z^{\star}$  by plugging  $\Phi^{\star}$  as follows:
$$
\boldsymbol {Z} ^ {\star} = \boldsymbol {\Phi} ^ {\star} \boldsymbol {\Sigma} ^ {- 1 / 2} = \left[ \begin{array}{c} \frac {1}{d _ {s}} \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {X} _ {1} ^ {(i,:)} \boldsymbol {\Sigma} ^ {- 1 / 2} \\ \vdots \\ \frac {1}{d _ {s}} \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {X} _ {N} ^ {(i,:)} \boldsymbol {\Sigma} ^ {- 1 / 2} \end{array} \right].
$$
Recall that in Assumption 4, we define the random vector  $\boldsymbol{x} = \boldsymbol{X}_n^{(i,:)\top}$  and write the covariance matrix for each feature of  $\pmb{x}$  as  $\pmb{\Sigma} = \mathbb{E}[\pmb{x}\pmb{x}^{\top}]$ . Here we further define the random vector  $\boldsymbol{z} = \boldsymbol{\Sigma}^{-1/2}\boldsymbol{x}$ . Clearly, the covariance matrix for each feature of  $\pmb{z}$  is  $\pmb{\Sigma}_{z} = \mathbb{E}[z\pmb{z}^{\top}] = \mathbb{I}_{d}$ . Therefore, by Proposition 4 of Zhu et al. [2022]), we have  $\lambda_{\min}(Z^{\star \top}Z^{\star}) \geq \Theta(N / d)$ , which finishes the proof.
Lemma 18. Suppose the number of samples  $N \geq \Omega(d^4)$  and the width  $d_m = \tilde{\Omega}(N^2\lambda_*^{-2})$ , where  $\lambda_* = \lambda_{\min}(\Phi^\star \Phi^{\star \top})$ , then under Assumption 4, with probability at least  $1 - \delta - e^{-d}$ , one has  $\lambda_0 = \lambda_{\min}(\mathbb{E}_{\boldsymbol{w} \sim \mathcal{N}(0, \eta_V \mathbb{I}_d)}[\sigma_r(\Phi \boldsymbol{w})\sigma_r(\Phi \boldsymbol{w})^T]) \geq \Theta(\eta_V N / d)$ .
Proof. By the Hermite expansion of  $\sigma_r$ , one has:
$$
\lambda_ {0} \geq \eta_ {V} \mu \left(\sigma_ {r}\right) ^ {2} \lambda_ {\min } \left(\Phi \Phi^ {\top}\right),
$$
where  $\mu (\sigma_r)$  is the 1-st Hermite coefficient of ReLU satisfying  $\mu (\sigma_r) > 0$ . By Lemma 16, when the width satisfies  $d_m = \tilde{\Omega} (N^2\lambda_*^{-2})$ , then with probability at least  $1 - \delta$ , one has  $\lambda_{\mathrm{min}}(\Phi \Phi^{\top})\geq \frac{3}{4}\lambda_{*}$ . Furthermore, by Lemma 17, when  $N\geq \Omega (d^{4})$  w.p. at least  $1 - e^{-d}$  one has  $\lambda_{*} = \Theta (N / d)$ . Thus, with probability at least  $1 - \delta -e^{-d}$ , one has:
$$
\lambda_ {0} \geq \frac {3}{4} \eta_ {V} \mu \left(\sigma_ {r}\right) ^ {2} \lambda_ {*} \geq \Theta \left(\eta_ {V} N / d\right)).
$$
Below, we provide a lower bound for the minimum eigenvalue of NTK, which plays a key role in analyzing convergence, generalization bounds, and memorization capacity [Arora et al., 2019a, Nguyen et al., 2021, Montanari and Zhong, 2022, Bombari et al., 2022].
To prove this, we need the following assumption.
Assumption 4. Let  $\pmb{x} = \sum_{i=1}^{d_s} (\pmb{X}^{(i,:)})^\top$ ,  $\pmb{\Sigma} = \mathbb{E}[\pmb{x}\pmb{x}^\top]$ , then we assume that  $\pmb{\Sigma}$  is positive definite, i.e.,  $\lambda_{\min}(\pmb{\Sigma}) \geq C_{\Sigma}$  for some positive constant  $C_{\Sigma}$ .
Remark: This assumption implies that the covariance matrix along each feature dimension is positive definite. In statistics and machine learning [Liu et al., 2021, Liang and Rakhlin, 2020, Vershynin, 2018, Hastie et al., 2022], the covariance of  $\pmb{x}$  is frequently assumed to be an identity matrix or positive definite matrix.
Lemma 19 (Minimum eigenvalue of limiting NTK). Under Assumptions 1 and 4 and scaling  $\tau_0 = d_m^{-1}$ , when  $N \geq \Omega(d^4)$ , the minimum eigenvalue of  $\pmb{K}$  can be lower bounded with probability at least  $1 - e^{-d}$  as:  $\lambda_{\min}(\pmb{K}) \geq \mu(\sigma_r)^2 \Omega(N / d)$ , where  $\mu(\sigma_r)$  represents the first Hermite coefficient of the ReLU activation function.
Proof. By the Hermite expansion of  $\sigma_r$ , we have:
$$
\begin{array}{l} \lambda_ {\min } (\boldsymbol {K}) > \lambda_ {\min } \left(\mathbb {E} _ {\boldsymbol {w} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I})} \left(\sigma_ {r} \left(\boldsymbol {\Phi} ^ {\star} \boldsymbol {w}\right) \sigma_ {r} \left(\boldsymbol {\Phi} ^ {\star} \boldsymbol {w}\right) ^ {\top}\right)\right) \\ = \lambda_ {\min } \left(\sum_ {s = 0} ^ {\infty} \mu_ {s} \left(\sigma_ {r}\right) ^ {2} \bigcirc_ {i = 1} ^ {s} \left(\Phi^ {\star} \Phi^ {\star^ {\top}}\right)\right) [ \text {N g u e n a n d M o n d e l l i , 2 0 2 0 , L e m m a D . 3} ] \\ \geq \mu (\sigma_ {r}) ^ {2} \lambda_ {\min } \left(\Phi^ {\star} \Phi^ {\star^ {\top}}\right). \\ \end{array}
$$
Note that when  $N \geq \Omega(d^4)$ , based on Lemma 17 and the fact that  $\Phi^\star \Phi^{\star \top}$  and  $\Phi^{\star \top} \Phi^\star$  share the same non-zero eigenvalues, with probability at least  $1 - e^{-d}$  one has:
$$
\lambda_ {\min } (\boldsymbol {K}) \geq \mu \left(\sigma_ {r}\right) ^ {2} \left(\frac {N}{d} - 9 N ^ {2 / 3} d ^ {1 / 3}\right) = \mu \left(\sigma_ {r}\right) ^ {2} \Theta (N / d),
$$
Remark: The minimum eigenvalue of the NTK plays an important role in the global convergence, similar to  $\alpha$  in Proposition 1. By defining  $\alpha \triangleq \sigma_{\mathrm{min}}(F_{\mathrm{pre}}^0)$ , under the over-parameterized regime with  $d_m \geq N$ , we have  $\alpha^2 = \lambda_{\mathrm{min}}((F_{\mathrm{pre}}^0)^\top F_{\mathrm{pre}}^0)$ , which is the exact minimum eigenvalue of the empirical NTK with respect to the weight in the output layer.
Proof. Before starting the proof, we introduce some useful lemmas that are used to analyze the randomness of initialized weight.
Lemma 20. Given an initial weight vector  $\boldsymbol{w} \in \mathbb{R}^{d_m}$  where each element is sampled independently from  $\mathcal{N}(0,1)$ , for any  $\tilde{\boldsymbol{w}}$  such that  $\| \tilde{\boldsymbol{w}} - \boldsymbol{w} \|_2 \leq R$ , with probability at least  $1 - 2 \exp(-d_m / 2)$ , one has:
$$
\left\| \tilde {\boldsymbol {w}} \right\| _ {2} \leq R + 3 \sqrt {d _ {m}}. \tag {69}
$$
Proof of Lemma 20. By triangle inequality and Lemma 6:
$$
\left\| \tilde {\boldsymbol {w}} \right\| _ {2} \leq \left\| \tilde {\boldsymbol {w}} - \boldsymbol {w} \right\| _ {2} + \left\| \boldsymbol {w} \right\| _ {2} \leq R + 3 \sqrt {d _ {m}}. \tag {70}
$$
Now we are ready to start our proof, in the following content, we avoid the tilde symbol  $(\tilde{\cdot})$  over the weight for simplicity. Because we aim to study the effect of width  $(d_m)$ , to simplify the proof, we set the embedding dimension of the input as one, we can write the output layer of the network into the following form:
$$
f (\boldsymbol {x}) = \tau_ {0} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\sigma_ {s} \left(\tau_ {1} x ^ {(i)} \boldsymbol {w} _ {Q} ^ {\top} \boldsymbol {w} _ {K} \boldsymbol {x} ^ {\top}\right) \left(\boldsymbol {w} _ {V} \boldsymbol {x} ^ {\top}\right)\right) ^ {\top} \boldsymbol {w} _ {O} \tag {71}
$$
The Hessian matrix  $H$  of the network can be written as the following structure:
$$
\boldsymbol {H} = \left( \begin{array}{l l l l} \boldsymbol {H} _ {Q, Q} & \boldsymbol {H} _ {Q, K} & \boldsymbol {H} _ {Q, V} & \boldsymbol {H} _ {Q, O} \\ \boldsymbol {H} _ {K, Q} & \boldsymbol {H} _ {K, K} & \boldsymbol {H} _ {K, V} & \boldsymbol {H} _ {K, O} \\ \boldsymbol {H} _ {V, Q} & \boldsymbol {H} _ {V, K} & \boldsymbol {H} _ {V, V} & \boldsymbol {H} _ {V, O} \\ \boldsymbol {H} _ {O, Q} & \boldsymbol {H} _ {O, K} & \boldsymbol {H} _ {O, V} & \boldsymbol {H} _ {O, O} \end{array} \right), \tag {72}
$$
where we partition the Hessian  $\pmb{H}$  of the network to each Hessian block e.g.,  $\pmb{H}_{Q,Q} = \frac{\partial^2 f}{\partial w_Q \partial w_Q}, \pmb{H}_{Q,K} = \frac{\partial^2 f}{\partial w_Q \partial w_K}$ , etc. Based on the triangle inequality of the norm, the spectral norm of the Hessian  $\pmb{H}$  can be upper bounded by the sum of the spectral norm of each Hessian block. We start by analyzing  $\pmb{H}_{Q,Q}$ . The time step  $t$  is hidden for simplification.
$$
\frac {\partial f (\boldsymbol {x})}{\partial w _ {Q} ^ {(j)}} = \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} x ^ {(i)} w _ {K} ^ {(j)} \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} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x}, \tag {73}
$$
where the Jacobian of Softmax is given by Lemma 11. Next, we calculate the second-order derivative. Each element of the Hessian  $H_{Q,Q}$  is
$$
H _ {Q, Q} ^ {(j, p)} = \frac {\partial^ {2} f (\pmb {x})}{\partial w _ {Q} ^ {(j)} \partial w _ {Q} ^ {(p)}} = \tau_ {0} ^ {2} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} w _ {K} ^ {(j)} w _ {K} ^ {(p)} (x ^ {(i)}) ^ {2} \pmb {w} _ {V} ^ {\top} \left(\pmb {w} _ {O} \circ \dot {\sigma_ {r}} \left(\pmb {w} _ {V} \pmb {x} ^ {\top} \beta_ {i}\right)\right) \pmb {x} ^ {\top} \pmb {\Gamma} _ {i} \pmb {x},
$$
where we denote by  $\Gamma_{i} = \left\{\mathrm{diag}\left(\beta_{i}\circ \boldsymbol {x} - \beta_{i}\boldsymbol{\beta}_{i}^{\top}\boldsymbol {x}\right) - \left(\mathrm{diag}\left(\beta_{i}\right) - \beta_{i}\boldsymbol{\beta}_{i}^{\top}\right)\boldsymbol{x}\boldsymbol{\beta}_{i}^{\top} - \right.$ $\beta_{i}\pmb{x}^{\top}\left(\mathrm{diag}\left(\beta_{i}\right) - \beta_{i}\beta_{i}^{\top}\right)\pmb{x}\Bigg\},$  o symbolizes Hadamard product. According to initialization of  $\pmb {w}_O,\pmb {w}_V,\pmb {w}_Q,\pmb {w}_K$  and Lemma 20, with probability at least  $1 - 8e^{-d_m / 2}$ , we have
$$
\left\| \boldsymbol {w} _ {V} \right\| _ {2} \leq 3 \sqrt {\eta_ {V} d _ {m}} + R, \quad \left\| \boldsymbol {w} _ {O} \right\| _ {2} \leq 3 \sqrt {\eta_ {O} d _ {m}} + R,
$$
$$
\left\| \boldsymbol {w} _ {Q} \right\| _ {2} \leq 3 \sqrt {\eta_ {Q} d _ {m}} + R, \quad \left\| \boldsymbol {w} _ {K} \right\| _ {2} \leq 3 \sqrt {\eta_ {K} d _ {m}} + R.
$$
The spectral norm of  $H_{Q,Q}$  can be bounded with probability at least  $1 - 8e^{-d_m / 2}$ :
$$
\begin{array}{l} \| \boldsymbol {H} _ {Q, Q} \| _ {2} = \| \tau_ {0} ^ {2} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} (x ^ {(i)}) ^ {2} \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} ^ {\top} \boldsymbol {\Gamma} _ {i} \boldsymbol {x} \boldsymbol {w} _ {K} \boldsymbol {w} _ {K} ^ {\top} \| _ {2} \\ \leq \left| \right. \tau_ {0} ^ {2} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(x ^ {(i)}\right) ^ {2} \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} ^ {\top} \boldsymbol {\Gamma} _ {i} \boldsymbol {x} \| \boldsymbol {w} _ {K} \boldsymbol {w} _ {K} ^ {\top} \| _ {2} \quad (\text {H o m o g e n e i t y o f n o r m}) \\ \leq \left| \tau_ {0} ^ {2} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(x ^ {(i)}\right) ^ {2} \boldsymbol {x} ^ {\top} \boldsymbol {\Gamma} _ {i} \boldsymbol {x} \right| \| \boldsymbol {w} _ {O} \| _ {2} \| \boldsymbol {w} _ {V} \| _ {2} \| \boldsymbol {w} _ {K} \| _ {2} ^ {2} \quad (\text {C a u c h y - S c h w a r z i n e q u a l i t y}) \\ \leq \tau_ {0} ^ {2} \tau_ {1} \left| \sum_ {i = 1} ^ {d _ {s}} \left(x ^ {(i)}\right) ^ {2} \boldsymbol {x} ^ {\top} \boldsymbol {\Gamma} _ {i} \boldsymbol {x} \right| \left(3 \sqrt {\eta_ {O} d _ {m}} + R\right) \left(3 \sqrt {\eta_ {V} d _ {m}} + R\right) \left(3 \sqrt {\eta_ {Q} d _ {m}} + R\right) \left(3 \sqrt {\eta_ {K} d _ {m}} + R\right) \quad (L e m m a 2 0) \\ = \mathcal {O} \left(1 / \sqrt {d _ {m}}\right), \\ \end{array}
$$
where the last equality is by the bound of  $\| \Gamma_i\| _2$  , by triangle inequality:
$$
\left\| \boldsymbol {\Gamma} _ {i} \right\| _ {2} \leq \left\| \operatorname {d i a g} \left(\boldsymbol {\beta} _ {i} \circ \boldsymbol {x} - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {x}\right) \right\| _ {2} + \left\| \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x} \boldsymbol {\beta} _ {i} ^ {\top} \right\| _ {2} + \left\| \boldsymbol {\beta} _ {i} \boldsymbol {x} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x} \right\| _ {2}. \tag {74}
$$
For the first part of Eq. (74), we have
$$
\begin{array}{l} \left\| \operatorname {d i a g} \left(\beta_ {i} \circ \boldsymbol {x} - \beta_ {i} \beta_ {i} ^ {\top} \boldsymbol {x}\right) \right\| _ {2} \leq \left\| \operatorname {d i a g} \left(\beta_ {i} \circ \boldsymbol {x} - \beta_ {i} \beta_ {i} ^ {\top} \boldsymbol {x}\right) \right\| _ {2} \leq \left\| \left(\beta_ {i} \circ \boldsymbol {x} - \beta_ {i} \beta_ {i} ^ {\top} \boldsymbol {x}\right) \right\| _ {\infty} \\ \leq \| (\boldsymbol {\beta} _ {i} \circ \boldsymbol {x}) \| _ {\infty} + \left\| \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {x} \right\| _ {\infty} \leq \| \boldsymbol {x} \| _ {\infty} + | \boldsymbol {\beta} _ {i} ^ {\top} \boldsymbol {x} | \leq \| \boldsymbol {x} \| _ {2} + \| \boldsymbol {\beta} _ {i} \| _ {2} \| \boldsymbol {x} \| _ {2} \leq 2 C _ {x}. \\ \end{array}
$$
For the second part of Eq. (74), we have
$$
\begin{array}{l} \left\| \left(\operatorname {d i a g} \left(\beta_ {i}\right) - \beta_ {i} \beta_ {i} ^ {\top}\right) \boldsymbol {x} \beta_ {i} ^ {\top} \right\| _ {2} \leq \left(\left\| \beta_ {i} \right\| _ {\infty} + \left\| \beta_ {i} \beta_ {i} ^ {\top} \right\| _ {2}\right) \left\| \boldsymbol {x} \beta_ {i} ^ {\top} \right\| _ {2} \\ \leq 2 \left\| \boldsymbol {x} \beta_ {i} ^ {\top} \right\| _ {2} \leq 2 \left\| \boldsymbol {x} \right\| _ {2} \left\| \beta_ {i} ^ {\top} \right\| _ {2} \leq 2 C _ {x}. \\ \end{array}
$$
For the third part of Eq. (74), we have
$$
\left\| \boldsymbol {\beta} _ {i} \boldsymbol {x} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x} \right\| _ {2} \leq \left\| \boldsymbol {\beta} _ {i} \boldsymbol {x} ^ {\top} \right\| _ {2} \left\| \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \right\| _ {2} \| \boldsymbol {x} \| _ {2}.
$$
Next, we analyze the Hessian  $H_{Q,K}$ , where each element is:
$$
H _ {Q, K} ^ {(j, p)} = \frac {\partial^ {2} f}{\partial w _ {Q} ^ {(j)} \partial w _ {K} ^ {(p)}} = \frac {\partial^ {2} f (\boldsymbol {x})}{\partial w _ {Q} ^ {(j)} \partial w _ {Q} ^ {(p)}} = \tau_ {0} ^ {2} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} w _ {K} ^ {(j)} w _ {Q} ^ {(p)} (x ^ {(i)}) ^ {2} \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} ^ {\top} \boldsymbol {\Gamma} _ {i} \boldsymbol {x}.
$$
Due to the symmetry, similar to  $H_{Q,Q}$ , the spectral norm of  $H_{Q,K}$  can be bounded with probability at least  $1 - 8e^{-d_m / 2}$ :  $\| H_{Q,K}\|_2 = \mathcal{O}(1 / \sqrt{d_m})$ .
Next, we analyze the Hessian  $H_{Q,V}$ , where each element is:
$$
H _ {Q, V} ^ {(j, p)} = \frac {\partial^ {2} f}{\partial w _ {Q} ^ {(j)} \partial w _ {V} ^ {(p)}} = \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} w _ {K} ^ {(j)} x ^ {(i)} w _ {O} ^ {(p)} \dot {\sigma} _ {r} \left(w _ {V} ^ {(p)} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i}\right) \boldsymbol {x} ^ {\top} \left(\operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x}.
$$
The spectral norm of  $H_{Q,V}$  can be bounded with probability at least  $1 - 8e^{-d_m / 2}$ :
$$
\begin{array}{l} \| \boldsymbol {H} _ {Q, V} \| _ {2} = \| \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} x ^ {(i)} \dot {\sigma_ {r}} \left(w _ {V} ^ {(p)} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i}\right) \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} \boldsymbol {w} _ {O} ^ {\top} \| _ {2} \\ \leq \left| \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} x ^ {(i)} \dot {\sigma} _ {r} \left(w _ {V} ^ {(p)} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i}\right) \boldsymbol {x} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x} \right| \| \boldsymbol {w} _ {K} \boldsymbol {w} _ {O} ^ {\top} \| _ {2} \quad (\text {H o m o g e n e i t y o f n o r m}) \\ \leq \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} | x _ {i} | \| \operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top} \| _ {2} \| \boldsymbol {x} \| _ {2} ^ {2} \| \boldsymbol {w} _ {K} \| _ {2} \| \boldsymbol {w} _ {O} \| _ {2} \quad \text {(C a u c h y - S c h w a r z i n e q u a l i t y)} \\ \leq \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} | x _ {i} | \| \operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top} \| _ {2} \| \boldsymbol {x} \| _ {2} ^ {2} \left(3 \sqrt {\eta_ {Q K} d _ {m}} + R\right) \left(3 \sqrt {\eta_ {O} d _ {m}} + R\right) \quad (L e m m a 2 0) \\ = \mathcal {O} \left(1 / d _ {m}\right), \\ \end{array}
$$
where the last step is by Weyl's inequality and the range of the output of Softmax (from zero to one):
$$
\left(\min  _ {j} \left(\beta_ {i} ^ {(j)}\right) + \| \boldsymbol {\beta} _ {i} \| _ {2} ^ {2}\right) ^ {2} \leq \| \operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \beta_ {i} \boldsymbol {\beta} _ {i} ^ {\top} \| _ {2} ^ {2} \leq \left(\max  _ {j} \left(\beta_ {i} ^ {(j)}\right) + \| \boldsymbol {\beta} _ {i} \| _ {2} ^ {2}\right) ^ {2} = \mathcal {O} (1).
$$
It suffices to bound  $\| \pmb{\beta}_i\|_2$ , i.e.,  $\| \sigma_s(x^{(i)}\pmb{x}\pmb{w}_Q^\top\pmb{w}_K)\|_2$ . Since the range of each element of the output of Softmax is from 0 to 1 and the sum of is one,  $\| \pmb{\beta}_i\| \leq 1$ .
Next, we analyze the Hessian  $H_{Q,O}$ , where each element is:
$$
H _ {Q, O} ^ {(j, p)} = \frac {\partial^ {2} f}{\partial w _ {Q} ^ {(j)} \partial w _ {V} ^ {(p)}} = \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} w _ {K} ^ {(j)} x ^ {(i)} w _ {V} ^ {(p)} \dot {\sigma_ {r}} \left(w _ {V} ^ {(p)} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i}\right) \boldsymbol {x} ^ {\top} \left(\mathrm {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x}.
$$
$$
\begin{array}{l} \| \boldsymbol {H} _ {Q, O} \| _ {2} = \| \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} x ^ {(i)} \dot {\sigma_ {r}} \left(w _ {V} ^ {(p)} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i}\right) \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} \boldsymbol {w} _ {V} ^ {\top} \| _ {2} \\ \leq \left| \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} x ^ {(i)} \dot {\sigma} _ {r} \left(w _ {V} ^ {(p)} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i}\right) \boldsymbol {x} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i}\right) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top}\right) \boldsymbol {x} \right| \| \boldsymbol {w} _ {K} \boldsymbol {w} _ {V} ^ {\top} \| _ {2} \quad (\text {H o m o g e n e i t y o f n o r m}) \\ \leq \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} | x _ {i} | \| \operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top} \| _ {2} \| \boldsymbol {x} \| _ {2} ^ {2} \| \boldsymbol {w} _ {K} \| _ {2} \| \boldsymbol {w} _ {V} \| _ {2} \quad \text {(C a u c h y - S c h w a r z i n e q u a l i t y)} \\ \leq \tau_ {0} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} | x _ {i} | \| \operatorname {d i a g} (\boldsymbol {\beta} _ {i}) - \boldsymbol {\beta} _ {i} \boldsymbol {\beta} _ {i} ^ {\top} \| _ {2} \| \boldsymbol {x} \| _ {2} ^ {2} \left(3 \sqrt {\eta_ {Q K} d _ {m}} + R\right) \left(3 \sqrt {\eta_ {V} d _ {m}} + R\right) \quad (L e m m a 2 0) \\ = \mathcal {O} \left(1 / d _ {m}\right). \\ \end{array}
$$
Next, we analyze the Hessian  $\pmb{H}_{O,V}$ . Each element of  $\pmb{H}_{O,V}$  is:
$$
H _ {O V} ^ {(p, j)} = \frac {\partial^ {2} f (\boldsymbol {x})}{\partial w _ {O} ^ {(p)} \partial w _ {V} ^ {(j)}} = \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \dot {\sigma_ {r}} \left(w _ {V} ^ {(p)} \beta_ {i} ^ {\top} \boldsymbol {x}\right) \boldsymbol {x} ^ {\top} \beta_ {i} \mathbb {1} _ {\{j = p \}}.
$$
Note that  $H_{O,V}$  is a diagonal matrix. Consequently, the spectral norm of  $H_{O,V}$  is:
$$
\| \boldsymbol {H} _ {O, V} \| _ {2} = \max  _ {j \in [ d _ {m} ]} | H _ {O, V} ^ {(j j)} | \leq | \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {x} ^ {\top} \boldsymbol {\beta} _ {i} | \leq \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \| \boldsymbol {x} \| _ {2} \| \boldsymbol {\beta} _ {i} \| _ {2} = \mathcal {O} (1 / d _ {m}).
$$
Next, since each element of  $H_{O,O}$  and  $H_{V,V}$  is zero, the corresponding spectral norm is zero. Lastly, due to the symmetry of  $\boldsymbol{w}_K$  and  $\boldsymbol{w}_Q$ , we can obtain the same bound for the corresponding Hessian block. Thus, the spectral norm of Hessian  $\pmb{H}$  is upper bounded by  $\mathcal{O}(1 / \sqrt{d_m})$ , which completes the proof.
(a) LeCun initialization,  $\gamma = 1$
(c) NTK initialization,  $\gamma = 1$