SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {W} _ {Q}} \ell (\boldsymbol {\theta}) \right\| _ {\mathrm {F}} = \left\| - \sum_ {n = 1} ^ {N} (f (\boldsymbol {X} _ {n}) - y _ {n}) \frac {\partial f (\boldsymbol {X} _ {n})}{\partial \boldsymbol {W} _ {Q}} \right\| _ {\mathrm {F}} \\ = \tau_ {0} \tau_ {1} \left\| \sum_ {n = 1} ^ {N} (f (\boldsymbol {X} _ {n}) - y _ {n}) \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {W} _ {k} \boldsymbol {X} _ {n} ^ {\top} \left(\mathrm {d i a g} (\boldsymbol {\beta} _ {i, n}) - \boldsymbol {\beta} _ {i, n} \boldsymbol {\beta} _ {i, n} ^ {\top}\right) \boldsymbol {X} _ {n} \boldsymbol {W} _ {V} ^ {\top} \left(\boldsymbol {w} _ {O} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n}\right)\right) \boldsymbol {X} _ {n} ^ {(i,:)} \right\| _ {\mathrm {F}} \\ \leq 2 \tau_ {0} \tau_ {1} \sum_ {n = 1} ^ {N} | (f (\boldsymbol {X} _ {n}) - y _ {n}) | d _ {s} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} \left(C _ {x} \sqrt {d _ {s}}\right) ^ {3} \\ \leq 2 \tau_ {0} \tau_ {1} \sqrt {N \sum_ {n = 1} ^ {N} | (f (\pmb {X} _ {n}) - y _ {n}) | ^ {2}} d _ {s} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} (C _ {x} \sqrt {d _ {s}}) ^ {3} \\ = 2 \rho \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \| \boldsymbol {f} - \boldsymbol {y} \| _ {2}. \tag {9} \\ \end{array}
$$
Similarly, for the gradient w.r.t  $\mathbf{W}_K$ , we have:
$$
\left\| \nabla_ {\boldsymbol {W} _ {K}} \ell (\boldsymbol {\theta}) \right\| _ {\mathrm {F}} \leq 2 \rho \tau_ {0} \bar {\lambda} _ {Q} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \| \boldsymbol {f} - \boldsymbol {y} \| _ {2}. \tag {10}
$$
Lemma 13 (Upper bound for the network function gradient norm). If  $\| \pmb{W}_Q^t\| _2\leq \bar{\lambda}_Q,\| \pmb{W}_K^t\| _2\leq$ $\bar{\lambda}_K,\| \pmb{W}_V^t\| _2\leq \bar{\lambda}_V,\| \pmb{w}_O^t\| _2\leq \bar{\lambda}_O$  , then one has:
$$
\left\| \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} ^ {t} \right\| _ {2} \leq c _ {2},
$$
$$
c _ {2} \triangleq \rho \sqrt {\bar {\lambda} _ {O} ^ {2} + \bar {\lambda} _ {V} ^ {2} + \left(2 \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2}\right) ^ {2} + \left(2 \tau_ {0} \bar {\lambda} _ {Q} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2}\right) ^ {2}}. \tag {11}
$$
Proof. To simplify the notation, in the proof below, we hide the index  $t$ . Firstly, note that:
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} \right\| _ {2} \leq \left\| \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} \right\| _ {\mathrm {F}} \\ = \sqrt {\sum_ {n = 1} ^ {N} \left(\left\| \nabla_ {\boldsymbol {w} _ {O}} f (\boldsymbol {X} _ {n} ; \boldsymbol {\theta}) \right\| _ {2} ^ {2} + \left\| \nabla_ {\boldsymbol {w} _ {Q}} f (\boldsymbol {X} _ {n} ; \boldsymbol {\theta}) \right\| _ {\mathrm {F}} ^ {2} + \left\| \nabla_ {\boldsymbol {w} _ {K}} f (\boldsymbol {X} _ {n} ; \boldsymbol {\theta}) \right\| _ {\mathrm {F}} ^ {2} + \left\| \nabla_ {\boldsymbol {W} _ {V}} f (\boldsymbol {X} _ {n} ; \boldsymbol {\theta}) \right\| _ {\mathrm {F}} ^ {2}\right)}. \tag {12} \\ \end{array}
$$
Then following the step as in Eqs. (7) to (10), each term can be bounded as follows:
$$
\begin{array}{l} \sum_ {n = 1} ^ {N} \left(\| \nabla_ {\boldsymbol {W} _ {V}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta}) \| _ {\mathrm {F}} ^ {2}\right) \leq (\rho \bar {\lambda} _ {O}) ^ {2}, \quad \sum_ {n = 1} ^ {N} \left(\| \nabla_ {\boldsymbol {w} _ {O}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta}) \| _ {2}\right) \leq (\rho \bar {\lambda} _ {V}) ^ {2}, \\ \sum_ {n = 1} ^ {N} \left(\left\| \nabla_ {\boldsymbol {w} _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta}) \right\| _ {\mathrm {F}}\right) \leq (2 \rho \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2}) ^ {2}, \quad \sum_ {n = 1} ^ {N} \left(\left\| \nabla_ {\boldsymbol {w} _ {K}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta}) \right\| _ {\mathrm {F}}\right) \leq (2 \rho \tau_ {0} \bar {\lambda} _ {Q} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2}) ^ {2}. \\ \end{array}
$$
Plugging these bounds back Eq. (12) finishes the proof.
Lemma 14 (Upper bound for the Lipschitzness of the network gradient). Suppose  $\max (\| \pmb{W}_Q^t\| _2,\left\| \pmb{W}_Q^{t'}\right\| _2)\leq \bar{\lambda}_Q,\max (\| \pmb{W}_K^t\| _2,\left\| \pmb{W}_K^{t'}\right\| _2)\leq \bar{\lambda}_K,\max (\| \pmb{W}_V^t\| _2,\left\| \pmb{W}_V^{t'}\right\| _2)\leq$ $\bar{\lambda}_V,\max (\| \pmb {w}_O^t\|_2,\left\| \pmb {w}_O^{t'}\right\|_2)\leq \bar{\lambda}_O$  , and define  $z\triangleq 2\tau_0C_x^2 d_s(\bar{\lambda}_Q + \bar{\lambda}_K)$  , then one has
$$
\left\| \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} ^ {t ^ {\prime}} - \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} ^ {t} \right\| _ {2} \leq c _ {3} \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \right\| _ {2}, \tag {13}
$$
$$
\begin{array}{l} (c _ {3}) ^ {2} \triangleq N (\tau_ {1} C _ {x} d _ {s} ^ {3 / 2} (1 + \bar {\lambda} _ {V} z)) ^ {2} \\ + N \left\{\tau_ {1} C _ {x} d _ {s} ^ {3 / 2} \left[ \bar {\lambda} _ {O} z + \bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} (1 + \bar {\lambda} _ {V} z) + 1 \right] \right\} ^ {2} \\ + N \left\{\tau_ {0} \tau_ {1} C _ {x} d _ {s} \left\{2 \bar {\lambda} _ {K} d _ {s} C _ {x} ^ {2} \left[ \bar {\lambda} _ {V} \left(\bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} (1 + \bar {\lambda} _ {V} z) + 1\right) + \bar {\lambda} _ {O} \right] + \bar {\lambda} _ {V} \bar {\lambda} _ {O} \left[ C _ {x} ^ {2} d _ {s} + 3 \bar {\lambda} _ {K} z \right] \right\} \right\} ^ {2} \\ + N \left\{\tau_ {0} \tau_ {1} C _ {x} d _ {s} \left\{2 \bar {\lambda} _ {Q} d _ {s} C _ {x} ^ {2} \left[ \bar {\lambda} _ {V} \left(\bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} (1 + \bar {\lambda} _ {V} z) + 1\right) + \bar {\lambda} _ {O} \right] + \bar {\lambda} _ {V} \bar {\lambda} _ {O} [ C _ {x} ^ {2} d _ {s} + 3 \bar {\lambda} _ {K} z ] \right\} \right\} ^ {2}. \tag {14} \\ \end{array}
$$
Proof. Firstly, note that:
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} ^ {t ^ {\prime}} - \nabla_ {\boldsymbol {\theta}} \boldsymbol {f} ^ {t} \right\| _ {2} ^ {2} \\ \leq \sum_ {n = 1} ^ {N} \left(\left\| \nabla_ {\boldsymbol {w} _ {O}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {w} _ {O}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {2} ^ {2} + \left\| \nabla_ {\boldsymbol {W} _ {V}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {W} _ {V}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} ^ {2}\right) \\ \left. + \left\| \nabla_ {\boldsymbol {W} _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {W} _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} ^ {2} + \left\| \nabla_ {\boldsymbol {W} _ {K}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {W} _ {K}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} ^ {2}\right). \tag {15} \\ \end{array}
$$
Then, we will prove each term separately. Regarding the first term in Eq. (15), we have
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {w} _ {O}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {w} _ {O}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {2} \\ = \left\| \boldsymbol {f} _ {\text {p r e}} ^ {t ^ {\prime}} - \boldsymbol {f} _ {\text {p r e}} ^ {t} \right\| _ {2} \\ \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) \\ \leq \tau_ {1} C _ {x} d _ {s} ^ {3 / 2} \left(1 + \bar {\lambda} _ {V} 2 \tau_ {0} C _ {x} ^ {2} d _ {s} \left(\bar {\lambda} _ {Q} + \bar {\lambda} _ {K}\right)\right) \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \right\| _ {2} \\ = \tau_ {1} C _ {x} d _ {s} ^ {3 / 2} \left(1 + \bar {\lambda} _ {V} z\right) \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \right\| _ {2}. \tag {16} \\ \end{array}
$$
where the first inequality is by Eq. (6), and in the last equality is by the definition of  $z$  in the lemma.
Regarding the second term in Eq. (15), we have:
$$
\begin{array}{l} \left\| \nabla_ {W _ {V}} f \left(\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}\right) - \nabla_ {W _ {V}} f \left(\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}\right) \right\| _ {\mathrm {F}} \\ = \tau_ {1} \left\| \sum_ {i = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right)\right) \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} ^ {\top} \boldsymbol {X} _ {n} - \sum_ {i = 1} ^ {d _ {s}} \left(\boldsymbol {w} _ {O} ^ {t} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right)\right) \boldsymbol {\beta} _ {i, n} ^ {t} ^ {\top} \boldsymbol {X} _ {n} \right\| _ {\mathrm {F}} \\ \leq \tau_ {1} C _ {x} \sqrt {d _ {s}} \sum_ {i = 1} ^ {d _ {s}} \left\| \left(\boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right)\right) \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} ^ {\top} - \left(\boldsymbol {w} _ {O} ^ {t} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right)\right) \boldsymbol {\beta} _ {i, n} ^ {t} ^ {\top} \right\| _ {\mathrm {F}} \\ \leq \tau_ {1} C _ {x} \sqrt {d _ {s}} \sum_ {i = 1} ^ {d _ {s}} \left[ \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right) \right\| _ {2} \left\| \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i, n} ^ {t} \right\| _ {2} + \left\| \boldsymbol {\beta} _ {i, n} ^ {t} \right\| _ {2} \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right) - \tilde {\boldsymbol {w}} _ {O} ^ {t} \circ \dot {\sigma_ {r}} \left(\tilde {\boldsymbol {W}} _ {V} ^ {t} \boldsymbol {X} _ {m} ^ {\top} \tilde {\boldsymbol {\beta}} _ {i, n} ^ {t}\right) \right\| _ {2} \right] \\ \leq \tau_ {1} C _ {x} \sqrt {d _ {s}} \sum_ {i = 1} ^ {d _ {s}} \left[ \bar {\lambda} _ {O} \left\| \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i, n} ^ {t} \right\| _ {2} + \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right) - \boldsymbol {w} _ {O} ^ {t} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right) \right\| _ {2} \right]. \tag {17} \\ \end{array}
$$
The term  $\left\| \beta_{i,n}^{t^{\prime}} - \beta_{i,n}^{t}\right\|_{2}$  can be bounded by Lemma 7 as follows:
$$
\left\| \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {\beta} _ {i, n} ^ {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) \leq z \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \right\| _ {2}. \tag {18}
$$
The term  $\left\| \pmb{w}_O^{t'} \circ \dot{\sigma}_r(\pmb{W}_V^{t'} \pmb{X}_n^\top \pmb{\beta}_{i,n}^{t'}) - \pmb{w}_O^t \circ \dot{\sigma}_r(\pmb{W}_V^t \pmb{X}_n^\top \pmb{\beta}_{i,n}^t)\right\|_2$  can be bounded by the triangle inequality, Cauchy-Schwarz inequality, and the same method in Eq. (6) as follows:
$$
\begin{array}{l} \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right) - \boldsymbol {w} _ {O} ^ {t} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right) \right\| _ {2} \\ \leq \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} \right\| _ {2} \left\| \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right) - \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right) \right\| _ {2} + \left\| \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right) \circ \left(\boldsymbol {w} _ {O} ^ {t ^ {\prime}} - \boldsymbol {w} _ {O} ^ {t}\right) \right\| _ {2} \\ \leq \bar {\lambda} _ {O} \left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t} \right\| _ {2} + \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} - \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} \\ \leq \bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} \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) + \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} - \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} \\ \leq \left[ \bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} \left(1 + \bar {\lambda} _ {V} z\right) + 1 \right] \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {\top} \right\| _ {2}. \tag {19} \\ \end{array}
$$
Plugging Eq. (18) and Eq. (19) back Eq. (17), we obtain:
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {W} _ {V}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {W} _ {V}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} \\ \leq \tau_ {1} C _ {x} d _ {s} ^ {3 / 2} \left[ \bar {\lambda} _ {O} z + \bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} \left(1 + \bar {\lambda} _ {V} z\right) + 1 \right] \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {\top} \right\| _ {2}. \tag {20} \\ \end{array}
$$
Regarding the third term in Eq. (15), let us denote by:
$$
\boldsymbol {U} _ {i, n} ^ {t ^ {\prime}} = \boldsymbol {W} _ {K} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \left(\mathrm {d i a g} (\boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}) - \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}} ^ {\top}\right) \boldsymbol {X} _ {n}, \quad \boldsymbol {h} _ {i, n} ^ {t ^ {\prime}} = \boldsymbol {W} _ {V} ^ {t ^ {\prime}} ^ {\top} \left(\boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma_ {r}} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right)\right),
$$
$$
\boldsymbol {U} _ {i, n} ^ {t} = \boldsymbol {W} _ {K} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \left(\operatorname {d i a g} \left(\boldsymbol {\beta} _ {i, n} ^ {t}\right) - \boldsymbol {\beta} _ {i, n} ^ {t} \boldsymbol {\beta} _ {i, n} ^ {t} ^ {\top}\right) \boldsymbol {X} _ {n}, \quad \boldsymbol {h} _ {i, n} ^ {t} = \boldsymbol {W} _ {V} ^ {t} ^ {\top} \left(\boldsymbol {w} _ {O} ^ {t} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right)\right).
$$
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {W} _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {W} _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} = \tau_ {0} \tau_ {1} \left\| \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {U} _ {i, n} ^ {t ^ {\prime}} \boldsymbol {h} _ {i, n} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {(i,:))} - \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {U} _ {i, n} ^ {t} \boldsymbol {h} _ {i, n} ^ {t} \boldsymbol {X} _ {n} ^ {(i,:)} \right\| _ {\mathrm {F}} \\ \leq \tau_ {0} \tau_ {1} C _ {x} \left\| \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {U} _ {i, n} ^ {t ^ {\prime}} \boldsymbol {h} _ {i, n} ^ {t ^ {\prime}} - \sum_ {i = 1} ^ {d _ {s}} \boldsymbol {U} _ {i, n} ^ {t} \boldsymbol {h} _ {i, n} ^ {t} \right\| _ {2} \leq \tau_ {0} \tau_ {1} C _ {x} \sum_ {i = 1} ^ {d _ {s}} \left(\left\| \boldsymbol {U} _ {i, n} ^ {t ^ {\prime}} \right\| _ {2} \left\| \boldsymbol {h} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {h} _ {i, n} ^ {t} \right\| _ {2} + \left\| \boldsymbol {h} _ {i, n} ^ {t} \right\| _ {2} \left\| \boldsymbol {U} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {U} _ {i, n} ^ {t} \right\| _ {2}\right). \tag {21} \\ \end{array}
$$
We bound each term separately:
$$
\left\| \boldsymbol {U} _ {i, n} ^ {t ^ {\prime}} \right\| _ {2} \leq \left\| \boldsymbol {W} _ {K} ^ {t ^ {\prime}} \right\| _ {2} \| \boldsymbol {X} _ {n} \| _ {2} ^ {2} \left\| \operatorname {d i a g} \left(\beta_ {i, n} ^ {t ^ {\prime}}\right) - \beta_ {i, n} ^ {t ^ {\prime}} \beta_ {i, n} ^ {t ^ {\prime} \top} \right\| _ {2} \leq 2 \bar {\lambda} _ {K} d _ {s} C _ {x} ^ {2}.
$$
$$
\left\| \boldsymbol {h} _ {i, n} ^ {t} \right\| _ {2} \leq \left\| \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} \left\| \boldsymbol {w} _ {O} ^ {t} \right\| _ {2} \leq \bar {\lambda} _ {V} \bar {\lambda} _ {O}.
$$
$$
\begin{array}{l} \left\| \boldsymbol {h} _ {i, n} ^ {t ^ {\prime}} - \boldsymbol {h} _ {i, n} ^ {t} \right\| _ {2} \leq \left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} \right\| _ {2} \left\| \boldsymbol {w} _ {O} ^ {t ^ {\prime}} \circ \dot {\sigma} _ {r} \left(\boldsymbol {W} _ {V} ^ {t ^ {\prime}} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t ^ {\prime}}\right) - \boldsymbol {w} _ {O} ^ {t} \circ \dot {\sigma} _ {r} \left(\tilde {\boldsymbol {W}} _ {V} ^ {t} \boldsymbol {X} _ {n} ^ {\top} \boldsymbol {\beta} _ {i, n} ^ {t}\right) \right\| _ {2} + \left\| \boldsymbol {W} _ {V} ^ {t ^ {\prime}} - \boldsymbol {W} _ {V} ^ {t} \right\| _ {2} \| \boldsymbol {w} _ {O} \| _ {2} \\ \leq \left[ \bar {\lambda} _ {V} \left(\bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} (1 + \bar {\lambda} _ {V} z) + 1\right) + \bar {\lambda} _ {O} \right] \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \right\| _ {2}. \\ \end{array}
$$
where the second inequality uses the result in Eq. (19).
where the last second inequality is by Lemma 7. Plugging back Eq. (21), we obtain:
$$
\begin{array}{l} \left\| \nabla_ {W _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {W _ {Q}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} \\ \leq \tau_ {0} \tau_ {1} C _ {x} d _ {s} \left\{2 \bar {\lambda} _ {K} d _ {s} C _ {x} ^ {2} \left[ \bar {\lambda} _ {V} \left(\bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} (1 + \bar {\lambda} _ {V} z) + 1\right) + \bar {\lambda} _ {O} \right] + \bar {\lambda} _ {V} \bar {\lambda} _ {O} \left[ C _ {x} ^ {2} d _ {s} + 3 \bar {\lambda} _ {K} z \right] \right\} \| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \| _ {2}. \tag {23} \\ \end{array}
$$
Similarly, the fourth term in Eq. (15) can be bounded by:
$$
\begin{array}{l} \left\| \nabla_ {\boldsymbol {W} _ {K}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t ^ {\prime}}) - \nabla_ {\boldsymbol {W} _ {K}} f (\boldsymbol {X} _ {n}; \boldsymbol {\theta} ^ {t}) \right\| _ {\mathrm {F}} \\ \leq \tau_ {0} \tau_ {1} C _ {x} d _ {s} \left\{2 \bar {\lambda} _ {Q} d _ {s} C _ {x} ^ {2} \left[ \bar {\lambda} _ {V} \left(\bar {\lambda} _ {O} C _ {x} \sqrt {d _ {s}} (1 + \bar {\lambda} _ {V} z) + 1\right) + \bar {\lambda} _ {O} \right] + \bar {\lambda} _ {V} \bar {\lambda} _ {O} \left[ C _ {x} ^ {2} d _ {s} + 3 \bar {\lambda} _ {K} z \right] \right\} \left\| \boldsymbol {\theta} ^ {t ^ {\prime}} - \boldsymbol {\theta} ^ {t} \right\| _ {2}. \tag {24} \\ \end{array}
$$
Plugging the upper bound for these four terms back Eq. (15) finishes the proof.
C.2 Proof of Proposition. 1
Proof. We can reformulate Eq. (1.2) as:
$$
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}\right) = \boldsymbol {w} _ {O} ^ {\top} \boldsymbol {f} _ {\text {p r e}},
$$
with  $\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}}$ . We show by induction for every  $t\geq 0$
$$
\left\{ \begin{array}{l} \left\| \boldsymbol {W} _ {Q} ^ {s} \right\| _ {2} \leq \bar {\lambda} _ {Q}, \left\| \boldsymbol {W} _ {K} ^ {s} \right\| _ {2} \leq \bar {\lambda} _ {K}, \quad s \in [ 0, t ], \\ \left\| \boldsymbol {W} _ {V} ^ {s} \right\| _ {2} \leq \bar {\lambda} _ {V}, \left\| \boldsymbol {w} _ {O} ^ {s} \right\| _ {2} \leq \bar {\lambda} _ {O}, \quad s \in [ 0, t ], \\ \sigma_ {\min } \left(\boldsymbol {F} _ {\text {p r e}} ^ {s}\right) \geq \frac {1}{2} \alpha , \quad s \in [ 0, t ], \\ \ell (\boldsymbol {\theta} ^ {s}) \leq \left(1 - \gamma \frac {\alpha^ {2}}{2}\right) ^ {s} \ell (\boldsymbol {\theta} ^ {0}), \quad s \in [ 0, t ]. \end{array} \right. \tag {25}
$$
It is clear that Eq. (25) holds for  $t = 0$ . Assume that Eq. (25) holds up to iteration  $t$ . By the triangle inequality, we have
$$
\begin{array}{l} \left\| \boldsymbol {W} _ {Q} ^ {t + 1} - \boldsymbol {W} _ {Q} ^ {0} \right\| _ {\mathrm {F}} \leq \sum_ {s = 0} ^ {t} \left\| \boldsymbol {W} _ {Q} ^ {s + 1} - \boldsymbol {W} _ {Q} ^ {s} \right\| _ {\mathrm {F}} = \gamma \sum_ {s = 0} ^ {t} \left\| \nabla_ {\boldsymbol {W} _ {Q}} \ell (\boldsymbol {\theta} ^ {s}) \right\| _ {\mathrm {F}} \\ \leq 2 \gamma \rho \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \sum_ {s = 0} ^ {t} \sqrt {2 \ell (\boldsymbol {\theta} ^ {s})} \leq 2 \gamma \rho \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \sum_ {s = 0} ^ {t} \left(1 - \gamma \frac {\alpha^ {2}}{2}\right) ^ {s / 2} \sqrt {2 \ell (\boldsymbol {\theta} ^ {0})}, \tag {26} \\ \end{array}
$$
where the second inequality follows from the upper bound of the gradient norm in Lemma 12, and the last inequality follows from induction assumption. Let  $u \coloneqq \sqrt{1 - \gamma\alpha^2 / 2}$ , we bound the RHS of the previous expression:
$$
\rho 2 \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \frac {2}{\alpha^ {2}} \left(1 - u ^ {2}\right) \frac {1 - u ^ {t + 1}}{1 - u} \sqrt {2 \ell (\boldsymbol {\theta} ^ {0})} \tag {27}
$$
$$
\leq \rho 2 \tau_ {0} \bar {\lambda} _ {K} \bar {\lambda} _ {V} \bar {\lambda} _ {O} d _ {s} C _ {x} ^ {2} \frac {4}{\alpha^ {2}} \sqrt {2 \ell (\boldsymbol {\theta} ^ {0})}, \quad (\text {s i n c e} u \in (0, 1)) \tag {28}
$$
$$
\leq C _ {Q}. \quad (\text {b y E q .} (2)). \tag {29}
$$
By Weyl's inequality, this implies:
$$
\left\| \boldsymbol {W} _ {Q} ^ {t + 1} \right\| _ {2} \leq \left\| \boldsymbol {W} _ {Q} ^ {0} \right\| _ {2} + C _ {Q} = \bar {\lambda} _ {Q}. \tag {30}
$$