SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Definition D.2 (Approximability by MLP). An operator  $P: \mathbb{R}^d \to \mathbb{R}^d$  is  $(\varepsilon, R, D, C)$ -approximable by MLP, if there exists a there exists a MLP  $\pmb{\theta}_{\mathrm{mlp}} = (\mathbf{W}_1, \mathbf{W}_2) \in \mathbb{R}^{D \times d} \times \mathbb{R}^{d \times D}$  with hidden dimension  $D$ ,  $\| \mathbf{W}_1 \|_{\mathrm{op}} + \| \mathbf{W}_2 \|_{\mathrm{op}} \leq C'$ , such that  $\sup_{\| \mathbf{w} \|_2 \leq R} \| P(\mathbf{w}) - \mathrm{MLP}_{\pmb{\theta}_{\mathrm{mlp}}}(\mathbf{w}) \|_2 \leq \varepsilon$ .
The definition above captures the proximal operator  $\mathbf{prox}_{\eta \mathcal{R}}$  for a broad class of regularizers, such as the (commonly-used)  $L_{1}$  and  $L_{2}$  regularizer listed in the following proposition, for all of which one can directly check that they can be exactly implemented by an MLP as stated below.
Proposition D.1 (Proximal operators for commonly-used regularizers). For regularizer  $\mathcal{R}$  in  $\{\lambda \| \cdot \| _1,\frac{\lambda}{2}\| \cdot \| _2^2,\mathbb{I}_{B_\infty (B)}(\cdot)\}$ , the operator  $\mathbf{prox}_{\eta \mathcal{R}}:\mathbb{R}^d\to \mathbb{R}^d$  is exactly approximable by MLP. More concretely, we have
1. For  $\mathcal{R} = \lambda \| \cdot \|_1$ ,  $\mathbf{prox}_{\eta \mathcal{R}}$  is  $(0, +\infty, 4d, 4 + 2\eta \lambda)$ -approximable by MLP.
2. For  $\mathcal{R} = \frac{\lambda}{2}\| \cdot \| _2^2$ ,  $\mathbf{prox}_{\eta \mathcal{R}}$  is  $(0, + \infty ,2d,2 + 2\eta \lambda)$ -approximable by MLP.
3. For  $\mathcal{R} = \mathbb{I}_{\mathsf{B}_{\infty}(B)}(\cdot)$ ,  $\mathbf{prox}_{\eta \mathcal{R}} = \mathrm{Proj}_{\mathsf{B}_{\infty}(B)}$  is  $(0, +\infty, 2d, 2 + 2B)$ -approximable by MLP.
Theorem D.2 (Convex ICPGD). Fix any  $B_w > 0$ ,  $L > 1$ ,  $\eta > 0$ , and  $\varepsilon + \varepsilon' \leq B_w / (2L)$ . Suppose that
1. The loss  $\ell (\cdot ,\cdot)$  is convex in the first argument;
2.  $\partial_s\ell$  is  $(\varepsilon ,R,M,C)$  -approximable by sum of relus with  $R = \max \{B_xB_w,B_y,1\}$
3.  $\mathcal{R}$  convex, and the proximal operator  $\mathbf{prox}_{\eta \mathcal{R}}(\mathbf{w})$  is  $(\eta \varepsilon', R', D', C')$ -approximable by MLP with  $R' = \sup_{\| \mathbf{w} \|_2 \leq B_w} \left\| \mathbf{w}_\eta^+ \right\|_2 + \eta \varepsilon$ .
Then there exists a transformer  $\mathrm{TF}_{\theta}$  with  $(L + 1)$  layers,  $\max_{\ell \in [L]}M^{(\ell)}\leq M$  heads within the first  $L$  layers,  $M^{(L + 1)} = 2$ , and hidden dimension  $D^{\prime}$  such that, for any input data  $(\mathcal{D},\mathbf{x}_{N + 1})$  such that
$$
\sup_{\| \mathbf{w}\|_{2}\leq B_{w}}\lambda_{\max}(\nabla^{2}\widehat{L}_{N}(\mathbf{w}))\leq 2 / \eta ,\qquad \exists \mathbf{w}^{\star}\in \operatorname *{arg  min}_{\mathbf{w}\in \mathbb{R}^{d}}\widehat{L}_{N}(\mathbf{w}) \text{such that}\| \mathbf{w}^{\star}\|_{2}\leq B_{w} / 2,
$$
$\mathrm{TF}_{\theta}(\mathbf{H}^{(0)})$  approximately implements (ICGD):
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): 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 {P G D}} ^ {\ell} \right\| _ {2} \leq \left(\varepsilon + \varepsilon^ {\prime}\right) \cdot \left(L \eta B _ {x}\right).
$$
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): 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;t_i]$ , where  $\widehat{y}_{N + 1} = \left\langle \widehat{\mathbf{w}}^L,\mathbf{x}_{N + 1}\right\rangle$  so that
$$
\left| \widehat {y} _ {N + 1} - \left\langle \mathbf {w} _ {\mathrm {P G D}} ^ {L}, \mathbf {x} _ {N + 1} \right\rangle \right| \leq (\varepsilon + \varepsilon^ {\prime}) \cdot (2 L \eta B _ {x} ^ {2}).
$$
Further, the weight matrices have norm bounds  $\| \pmb {\theta}\| \leq 3 + R + 2\eta C + C^{\prime}$
The proof of Theorem D.2 is essentially similar to the proof of Theorem D.1, using the following generalized version of Lemma D.1.
Lemma D.2 (Composition of error for approximating convex PGD). Suppose  $f: \mathbb{R}^d \to \mathbb{R}$  is a convex function and  $\mathcal{R}$  is a convex regularizer. Let  $\mathbf{w}^\star \in \arg \min_{\mathbf{w} \in \mathbb{R}^d} f(\mathbf{w}) + \mathcal{R}(\mathbf{w})$ ,  $R \geq 2\|\mathbf{w}^\star\|_2$ , and assume that  $\nabla f$  is  $L_f$ -smooth on  $\mathsf{B}_2^d(R)$ . Let sequences  $\{\widehat{\mathbf{w}}^\ell\}_{\ell \geq 0} \subset \mathbb{R}^d$  and  $\{\mathbf{w}_{\mathrm{GD}}^\ell\}_{\ell \geq 0} \subset \mathbb{R}^d$  be given by  $\widehat{\mathbf{w}}^0 = \mathbf{w}_{\mathrm{GD}}^0 = \mathbf{0}$ ,
$$
\left\{ \begin{array}{l} \widehat {\mathbf {w}} ^ {\ell + 1} = \mathbf {p r o x} _ {\eta \mathcal {R}} \big (\widehat {\mathbf {w}} ^ {\ell} - \eta \nabla f (\widehat {\mathbf {w}} ^ {\ell}) \big) + \varepsilon^ {\ell}, \qquad \| \varepsilon^ {\ell} \| _ {2} \leq \varepsilon , \\ \mathbf {w} _ {\mathrm {G D}} ^ {\ell + 1} = \mathbf {p r o x} _ {\eta \mathcal {R}} \big (\mathbf {w} _ {\mathrm {G D}} ^ {\ell} - \eta \nabla f (\mathbf {w} _ {\mathrm {G D}} ^ {\ell}) \big), \end{array} \right.
$$
The proof of the above lemma is done by utilizing the non-expansiveness of the PGD operator  $\mathbf{w} \mapsto \mathbf{prox}_{\eta \mathcal{R}}(\mathbf{w} - \eta \nabla f(\mathbf{w}))$  and otherwise following the same arguments as for Lemma D.1.
D.3 Gradient descent on two-layer neural networks
We now move beyond the convex setting by showing that transformers can implement gradient descent on two-layer neural networks in context.
Suppose that the prediction function  $\operatorname{pred}(\mathbf{x};\mathbf{w}) := \sum_{k=1}^{K} u_k r(\mathbf{v}_k^\top \mathbf{x})$  is given by a two-layer neural network, parameterized by  $\mathbf{w} = [\mathbf{v}_k;u_k]_{k\in [K]} \in \mathbb{R}^{K(d+1)}$ . Consider the empirical risk minimization problem:
$$
\min  _ {\mathbf {w} \in \mathcal {W}} \widehat {L} _ {N} (\mathbf {w}) := \frac {1}{2 N} \sum_ {i = 1} ^ {N} \ell \left(\operatorname {p r e d} \left(\mathbf {x} _ {i}; \mathbf {w}\right), y _ {i}\right) = \frac {1}{2 N} \sum_ {i = 1} ^ {N} \ell \left(\sum_ {k = 1} ^ {K} u _ {k} r \left(\mathbf {v} _ {k} ^ {\top} \mathbf {x} _ {i}\right), y _ {i}\right), \tag {13}
$$
where  $\mathcal{W}$  is a bounded domain. For the sake of simplicity, in the following discussion we assume that  $\mathrm{Proj}_{\mathcal{W}}$  can be exactly implemented by a MLP layer (e.g.  $\mathcal{W} = \mathsf{B}_{\infty}(R_w)$  for some  $R_{w} > 0$ ).
Theorem D.3 (Approximate ICGD on two-layer NNs). Fix any  $B_v, B_u > 0$ ,  $L \geq 1$ ,  $\eta > 0$ , and  $\varepsilon > 0$ . Suppose that
1. Both the activation function  $r$  and the loss function  $\ell$  is  $C^4$ -smooth;
2.  $\mathcal{W}$  is a closed domain such that  $\mathcal{W} \subset \{\mathbf{w} = [\mathbf{v}_k; u_k]_{k \in [K]} \in \mathbb{R}^{K(d + 1)} : \| \mathbf{v}_k \|_2 \leq B_v, |u_k| \leq B_u\}$ , and  $\mathrm{Proj}_{\mathcal{W}} = \mathrm{MLP}_{\theta_{\mathrm{mlp}}}$  for some MLP layer  $\theta_{\mathrm{mlp}}$  with hidden dimension  $D_w$  and  $\left\| \theta_{\mathrm{mlp}} \right\| \leq C_w$ ;
Then there exists a  $(2L)$ -layer transformer  $\mathrm{TF}_{\theta}$  with
$$
\max _ {\ell \in [ 2 L ]} M ^ {(\ell)} \leq \widetilde {\mathcal {O}} (\varepsilon^ {- 2}), \quad \max _ {\ell \in [ 2 L ]} D ^ {(\ell)} \leq \widetilde {\mathcal {O}} (\varepsilon^ {- 2}) + D _ {w}, \quad \| \pmb {\theta} \| \leq \mathcal {O} (1 + \eta) + C _ {w},
$$
where  $\mathcal{O}(\cdot)$  hides the constants that depend on  $K$ , the radius parameters  $B_{x}, B_{y}, B_{u}, B_{v}$  and the smoothness of  $r$  and  $\ell$ , such that for any input data  $(\mathcal{D}, \mathbf{x}_{N+1})$  such that input sequence  $\mathbf{H}^{(0)} \in \mathbb{R}^{D \times (N+1)}$  takes form (3),  $\mathrm{TF}_{\theta}(\mathbf{H}^{(0)})$  approximately implements in-context gradient descent on risk (13): For every  $\ell \in [L]$ , the 2 $\ell$ -th layer's output  $\mathbf{h}_i^{(2\ell)} = [\mathbf{x}_i; y_i'; \widehat{\mathbf{w}}^{\ell}; \mathbf{0}; 1; t_i]$  for every  $i \in [N+1]$ , and
$$
\widehat {\mathbf {w}} ^ {\ell} = \operatorname {P r o j} _ {\mathcal {W}} \left(\widehat {\mathbf {w}} ^ {\ell - 1} - \eta \left(\nabla \widehat {L} _ {N} \left(\widehat {\mathbf {w}} ^ {\ell - 1}\right) + \varepsilon^ {\ell - 1}\right)\right), \quad \widehat {\mathbf {w}} ^ {0} = \mathbf {0}, \tag {14}
$$
where  $\left\| \varepsilon^{\ell -1}\right\| _2\leq \varepsilon$  is an error term.
As a direct corollary, the transformer constructed above can approximate the true gradient descent trajectory  $\{\mathbf{w}_{\mathrm{GD}}^{\ell}\}_{\ell \geq 0}$  on (16), defined as  $\mathbf{w}_{\mathrm{GD}}^{0} = \mathbf{0}$  and  $\mathbf{w}_{\mathrm{GD}}^{\ell +1} = \mathbf{w}_{\mathrm{GD}}^{\ell} - \eta \nabla \widehat{L}_N(\mathbf{w}_{\mathrm{GD}}^{\ell})$  for all  $\ell \geq 0$ .
Corollary D.1 (Approximating multi-step ICGD on two-layer NNs). For any  $L \geq 1$ , under the same setting as Theorem D.3, the  $(2L)$ -layer transformer  $\mathrm{TF}_{\theta}$  there approximates the true gradient descent trajectory  $\{\mathbf{w}_{\mathrm{GD}}^{\ell}\}_{\ell \geq 0}$ : For the intermediate iterates  $\{\widehat{\mathbf{w}}^{\ell}\}_{\ell \in [L]}$  considered therein, we have
$$
\left\| \widehat {\mathbf {w}} ^ {\ell} - \mathbf {w} _ {\mathrm {G D}} ^ {\ell} \right\| _ {2} \leq L _ {f} ^ {- 1} (1 + \eta L _ {f}) ^ {\ell} \varepsilon ,
$$
where  $L_{f} = \sup_{\mathbf{w}\in \mathcal{W}}\left\| \nabla^{2}\widehat{L}_{N}(\mathbf{w})\right\|_{\mathrm{op}}$  denotes the smoothness of  $\widehat{L}_N$  within  $\mathcal{W}$
Remark on error accumulation Note that in Corollary D.1, the error accumulates exponentially in  $\ell$  rather than linearly as in Theorem D.1. This is as expected, since gradient descent on non-convex objectives is inherently unstable at a high level (a slight error added upon each step may result in a drastically different trajectories); technically, this happens as the stability-like property Lemma D.1 no longer holds for the non-convex case.
Corollary D.1 is a simple implication of Theorem D.3 and Part (b) of the following convergence and trajectory closeness result for inexact gradient descent. For any closed convex set  $\mathcal{W} \subset \mathbb{R}^d$ , any function  $f: \mathcal{W} \to \mathbb{R}$ , and any initial point  $\mathbf{w} \in \mathcal{W}$ , let
$$
\mathrm {G} _ {\mathcal {W}, \eta} ^ {f} (\mathbf {w}) := \frac {\mathbf {w} - \operatorname {P r o j} _ {\mathcal {W}} (\mathbf {w} - \eta \nabla f (\mathbf {w}))}{\eta}
$$
denote the gradient mapping at  $\mathbf{w}$  with step size  $\eta$ , a standard measure of stationarity in constrained optimization [63]. Note that  $\mathsf{G}_{\mathcal{W},\eta}^{f}(\mathbf{w}) = \nabla f(\mathbf{w})$  when  $\mathbf{w} - \eta \nabla f(\mathbf{w}) \in \mathcal{W}$  (so that the projection does not take effect).
Lemma D.3 (Convergence and trajectory closeness of inexact GD). Suppose  $f: \mathcal{W} \to \mathbb{R}$ , where  $\mathcal{W} \subset \mathbb{R}^d$  is a convex closed domain and  $\nabla f$  is  $L_f$ -Lipschitz on  $\mathcal{W}$ . Let sequence  $\{\widehat{\mathbf{w}}^\ell\}_{\ell \geq 0} \subset \mathbb{R}^d$  be given by  $\widehat{\mathbf{w}}^0 = \mathbf{w}^0$ ,
$$
\widehat {\mathbf {w}} ^ {\ell + 1} = \operatorname {P r o j} _ {\mathcal {W}} \left(\widehat {\mathbf {w}} ^ {\ell} - \eta (\nabla f (\widehat {\mathbf {w}} ^ {\ell}) + \varepsilon^ {\ell})\right), \quad \| \varepsilon^ {\ell} \| _ {2} \leq \varepsilon ,
$$
for all  $\ell \geq 0$ . Then the following holds.
(a) As long as  $\eta \leq 1 / L_f$ , for all  $L \geq 1$
$$
\min  _ {\ell \in [ L - 1 ]} \left\| \mathsf {G} _ {\mathcal {W}, \eta} ^ {f} (\widehat {\mathbf {w}} ^ {\ell}) \right\| _ {2} ^ {2} \leq \frac {1}{L} \sum_ {\ell = 0} ^ {L - 1} \left\| \mathsf {G} _ {\mathcal {W}, \eta} ^ {f} (\widehat {\mathbf {w}} ^ {\ell}) \right\| _ {2} ^ {2} \leq \frac {8 (f (\mathbf {w} ^ {0}) - \inf  _ {\mathbf {w} \in \mathcal {W}} f (\mathbf {w}))}{\eta L} + 1 0 \varepsilon^ {2}.
$$
(b) Let the sequences  $\{\mathbf{w}_{\mathrm{GD}}^{\ell}\}_{\ell \geq 0} \subset \mathbb{R}^{d}$  and be given by  $\mathbf{w}_{\mathrm{GD}}^{0} = \mathbf{w}^{0}$  and  $\mathbf{w}_{\mathrm{GD}}^{\ell +1} = \operatorname{Proj}_{\mathcal{W}}(\mathbf{w}_{\mathrm{GD}}^{\ell} - \eta \nabla f(\mathbf{w}_{\mathrm{GD}}^{\ell}))$ . Then it holds that
$$
\left\| \widehat {\mathbf {w}} ^ {\ell} - \mathbf {w} _ {\mathrm {G D}} ^ {\ell} \right\| _ {2} \leq L _ {f} ^ {- 1} (1 + \eta L _ {f}) ^ {\ell} \varepsilon , \quad \forall \ell \geq 0.
$$
E Proofs for Section D
E.1 Approximating a single GD step
Proposition E.1 (Approximating a single GD step by a single attention layer). Let  $\ell(\cdot, \cdot): \mathbb{R}^2 \to \mathbb{R}$  be a loss function such that  $\partial_1\ell$  is  $(\varepsilon, R, M, C)$ -approximable by sum of relus with  $R = \max\{B_xB_w, B_y, 1\}$ . Let  $\widehat{L}_N(\mathbf{w}) \coloneqq \frac{1}{N}\sum_{i=1}^{N}\ell(\mathbf{w}^\top\mathbf{x}_i, y_i)$  denote the empirical risk with loss function  $\ell$  on dataset  $\{(\mathbf{x}_i, y_i)\}_{i \in [N]}$ .
Then, for any  $\varepsilon > 0$ , there exists an attention layer  $\pmb{\theta} = \{(\mathbf{Q}_m, \mathbf{K}_m, \mathbf{V}_m)\}_{m \in [M]}$  with  $M$  heads such that, for any input sequence that takes form  $\mathbf{h}_i = [\mathbf{x}_i; y_i'; \mathbf{w}; \mathbf{0}_{D-2d-3}; 1; t_i]$  with  $\|\mathbf{w}\|_2 \leq B_w$ , it gives output  $\widetilde{\mathbf{h}}_i = [\mathrm{Attn}_{\pmb{\theta}}(\mathbf{H})]_i = [\mathbf{x}_i; y_i'; \widetilde{\mathbf{w}}; \mathbf{0}_{D-2d-3}; 1; t_i]$  for all  $i \in [N+1]$ , where
$$
\left\| \widetilde {\mathbf {w}} - (\mathbf {w} - \eta \nabla \widehat {L} _ {N} (\mathbf {w})) \right\| _ {2} \leq \varepsilon \cdot (\eta B _ {x}).
$$
Further,  $\| \pmb {\theta}\| \leq 2 + R + 2\eta C$
Proof of Proposition E.1. As  $\partial_s\ell$  is  $(\varepsilon, R, M, C)$ -approximable by sum of relus, there exists a function  $f: [-R, R]^2 \to \mathbb{R}$  of form
$$
f (s, t) = \sum_ {m = 1} ^ {M} c _ {m} \sigma \left(a _ {m} s + b _ {m} t + d _ {m}\right) \text {w i t h} \sum_ {m = 1} ^ {M} | c _ {m} | \leq C, | a _ {m} | + | b _ {m} | + | d _ {m} | \leq 1, \forall m \in [ M ],
$$
such that  $\sup_{(s,t)\in [-R,R]^2}|f(s,t) - \partial_s\ell (s,t)|\leq \varepsilon$
Next, for every  $m\in [M]$ , we define matrices  $\mathbf{Q}_m,\mathbf{K}_m,\mathbf{V}_m\in \mathbb{R}^{D\times D}$  such that
$$
\mathbf {Q} _ {m} \mathbf {h} _ {i} = \left[ \begin{array}{c} a _ {m} \mathbf {w} \\ b _ {m} \\ d _ {m} \\ - 2 \\ \mathbf {0} \end{array} \right], \quad \mathbf {K} _ {m} \mathbf {h} _ {j} = \left[ \begin{array}{c} \mathbf {x} _ {j} \\ y _ {j} ^ {\prime} \\ 1 \\ R (1 - t _ {j}) \\ \mathbf {0} \end{array} \right], \quad \mathbf {V} _ {m} \mathbf {h} _ {j} = - \frac {(N + 1) \eta c _ {m}}{N} \cdot \left[ \begin{array}{c} \mathbf {0} _ {d} \\ 0 \\ \mathbf {x} _ {j} \\ \mathbf {0} _ {D - 2 d - 1} \end{array} \right]
$$
for all  $i,j\in [N + 1]$ . As the input has structure  $\mathbf{h}_i = [\mathbf{x}_i;y_i';\mathbf{w};\mathbf{0}_{D - 2d - 3};1;t_i]$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max  _ {m \in [ M ]} \| \mathbf {Q} _ {m} \| _ {\mathrm {o p}} \leq 3, \quad \max  _ {m \in [ M ]} \| \mathbf {K} _ {m} \| _ {\mathrm {o p}} \leq 2 + R, \quad \sum_ {m \in [ M ]} \| \mathbf {V} _ {m} \| _ {\mathrm {o p}} \leq 2 \eta C.
$$
Consequently,  $\| \pmb {\theta}\| \leq 2 + R + 2\eta C$
Now, for every  $i,j\in [N + 1]$  , we have
$$
\begin{array}{l} \sigma \left(\left\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \right\rangle\right) = \sigma \left(a _ {m} \mathbf {w} ^ {\top} \mathbf {x} _ {j} + b _ {m} (1 - t _ {j}) y _ {j} + d _ {m} - 2 R t _ {j}\right) \\ = \sigma \left(a _ {m} \mathbf {w} ^ {\top} \mathbf {x} _ {j} + b _ {m} y _ {j} + d _ {m}\right) 1 \left\{t _ {j} = 1 \right\}, \\ \end{array}
$$
where the last equality follows from the bound
$$
\left| a _ {m} \mathbf {w} ^ {\top} \mathbf {x} _ {j} + b _ {m} \left(1 - t _ {j}\right) y _ {j} + d _ {m} \right| \leq \left| a _ {m} \right| B _ {x} B _ {w} + R \leq 2 R, \tag {15}
$$
so that the above relu equals 0 if  $t_j \leq 0$ . Therefore,
$$
\begin{array}{l} \sum_ {m = 1} ^ {M} \sigma (\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \rangle) \mathbf {V} _ {m} \mathbf {h} _ {j} \\ = \left(\sum_ {m = 1} ^ {M} c _ {m} \sigma \left(a _ {m} \mathbf {w} ^ {\top} \mathbf {x} _ {j} + b _ {m} y _ {j} + d _ {m}\right)\right) \cdot \frac {- (N + 1) \eta}{N} 1 \{t _ {j} = 0 \} [ \mathbf {0} _ {d + 1}; \mathbf {x} _ {j}; \mathbf {0} _ {2} ] \\ = f \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j}, y _ {j}\right) \cdot \frac {- (N + 1) \eta}{N} 1 \{t _ {j} = 0 \} \left[ \mathbf {0} _ {d + 1}; \mathbf {x} _ {j}; \mathbf {0} _ {D - 2 d - 1} \right]. \\ \end{array}
$$
Thus letting the attention layer  $\pmb{\theta} = \{(\mathbf{V}_m, \mathbf{Q}_m, \mathbf{K}_m)\}_{m \in [M]}$ , we have
$$
\begin{array}{l} \widetilde {\mathbf {h}} _ {i} = \left[ \operatorname {A t t n} _ {\boldsymbol {\theta}} (\mathbf {H}) \right] _ {i} = \mathbf {h} _ {i} + \frac {1}{N + 1} \sum_ {j = 1} ^ {N + 1} \sum_ {m = 1} ^ {M} \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} \\ = \mathbf {h} _ {i} - \frac {\eta}{N} \sum_ {j = 1} ^ {N} f \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j}, y _ {j}\right) \left[ \mathbf {0} _ {d + 1}; \mathbf {x} _ {j}; \mathbf {0} _ {2} \right] \\ = \left[ \mathbf {x} _ {i}; y _ {i}; \mathbf {w}; 1; t _ {i} \right] \underbrace {- \frac {\eta}{N} \sum_ {j = 1} ^ {N} \partial_ {s} \ell \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j} , y _ {j}\right) \left[ \mathbf {0} _ {d + 1} ; \mathbf {x} _ {j} ; \mathbf {0} _ {D - 2 d - 1} \right]} _ {\left[ \mathbf {0} _ {d + 1}; - \eta \nabla \widehat {L} _ {N} (\mathbf {w}); \mathbf {0} _ {D - 2 d - 1} \right]} + \left[ \mathbf {0} _ {d + 1}; \varepsilon ; \mathbf {0} _ {D - 2 d - 1} \right] \\ = \left[ \mathbf {x} _ {i}; y _ {i}; \mathbf {w} _ {\eta} ^ {+} + \varepsilon ; \mathbf {0} _ {D - 2 d - 3}; 1; t _ {i} \right], \\ \end{array}
$$
where the error vector  $\varepsilon \in \mathbb{R}^d$  satisfies
$$
\left\| \boldsymbol {\varepsilon} \right\| _ {2} = \left\| - \frac {\eta}{N} \sum_ {j = 1} ^ {N} \left(f \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j}, y _ {j}\right) - \partial_ {s} \ell \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j}, y _ {j}\right)\right) \mathbf {x} _ {j} \right\| _ {2}
$$
$$
\begin{array}{l} \leq \frac {\eta}{N} \sum_ {j = 1} ^ {N} \left| f \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j}, y _ {j}\right) - \partial_ {s} \ell \left(\mathbf {w} ^ {\top} \mathbf {x} _ {j}, y _ {j}\right) \right| \cdot \| \mathbf {x} _ {j} \| _ {2} \\ \leq \frac {\eta}{N} \cdot N \cdot \varepsilon \cdot B _ {x} = \varepsilon \cdot (\eta B _ {x}). \\ \end{array}
$$
E.2 Proof of Theorem D.1
We first prove part (a), which requires constructing the first  $L$  layers of  $\theta$ . Note that by our precondition  $L \leq B_w / (2\varepsilon)$ .
By our precondition, the partial derivative of the loss  $\partial_s\ell$  is  $(\varepsilon ,R,M,C)$ -approximable by sum of relus. Therefore we can apply Proposition E.1 to obtain that, there exists a single attention layer  $\pmb{\theta}^{(1)} = \{(\mathbf{Q}_m,\mathbf{K}_m,\mathbf{V}_m)\}_{m\in [M]}$  with  $M$  heads (and norm bounds specified in Proposition E.1), such that for any  $\mathbf{w}$  with  $\| \mathbf{w}\| _2\leq B_w$ , the attention layer Attn  $\pmb{\theta}^{(1)}$  maps the input  $\mathbf{h}_i = [\mathbf{x}_i;y_i';\mathbf{w};\mathbf{0}_{D - 2d - 3};1;t_i]$  to output  $\mathbf{h}_i' = [\mathbf{x}_i;y_i';\widehat{\mathbf{w}};\mathbf{0}_{D - 2d - 3};1;t_i]$  for all  $i\in [N + 1]$ , where
$$
\left\| \widehat {\mathbf {w}} - \left(\mathbf {w} - \eta \nabla \widehat {L} _ {N} (\mathbf {w})\right) \right\| _ {2} \leq \varepsilon \cdot \left(\eta B _ {x}\right) =: \varepsilon^ {\prime}.
$$
Consider the  $L$ -layer transformer  $\pmb{\theta}^{1:L} = (\pmb{\theta}^{(1)}, \dots, \pmb{\theta}^{(1)})$  which stacks the same attention layer  $\pmb{\theta}^{(1)}$  for  $L$  times, and for the given input  $\mathbf{h}_i^{(0)} = [\mathbf{x}_i; y_i'; \mathbf{w}^0; \mathbf{0}_{D-2d-3}; 1; t_i]$ , its  $\ell$ -th layer's output  $\mathbf{h}_i^{(\ell)} = [\mathbf{x}_i; y_i'; \widehat{\mathbf{w}}^\ell; \mathbf{0}_{D-2d-3}; 1; t_i]$ .
We now inductively show that  $\left\| \widehat{\mathbf{w}}^{\ell}\right\|_{2} \leq B_{w}$  and  $\left\| \widehat{\mathbf{w}}^{\ell} - \mathbf{w}_{\mathrm{GD}}^{\ell}\right\|_{2} \leq \ell \varepsilon$  for all  $\ell \in [L]$ . The base case of  $\ell = 0$  is trivial. Suppose the claim holds for  $\ell$ . Then for  $\ell + 1 \leq L \leq B_{w} / (2\varepsilon)$ , the sequence  $\{\widehat{\mathbf{w}}^{i}\}_{i \leq \ell + 1}$  and  $\{\mathbf{w}_{\mathrm{GD}}^{i}\}_{i \leq \ell + 1}$  satisfies the precondition of the error composition lemma (Lemma D.1) with error bound  $\varepsilon$ , from which we obtain  $\left\| \widehat{\mathbf{w}}^{\ell + 1}\right\|_{2} \leq B_{w}$  and
$$
\left\| \widehat {\mathbf {w}} ^ {\ell + 1} - \mathbf {w} _ {\mathrm {G D}} ^ {\ell + 1} \right\| _ {2} \leq (\ell + 1) \varepsilon^ {\prime}.
$$
This finishes the induction, and gives the following approximation guarantee for all  $\ell \in [L]$ :
$$
\left\| \widehat {\mathbf {w}} ^ {\ell} - \mathbf {w} _ {\mathrm {G D}} ^ {\ell} \right\| _ {2} \leq \ell \varepsilon^ {\prime} \leq \varepsilon \cdot (L \eta B _ {x}),
$$
which proves part (a).
We now prove part (b), which requires constructing the last attention layer  $\pmb{\theta}^{(L + 1)}$ . Recall  $\mathbf{h}_i^{(L)} = [\mathbf{x}_i;y_i';\hat{\mathbf{w}}^L;\mathbf{0}_{D - 2d - 3};1;t_i]$  for all  $i\in [N + 1]$ . We construct a 2-head attention layer  $\pmb{\theta}^{(L + 1)} = \{(\mathbf{Q}_m^{(L + 1)},\mathbf{K}_m^{(L + 1)},\mathbf{V}_m^{(L + 1)})\}_{m = 1,2}$  such that for every  $i,j\in [N + 1]$ ,
$$
\mathbf {Q} _ {1} ^ {(L + 1)} \mathbf {h} _ {i} ^ {(L)} = [ \mathbf {x} _ {i}; \mathbf {0} _ {D - d} ], \quad \mathbf {K} _ {1} ^ {(L + 1)} \mathbf {h} _ {j} ^ {(L)} = [ \widehat {\mathbf {w}} ^ {L}; \mathbf {0} _ {D - d} ], \quad \mathbf {V} _ {1} ^ {(L + 1)} \mathbf {h} _ {j} ^ {(L)} = [ \mathbf {0} _ {d}; 1; \mathbf {0} _ {D - d - 1} ],
$$
$$
\mathbf {Q} _ {2} ^ {(L + 1)} \mathbf {h} _ {i} ^ {(L)} = [ \mathbf {x} _ {i}; \mathbf {0} _ {D - d} ], \quad \mathbf {K} _ {2} ^ {(L + 1)} \mathbf {h} _ {j} ^ {(L)} = [ - \widehat {\mathbf {w}} ^ {L}; \mathbf {0} _ {D - d} ], \quad \mathbf {V} _ {2} ^ {(L + 1)} \mathbf {h} _ {j} ^ {(L)} = [ \mathbf {0} _ {d}; - 1; \mathbf {0} _ {D - d - 1} ].
$$
Note that the weight matrices have norm bound
$$
\max _ {i = 1, 2} \left\| \mathbf {Q} _ {i} ^ {(L + 1)} \right\| _ {\mathrm {o p}} \leq 1, \quad \max _ {i = 1, 2} \left\| \mathbf {K} _ {i} ^ {(L + 1)} \right\| _ {\mathrm {o p}} \leq 1, \quad \sum_ {i = 1} ^ {2} \left\| \mathbf {V} _ {i} ^ {(L + 1)} \right\| _ {\mathrm {o p}} \leq 2.
$$