SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbf {h} _ {N + 1} ^ {(L + 1)} = \mathbf {h} _ {N + 1} ^ {(L)} + \frac {1}{N + 1} \sum_ {j = 1} ^ {N + 1} \sum_ {m = 1} ^ {2} \sigma \left(\left\langle \mathbf {Q} ^ {(L + 1)} \mathbf {h} _ {N + 1} ^ {(L)}, \mathbf {K} ^ {(L + 1)} \mathbf {h} _ {j} ^ {(L)} \right\rangle\right) \mathbf {V} ^ {(L + 1)} \mathbf {h} _ {j} ^ {(L)} \\ = \left[ \mathbf {x} _ {i}; 0; \widehat {\mathbf {w}} ^ {L}; \mathbf {0} _ {D - 2 d - 3}; 1; 1 \right] + \left(\sigma \left(\left\langle \widehat {\mathbf {w}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle\right) - \sigma \left(- \left\langle \widehat {\mathbf {w}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle\right)\right) \cdot \left[ \mathbf {0} _ {d}; 1; \mathbf {0} _ {D - d - 1} \right] \\ \stackrel {(i)} {=} \left[ \mathbf {x} _ {i}; 0; \widehat {\mathbf {w}} ^ {L}; \mathbf {0} _ {D - 2 d - 3}; 1; 1 \right] + \left[ \mathbf {0} _ {d}; \left\langle \widehat {\mathbf {w}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle ; \mathbf {0} _ {D - d - 1} \right] \\ = [ \mathbf {x} _ {i}; \underbrace {\left\langle \widehat {\mathbf {w}} ^ {L} , \mathbf {x} _ {N + 1} \right\rangle} _ {\widehat {y} _ {N + 1}}; \widehat {\mathbf {w}} ^ {L}; \mathbf {0} _ {D - 2 d - 3}; 1; 1 ], \\ \end{array}
$$
Above, (i) uses the identity  $t = \sigma(t) - \sigma(-t)$ . Further by part (a) we have
$$
\left| \widehat {y} _ {N + 1} - \left\langle \mathbf {w} _ {\mathrm {G D}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle \right| = \left| \left\langle \widehat {\mathbf {w}} ^ {L} - \mathbf {w} _ {\mathrm {G D}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle \right| \leq \varepsilon \cdot (L \eta B _ {x} ^ {2}).
$$
This proves part (b), and also finishes the proof Theorem D.1 where the overall  $(L + 1)$ -layer attention-only transformer is given by  $\mathrm{TF}_{\theta}^{0}$  with
$$
\boldsymbol {\theta} = (\underbrace{\boldsymbol{\theta}^{(1)},\ldots,\boldsymbol{\theta}^{(1)}}_{L \text{times}},\boldsymbol{\theta}^{(L + 1)}).
$$
E.3 Proof of Lemma D.1
As  $f$  is a convex,  $L_{f}$  smooth function on  $\mathsf{B}_2^d (R)$ , the mapping  $\mathcal{T}_{\eta}:\mathbf{w}\mapsto \mathbf{w} - \eta \nabla f(\mathbf{w})$  is non-expansive in  $\| \cdot \| _2$ : Indeed, for any  $\mathbf{w},\mathbf{w}'\in \mathsf{B}_2^d (R)$  we have
$$
\begin{array}{l} \left\| \mathcal {T} _ {\eta} (\mathbf {w}) - \mathcal {T} _ {\eta} (\mathbf {w} ^ {\prime}) \right\| _ {2} = \left\| \mathbf {w} - \eta \nabla f (\mathbf {w}) - \left(\mathbf {w} ^ {\prime} - \eta \nabla f (\mathbf {w} ^ {\prime})\right) \right\| _ {2} ^ {2} \\ = \| \mathbf {w} - \mathbf {w} ^ {\prime} \| _ {2} ^ {2} - 2 \eta \left\langle \mathbf {w} - \mathbf {w} ^ {\prime}, \nabla f (\mathbf {w}) - \nabla f (\mathbf {w} ^ {\prime}) \right\rangle + \eta^ {2} \| \nabla f (\mathbf {w}) - \nabla f (\mathbf {w} ^ {\prime}) \| _ {2} ^ {2} \\ \stackrel {(i)} {\leq} \| \mathbf {w} - \mathbf {w} ^ {\prime} \| _ {2} ^ {2} - \left(2 \eta / L _ {f} - \eta^ {2}\right) \| \nabla f (\mathbf {w}) - \nabla f (\mathbf {w} ^ {\prime}) \| _ {2} ^ {2} \stackrel {(i i)} {\leq} \| \mathbf {w} - \mathbf {w} ^ {\prime} \| _ {2} ^ {2}. \\ \end{array}
$$
Above, (i) uses the property  $\langle \mathbf{w} - \mathbf{w}', \nabla f(\mathbf{w}) - \nabla f(\mathbf{w}') \rangle \geq \frac{1}{L_f} \| \nabla f(\mathbf{w}) - \nabla f(\mathbf{w}') \|_2^2$  for smooth convex functions [63, Theorem 2.1.5]; (ii) uses the precondition that  $\eta \leq 2 / L_f$ .
The lemma then follows directly by induction on  $L$ . The base case of  $L = 0$  follows directly by assumption that  $\widehat{\mathbf{w}}^0 = \mathbf{w}_{\mathrm{GD}}^0 \in \mathsf{B}_2^d (R / 2)$ . Suppose the claim holds for iterate  $L$ . For iterate  $L + 1 \leq R / (2\varepsilon)$ , we have
$$
\begin{array}{l} \left\| \widehat {\mathbf {w}} ^ {L + 1} - \mathbf {w} _ {\mathrm {G D}} ^ {L + 1} \right\| _ {2} = \left\| \mathcal {T} _ {\eta} (\widehat {\mathbf {w}} ^ {L}) + \varepsilon^ {L} - \mathcal {T} _ {\eta} (\mathbf {w} _ {\mathrm {G D}} ^ {L}) \right\| _ {2} \\ \leq \left\| \mathcal {T} _ {\eta} \left(\widehat {\mathbf {w}} ^ {L}\right) - \mathcal {T} _ {\eta} \left(\mathbf {w} _ {\mathrm {G D}} ^ {L}\right) \right\| _ {2} + \left\| \varepsilon^ {L} \right\| _ {2} \\ \stackrel {(i)} {\leq} \left\| \widehat {\mathbf {w}} ^ {L} - \mathbf {w} _ {\mathrm {G D}} ^ {L} \right\| _ {2} + \varepsilon \stackrel {(i i)} {\leq} (L + 1) \varepsilon . \\ \end{array}
$$
Above, (i) uses the non-expansiveness, and (ii) uses the inductive hypothesis. Similarly, by our assumption  $\mathbf{w}^{\star} = \mathcal{T}_{\eta}(\mathbf{w}^{\star})$
$$
\left\| \widehat {\mathbf {w}} ^ {L + 1} - \mathbf {w} ^ {\star} \right\| _ {2} = \left\| \mathcal {T} _ {\eta} (\widehat {\mathbf {w}} ^ {L}) + \varepsilon^ {L} - \mathcal {T} _ {\eta} (\mathbf {w} ^ {\star}) \right\| _ {2} \leq \left\| \widehat {\mathbf {w}} ^ {L} - \mathbf {w} ^ {\star} \right\| _ {2} + \left\| \varepsilon^ {L} \right\| _ {2} \leq \frac {R}{2} + (L + 1) \varepsilon \leq R.
$$
This finishes the induction.
E.4 Convex ICGD with  $\ell_2$  regularization
In the same setting as Theorem D.1, consider the ICGD dynamics over an  $\ell_2$ -regularized empirical risk:
$$
\mathbf {w} _ {\mathrm {G D}} ^ {t + 1} := \mathbf {w} _ {\mathrm {G D}} ^ {t} - \eta \nabla \widehat {L} _ {N} ^ {\lambda} \left(\mathbf {w} _ {\mathrm {G D}} ^ {t}\right) \quad (\text {I C G D -} \ell_ {2})
$$
with initialization  $\mathbf{w}_{\mathrm{GD}}^0 \in \mathbb{R}^d$  and learning rate  $\eta > 0$ , where  $\widehat{L}_N^\lambda (\mathbf{w}) \coloneqq \widehat{L}_N(\mathbf{w}) + \frac{\lambda}{2}\| \mathbf{w}\|_2^2$  denotes the  $\ell_2$ -regularized empirical risk.
Corollary E.1 (Convex ICGD with  $\ell_2$  regularization). Fix any  $B_w > 0$ ,  $L > 1$ ,  $\eta > 0$ , and  $\varepsilon < B_x B_w$ . Suppose the loss  $\ell(\cdot, \cdot)$  is convex in the first argument, and  $\partial_s \ell$  is  $(\varepsilon, R, M, C)$ -approximable by sum of relus with  $R = \max \{B_x B_w, B_y, 1\}$ .
Then, there exists an attention-only transformer  $\mathrm{TF}_{\pmb{\theta}}^{0}$  with  $(L + 1)$  layers,  $\max_{\ell \in [L]}M^{(\ell)}\leq M + 1$  heads within the first  $L$  layers, and  $M^{(L + 1)} = 2$  such that for any input data  $(\mathcal{D},\mathbf{x}_{N + 1})$  with
$$
\sup_{\| \mathbf{w}\|_{2}\leq B_{w}}\lambda_{\max}(\nabla^{2}\widehat{L}_{N}^{\lambda}(\mathbf{w}))\leq 2\eta^{-1},\quad \exists \mathbf{w}^{\star}\in \operatorname *{arg  min}_{\mathbf{w}\in \mathbb{R}^{d}}\widehat{L}_{N}^{\lambda}(\mathbf{w})\text{such that}\| \mathbf{w}^{\star}\|_{2}\leq B_{w} / 2,
$$
$\mathrm{TF}_{\pmb{\theta}}^{0}(\mathbf{H}^{(0)})$  approximately implements (ICGD-  $\ell_2$ ):
1. (Parameter space) For every  $\ell \in [L]$ , the  $\ell$ -th layer's output  $\mathbf{H}^{(\ell)} = \mathrm{TF}_{\pmb{\theta}^{(1:\ell)}}(\mathbf{H}^{(0)})$  approximates  $\ell$  steps of (ICGD- $\ell_2$ ): We have  $\mathbf{h}_i^{(\ell)} = [\mathbf{x}_i; y_i'; \widehat{\mathbf{w}}^\ell; \mathbf{0}_{D-2d-3}; 1; t_i]$  for every  $i \in [N+1]$ , where
$$
\left\| \widehat {\mathbf {w}} ^ {\ell} - \mathbf {w} _ {\mathrm {G D}} ^ {\ell} \right\| _ {2} \leq \varepsilon \cdot (2 L \eta B _ {x}).
$$
2. (Prediction space) The final output  $\mathbf{H}^{(L + 1)} = \mathrm{TF}_{\pmb{\theta}}(\mathbf{H}^{(0)})$  approximates the prediction of  $L$  steps of  $(ICGD - \ell_2)$ : We have  $\mathbf{h}_{N + 1}^{(L + 1)} = [\mathbf{x}_{N + 1};\widehat{y}_{N + 1};\widehat{\mathbf{w}}^L;\mathbf{0}_{D - 2d - 3};1;0]$ , where
$$
\left| \widehat {y} _ {N + 1} - \left\langle \mathbf {w} _ {\mathrm {G D}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle \right| \leq \varepsilon \cdot (2 L \eta B _ {x} ^ {2}).
$$
Further, the transformer admits norm bound  $\| \pmb {\theta}\| \leq 2 + R + (2C + \lambda)\eta$
Proof. This construction is the same as in the proof of Theorem D.1, except that within each layer  $\ell \in [L]$ , we add one more attention head  $(\mathbf{Q}^{(\ell)}, \mathbf{K}^{(\ell)}, \mathbf{V}^{(\ell)}) \subset \mathbb{R}^{D \times D}$  which when acting on its input  $\mathbf{h}_i^{(\ell-1)} = [*\;*\; \widehat{\mathbf{w}}^{\ell-1}; 1; *]$  gives
$$
\mathbf {Q} ^ {(\ell)} \mathbf {h} _ {i} ^ {(\ell - 1)} = \left[ \begin{array}{c} 1 \\ \mathbf {0} _ {D - 1} \end{array} \right], \quad \mathbf {K} ^ {(\ell)} \mathbf {h} _ {j} ^ {(\ell - 1)} = \left[ \begin{array}{c} 1 \\ \mathbf {0} _ {D - 1} \end{array} \right], \quad \mathbf {V} ^ {(\ell)} \mathbf {h} _ {j} ^ {(\ell - 1)} = \left[ \begin{array}{c} \mathbf {0} _ {d + 1} \\ - \eta \lambda \widehat {\mathbf {w}} ^ {\ell - 1} \\ \mathbf {0} _ {2} \end{array} \right]
$$
for all  $i,j\in [N + 1]$ . Note that  $\left\| \mathbf{Q}^{(\ell)}\right\|_{\mathrm{op}} = \left\| \mathbf{K}^{(\ell)}\right\|_{\mathrm{op}} = 1$ , and  $\left\| \mathbf{V}^{(\ell)}\right\|_{\mathrm{op}} = \eta \lambda$ . Further, it is straightforward to check that the output of this attention head on every  $\mathbf{h}_i^{(\ell)}$  is
$$
\frac {1}{N + 1} \sum_ {j = 1} ^ {N + 1} \sigma \Big (\big \langle \mathbf {Q} ^ {(\ell)} \mathbf {h} _ {i} ^ {(\ell - 1)}, \mathbf {K} ^ {(\ell)} \mathbf {h} _ {j} ^ {(\ell - 1)} \big \rangle \Big) \mathbf {V} ^ {(\ell)} \mathbf {h} _ {j} ^ {(\ell - 1)} = \left[ \begin{array}{c} \mathbf {0} _ {d + 1} \\ - \eta \lambda \widehat {\mathbf {w}} ^ {\ell - 1} \\ \mathbf {0} _ {2} \end{array} \right].
$$
Adding this onto the original output of the  $\ell$ -th layer exactly implements the gradient of the regularizer  $\mathbf{w} \mapsto \frac{\lambda}{2} \| \mathbf{w} \|_2^2$ . The rest of the proof follows by repeating the argument of Theorem D.1, and combining the norm bound for the additional attention head here with the norm bound therein.
E.5 Proof of Theorem D.3
We only need to prove the following single-step version of Theorem D.3.
Proposition E.2. Under the assumptions of Theorem D.3, there exists a 2-layer transformer  $\mathrm{TF}_{\theta}$  with the same bounds on the number of heads, hidden dimension and the norm, such that for any input data  $(\mathcal{D},\mathbf{x}_{N + 1})$  and any  $\mathbf{w}\in \mathbb{R}^d$ ,  $\mathrm{TF}_{\theta}$  maps
$$
\mathbf {h} _ {i} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; \mathbf {w}; 0; 1; t _ {i} ] \quad \rightarrow \quad \mathbf {h} _ {i} ^ {\prime} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; \mathbf {w} _ {\eta} ^ {+}; 0; 1; t _ {i} ],
$$
$$
\mathbf {w} _ {\eta} ^ {+} = \operatorname {P r o j} _ {\mathcal {W}} \left(\mathbf {w} - \eta \nabla \widehat {L} _ {N} (\mathbf {w}) + \varepsilon (\mathbf {w})\right), \quad \| \varepsilon (\mathbf {w}) \| _ {2} \leq \eta \varepsilon .
$$
Before we present the formal (and technical) proof of Proposition E.2, we first provide some intuitions. To begin with, we first note that
$$
\nabla_ {\mathbf {w}} \widehat {L} _ {N} (\mathbf {w}) = \frac {1}{N} \sum_ {i = 1} ^ {N} \partial_ {1} \ell \left(\operatorname {p r e d} \left(\mathbf {x} _ {i}; \mathbf {w}\right), y _ {i}\right) \cdot \nabla_ {\mathbf {w}} \operatorname {p r e d} \left(\mathbf {x} _ {i}; \mathbf {w}\right), \tag {16}
$$
where  $\partial_1\ell$  is the partial derivative of  $\ell$  with respect to the first component, and
$$
\nabla_ {\mathbf {w}} \operatorname {p r e d} \left(\mathbf {x} _ {i}; \mathbf {w}\right) = \left[ \begin{array}{c} u _ {1} \cdot r ^ {\prime} \left(\left\langle \mathbf {v} _ {1}, \mathbf {x} _ {i} \right\rangle\right) \cdot \mathbf {x} _ {i} \\ r \left(\left\langle \mathbf {v} _ {1}, \mathbf {x} _ {i} \right\rangle\right) \\ \vdots \\ u _ {K} \cdot r ^ {\prime} \left(\left\langle \mathbf {v} _ {K}, \mathbf {x} _ {i} \right\rangle\right) \cdot \mathbf {x} _ {i} \\ r \left(\left\langle \mathbf {v} _ {K}, \mathbf {x} _ {i} \right\rangle\right) \end{array} \right] \in \mathbb {R} ^ {K (d + 1)}. \tag {17}
$$
Therefore, the basic idea is that we can use an attention layer to approximate  $(\mathbf{x}_i,\mathbf{w})\mapsto \operatorname {pred}(\mathbf{x}_i;\mathbf{w})$  then use an MLP layer to implement  $(\mathrm{pred}(\mathbf{x}_i;\mathbf{w}),y_i',t_i)\mapsto 1\{i < N + 1\} \cdot \partial_1\ell (\mathrm{pred}(\mathbf{x}_i;\mathbf{w}),y_i)$  and then use an attention layer to compute the gradient descent step  $\mathbf{w}\mapsto \mathbf{w} - \eta \nabla L_N(\mathbf{w})$  , and finally use an MLP layer to implement the projection into  $\mathcal{W}$
Based on the observations above, we now present the proof of Proposition E.2.
Proof of Proposition E.2. We write  $D_0 = d + 1 + K(d + 1)$  be the length of the vector  $[\mathbf{x}_i;y_i;\mathbf{w}]$ . We also define
$$
B _ {r} := \max _ {| t | \leq B _ {x} B _ {u}} | r (t) |, \qquad B _ {g} := \max _ {| t | \leq K B _ {r}, | y | \leq B _ {y}} | \partial_ {t} \ell (t, y) |.
$$
Let us fix  $\varepsilon_r, \varepsilon_p, \varepsilon_\ell > 0$  that will be specified later in proof (see (18)). By our assumption and Proposition B.1, the following facts hold.
(1) The function  $r(t)$  is  $(\varepsilon_r,R_1,M_1,C_1)$  for  $R_{1} = \max \{B_{x}B_{u},1\}$ ,  $M_1\leq \widetilde{\mathcal{O}}\left(C_1^2\varepsilon_r^{-2}\right)$ , where  $C_1$  depends only on  $R_{1}$  and the  $C^2$ -smoothness of  $r$ . Therefore, there exists
$$
\bar {r} (t) = \sum_ {m = 1} ^ {M} c _ {m} ^ {1} \sigma (\left\langle \mathbf {a} _ {m} ^ {1}, [ t; 1 ] \right\rangle) \quad \text {w i t h} \quad \sum_ {m = 1} ^ {M} \left| c _ {m} ^ {1} \right| \leq C _ {1}, \left\| \mathbf {a} _ {m} ^ {1} \right\| _ {1} \leq 1, \forall m \in [ M _ {1} ],
$$
such that  $\sup_{t\in [-R_1,R_1]}\left|r(t) - \overline{r} (t)\right|\leq \varepsilon_r$
(2) The function  $(t,y)\mapsto \partial_1\ell (t,y)$  is  $(\varepsilon_{\ell},R_{2},M_{2},C_{2})$  for  $R_{2} = \max \left\{KB_{r},B_{y},1\right\} M_{2}\leq \tilde{\mathcal{O}}\left(C_2^2\varepsilon_\ell^{-2}\right)$ , where  $C_2$  depends only on  $R_{2}$  and the  $C^3$ -smoothness of  $\partial_1\ell$ . Therefore, there exists
$$
g (t, y) = \sum_ {m = 1} ^ {M} c _ {m} ^ {2} \sigma (\langle \mathbf {a} _ {m} ^ {2}, [ t; y; 1 ] \rangle) \quad \mathrm {w i t h} \quad \sum_ {m = 1} ^ {M} | c _ {m} ^ {2} | \leq C _ {2}, \left\| \mathbf {a} _ {m} ^ {2} \right\| _ {1} \leq 1, \forall m \in [ M _ {2} ],
$$
such that  $\sup_{(t,y)\in [-R_2,R_2]^2}|g(t,y) - \partial_1\ell (t,y)|\leq \varepsilon_\ell$
(3) The function  $(s,t)\mapsto s\cdot r^{\prime}(t)$  is  $(\varepsilon_{p},R_{3},M_{3},C_{3})$  for  $R_{3} = \max \{B_{x}B_{u},B_{g}B_{u},1\}$ ,  $M_3\leq \widetilde{\mathcal{O}}\left(C_3^2\varepsilon_p^{-2}\right)$ , where  $C_3$  depends only on  $R_{3}$  and the  $C^3$ -smoothness of  $r^\prime$ . Therefore, there exists
$$
P (s, t) = \sum_ {m = 1} ^ {M} c _ {m} ^ {3} \sigma (\left\langle \mathbf {a} _ {m} ^ {3}, [ s; t; 1 ] \right\rangle) \quad \text {w i t h} \quad \sum_ {m = 1} ^ {M} \left| c _ {m} ^ {3} \right| \leq C _ {3}, \left\| \mathbf {a} _ {m} ^ {3} \right\| _ {1} \leq 1, \forall m \in [ M _ {3} ],
$$
such that  $\sup_{(s,t)\in [-R_3,R_3]^2}|P(s,t) - s\cdot r'(t)|\leq \varepsilon_p$
In the following, we proceed to construct the desired transformer step by step.
Step 1: construction of  $\pmb{\theta}_{\mathrm{attn}}^{(1)}$ . We consider the matrices  $\{\mathbf{Q}_{k,m}^{(1)}, \mathbf{K}_{k,m}^{(1)}, \mathbf{V}_{k,m}^{(1)}\}_{k \in [K], m \in [M_1]}$  so that for all  $i, j \in [N + 1]$ , we have
$$
\mathbf {Q} _ {k, m} ^ {(1)} \mathbf {h} _ {i} = \left[ \begin{array}{c} \mathbf {a} _ {m} ^ {1} [ 1 ] \cdot \mathbf {x} _ {i} \\ \mathbf {a} _ {m} ^ {1} [ 2 ] \\ \mathbf {0} \end{array} \right], \quad \mathbf {K} _ {k, m} ^ {(1)} \mathbf {h} _ {j} = \left[ \begin{array}{c} \mathbf {v} _ {k} \\ 1 \\ \mathbf {0} \end{array} \right], \quad \mathbf {V} _ {k, m} ^ {(1)} \mathbf {h} _ {j} = c _ {m} ^ {1} \cdot u _ {k} \mathbf {e} _ {D _ {0} + 1}.
$$
As the input has structure  $\mathbf{h}_i = [\mathbf{x}_i; y_i'; \mathbf{w}; 0; 1; t_i]$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max  _ {k, m} \left\| \mathbf {Q} _ {k, m} ^ {(1)} \right\| _ {\mathrm {o p}} \leq 1, \quad \max  _ {k, m} \left\| \mathbf {K} _ {k, m} ^ {(1)} \right\| _ {\mathrm {o p}} \leq 1, \quad \sum_ {k, m} \left\| \mathbf {V} _ {k, m} ^ {(1)} \right\| _ {\mathrm {o p}} \leq C _ {1}.
$$
A simple calculation shows that
$$
\sum_ {m \in [ M _ {1} ], k \in [ K ]} \sigma \left(\left\langle \mathbf {Q} _ {k, m} ^ {(1)} \mathbf {h} _ {i}, \mathbf {K} _ {k, m} ^ {(1)} \mathbf {h} _ {j} \right\rangle\right) \mathbf {V} _ {k, m} ^ {(1)} \mathbf {h} _ {j} = \sum_ {k = 1} ^ {K} u _ {k} \bar {r} (\left\langle \mathbf {v} _ {k}, \mathbf {x} _ {i} \right\rangle) \cdot \mathbf {e} _ {D _ {0} + 1}.
$$
For simplicity, we denote  $\overline{\mathrm{pred}} (\mathbf{x};\mathbf{w}):= \sum_{k = 1}^{K}u_k\overline{r} (\langle \mathbf{v}_k,\mathbf{x}\rangle)$  in the following analysis. Thus, letting the attention layer  $\pmb{\theta}_{\mathrm{attn}}^{(1)} = \{(\mathbf{V}_{k,m}^{(1)},\mathbf{Q}_{k,m}^{(1)},\mathbf{K}_{k,m}^{(1)})\}_{(k,m)}$  , we have
$$
\operatorname {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {(1)}}: \mathbf {h} _ {i} \mapsto \mathbf {h} _ {i} ^ {(0. 5)} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; \mathbf {w}; \overline {{\operatorname* {p r e d}}} (\mathbf {x} _ {i}; \mathbf {w}); \mathbf {0}; 1; t _ {i} ].
$$
Step 2: construction of  $\pmb{\theta}_{\mathrm{mlp}}^{(1)}$ . We pick matrices  $\mathbf{W}_1, \mathbf{W}_2$  so that  $\mathbf{W}_1$  maps
$$
\mathbf {W} _ {1} \mathbf {h} _ {i} ^ {(0. 5)} = \left[ \mathbf {a} _ {m} ^ {2} [ 1 ] \cdot \overline {{\mathrm {p r e d}}} (\mathbf {x} _ {i}; \mathbf {w}) + \mathbf {a} _ {m} ^ {2} [ 2 ] \cdot y _ {i} ^ {\prime} + \mathbf {a} _ {m} ^ {2} [ 3 ] - R _ {2} (1 - t _ {i}) \right] _ {m \in [ M _ {2} ]} \in \mathbb {R} ^ {M _ {2}},
$$
and  $\mathbf{W}_2\in \mathbb{R}^{D\times M_3}$  with entries being  $(\mathbf{W}_2)_{(j,m)} = c_m^2 1\{j = D_0 + 2\}$ . It is clear that  $\| \mathbf{W}_1\|_{\mathrm{op}}\leq R_2 + 1$ ,  $\| \mathbf{W}_2\|_{\mathrm{op}}\leq C_2$ . Then we have
$$
\begin{array}{l} \mathbf {W} _ {2} \sigma (\mathbf {W} _ {1} \mathbf {h} _ {i} ^ {(0. 5)}) = \sum_ {m \in [ M _ {3} ]} \sigma \big (\left\langle \mathbf {a} _ {m} ^ {2}, [ \overline {{\mathrm {p r e d}}} (\mathbf {x} _ {i}; \mathbf {w}); y _ {i} ^ {\prime}; 1 ] \right\rangle - R _ {2} (1 - t _ {j}) \big) \cdot c _ {m} ^ {2} \mathbf {e} _ {D _ {0} + 2} \\ = 1 \{t _ {j} = 1 \} \cdot g (\overline {{\mathrm {p r e d}}} (\mathbf {x} _ {i}; \mathbf {w}), y _ {i} ^ {\prime}) \cdot \mathbf {e} _ {D _ {0} + 2}. \\ \end{array}
$$
In the following, we abbreviate  $g_{i} = 1\{t_{j} = 1\} \cdot g(\overline{\mathrm{pred}} (\mathbf{x}_{i};\mathbf{w}),y_{i}^{\prime})$ . Hence,  $\theta_{\mathrm{m1p}}$  maps
$$
\mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}}}: \mathbf {h} _ {i} ^ {(0. 5)} \mapsto \mathbf {h} _ {i} ^ {(1)} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; \mathbf {w}; \overline {{\operatorname {p r e d}}} (\mathbf {x} _ {i}; \mathbf {w}); g _ {i}; \mathbf {0}; 1; t _ {i} ].
$$
By the definition of the function  $g$ , for each  $i \in [N]$
$$
\left| g _ {i} - \partial_ {1} \ell \left(\operatorname {p r e d} \left(\mathbf {x} _ {i}; \mathbf {w}\right), y _ {i}\right) \right| \leq \varepsilon_ {\ell} + B _ {u} L _ {\ell} \varepsilon_ {r},
$$
where  $L_{\ell} \coloneqq \max_{|t|\leq KB_r, |y|\leq B_y}\left|\partial_{tt}^2\ell(t,y)\right|$  is the smoothness of  $\partial_1\ell$ . Also,  $g_{N+1} = 0$  by definition.
Step 3: construction of  $\pmb{\theta}_{\mathrm{attn}}^{(2)}$ . We consider the matrices  $\{\mathbf{Q}_{k,1,m}^{(2)}, \mathbf{K}_{k,1,m}^{(2)}, \mathbf{V}_{k,1,m}^{(2)}\}_{k \in [K], m \in [M_3]}$  so that for all  $i, j \in [N + 1]$ , we have
$$
\mathbf {Q} _ {k, 1, m} ^ {(2)} \mathbf {h} _ {i} ^ {(1)} = \left[ \begin{array}{c} \mathbf {a} _ {m} ^ {3} [ 1 ] \cdot u _ {k} \\ \mathbf {a} _ {m} ^ {3} [ 2 ] \cdot \mathbf {v} _ {k} \\ \mathbf {a} _ {m} ^ {3} [ 3 ] \\ \mathbf {0} \end{array} \right], \qquad \mathbf {K} _ {k, 1, m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} g _ {j} \\ \mathbf {x} _ {j} \\ 1 \\ \mathbf {0} \end{array} \right], \qquad \mathbf {V} _ {k, 1, m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = - \frac {(N + 1) \eta c _ {m} ^ {3}}{N} \cdot \left[ \begin{array}{c} \mathbf {0} _ {k (d + 1)} \\ \mathbf {x} _ {j} \\ \mathbf {0} \end{array} \right].
$$
We further consider the matrices  $\{\mathbf{Q}_{k,2,m}^{(2)},\mathbf{K}_{k,2,m}^{(2)},\mathbf{V}_{k,2,m}^{(2)}\}_{k\in [K],m\in [M_1]}$  so that for all  $i,j\in [N + 1]$ , we have
$$
\mathbf {Q} _ {k, 2, m} ^ {(2)} \mathbf {h} _ {i} ^ {(1)} = \left[ \begin{array}{c} \mathbf {a} _ {m} ^ {1} [ 1 ] \cdot \mathbf {v} _ {k} \\ \mathbf {a} _ {m} ^ {1} [ 2 ] \\ \mathbf {0} \end{array} \right], \qquad \mathbf {K} _ {k, 2, m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = \left[ \begin{array}{c} \mathbf {x} _ {j} \\ 1 \\ \mathbf {0} \end{array} \right], \qquad \mathbf {V} _ {k, 2, m} ^ {(2)} \mathbf {h} _ {j} ^ {(1)} = - \frac {(N + 1) \eta c _ {m} ^ {1}}{N} \cdot \left[ \begin{array}{c} \mathbf {0} _ {k (d + 1) + d} \\ g _ {j} \\ \mathbf {0} \end{array} \right].
$$
By the structure of the input  $\mathbf{h}_i^{(1)}$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max _ {(k, w, m)} \left\| \mathbf {Q} _ {k, w, m} ^ {(2)} \right\| _ {\mathrm {o p}} \leq 1, \quad \max _ {(k, w, m)} \left\| \mathbf {K} _ {k, w, m} ^ {(2)} \right\| _ {\mathrm {o p}} \leq 1, \quad \sum_ {(k, w, m)} \left\| \mathbf {V} _ {k, w, m} ^ {(2)} \right\| _ {\mathrm {o p}} \leq 2 \eta C _ {1} + 2 \eta C _ {3}.
$$
Furthermore, a simple calculation shows that
$$
\mathbf {g} (\mathbf {w}) =: \frac {1}{N + 1} \sum_ {i = 1} ^ {N + 1} \sum_ {(k, w, m)} \sigma \Big (\Big \langle \mathbf {Q} _ {k, w, m} ^ {(2)} \mathbf {h} _ {i}, \mathbf {K} _ {k, w, m} ^ {(2)} \mathbf {h} _ {j} \Big \rangle \Big) \mathbf {V} _ {k, w, m} ^ {(2)} \mathbf {h} _ {j} = - \frac {\eta}{N} \sum_ {j = 1} ^ {N + 1} \left[ \begin{array}{c} \mathbf {0} _ {d + 1} \\ P (u _ {1} g _ {j}, \langle \mathbf {v} _ {1}, \mathbf {x} _ {j} \rangle) \cdot \mathbf {x} _ {j} \\ \overline {{r}} (\langle \mathbf {v} _ {1}, \mathbf {x} _ {j} \rangle) \cdot g _ {j} \\ \vdots \\ P (u _ {K} g _ {j}, \langle \mathbf {v} _ {K}, \mathbf {x} _ {j} \rangle) \cdot \mathbf {x} _ {j} \\ \overline {{r}} (\langle \mathbf {v} _ {K}, \mathbf {x} _ {j} \rangle) \cdot g _ {j} \\ \mathbf {0} \end{array} \right],
$$
where the summation is taken over all possibilities of the tuple  $(k, w, m)$ , i.e. over the union of  $[K] \times \{1\} \times [M_3]$  and  $[K] \times \{2\} \times [M_1]$ .
By our definition, we have  $|P(s,t) - sr'(t)| \leq \varepsilon_p$  for all  $s,t \in [-R_3,R_3]$ . Therefore, for each  $i \in [N]$ ,  $k \in [K]$ ,
$$
\begin{array}{l} | P (u _ {k} g _ {j}, \langle \mathbf {v} _ {k}, \mathbf {x} _ {j} \rangle) - \partial_ {1} \ell (\operatorname * {p r e d} (\mathbf {x} _ {j}; \mathbf {w}), y _ {j}) \cdot u _ {k} \cdot r ^ {\prime} (\langle \mathbf {v} _ {k}, \mathbf {x} _ {j} \rangle) | \leq \varepsilon_ {p} + | g _ {j} - \partial_ {1} \ell (\operatorname * {p r e d} (\mathbf {x} _ {i}; \mathbf {w}), y _ {i}) | \cdot | u _ {k} | \cdot | r ^ {\prime} (\langle \mathbf {v} _ {k}, \mathbf {x} _ {j} \rangle) | \\ \leq \varepsilon_ {p} + B _ {u} L _ {r} (\varepsilon_ {\ell} + B _ {u} L _ {\ell} \varepsilon_ {r}), \\ \end{array}
$$
where  $L_{r} \coloneqq \max_{|t| \leq B_{x}B_{u}} |r'(t)|$  is the upper bound of  $r'$ . Similarly, for each  $i \in [N]$ ,  $k \in [K]$ , we have
$$
\left| \bar {r} \left(\left\langle \mathbf {v} _ {k}, \mathbf {x} _ {j} \right\rangle\right) \cdot g _ {j} - r \left(\left\langle \mathbf {v} _ {k}, \mathbf {x} _ {j} \right\rangle\right) \cdot \partial_ {1} \ell \left(\operatorname {p r e d} \left(\mathbf {x} _ {j}; \mathbf {w}\right), y _ {j}\right) \right| \leq 2 B _ {g} \varepsilon_ {r} + 2 B _ {r} \left(\varepsilon_ {\ell} + B _ {u} L _ {\ell} ^ {2} \varepsilon_ {r}\right).
$$
As for the case  $i = N + 1$ , we have  $g_{N + 1} = 0$  and  $|P(u_kg_{N + 1},\langle \mathbf{v}_k,\mathbf{x}_{N + 1}\rangle)|\leq \varepsilon_p$  for each  $k\in [K]$  by definition. Combining these estimations and using (16) and (17), we can conclude that