SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\frac {d f \left(\boldsymbol {\theta} ^ {t}\right)}{d t} = - \boldsymbol {K} ^ {t} \left(f \left(\boldsymbol {\theta} ^ {t}\right) - \boldsymbol {y}\right), \tag {5}
$$
where  $\pmb{K}^{t} = \left(\frac{\partial f(\pmb{\theta}^{t})}{\partial\pmb{\theta}}\right)\left(\frac{\partial f(\pmb{\theta}^{t})}{\partial\pmb{\theta}}\right)^{\top}\in \mathbb{R}^{N\times N}$ . Jacot et al. [2018], Arora et al. [2019b] have shown that for fully-connected neural networks, under the infinite-width setting and proper initialization,  $\pmb{K}^{t}$  will be stable during training and  $\pmb{K}^{0}$  will converge to a fixed matrix  $\pmb{K}^{\star}\in \mathbb{R}^{N\times N}$ , where  $(\pmb{K}^{\star})^{(ij)} = K^{\star}(\pmb{X}_i,\pmb{X}_j)$  is the NTK value for the inputs  $\pmb{X}_i$  and  $\pmb{X}_j$ . Then, we rewrite Eq. (5) as:
$$
\frac {d f (\boldsymbol {\theta} ^ {t})}{d t} = - \boldsymbol {K} ^ {\star} (f (\boldsymbol {\theta} ^ {t}) - \boldsymbol {y}).
$$
This implies the network output for any  $X \in \mathbb{R}^{d_s \times d}$  can be calculated by the kernel regression predictor with the associated NTK:
$$
f (\boldsymbol {X}) = \left(K ^ {\star} \left(\boldsymbol {X}, \boldsymbol {X} _ {1}\right), \dots , K ^ {\star} \left(\boldsymbol {X}, \boldsymbol {X} _ {N}\right)\right) \cdot \left(\boldsymbol {K} ^ {\star}\right) ^ {- 1} \boldsymbol {y},
$$
where  $K^{\star}(\pmb{X},\pmb{X}_n)$  is the kernel value between test data  $\pmb{X}$  and training data  $\pmb{X}_n$ .
B.2 Preliminary on Sub-Exponential random variables
Below, we overview the definition of a sub-exponential random variable and few related lemma based on Wainwright [2019]. A random variable  $X$  with mean  $\mu$  is called sub-exponential random variable if there exist non-negative parameters  $(\nu, \alpha)$  such that
$$
\mathbb {E} \left[ e ^ {\lambda (X - \mu)} \right] \leq e ^ {\frac {\nu^ {2} \lambda^ {2}}{2}} \quad \text {f o r a l l} | \lambda | <   \frac {1}{\alpha},
$$
and we denote by  $X\sim SE(\nu ,\alpha)$
Lemma 2. The multiplication of a scalar  $s \in \mathbb{R}^+$  and a sub-exponential random variable  $X \sim SE(\nu, \alpha)$  is still a sub-exponential random variable:  $sX \sim SE(s\nu, s\alpha)$ .
Lemma 3. Given a set of independent sub-exponential random variables  $X_{i}\sim (\nu_{i},\alpha_{i})$  for  $i\in$ $1,\ldots ,N$  , then  $\sum_{i = 1}^{N}X_{i}\sim (\sqrt{\sum_{i = 1}^{N}\nu_{i}^{2}},\max_{i}\alpha_{i})$
Lemma 4. Given a sub-exponential random variable  $X \sim SE(\nu, \alpha)$  with mean  $\mu$ , the following inequality holds with probability at least  $1 - \delta$ :
$$
| X - \mu | <   \max \left(\nu \sqrt {2 \log {\frac {2}{\delta}}}, 2 \alpha \log {\frac {2}{\delta}}\right).
$$
B.3 Related work on over-parameterization for convergence analysis
Recent empirical observation shows that neural networks can fit arbitrary labels with zero training loss when applying the first-order methods, e.g., gradient descent (GD) [Zhang et al., 2017]. Due to the highly non-convex and non-smooth instinct of the neural network, a large body of work have attempted to explain such a phenomenon. Early work studied the convergence of stochastic gradient descent (SGD) for training two-layer over-parameterized ReLU neural network with cross-entropy loss [Li and Liang, 2018]. Du et al. [2019b] show that only training the first layer of two-layer ReLU network with square loss by GD can lead to global convergence under the assumption that the Gram matrix is positive definite. Du et al. [2019a] extend the result to deep neural network with smooth activation function and shows that the convergence is guaranteed when the widths of all the hidden layers scale in  $\Omega(N^4)$ , where  $N$  is the number of data points. Meanwhile, Zou and Gu [2019] prove that the condition for the convergence of GD for deep ReLU network is  $\Omega(N^8)$ , which improves upon Allen-Zhu et al. [2019a] that show the result of  $\Omega(N^{24})$ . Allen-Zhu et al. [2019a] also provide several convergence analyses under the setting of SGD and various loss functions. The assumption regarding the positive definiteness of the Gram matrix made in Du et al. [2019b] has been rigorously proved in Nguyen et al. [2021]. This facilitates Nguyen [2021] to demonstrate that deep ReLU network under LeCUN initialization with width in the order  $\Omega(N^3)$  is enough for global convergence. Recent breakthrough [Bombari et al., 2022] improves previous results by showing that sub-linear layer widths suffice for deep neural network with smooth activation function.
Table 3: Over-parameterization conditions for the convergence analysis of neural network under gradient descent training with squared loss.  $L$  is the depth of the network.
Model	Depth	Initialization	Activation	Width
Allen-Zhu et al. [2019a]	FCNN/CNN	Deep	NTK	ReLU	Ω(N24L12)
Du et al. [2019a]	FCNN/CNN	Deep	NTK	Smooth	Ω(N42O(L))
Oymak and Soltanolkotabi [2020]	FCNN	Shallow	Standard Gaussian	ReLU	Ω(N2)
Zou and Gu [2019]	FCNN	Deep	He	ReLU	Ω(N8L12)
Du et al. [2019b]	FCNN	Shallow	NTK	ReLU	Ω(N6)
Nguyen [2021]	FCNN	Deep	LeCun	ReLU	Ω(N3)
Chen et al. [2021]	FCNN	Deep	NTK	ReLU	Ω(L22)
Song et al. [2021]	FCNN	Shallow	He/Lecun	Smooth	Ω(N3/2)
Bombari et al. [2022]	FCNN	Deep	He/LeCun	Smooth	Ω(√N)
Allen-Zhu et al. [2019b]	RNN	-	NTK	ReLU	Ω(Nc), c > 1
Hron et al. [2020]	Transformer	Deep	NTK	ReLU	-
Yang [2020]	Transformer	Deep	NTK	Softmax+ReLU	-
Our	Transformer	Shallow	Table 1	Softmax+ReLU	Ω(N)
C Proof for convergence analysis
This section is developed for the proof of the convergence result and we outline the flowchart below: Specifically, in Appendix C.1, we provide some auxiliary lemmas. Lemma 5 and Lemma 6 show that the norm of the parameters can be bounded with high probability at initialization. Lemmas 7, 9 and 10 present that the network output and the output of softmax between two adjacent time steps can be upper bounded. The Lipschitzness of network gradient and its norm is bounded in Lemmas 13 and 14. In Appendix C.2, we prove the convergence of the general cases, i.e., Proposition 1. In Appendix C.3 and C.4 we present the proof for  $d_m^{-1/2}$  and  $d_m^{-1}$  scaling.
C.1 Auxiliary lemmas
Lemma 5 (Corollary 5.35 of Vershynin [2012]). For a weight matrix  $\mathbf{W} \in \mathbb{R}^{d_1 \times d_2}$  where each element is sampled independently from  $\mathcal{N}(0,1)$ , for every  $\zeta \geq 0$ , with probability at least  $1 - 2\exp(-\zeta^2 / 2)$  one has:
$$
\sqrt {d _ {1}} - \sqrt {d _ {2}} - \zeta \leq \sigma_ {\min } (\boldsymbol {W}) \leq \sigma_ {\max } (\boldsymbol {W}) \leq \sqrt {d _ {1}} + \sqrt {d _ {2}} + \zeta ,
$$
where  $\sigma_{\mathrm{max}}(\mathbf{W})$  and  $\sigma_{\mathrm{min}}(\mathbf{W})$  represents the maximum and minimum singular value of  $\mathbf{W}$ , respectively.
Lemma 6 (Upper bound of spectral norms of initial weight). For a weight matrix  $\mathbf{W} \in \mathbb{R}^{d_m \times d}$  where  $d_m > d$ , each element is sampled independently from  $\mathcal{N}(0,1)$ , with probability at least  $1 - 2\exp(-d_m/2)$ , one has:
$$
\left\| \boldsymbol {W} \right\| _ {2} \leq 3 \sqrt {d _ {m}}.
$$
Proof of Lemma 6. Following Lemma 5, one has:
$$
\| \boldsymbol {W} \| _ {2} \leq \sqrt {d _ {m}} + \sqrt {d} + \zeta .
$$
Letting  $\zeta = \sqrt{d_m}$  and using the fact that  $d_{m} > d$  complete the proof.
Lemma 7. Recall from Table 2,  $\beta_{i} = \sigma_{s}\left(\tau_{0}\pmb{X}^{(i,:)}\pmb{W}_{Q}^{\top}\pmb{W}_{K}\pmb{X}^{\top}\right)^{\top}$  is the  $i$ -th row of the output of Softmax, if  $\max (\| \pmb{W}_Q^t\| _2,\| \pmb{W}_Q^{t'}\| _2)\leq \bar{\lambda}_Q,\max (\| \pmb{W}_K^t\| _2,\| \pmb{W}_K^{t'}\| _2)\leq \bar{\lambda}_K$  then its difference
between  $t^\prime$  step and  $t$  step has the following upper bound:
$$
\left\| \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i} ^ {t} \right\| _ {2} \leq 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} + \bar {\lambda} _ {K} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t} \right\| _ {2}\right).
$$
$$
\begin{array}{l} \left\| \beta_ {i} ^ {t ^ {\prime}} - \beta_ {i} ^ {t} \right\| _ {2} \\ \leq \left\| \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i} ^ {t} \right\| _ {1} \\ \leq 2 \left\| \tau_ {0} \boldsymbol {X} ^ {(i,:)} \boldsymbol {W} _ {Q} ^ {t ^ {\prime} \top} \boldsymbol {W} _ {K} ^ {t ^ {\prime}} \boldsymbol {X} ^ {\top} - \tau_ {0} \boldsymbol {X} ^ {(i,:)} \boldsymbol {W} _ {Q} ^ {t \top} \boldsymbol {W} _ {K} ^ {t} \boldsymbol {X} ^ {\top} \right\| _ {\infty} (\text {B y C o r o l l a r y A . 7 i n E d e l m a n e t a l . [ 2 0 2 2 ]}) \\ = 2 \max _ {j} | \tau_ {0} \boldsymbol {X} ^ {(i,:)} \boldsymbol {W} _ {Q} ^ {t ^ {\prime} \top} \boldsymbol {W} _ {K} ^ {t ^ {\prime}} \boldsymbol {X} ^ {(j,:) \top} - \tau_ {0} \boldsymbol {X} ^ {(i,:)} \boldsymbol {W} _ {Q} ^ {t \top} \boldsymbol {W} _ {K} ^ {t} \boldsymbol {X} ^ {(j,:) \top} | \\ \leq 2 \max  _ {j} \left(\tau_ {0} \left\| \boldsymbol {X} ^ {(i,:)} \right\| _ {2} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime} \top} \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t \top} \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} \left\| \boldsymbol {X} ^ {(j,:)} \right\| _ {2}\right) (\text {B y C a u c h y - S c h w a r z i n e q u a l i t y}) \\ \leq 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime} \top} \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t \top} \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} (\text {B y A s s u m p t i o n 1}) \\ \leq 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} \right\| _ {2} \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} + \left\| \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t} \right\| _ {2}\right) (\text {B y C a u c h y - S c h w a r z i n e q u a l i t y , T r i a n g l e i n e q u a l i t y}) \\ \leq 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} + \bar {\lambda} _ {K} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t} \right\| _ {2}\right). \\ \end{array}
$$
Lemma 8 (Upper bound the Euclidean norm of the output of softmax). Suppose  $\mathbf{v} = \operatorname{Softmax}(\mathbf{u}) \in \mathbb{R}^{d_s}$ , then one has:  $1 / \sqrt{d_s} \leq \| \mathbf{v} \|_2 \leq 1$ .
Proof. By the inequality of arithmetic and geometric mean:
$$
\frac {1}{\sqrt {d _ {s}}} = \frac {\sum_ {i = 1} ^ {d _ {s}} \boldsymbol {v} ^ {(i)}}{\sqrt {d _ {s}}} \leq \| \boldsymbol {v} \| _ {2} = \sqrt {\sum_ {i = 1} ^ {d _ {s}} \left((\boldsymbol {v} ^ {(i)}) ^ {2}\right)} \leq \sqrt {\left(\sum_ {i = 1} ^ {d _ {s}} (\boldsymbol {v} ^ {(i)})\right) ^ {2}} = 1.
$$
Lemma 9. If at  $t'$  and  $t$  step,  $\max(\| W_{Q}^{t}\|_{2},\left\| W_{Q}^{t^{\prime}}\right\|_{2}) \leq \bar{\lambda}_{Q}, \max(\| W_{K}^{t}\|_{2},\left\| W_{K}^{t^{\prime}}\right\|_{2}) \leq \bar{\lambda}_{K}, \max(\| W_{V}^{t}\|_{2},\left\| W_{V}^{t^{\prime}}\right\|_{2}) \leq \bar{\lambda}_{V}$ , then the difference between the output of the last hidden layer at  $t'$  step and  $t$  can be upper bounded by:
$$
\left\| \pmb {F} _ {p r e} ^ {t ^ {\prime}} - \pmb {F} _ {p r e} ^ {t} \right\| _ {\mathrm {F}} \leq \rho \left(\left\| \pmb {W} _ {V} ^ {t ^ {\prime}} - \pmb {W} _ {V} ^ {t} \right\| _ {2} + \bar {\lambda} _ {V} 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} \left\| \pmb {W} _ {K} ^ {t ^ {\prime}} - \pmb {W} _ {K} ^ {t} \right\| _ {2} + \bar {\lambda} _ {K} \left\| \pmb {W} _ {Q} ^ {t ^ {\prime}} - \pmb {W} _ {Q} ^ {t} \right\| _ {2}\right)\right).
$$
Proof. Note that  $\pmb{F}_{\mathrm{pre}} \in \mathbb{R}^{N \times d_m}$  is the output of the last hidden layer given  $\{X_n\}_{n=1}^N$  sample as defined in Table 2. We firstly analyze each sample, i.e.,  $\pmb{f}_{\mathrm{pre}} \in \mathbb{R}^{d_m}$  and we drop the index  $n$  for simplification.
$$
\begin{array}{l} \left\| \boldsymbol {f} _ {\text {p r e}} ^ {t ^ {\prime}} - \boldsymbol {f} _ {\text {p r e}} ^ {t} \right\| _ {2} = \tau_ {1} \left\| \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}}\right) - \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} ^ {t}\right) \right\| _ {2} \\ \leq \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} ^ {t} \right\| _ {2} (\text {B y L i p s c h i t z c o n t i n u i t y o f} \sigma_ {r}) \\ \leq \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(\left\| \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t}\right) \boldsymbol {X} ^ {\top} \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} \right\| _ {2} + \left\| \boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} \left(\boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i} ^ {t}\right) \right\| _ {2}\right) (\text {B y T r i a n g l e i n e q u a l i t y}) \\ \leq \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \left(\left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} \| \boldsymbol {X} \| _ {2} \left\| \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} \right\| _ {2} + \left\| \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} \| \boldsymbol {X} \| _ {2} \left\| \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i} ^ {t} \right\| _ {2}\right) (\text {B y 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_ {1} C _ {x} \sqrt {d _ {s}} \sum_ {i = 1} ^ {d _ {s}} \left(\left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} + \bar {\lambda} _ {V} \left\| \boldsymbol {\beta} _ {i} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i} ^ {t} \right\| _ {2}\right) (\text {B y}) \\ \leq \tau_ {1} C _ {x} d _ {s} ^ {3 / 2} \left(\left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} + \bar {\lambda} _ {V} 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} + \bar {\lambda} _ {K} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t} \right\| _ {2}\right)\right) \tag {6} \\ \end{array}
$$
Next, we bound the difference given  $N$  data sample:
$$
\begin{array}{l} \left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} - \boldsymbol {F} _ {\text {p r e}} ^ {t} \right\| _ {\mathrm {F}} \leq \sqrt {N} \tau_ {1} C _ {x} d _ {s} ^ {3 / 2} \left(\left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} + \bar {\lambda} _ {V} 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} + \bar {\lambda} _ {K} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t} \right\| _ {2}\right)\right) \\ = \rho \left(\left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} + \bar {\lambda} _ {V} 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} - \boldsymbol {W} _ {K} ^ {t} \right\| _ {2} + \bar {\lambda} _ {K} \left\| \boldsymbol {W} _ {Q} ^ {t ^ {\prime}} - \boldsymbol {W} _ {Q} ^ {t} \right\| _ {2}\right)\right), \\ \end{array}
$$
where the last equality is by the definition of  $\rho$  in Proposition. 1.
Lemma 10. If at  $t'$  and  $t$  step,  $\max(\| \boldsymbol{W}_Q^t \|_2, \| \boldsymbol{W}_Q^{t'} \|_2) \leq \bar{\lambda}_Q, \max(\| \boldsymbol{W}_K^t \|_2, \| \boldsymbol{W}_K^{t'} \|_2) \leq \bar{\lambda}_K, \max(\| \boldsymbol{W}_V^t \|_2, \| \boldsymbol{W}_V^{t'} \|_2) \leq \bar{\lambda}_V, \max(\| \boldsymbol{w}_O^t \|_2, \| \boldsymbol{w}_O^{t'} \|_2) \leq \bar{\lambda}_O$ , then the difference between the network output at  $t'$  step and  $t$  step can be upper bounded by:
$$
\left\| \boldsymbol {f} ^ {t ^ {\prime}} - \boldsymbol {f} ^ {t} \right\| _ {2} \leq \rho \bar {\lambda} _ {V} \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} - \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} + \bar {\lambda} _ {O} \left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} - \boldsymbol {F} _ {\text {p r e}} ^ {t} \right\| _ {\mathrm {F}}.
$$
$$
\left\| \boldsymbol {f} ^ {t ^ {\prime}} - \boldsymbol {f} ^ {t} \right\| _ {2} = \left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} \boldsymbol {w} _ {O} ^ {t ^ {\prime}} - \boldsymbol {F} _ {\text {p r e}} ^ {t} \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} \leq \left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} \right\| _ {2} \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} - \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} + \left\| \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} \left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} - \boldsymbol {F} _ {\text {p r e}} ^ {t} \right\| _ {2},
$$
where we use triangle inequality. Then, the first term of the RHS can be bounded by:
$$
\left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} \right\| _ {2} \leq \left\| \boldsymbol {F} _ {\text {p r e}} ^ {t ^ {\prime}} \right\| _ {\mathrm {F}} = \left\| \left[ \begin{array}{c} \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {1} ^ {\top} \boldsymbol {\beta} _ {i, 1} ^ {t ^ {\prime}}\right) \\ \vdots \\ \tau_ {1} \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {N} ^ {\top} \boldsymbol {\beta} _ {i, N} ^ {t ^ {\prime}}\right) \end{array} \right] \right\| _ {\mathrm {F}} \leq \tau_ {1} \sqrt {N} d _ {s} ^ {3 / 2} C _ {x} \left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} \right\| _ {2} = \rho \left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} \right\| _ {2},
$$
Lemma 11 (Jacobian of Softmax). Suppose  $\pmb{v} = \mathrm{Softmax}(\pmb{u}) \in \mathbb{R}^{d_s}$ , then  $\frac{\partial\pmb{v}}{\partial\pmb{u}} = \mathrm{diag}(\pmb{v}) - \pmb{v}\pmb{v}^{\top}$ .
Proof. We can reformulate  $\pmb{v}$  as:  $\pmb{v} = \begin{bmatrix} \frac{\exp(\pmb{u}^{(1)})}{\sum_{i=1}^{ds} \exp(\pmb{u}^{(i)})} \\ \vdots \\ \frac{\exp(\pmb{u}^{(ds)})}{\sum_{i=1}^{ds} \exp(\pmb{u}^{(i)})} \end{bmatrix}$ .
$$
\begin{array}{l} \frac {\partial \boldsymbol {v} ^ {(j)}}{\partial \boldsymbol {u} ^ {(k)}} = \frac {\partial \frac {\exp \left(\boldsymbol {u} ^ {(j)}\right)}{\sum_ {i = 1} ^ {d _ {s}} \exp \left(\boldsymbol {u} ^ {(i)}\right)}}{\partial \boldsymbol {u} ^ {(k)}} = \left\{ \begin{array}{l l} \frac {- \exp \left(\boldsymbol {u} ^ {(j)}\right) - \exp \left(\boldsymbol {u} ^ {(k)}\right)}{\left(\sum_ {i = 1} ^ {d _ {s}} \exp \left(\boldsymbol {u} ^ {(i)}\right)\right) ^ {2}} & \text {i f} j \neq k \\ \frac {\exp \left(\boldsymbol {u} ^ {(k)}\right) \sum_ {i = 1} ^ {d _ {s}} \exp \left(\boldsymbol {u} ^ {(i)}\right) - \left(\exp \left(\boldsymbol {u} ^ {(k)}\right)\right) ^ {2}}{\left(\sum_ {i = 1} ^ {d _ {s}} \exp \left(\boldsymbol {u} ^ {(i)}\right)\right) ^ {2}} & \text {i f} j = k \end{array} \right. \\ = \left\{ \begin{array}{l l} - \boldsymbol {v} ^ {(j)} \boldsymbol {v} ^ {(k)} & \text {i f} j \neq k \\ \boldsymbol {v} ^ {(k)} - \boldsymbol {v} ^ {(j)} \boldsymbol {v} ^ {(k)} & \text {i f} j = k \end{array} \right. \\ \end{array}
$$
$$
\frac {\partial \boldsymbol {v}}{\partial \mathbf {u}} = \operatorname {d i a g} (\boldsymbol {v}) - \boldsymbol {v} \boldsymbol {v} ^ {\top}.
$$
Lemma 12 (Upper bound for the loss gradient norm). If  $\left\| \boldsymbol{W}_Q^t \right\|_2 \leq \bar{\lambda}_Q, \left\| \boldsymbol{W}_K^t \right\|_2 \leq \bar{\lambda}_K, \left\| \boldsymbol{W}_V^t \right\|_2 \leq \bar{\lambda}_V, \left\| \boldsymbol{w}_O^t \right\|_2 \leq \bar{\lambda}_O$ , then the gradient norm with respect to  $\boldsymbol{W}_Q, \boldsymbol{W}_K, \boldsymbol{W}_V, \boldsymbol{w}_O$  can be upper bounded by:
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {W} _ {Q}} \ell (\boldsymbol {\theta} ^ {\top}) \right\| _ {\mathrm {F}} \leq 2 \rho \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \| \boldsymbol {f} ^ {t} - \boldsymbol {y} \| _ {2}, \quad \left\| \nabla_ {\boldsymbol {W} _ {K}} \ell (\boldsymbol {\theta} ^ {\top}) \right\| _ {\mathrm {F}} \leq 2 \rho \tau_ {0} \bar {\lambda} _ {Q} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \| \boldsymbol {f} ^ {t} - \boldsymbol {y} \| _ {2}, \\ \left\| \nabla_ {\boldsymbol {W} _ {V}} \ell (\boldsymbol {\theta} ^ {\top}) \right\| _ {\mathrm {F}} \leq \rho \bar {\lambda} _ {O} \left\| \boldsymbol {f} ^ {t} - \boldsymbol {y} \right\| _ {2}, \quad \left\| \nabla_ {\boldsymbol {w} _ {O}} \ell (\boldsymbol {\theta} ^ {\top}) \right\| _ {2} \leq \rho \bar {\lambda} _ {V} \left\| \boldsymbol {f} ^ {t} - \boldsymbol {y} \right\| _ {2}. \\ \end{array}
$$
Proof. To simplify the notation, in the proof below, we hide the index  $t$ . Firstly, consider the gradient w.r.t  $W_V$ ,
$$
\begin{array}{l} \| \nabla_ {\boldsymbol {W} _ {V}} \ell (\boldsymbol {\theta}) \| _ {\mathrm {F}} = \left\| - \sum_ {n = 1} ^ {N} (f (\boldsymbol {X} _ {n}) - y _ {n}) \frac {\partial f (\boldsymbol {X} _ {n})}{\partial \boldsymbol {W} _ {V}} \right\| _ {\mathrm {F}} \\ = \tau_ {1} \left\| \sum_ {n = 1} ^ {N} (f (\boldsymbol {X} _ {n}) - y _ {n}) \sum_ {i = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n}\right)\right) \boldsymbol {\beta} _ {i, n} ^ {\top} \boldsymbol {X} _ {n} \right\| _ {\mathrm {F}} \leq \tau_ {1} \sum_ {n = 1} ^ {N} | (f (\boldsymbol {X} _ {n}) - y _ {n}) | d _ {s} ^ {3 / 2} \bar {\lambda} _ {O} C _ {x} \\ \leq \tau_ {1} \sqrt {N \sum_ {n = 1} ^ {N} \left| \left(f \left(\boldsymbol {X} _ {n}\right) - y _ {n}\right) \right| ^ {2}} d _ {s} ^ {3 / 2} \bar {\lambda} _ {O} C _ {x} = \tau_ {1} d _ {s} ^ {3 / 2} \bar {\lambda} _ {O} C _ {x} \sqrt {N} \| \boldsymbol {f} - \boldsymbol {y} \| _ {2} = \rho \bar {\lambda} _ {O} \| \boldsymbol {f} - \boldsymbol {y} \| _ {2}, \tag {7} \\ \end{array}
$$
where the last equality is by the definition of  $\rho$  in Proposition. 1. Next, consider the gradient w.r.t  $w_{O}$
$$
\begin{array}{l} \| \nabla_ {\pmb {w} _ {O}} \ell (\pmb {\theta}) \| _ {2} = \left\| - \sum_ {n = 1} ^ {N} (f (\pmb {X} _ {n}) - y _ {n}) \frac {\partial f (\pmb {X} _ {n})}{\partial \pmb {w} _ {O}} \right\| _ {2} = \tau_ {1} \left\| \sum_ {n = 1} ^ {N} (f (\pmb {X} _ {n}) - y _ {n}) \sum_ {i = 1} ^ {d _ {s}} \sigma_ {r} \left(\pmb {W} _ {V} \pmb {X} _ {n} ^ {\top} \pmb {\beta} _ {i, n}\right) \right\| _ {2} \\ \leq \tau_ {1} \sum_ {n = 1} ^ {N} | (f (\boldsymbol {X} _ {n}) - y _ {n}) | d _ {s} ^ {3 / 2} \bar {\lambda} _ {V} C _ {x} \leq \tau_ {1} \sqrt {N \sum_ {n = 1} ^ {N} | (f (\boldsymbol {X} _ {n}) - y _ {n}) | ^ {2}} d _ {s} ^ {3 / 2} \bar {\lambda} _ {V} C _ {x} = \rho \bar {\lambda} _ {V} \| \boldsymbol {f} - \boldsymbol {y} \| _ {2}, \tag {8} \\ \end{array}
$$
Next, consider the gradient w.r.t  $W_{Q}$