SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We draw ICL instances  $\mathbf{Z} := (\mathbf{H}, y_{N+1}) = (\mathcal{D}, (\mathbf{x}_{N+1}, y_{N+1}))$  from a (meta-)distribution denoted as  $\pi$ , which first sample an in-context data distribution  $\mathsf{P} \sim \pi$ , then sample iid examples  $(\mathbf{x}_i, y_i)_{i=1}^{N+1} \stackrel{\mathrm{iid}}{\sim} \mathsf{P}^{\otimes(N+1)}$  and form  $\mathcal{D} = \{(\mathbf{x}_i, y_i)\}_{i \in [N]}$ . Our pretraining loss is the average ICL loss on  $n$  pretraining instances  $\mathbf{Z}^{(1:n)} \stackrel{\mathrm{iid}}{\sim} \pi$ , and we consider the corresponding test ICL loss on a new test instance:
$$
\widehat {L} _ {\mathrm {i c l}} (\boldsymbol {\theta}) := \frac {1}{n} \sum_ {j = 1} ^ {n} \ell_ {\mathrm {i c l}} (\boldsymbol {\theta}; \mathbf {Z} ^ {j}), \quad L _ {\mathrm {i c l}} (\boldsymbol {\theta}) := \mathbb {E} _ {\mathbf {P} \sim \pi , \mathbf {Z} \sim \mathsf {P} ^ {\otimes (N + 1)}} [ \ell_ {\mathrm {i c l}} (\boldsymbol {\theta}; \mathbf {Z}) ].
$$
Our pretraining algorithm is to solve a standard constrained empirical risk minimization (ERM) problem over transformers with  $L$  layers,  $M$  heads, and norm bound  $B$  (recall the definition of the  $\| \cdot \|$  norm in (2)):
$$
\begin{array}{l} \widehat{\boldsymbol{\theta}} := \operatorname *{arg  min}_{\boldsymbol {\theta}\in \Theta_{L,M,D^{\prime},B}}\widehat{L}_{\mathrm{icl}}(\boldsymbol {\theta}), \\ \Theta_ {L, M, D ^ {\prime}, B} := \left\{\boldsymbol {\theta} = (\boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {(1: L)}, \boldsymbol {\theta} _ {\mathrm {m l p}} ^ {(1: L)}): \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq M, \max  _ {\ell \in [ L ]} D ^ {(\ell)} \leq D ^ {\prime}, \| \boldsymbol {\theta} \| \leq B \right\}. \tag {TF-ERM} \\ \end{array}
$$
Generalization guarantee By standard uniform concentration analysis via chaining arguments (Proposition B.4; see also [87, Chapter 5] for similar arguments), we have the following excess loss guarantee for (TF-ERM). The proof can be found in Appendix L.2.
Theorem K.1 (Generalization for pretraining). With probability at least  $1 - \xi$  (over the pretraining instances  $\{\mathbf{Z}^j\}_{j \in [n]}$ ), the solution  $\widehat{\pmb{\theta}}$  to (TF-ERM) satisfies
$$
L _ {\mathrm {i c l}} (\widehat {\boldsymbol {\theta}}) \leq \inf  _ {\boldsymbol {\theta} \in \Theta_ {L, M, D ^ {\prime}, B}} L _ {\mathrm {i c l}} (\boldsymbol {\theta}) + \mathcal {O} \left(B _ {y} ^ {2} \sqrt {\frac {L ^ {2} (M D ^ {2} + D D ^ {\prime}) \iota + \log (1 / \xi)}{n}}\right),
$$
where  $\iota = \log (2 + \max \{B,\mathsf{R},B_y\})$  is a log factor.
K.2 Examples of pretraining for in-context regression problems
In Theorem K.1, the comparator  $\inf_{\theta \in \Theta_{L,M,D',B}} L_{\mathrm{icl}}(\theta)$  is simply the smallest expected ICL loss for ICL instances drawn from  $\pi$ , among all transformers within the norm ball  $\Theta_{L,M,D',B}$ . Using our constructions in Section 3 & 4, we show that this comparator loss is small on various (metadistribution  $\pi$ 's, by which we obtain end-to-end guarantees for pretraining transformers with small ICL loss at test time. Here we showcase this argument on several representative regression problems.
Linear regression For any in-context data distribution  $\mathsf{P}$ , let  $\mathbf{w}_{\mathsf{P}}^{\star} \coloneqq \mathbb{E}_{\mathsf{P}}[\mathbf{x}\mathbf{x}^{\top}]^{-1}\mathbb{E}_{\mathsf{P}}[\mathbf{x}y]$  denote the best linear predictor for  $\mathsf{P}$ . We show that with mild choices of  $L, M, B$ , the learned transformer can perform in-context linear regression with near-optimal statistical power, in that on the sampled  $\mathsf{P} \sim \pi$  and ICL instance  $\{(x_{i}, y_{i})\}_{i \in [N + 1]}\stackrel{\mathrm{id}}{\sim} \mathsf{P}$ , it competes with the best linear predictor  $\mathbf{w}_{\mathsf{P}}^{\star}$  for this particular  $\mathsf{P}$ . The proof follows directly by on combining Corollary 5 with Theorem K.1, and can be found in Appendix L.3.
Theorem K.2 (Pretraining transformers for in-context linear regression). Suppose  $\mathsf{P} \sim \pi$  is almost surely well-posed for in-context linear regression (Assumption A) with the canonical parameters. Then, for  $N \geq \widetilde{\mathcal{O}}(d)$ , with probability at least  $1 - \xi$  (over the training instances  $\mathbf{Z}^{(1:n)}$ ), the solution
$\widehat{\pmb{\theta}}$  of (TF-ERM) with  $L = \mathcal{O}(\kappa \log (\kappa N / \sigma))$  layers,  $M = 3$  heads,  $D^{\prime} = 0$  (attention-only), and  $B = \mathcal{O}(\sqrt{\kappa d})$  achieves small excess ICL risk over  $\mathbf{w}_{\mathsf{P}}^{\star}$ :
$$
L _ {\mathrm {i c l}} (\widehat {\pmb {\theta}}) - \mathbb {E} _ {\mathsf {P} \sim \pi} \mathbb {E} _ {(\mathbf {x}, y) \sim \mathsf {P}} \left[ \frac {1}{2} (y - \langle \mathbf {w} _ {\mathsf {P}} ^ {\star}, \mathbf {x} \rangle) ^ {2} \right] \leq \widetilde {\mathcal {O}} \left(\sqrt {\frac {\kappa^ {2} d ^ {2} + \log (1 / \xi)}{n}} + \frac {d \sigma^ {2}}{N}\right),
$$
where  $\tilde{\mathcal{O}} (\cdot)$  only hides polylogarithmic factors in  $\kappa ,N,1 / \sigma$
To our best knowledge, Theorem K.2 offers the first end-to-end result for pretraining a transformer to perform in-context linear regression with explicit excess loss bounds. The  $\widetilde{\mathcal{O}} (\sqrt{\kappa^2d^2 / n})$  term originates from the generalization of pretraining (Theorem K.1), where as the  $\widetilde{\mathcal{O}} (d\sigma^2 /N)$  term agrees with the standard fast rate for the excess loss of linear regression [38]. Further, as long as  $n\geq \widetilde{\mathcal{O}} (\kappa^2 N / \sigma^2)$ , the excess risk achieves the optimal rate  $\widetilde{\mathcal{O}} (d\sigma^2 /N)$  (up to log factors).
Additional examples By similar arguments as in the proof of Theorem K.2, we can directly turn most of our other expressivity results into results on the pretrained transformers. Here we present three such additional examples (proofs in Appendix L.4-L.6). The first example is for the sparse linear regression problem considered in Theorem 8.
Theorem K.3 (Pretraining transformers for in-context sparse linear regression). Suppose each  $\mathsf{P} \sim \pi$  is almost surely an instance of the sparse linear model specified in Theorem 8 with parameters  $B_w^{\star}$  and  $\sigma$ . Suppose  $N \geq \widetilde{\mathcal{O}}(s \log((d \vee N) / \sigma))$  and let  $\kappa := B_w^{\star} / \sigma$ .
Then with probability at least  $1 - \xi$  (over the training instances  $\mathbf{Z}^{(1:n)}$ ), the solution  $\widehat{\theta}$  of (TF-ERM) with  $L = \widetilde{\mathcal{O}} (\kappa^2 (1 + d / N))$  layers,  $M = 2$  heads,  $D^{\prime} = 2d$ , and  $B = \tilde{\mathcal{O}} (\mathrm{poly}(d,B_{w}^{\star},\sigma))$  achieves small excess ICL risk:
$$
L _ {\mathrm {i c l}} (\widehat {\pmb {\theta}}) - \sigma^ {2} \leq \widetilde {\mathcal {O}} \left(\sqrt {\frac {\kappa^ {4} d ^ {2} (1 + d / N) ^ {2} + \log (1 / \xi)}{n}} + \sigma^ {2} \frac {s \log d}{N}\right),
$$
where  $\tilde{\mathcal{O}} (\cdot)$  only hides polylogarithmic factors in  $d,N,1 / \sigma$
Our next example is for the problem of noisy linear regression with mixed noise levels considered in Theorem 12 and Theorem J.1. There, the constructed transformer uses the post-ICL validation mechanism to perform ridge regression with an adaptive regularization strength depending on the particular input sequence.
Theorem K.4 (Pretraining transformers for in-context noisy linear regression with algorithm selection). Suppose  $\pi$  is the data generating model (noisy linear model with mixed noise levels) considered in Theorem J.1, with  $\sigma_{\max} \leq \mathcal{O}(1)$ . Let  $N \geq d / 10$ .
Then, with probability at least  $1 - \xi$  (over the training instances  $\mathbf{Z}^{(1:n)}$ ), the solution  $\widehat{\pmb{\theta}}$  of (TF-ERM) with input dimension  $D = \Theta(dK)$ ,  $L = \mathcal{O}(\sigma_{\min}^{-2}\log(N / \sigma_{\min}))$  layers,  $M = \mathcal{O}(K)$  heads,  $D' = \mathcal{O}(K^2)$ , and  $B = \mathcal{O}(\mathrm{poly}(K,\sigma_{\min}^{-1},d,N))$  achieves small excess ICL risk:
$$
L _ {\mathrm {i c l}} (\widehat {\boldsymbol {\theta}}) - \text {B a y e s R i s k} _ {\pi} \leq \widetilde {\mathcal {O}} \left(\sqrt {\frac {\sigma_ {\operatorname* {m i n}} ^ {- 4} K ^ {3} d ^ {2} + \log (1 / \xi)}{n}} + \frac {\sigma_ {\operatorname* {m a x}} ^ {2}}{\sigma_ {\operatorname* {m i n}} ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right),
$$
where  $\widetilde{\mathcal{O}} (\cdot)$  only hides polylogarithmic factors in  $d,N,K,1 / \sigma_{\mathrm{min}}$
Our final example is for in-context logistic regression. For simplicity we consider the realizable case.
Theorem K.5 (Pretraining transformers for in-context logistic regression; square loss guarantee). Suppose for  $\mathsf{P} \sim \pi$ ,  $\mathsf{P}$  is almost surely a realizable logistic model (i.e.  $\mathsf{P} = \mathsf{P}_{\beta}^{\log}$  with  $\| \beta \|_2 \leq B_w^\star$  as in Corollary G.1). Suppose that  $B_w^\star = \mathcal{O}(1)$  and  $N \geq \mathcal{O}(d)$ .
Then, with probability at least  $1 - \xi$  (over the training instances  $\mathbf{Z}^{(1:n)})$ , the solution  $\widehat{\pmb{\theta}}$  of (TF-ERM) with  $L = \mathcal{O}(\log(N))$  layers,  $M = \widetilde{\mathcal{O}}(d^3N)$  heads,  $D' = 0$ , and  $B = \mathcal{O}(\mathrm{poly}(d,N))$  achieves small excess ICL risk:
$$
L _ {\mathrm {i c l}} (\widehat {\boldsymbol {\theta}}) - \mathbb {E} _ {\mathsf {P} _ {\beta} ^ {\log} \sim \pi} \mathbb {E} _ {(\mathbf {x}, y) \sim \mathsf {P} _ {\beta} ^ {\log}} \left[ \frac {1}{2} (y - \sigma_ {\log} (\langle \boldsymbol {\beta}, \mathbf {x} \rangle)) ^ {2} \right] \leq \widetilde {\mathcal {O}} \left(\sqrt {\frac {d ^ {5} N + \log (1 / \xi)}{n}} + \frac {d}{N}\right),
$$
where  $\widetilde{\mathcal{O}} (\cdot)$  only hides polylogarithmic factors in  $d,N$
Remark on generality of transformer All results above are established by the expressivity results in Section 3 & 4 for transformers to implement various ICL procedures (such as least squares, Lasso, GLM, and ridge regression with in-context algorithm selection), combined with the generalization bound (Theorem K.1). However, the transformer itself was not specified to encode any actual structure about the problem at hand in any result above, other than having sufficiently large number of layers, number of heads, and weight norms, which illustrates the flexibility of the transformer architecture.
L Proofs for Section K
L.1 Lipschitzness of transformers
For any  $p \in [1, \infty]$ , let  $\| \mathbf{H} \|_{2,p} \coloneqq (\sum_{i=1}^{N} \| \mathbf{h}_i \|_2^p)^{1/p}$  denote the column-wise  $(2, p)$ -norm of  $\mathbf{H}$ . For any radius  $R > 0$ , we denote  $\mathcal{H}_{\mathbb{R}} \coloneqq \{\mathbf{H} : \| \mathbf{H} \|_{2,\infty} \leq R\}$  be the ball of radius  $R$  under norm  $\| \cdot \|_{2,\infty}$ .
Lemma L.1. For a single MLP layer  $\theta_{\mathrm{mlp}} = (\mathbf{W}_1,\mathbf{W}_2)$ , we introduce its norm (as in (2))
$$
\left\| \left. \boldsymbol {\theta} _ {\mathrm {m l p}} \right\| \right\| = \left\| \mathbf {W} _ {1} \right\| _ {\mathrm {o p}} + \left\| \mathbf {W} _ {2} \right\| _ {\mathrm {o p}}.
$$
For any fixed hidden dimension  $D'$ , we consider
$$
\Theta_ {\mathrm {m l p}, B} := \left\{\boldsymbol {\theta} _ {\mathrm {m l p}}: \| \boldsymbol {\theta} _ {\mathrm {m l p}} \| \leq B \right\}.
$$
Then for  $\mathbf{H} \in \mathcal{H}_{\mathbb{R}}$ ,  $\pmb{\theta}_{\mathrm{m1p}} \in \Theta_{\mathrm{m1p},B}$ , the function  $(\pmb{\theta}_{\mathrm{m1p}},\mathbf{H}) \mapsto \mathrm{MLP}_{\pmb{\theta}_{\mathrm{m1p}}}(\mathbf{H})$  is (BR)-Lipschitz w.r.t.  $\pmb{\theta}_{\mathrm{m1p}}$  and  $(1 + B^{2})$ -Lipschitz w.r.t.  $\mathbf{H}$ .
Proof. Recall that by our definition, for the parameter  $\theta_{\mathrm{mlp}} = (\mathbf{W}_1,\mathbf{W}_2)\in \Theta_{\mathrm{mlp},B}$  and the input  $\mathbf{H} = [\mathbf{h}_i]\in \mathbb{R}^{D\times N}$ , the output  $\mathrm{MLP}_{\theta_{\mathrm{mlp}}}(\mathbf{H}) = \mathbf{H} + \mathbf{W}_2\sigma (\mathbf{W}_1\mathbf{H}) = [\mathbf{h}_i + \mathbf{W}_2\sigma (\mathbf{W}_1\mathbf{h}_i)]_i$ . Therefore, for  $\theta_{\mathrm{mlp}}^{\prime} = (\mathbf{W}_{1}^{\prime},\mathbf{W}_{2}^{\prime})\in \Theta_{\mathrm{mlp},B}$ , we have
$$
\begin{array}{l} \left\| \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}}} (\mathbf {H}) - \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}} ^ {\prime}} (\mathbf {H}) \right\| _ {2, \infty} \\ = \max _ {i} \| \mathbf {W} _ {2} \sigma (\mathbf {W} _ {1} \mathbf {h} _ {i}) - \mathbf {W} _ {2} ^ {\prime} \sigma (\mathbf {W} _ {1} ^ {\prime} \mathbf {h} _ {i}) \| _ {2} \\ = \max  _ {i} \| \left(\mathbf {W} _ {2} - \mathbf {W} _ {2} ^ {\prime}\right) \sigma \left(\mathbf {W} _ {1} \mathbf {h} _ {i}\right) + \mathbf {W} _ {2} ^ {\prime} \left(\sigma \left(\mathbf {W} _ {1} \mathbf {h} _ {i}\right) - \sigma \left(\mathbf {W} _ {1} ^ {\prime} \mathbf {h} _ {i}\right)\right) \| _ {2} \\ \leq \max _ {i} \| \mathbf {W} _ {2} - \mathbf {W} _ {2} ^ {\prime} \| _ {\mathrm {o p}} \| \sigma (\mathbf {W} _ {1} \mathbf {h} _ {i}) \| _ {2} + \| \mathbf {W} _ {2} ^ {\prime} \| _ {\mathrm {o p}} \| \sigma (\mathbf {W} _ {1} \mathbf {h} _ {i}) - \sigma (\mathbf {W} _ {1} ^ {\prime} \mathbf {h} _ {i}) \| _ {2} \\ \leq \max _ {i} \| \mathbf {W} _ {2} - \mathbf {W} _ {2} ^ {\prime} \| _ {\mathrm {o p}} \| \mathbf {W} _ {1} \mathbf {h} _ {i} \| _ {2} + \| \mathbf {W} _ {2} ^ {\prime} \| _ {\mathrm {o p}} \| \mathbf {W} _ {1} \mathbf {h} _ {i} - \mathbf {W} _ {1} ^ {\prime} \mathbf {h} _ {i} \| _ {2} \\ \leq B R \left\| \mathbf {W} _ {2} - \mathbf {W} _ {2} ^ {\prime} \right\| _ {\mathrm {o p}} + B R \left\| \mathbf {W} _ {1} - \mathbf {W} _ {1} ^ {\prime} \right\| _ {\mathrm {o p}}, \\ \end{array}
$$
where the second inequality follows from the 1-Lipschitznees of  $\sigma = [\cdot]_{+}$ . Similarly, for  $\mathbf{H}' = [\mathbf{h}_i'] \in \mathbb{R}^{D \times N}$ ,
$$
\begin{array}{l} \left\| \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m 1 p}}} (\mathbf {H}) - \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m 1 p}}} (\mathbf {H} ^ {\prime}) \right\| _ {2, \infty} = \max _ {i} \| \mathbf {h} _ {i} + \mathbf {W} _ {1} \sigma (\mathbf {W} _ {2} \mathbf {h} _ {i}) - \mathbf {h} _ {i} ^ {\prime} - \mathbf {W} _ {1} \sigma (\mathbf {W} _ {2} \mathbf {h} _ {i} ^ {\prime}) \| _ {2} \\ \leq \left\| \mathbf {H} - \mathbf {H} ^ {\prime} \right\| _ {2, \infty} + \max  _ {i} \left\| \mathbf {W} _ {1} \left(\sigma \left(\mathbf {W} _ {2} \mathbf {h} _ {i}\right) - \sigma \left(\mathbf {W} _ {2} \mathbf {h} _ {i} ^ {\prime}\right)\right) \right\| _ {2} \\ \leq \left\| \mathbf {H} - \mathbf {H} ^ {\prime} \right\| _ {2, \infty} + \max _ {i} B \left\| \sigma (\mathbf {W} _ {2} \mathbf {h} _ {i}) - \sigma (\mathbf {W} _ {2} \mathbf {h} _ {i} ^ {\prime}) \right\| _ {2} \\ \leq \| \mathbf {H} - \mathbf {H} ^ {\prime} \| _ {2, \infty} + B ^ {2} \| \mathbf {H} - \mathbf {H} ^ {\prime} \| _ {2, \infty}. \\ \end{array}
$$
Lemma L.2. For a single attention layer  $\theta_{\mathrm{attn}} = \{(\mathbf{V}_m,\mathbf{Q}_m,\mathbf{K}_m)\}_{m\in [M]}\subset \mathbb{R}^{D\times D}$ , we introduce its norm (as in (2))
$$
\| \boldsymbol {\theta} _ {\mathrm {a t t n}} \| := \max  _ {m \in [ M ]} \left\{\| \mathbf {Q} _ {m} \| _ {\mathrm {o p}}, \| \mathbf {K} _ {m} \| _ {\mathrm {o p}} \right\} + \sum_ {m = 1} ^ {M} \| \mathbf {V} _ {m} \| _ {\mathrm {o p}}.
$$
For any fixed dimension  $D$ , we consider
$$
\Theta_ {\mathrm {a t t n}, B} := \left\{\theta_ {\mathrm {a t t n}}: \| \theta_ {\mathrm {a t t n}} \| \leq B \right\}.
$$
Then for  $\mathbf{H} \in \mathcal{H}_{\mathbb{R}}$ ,  $\theta_{\mathrm{attn}} \in \Theta_{\mathrm{attn},B}$ , the function  $(\theta_{\mathrm{attn}},\mathbf{H}) \mapsto \mathrm{Attn}_{\theta_{\mathrm{attn}}}(\mathbf{H})$  is  $(B^{2}\mathsf{R}^{3})$ -Lipschitz w.r.t.  $\theta_{\mathrm{attn}}$  and  $(1 + B^{3}\mathsf{R}^{2})$ -Lipschitz w.r.t.  $\mathbf{H}$ .
Proof. Recall that by our definition, for the parameter  $\theta_{\mathrm{attn}} = \{(\mathbf{V}_m,\mathbf{Q}_m,\mathbf{K}_m)\}_{m\in [M]}\in \Theta_{\mathrm{attn},B}$  and the input  $\mathbf{H} = [\mathbf{h}_i]\in \mathbb{R}^{D\times N}$ , the output  $\mathrm{Attn}\theta_{\mathrm{attn}}(\mathbf{H}) = [\widetilde{\mathbf{h}}_i]$  is given by
$$
\widetilde {\mathbf {h}} _ {i} = \mathbf {h} _ {i} + \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \sigma (\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \rangle) \cdot \mathbf {V} _ {m} \mathbf {h} _ {j}.
$$
Now, for  $\theta_{\mathrm{attn}}^{\prime} = \left\{(\mathbf{V}_{m}^{\prime},\mathbf{Q}_{m}^{\prime},\mathbf{K}_{m}^{\prime})\right\}_{m\in [M]}$ , we consider
$$
\widetilde {\mathbf {h}} _ {i} ^ {\prime} = \left[ \mathrm {A t t n} _ {\theta_ {\mathrm {a t t n}} ^ {\prime}} (\mathbf {H}) \right] _ {i} = \mathbf {h} _ {i} + \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \sigma (\langle \mathbf {Q} _ {m} ^ {\prime} \mathbf {h} _ {i}, \mathbf {K} _ {m} ^ {\prime} \mathbf {h} _ {j} \rangle) \cdot \mathbf {V} _ {m} ^ {\prime} \mathbf {h} _ {j}, \qquad \forall i \in [ N ].
$$
Clearly  $\left\| \mathrm{Att}\mathbf{n}_{\pmb{\theta}_{\mathrm{attn}}}(\mathbf{H}) - \mathrm{Att}\mathbf{n}_{\pmb{\theta}_{\mathrm{attn}}^{\prime}}(\mathbf{H})\right\|_{2,\infty} = \max_{i}\left\| \widetilde{\mathbf{h}}_{i} - \widetilde{\mathbf{h}}_{i}^{\prime}\right\|_{2}$ . For any  $i\in [N]$ , we have
where the second inequality uses the definition of operator norm, the third inequality follows from the triangle inequality, the forth inequality is because  $\| \mathbf{Q}_m\mathbf{h}_i\| _2\leq BR,\| \mathbf{K}_m\mathbf{h}_j\| _2\leq BR$  and  $\sigma$  is 1-Lipschitz. This completes the proof the Lipschitzness w.r.t.  $\theta_{\mathrm{attn}}$
Similarly, we consider  $\mathbf{H}' = [\mathbf{h}_i']$ , and
$$
\widetilde {\mathbf {h}} _ {i} ^ {\prime} = \left[ \mathrm {A t t n} _ {\theta_ {\mathrm {a t t n}} ^ {\prime}} (\mathbf {H}) \right] _ {i} = \mathbf {h} _ {i} ^ {\prime} + \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \sigma \big (\big \langle \mathbf {Q} _ {m} \mathbf {h} _ {i} ^ {\prime}, \mathbf {K} _ {m} \mathbf {h} _ {j} ^ {\prime} \big \rangle \big) \cdot \mathbf {V} _ {m} \mathbf {h} _ {j} ^ {\prime}, \qquad \forall i \in [ N ].
$$
By definition, we can similarly bound
$$
\begin{array}{l} \left\| \left(\widetilde {\mathbf {h}} _ {i} ^ {\prime} - \mathbf {h} _ {i} ^ {\prime}\right) - \left(\widetilde {\mathbf {h}} _ {i} - \mathbf {h} _ {i}\right) \right\| _ {2} \\ = \left\| \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \left[ \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) \mathbf {V} _ {m} \mathbf {h} _ {j} - \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i} ^ {\prime}, \mathbf {K} _ {m} \mathbf {h} _ {j} ^ {\prime} \right\rangle\right) \mathbf {V} _ {m} \mathbf {h} _ {j} ^ {\prime} \right] \right\| _ {2} \\ \leq \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \| \mathbf {V} _ {m} \| _ {\mathrm {o p}} \left\| \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) \mathbf {h} _ {j} - \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i} ^ {\prime}, \mathbf {K} _ {m} \mathbf {h} _ {j} ^ {\prime} \right\rangle\right) \mathbf {h} _ {j} ^ {\prime} \right\| _ {2} \\ \end{array}
$$
$$
\begin{array}{l} \leq \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \| \mathbf {V} _ {m} \| _ {\mathrm {o p}} \left\{\left| \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) \right| \cdot \left\| \mathbf {h} _ {j} - \mathbf {h} _ {j} ^ {\prime} \right\| _ {2} \right. \\ + \left| \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) - \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i} ^ {\prime}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) \right| \cdot \left\| \mathbf {h} _ {j} ^ {\prime} \right\| _ {2} \\ \left. + \left| \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i} ^ {\prime}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) - \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i} ^ {\prime}, \mathbf {K} _ {m} \mathbf {h} _ {j} ^ {\prime} \right\rangle\right) \right| \cdot \left\| \mathbf {h} _ {j} ^ {\prime} \right\| _ {2} \right\} \\ \leq \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \| \mathbf {V} _ {m} \| _ {\mathrm {o p}} \cdot 3 \| \mathbf {Q} _ {m} \| _ {\mathrm {o p}} \| \mathbf {K} _ {m} \| _ {\mathrm {o p}} R ^ {2} \left\| \mathbf {h} _ {j} - \mathbf {h} _ {j} ^ {\prime} \right\| _ {2} \\ \leq \mathsf {R} ^ {2} \| \mathbf {H} - \mathbf {H} ^ {\prime} \| _ {2, \infty} \cdot 3 \max  _ {m \in [ M ]} \| \mathbf {Q} _ {m} \| _ {\mathrm {o p}} \| \mathbf {K} _ {m} \| _ {\mathrm {o p}} \cdot \sum_ {m = 1} ^ {M} \| \mathbf {V} _ {m} \| _ {\mathrm {o p}} \\ \leq B ^ {3} \mathsf {R} ^ {2} \| \mathbf {H} - \mathbf {H} ^ {\prime} \| _ {2, \infty}, \\ \end{array}
$$
where the last inequality uses  $\| \theta_{\mathrm{attn}}\| \leq B$  and the AM-GM inequality. This completes the proof the Lipschitzness w.r.t. H.
Corollary L.1. For a fixed number of heads  $M$  and hidden dimension  $D'$ , we consider
$$
\Theta_ {\mathrm {T F}, 1, B} = \left\{\boldsymbol {\theta} = \left(\boldsymbol {\theta} _ {\mathrm {a t t n}}, \boldsymbol {\theta} _ {\mathrm {m l p}}\right): M \text {h e a d s , h i d d e n d i m e n s i o n} D ^ {\prime}, \| \boldsymbol {\theta} \| \leq B \right\}.
$$
Then for the function  $\mathrm{TF}^{\mathsf{R}}$  given by
$$
\mathrm {T F} ^ {\mathsf {R}}: (\boldsymbol {\theta}, \mathbf {H}) \mapsto \operatorname {c l i p} _ {\mathsf {R}} \left(\operatorname {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}}} \left(\operatorname {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}}} (\mathbf {H})\right)\right), \quad \boldsymbol {\theta} \in \Theta_ {\mathrm {T F}, 1, B}, \mathbf {H} \in \mathcal {H} _ {\mathsf {R}}
$$
$\mathrm{TF}^{\mathsf{R}}$  is  $B_{\Theta}$ -Lipschitz w.r.t  $\pmb{\theta}$  and  $L_{H}$ -Lipschitz w.r.t.  $\mathbf{H}$ , where  $B_{\Theta} := B\mathsf{R}(1 + B\mathsf{R}^{2} + B^{3}\mathsf{R}^{2})$  and  $B_{H} := (1 + B^{2})(1 + B^{2}\mathsf{R}^{3})$ .
Proof. For any  $\pmb{\theta} = (\pmb{\theta}_{\mathrm{attn}}, \pmb{\theta}_{\mathrm{mlp}})$ ,  $\mathbf{H} \in \mathcal{H}_{\mathbb{R}}$ , and  $\theta' = (\theta_{\mathrm{attn}}', \theta_{\mathrm{mlp}}')$ , we have
$$
\begin{array}{l} \left\| \mathrm {T F} _ {\boldsymbol {\theta}} (\mathbf {H}) - \mathrm {T F} _ {\theta^ {\prime}} (\mathbf {H}) \right\| _ {2, \infty} \leq \left\| \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}}} \left(\operatorname {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}}} (\mathbf {H})\right) - \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}}} \left(\operatorname {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {\prime}} (\mathbf {H})\right) \right\| _ {2, \infty} \\ + \left\| \operatorname {M L P} _ {\theta_ {\mathrm {m l p}}} \left(\operatorname {A t t n} _ {\theta_ {\mathrm {a t t n}} ^ {\prime}} (\mathbf {H})\right) - \operatorname {M L P} _ {\theta_ {\mathrm {m l p}} ^ {\prime}} \left(\operatorname {A t t n} _ {\theta_ {\mathrm {a t t n}} ^ {\prime}} (\mathbf {H})\right) \right\| _ {2, \infty} \\ \leq \left(1 + B ^ {2}\right) \left\| \operatorname {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}}} (\mathbf {H}) - \operatorname {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {\prime}} (\mathbf {H}) \right\| _ {2, \infty} + B \overline {{\mathsf {R}}} \left\| \boldsymbol {\theta} _ {\mathrm {m l p}} - \boldsymbol {\theta} _ {\mathrm {m l p}} ^ {\prime} \right\| \\ \leq \left(1 + B ^ {2}\right) B ^ {2} \mathsf {R} ^ {3} \left\| \boldsymbol {\theta} _ {\mathrm {a t t n}} - \boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {\prime} \right\| + B \overline {{\mathsf {R}}} \left\| \boldsymbol {\theta} _ {\mathrm {m l p}} - \boldsymbol {\theta} _ {\mathrm {m l p}} ^ {\prime} \right\| \\ \leq B _ {\Theta} \left\| \boldsymbol {\theta} - \theta^ {\prime} \right\|, \\ \end{array}
$$