SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where the last inequality is because the marginal distribution of  $y_{N + 1}$  is simply  $\mathsf{N}(0,\sigma^2 +\| \mathbf{w}_\star \| _2^2)$ . Combining these yields
$$
\mathbb {E} \left[ 1 \{\mathcal {E} ^ {c} \} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \mid \mathcal {D} \right] \leq \mathcal {O} \left(\frac {\delta_ {1} ^ {5} B _ {y} ^ {2}}{N ^ {5}}\right) \leq \mathcal {O} \left(\frac {\delta_ {1} ^ {5} \left(\left(B _ {w} ^ {\star}\right) ^ {2} + \sigma^ {2}\right) \log \left(1 / \delta_ {1}\right)}{N ^ {4}}\right).
$$
Therefore, choosing  $\delta_1 = \min \{\delta, \frac{\sigma}{B_w^*}\}$  is enough for our purpose, and under such choice of  $\delta_1$ ,
$$
\mathbb {E} \left[ 1 \{\mathcal {E} ^ {c} \} (\hat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \mid \mathcal {D} \right] \leq \mathcal {O} \left(\frac {\sigma^ {2}}{N ^ {4}}\right).
$$
Conclusion. Combining the inequalities above, we can conclude that under  $\mathcal{E}_0$
$$
\mathbb {E} \left[ \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \mid \mathcal {D} \right] \leq \sigma^ {2} \left[ 1 + \mathcal {O} \left(\frac {s \log (d / \delta)}{N}\right) \right].
$$
It remains to show that  $\mathbb{P}(\mathcal{E}_0) \geq 1 - \delta$ . By Proposition H.2, Lemma B.2 and Lemma B.4, we have
$$
\mathbb {P} (\mathcal {E} _ {w}) \leq 3 \exp (- N / 3 2), \qquad \mathbb {P} (\mathcal {E} _ {r}) \leq \frac {\delta}{2}, \qquad \mathbb {P} (\mathcal {E} _ {b}) \leq \frac {\delta}{4}.
$$
Therefore, as long as  $N \geq 32\log(12/\delta)$ , we have  $\mathbb{P}(\mathcal{E}_0) \geq 1 - \delta$ . This completes the proof.
We also remark that in the construction above,
$$
R = \mathcal {O} \left(\left(B _ {w} ^ {\star} + \sigma\right) \sqrt {d} \log \left(N \cdot \left(1 + B _ {w} ^ {\star} / \sigma\right)\right)\right),
$$
which would be useful for bounding  $\| \theta \|$
I Proofs for Section 4
I.1 Proof of Proposition 10
We begin by restating Proposition 10 into the following version, which contains additional size bounds on  $\theta$ .
Theorem I.1 (Full statement of Proposition 10). Suppose that for
$$
\widehat {L} _ {\text {v a l}} (f) := \frac {1}{| \mathcal {D} _ {\text {v a l}} |} \sum_ {(\mathbf {x} _ {i}, y _ {i}) \in \mathcal {D} _ {\text {v a l}}} \ell (f (\mathbf {x} _ {i}), y _ {i}),
$$
$\ell (\cdot ,\cdot)$  is  $(\gamma /3,R,M,C)$ -approximable by sum of relus (Definition D.1). Then there exists a 3-layer transformer  $\mathrm{TF}_{\theta}$  with
$$
\max _ {\ell \in [ 3 ]} M ^ {(\ell)} \leq (M + 3) K, \quad \max _ {\ell \in [ 3 ]} D ^ {(\ell)} \leq K ^ {2} + K + 1, \quad \| \boldsymbol {\theta} \| \leq \frac {2 N K C}{| \mathcal {D} _ {\mathrm {v a l}} |} + 3 \gamma^ {- 1} + 7 K R.
$$
that maps
$$
\mathbf {h} _ {i} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \mathbf {0} _ {K + 1}; 1; t _ {i} ] \quad \rightarrow \quad \mathbf {h} _ {i} ^ {\prime} = [ *; \widehat {f} (\mathbf {x} _ {i}); 1; t _ {i} ], i \in [ N + 1 ],
$$
$$
L (\widehat {f}) \leq \min  _ {k _ {\star} \in [ K ]} L \left(f _ {k _ {\star}}\right) + \max  _ {k \in [ K ]} \left| \widehat {L} _ {\text {v a l}} \left(f _ {k}\right) - L \left(f _ {k}\right) \right| + \gamma .
$$
To prove Theorem I.1, we first state and prove the following two propositions.
Proposition I.1 (Evaluation layer). There exists a 1-layer transformer  $\mathrm{TF}_{\theta}$  with  $MK$  heads and  $\|\pmb{\theta}\| \leq 3R + 2NKC / |\mathcal{D}_{\mathrm{val}}|$  such that for all  $\mathbf{H}$  such that  $\max_i\{|y_i^{\prime}|\} \leq R, \max_{i,k}\{|f_k(\mathbf{x}_i)|\} \leq R,$ $\mathrm{TF}_{\theta}$  maps
$$
\begin{array}{l} \mathbf {h} _ {i} = \left[ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \mathbf {0} _ {K + 1}; 1; t _ {i} \right] \\ \rightarrow \quad \mathbf {h} _ {i} ^ {\prime} = \left[ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} \left(\mathbf {x} _ {i}\right); \dots ; f _ {K} \left(\mathbf {x} _ {i}\right); \widetilde {L} _ {\text {v a l}} \left(f _ {1}\right); \dots ; \widetilde {L} _ {\text {v a l}} \left(f _ {K}\right); 0; 1; t _ {i} \right], \quad i \in [ N + 1 ], \\ \end{array}
$$
where  $\widetilde{L}_{\mathrm{val}}(\cdot)$  is a functional such that  $\max_k\left|\widetilde{L}_{\mathrm{val}}(f_k) - \widehat{L}_{\mathrm{val}}(f_k)\right|\leq \varepsilon$
Proof of Proposition I.1. As  $\ell$  is  $(\varepsilon, R, M, C)$ -approximable by sum of relus, there exists a function  $g: \mathbb{R}^2 \to \mathbb{R}$  of form
$$
g (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}|g(s,t) - \ell (s,t)|\leq \varepsilon$  . We define
$$
\widetilde {L} _ {\mathrm {v a l}} (f) := \frac {1}{| \mathcal {D} _ {\mathrm {v a l}} |} \sum_ {(\mathbf {x} _ {i}, y _ {i}) \in \mathcal {D} _ {\mathrm {v a l}}} g (f (\mathbf {x} _ {i}), y _ {i}),
$$
Next, for every  $m \in [M]$  and  $k \in [K]$ , we define matrices  $\mathbf{Q}_{m,k}, \mathbf{K}_{m,k}, \mathbf{V}_{m,k} \in \mathbb{R}^{D \times D}$  such that for all  $i,j \in [N + 1]$ ,
$$
\mathbf {Q} _ {m, k} \mathbf {h} _ {i} = \left[ \begin{array}{c} a _ {m} \\ b _ {m} \\ d _ {m} \\ - 2 \\ \mathbf {0} \end{array} \right], \quad \mathbf {K} _ {m, k} \mathbf {h} _ {j} = \left[ \begin{array}{c} f _ {k} (\mathbf {x} _ {j}) \\ y _ {j} \\ 1 \\ R (1 + t _ {j}) \\ \mathbf {0} \end{array} \right], \quad \mathbf {V} _ {m, k} \mathbf {h} _ {j} = \frac {(N + 1) c _ {m}}{| \mathcal {D} _ {\mathrm {v a l}} |} \cdot \mathbf {e} _ {D - (K - k) - 3}
$$
where  $\mathbf{e}_s\in \mathbb{R}^D$  is the vector with  $s$ -th entry being 1 and others being 0. As the input has structure  $\mathbf{h}_i = [\mathbf{x}_i;y_i';*;f_1(\mathbf{x}_i);\dots ;f_K(\mathbf{x}_i);\mathbf{0}_{K + 1};1;t_i]$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max _ {m \in [ M ], k \in [ K ]} \| \mathbf {Q} _ {m, k} \| _ {\mathrm {o p}} \leq 3, \quad \max _ {m \in [ M ], k \in [ K ]} \| \mathbf {K} _ {m, k} \| _ {\mathrm {o p}} \leq 2 + R, \quad \sum_ {m \in [ M ], k \in [ K ]} \| \mathbf {V} _ {m, k} \| _ {\mathrm {o p}} \leq \frac {K (N + 1) C}{| \mathcal {D} _ {\mathrm {v a l}} |}.
$$
$$
\begin{array}{l} \sigma \left(\left\langle \mathbf {Q} _ {m, k} \mathbf {h} _ {i}, \mathbf {K} _ {m, k} \mathbf {h} _ {j} \right\rangle\right) = \sigma \left(a _ {m} f _ {k} \left(\mathbf {x} _ {j}\right) + b _ {m} y _ {j} + d _ {m} - 2 R \left(1 + t _ {j}\right)\right) \\ = \sigma \left(a _ {m} \mathbf {w} ^ {\top} \mathbf {x} _ {j} + b _ {m} y _ {j} + d _ {m}\right) 1 \{t _ {j} = - 1 \}, \\ \end{array}
$$
where the last equality follows from the bound  $|a_{m}f_{k}(\mathbf{x}_{j}) + b_{m}y_{j} + d_{m}| \leq R(|a_{m}| + |b_{m}|) + d_{m} \leq 2R$ , so that the above relu equals 0 if  $t_j \leq 0$ . Therefore, for each  $i \in [N + 1]$  and  $k \in [K]$ ,
$$
\begin{array}{l} \sum_ {m = 1} ^ {M} \sigma (\langle \mathbf {Q} _ {m, k} \mathbf {h} _ {i}, \mathbf {K} _ {m, k} \mathbf {h} _ {j} \rangle) \mathbf {V} _ {m, k} \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)}{| \mathcal {D} _ {\mathrm {v a l}} |} 1 \{t _ {j} = - 1 \} \mathbf {e} _ {D - (K - k) - 3} \\ = g \left(f _ {k} \left(\mathbf {x} _ {j}\right), y _ {j}\right) \cdot \frac {(N + 1)}{\left| \mathcal {D} _ {\text {v a l}} \right|} 1 \left\{t _ {j} = - 1 \right\} \mathbf {e} _ {D - (K - k) - 3}. \\ \end{array}
$$
Thus letting the attention layer  $\pmb{\theta} = \{(\mathbf{V}_{m,k},\mathbf{Q}_{m,k},\mathbf{K}_{m,k})\}_{(m,k)\in [M]\times [K]}$ , we have
$$
\widetilde {\mathbf {h}} _ {i} = [ \operatorname {A t t n} _ {\boldsymbol {\theta}} (\mathbf {H}) ] _ {i} = \mathbf {h} _ {i} + \frac {1}{N + 1} \sum_ {j = 1} ^ {N + 1} \sum_ {m, k} \sigma (\langle \mathbf {Q} _ {m, k} \mathbf {h} _ {i}, \mathbf {K} _ {m, k} \mathbf {h} _ {j} \rangle) \mathbf {V} _ {m, k} \mathbf {h} _ {j}
$$
$$
\begin{array}{l} = \mathbf {h} _ {i} + \frac {1}{\left| \mathcal {D} _ {\mathrm {v a l}} \right|} \sum_ {j = 1} ^ {N + 1} \sum_ {k = 1} ^ {K} g \left(f _ {k} \left(\mathbf {x} _ {j}\right), y _ {j}\right) \cdot 1 \left\{t _ {j} = - 1 \right\} \mathbf {e} _ {D - (K - k) - 3} \\ = \mathbf {h} _ {i} + \sum_ {k = 1} ^ {K} \left(\frac {1}{\left| \mathcal {D} _ {\mathrm {v a l}} \right|} \sum_ {\left(\mathbf {x} _ {j}, y _ {j}\right) \in \mathcal {D} _ {\mathrm {v a l}}} g \left(f _ {k} \left(\mathbf {x} _ {j}\right), y _ {j}\right)\right) \mathbf {e} _ {D - (K - k) - 3} \\ = \mathbf {h} _ {i} + \sum_ {k = 1} ^ {K} \widetilde {L} _ {\text {v a l}} (f _ {k}) \cdot \mathbf {e} _ {D - (K - k) - 3} \\ = \left[ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} \left(\mathbf {x} _ {i}\right); \dots ; f _ {K} \left(\mathbf {x} _ {i}\right); \mathbf {0} _ {K + 1}; 1; t _ {i} \right] + \left[ \mathbf {0} _ {D - K - 3}; \widetilde {L} _ {\text {v a l}} \left(f _ {1}\right); \dots ; \widetilde {L} _ {\text {v a l}} \left(f _ {K}\right); 0; 0; 0 \right] \\ = \left[ \begin{array}{l} \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \widetilde {L} _ {\text {v a l}} (f _ {1}); \dots ; \widetilde {L} _ {\text {v a l}} (f _ {K}); 0; 1; t _ {i} \end{array} \right], \quad i \in [ N + 1 ]. \\ \end{array}
$$
Proposition I.2 (Selection layer). There exists a 3-layer transformer  $\mathrm{TF}_{\theta}$  with
$$
\max _ {\ell \in [ 3 ]} M ^ {(\ell)} \leq 2 K + 2, \quad \max _ {\ell \in [ 3 ]} D ^ {(\ell)} \leq K ^ {2} + K + 1, \quad \| \boldsymbol {\theta} \| \leq \gamma^ {- 1} + 3 K R + 2.
$$
such that  $\mathrm{TF}_{\theta}$  maps
$$
\begin{array}{l} \mathbf {h} _ {i} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \mathbb {L} _ {1}; \dots ; \mathbb {L} _ {K}; 0; 1; t _ {i} ] \\ \rightarrow \quad \mathbf {h} _ {i} ^ {\prime} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); *; \dots ; *; \widehat {f} (\mathbf {x} _ {i}); 1; t _ {i} ], \qquad i \in [ N + 1 ], \\ \end{array}
$$
where  $\widehat{f} = \sum_{k=1}^{K} \lambda_k f_k$  is an aggregated predictor, where the weights  $\lambda_1, \dots, \lambda_K \geq 0$  are functions only on  $\mathbb{L}_1, \dots, \mathbb{L}_k$  such that
$$
\sum_ {k = 1} ^ {K} \lambda_ {k} = 1, \quad \lambda_ {k} > 0 \text {o n l y i f} \mathbb {L} _ {k} \leq \min  _ {k ^ {\star} \in [ K ]} \mathbb {L} _ {k ^ {\star}} + \gamma .
$$
Proof of Proposition I.2. We construct a  $\pmb{\theta}$  which is a composition of 2 MLP layers followed by an attention layer  $(\pmb{\theta}_{\mathrm{mlp}}^{(1)},\pmb{\theta}_{\mathrm{mlp}}^{(2)},\pmb{\theta}_{\mathrm{attn}}^{(3)})$
Step 1: construction of  $\theta_{\mathrm{mlp}}^{(1)}$ . We consider matrix  $\mathbf{W}_1^{(1)}$  that maps
$$
\begin{array}{l} \mathbf {h} = \left[ * _ {D - K - 3}; \mathbb {L} _ {1}; \dots ; \mathbb {L} _ {K}; *; *; * \right] \\ \mapsto \mathbf {W} _ {1} ^ {(1)} \mathbf {h} = \left[ \mathbb {L} _ {1} - \mathbb {L} _ {2}; \dots ; \mathbb {L} _ {1} - \mathbb {L} _ {K}; \dots ; \mathbb {L} _ {K} - \mathbb {L} _ {K - 1}; \mathbb {L} _ {1}; - \mathbb {L} _ {1}; \dots ; \mathbb {L} _ {K}; - \mathbb {L} _ {K} \right], \\ \end{array}
$$
i.e.  $\mathbf{W}_1^{(1)}\mathbf{h}$  is a  $K^2 + K$  dimensional vector so that its entry contains  $\{\mathbb{L}_k - \mathbb{L}_l\}_{k,l\in [K]}$  and  $\{\mathbb{L}_k, - \mathbb{L}_k\}_{k\in [K]}$ . Clearly, such  $\mathbf{W}_1^{(1)}$  exists and can be chosen so that  $\left\| \mathbf{W}_1^{(1)}\right\|_{\mathrm{op}}\leq 2K$ . We then consider a matrix  $\mathbf{W}_2^{(1)}$  that maps
$$
\sigma \left(\mathbf {W} _ {1} ^ {(1)} \mathbf {h}\right) \mapsto \mathbf {W} _ {2} ^ {(1)} \sigma \left(\mathbf {W} _ {1} ^ {(1)} \mathbf {h}\right) = \left[ \mathbf {0} _ {D - K - 3}; c _ {1} - \mathbb {L} _ {1}; \dots ; c _ {K} - \mathbb {L} _ {K}; \mathbf {0} _ {3} \right] \in \mathbb {R} ^ {D},
$$
where  $c_{k} = c_{k}(\mathbb{L})\coloneqq \sum_{l\neq k}\sigma (\mathbb{L}_{k} - \mathbb{L}_{l})$  . Notice that
$$
c _ {k} - \mathbb {L} _ {k} = - \sigma (\mathbb {L} _ {k}) + \sigma (- \mathbb {L} _ {k}) + \sum_ {l \neq k} \sigma (\mathbb {L} _ {k} - \mathbb {L} _ {l}),
$$
and hence such  $\mathbf{W}_2^{(1)}$  exists and can be chosen so that  $\left\| \mathbf{W}_2^{(1)}\right\|_{\mathrm{op}}\leq K + 1$ . We set  $\pmb{\theta}_{\mathrm{mlp}}^{(1)} = (\mathbf{W}_1^{(1)},\mathbf{W}_2^{(1)})$ , then  $\mathrm{MLP}_{\pmb{\theta}_{\mathrm{mlp}}^{(1)}}$  maps  $\mathbf{h}_i$  to
$$
\mathbf {h} _ {i} ^ {(1)} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); c _ {1}; \dots ; c _ {K}; 0; 1; t _ {i} ].
$$
The basic property of  $\{c_k\}_{k\in [K]}$  is that, if  $c_{k}\leq \gamma$  , then  $\mathbb{L}_k\leq \min_{k^{\star}\in [K]}\mathbb{L}_{k^{\star}} + \gamma$
Step 2: construction of  $\theta_{\mathrm{mlp}}^{(2)}$ . We consider matrix  $\mathbf{W}_1^{(2)}$  that maps
$$
\mathbf {h} = \left[ * _ {D - K - 3}; c _ {1}; \dots ; c _ {K}; *; 1; * \right]
$$
$$
\mapsto \mathbf {W} _ {1} ^ {(2)} \mathbf {h} = \left[ 1 - \gamma^ {- 1} c _ {1}; c _ {1}; - c _ {1}; \dots ; 1 - \gamma^ {- 1} c _ {K}; c _ {K}; - c _ {K} \right] \in \mathbb {R} ^ {3 K},
$$
and  $\mathbf{W}_1^{(2)}$  can be chosen so that  $\left\| \mathbf{W}_1^{(2)}\right\|_{\mathrm{op}}\leq K + 1 + \gamma^{-1}$ . We then consider a matrix  $\mathbf{W}_2^{(2)}$  that maps
$$
\sigma (\mathbf {W} _ {1} ^ {(2)} \mathbf {h}) \mapsto \mathbf {W} _ {2} ^ {(2)} \sigma (\mathbf {W} _ {1} ^ {(1)} \mathbf {h}) = [ \mathbf {0} _ {D - K - 3}; \sigma (1 - \gamma^ {- 1} c _ {1}) - c _ {1}; \dots ; \sigma (1 - \gamma^ {- 1} c _ {K}) - c _ {K}; \mathbf {0} _ {3} ] \in \mathbb {R} ^ {D},
$$
which exists and can be chosen so that  $\left\| \mathbf{W}_2^{(1)}\right\|_{\mathrm{op}}\leq 2$  . We set  $\theta_{\mathrm{mlp}}^{(2)} = (\mathbf{W}_1^{(2)},\mathbf{W}_2^{(2)})$  , then  $\mathrm{MLP}_{\theta_{\mathrm{mlp}}^{(2)}}$  maps  $\mathbf{h}_i^{(1)}$  to
$$
\mathbf {h} _ {i} ^ {(2)} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); u _ {1}; \dots ; u _ {K}; 0; 1; t _ {i} ],
$$
where  $u_{k} = \sigma (1 - \gamma^{-1}c_{k})\forall k\in [K]$ . Clearly,  $u_{k}\in [0,1]$ , and  $u_{k} > 0$  if and only if  $c_{k}\leq \gamma$ .
Step 3: construction of  $\theta_{\mathrm{attn}}^{(3)}$ . We define
$$
\lambda_ {1} = 1 - \sigma (1 - u _ {1}), \quad \lambda_ {k} = \sigma (1 - u _ {1} - \dots - u _ {k - 1}) - \sigma (1 - u _ {1} - \dots - u _ {k}) \forall k \geq 2.
$$
Clearly,  $\lambda_{k}\geq 0$  , and  $\sum_{k}\lambda_{k} = 1$  . Further,
$$
\lambda_ {k} > 0 \Rightarrow u _ {k} > 0 \Rightarrow c _ {k} \leq \gamma \Rightarrow \mathbb {L} _ {k} \leq \min  _ {k ^ {\star} \in [ K ]} \mathbb {L} _ {k ^ {\star}} + \gamma .
$$
Therefore, it remains to construct  $\theta_{\mathrm{attn}}^{(3)}$  that implements  $\widehat{f} = \sum_{k=1}^{K} \lambda_k f_k$  based on  $[\mathbf{h}_i^{(2)}]_i$ . Notice that
$$
\begin{array}{l} \widehat {f} (\mathbf {x} _ {i}) = \sigma (1) \cdot f _ {1} (\mathbf {x} _ {i}) + \sum_ {k = 1} ^ {K - 1} \sigma \left(1 - u _ {1} - \dots - u _ {k - 1}\right) \cdot \left(f _ {k} (\mathbf {x} _ {i}) - f _ {k - 1} (\mathbf {x} _ {i})\right) \tag {41} \\ - \sigma (1 - u _ {1} - \dots - u _ {K}) \cdot f _ {K} (\mathbf {x} _ {i}), \\ \end{array}
$$
and hence we construct  $\theta_{\mathrm{attn}}^{(3)}$  as follows: for every  $k \in [K + 1]$  and  $w \in \{0,1\}$ , we define matrices  $\mathbf{Q}_{k,w}, \mathbf{K}_{k,w}, \mathbf{V}_{k,w} \in \mathbb{R}^{D \times D}$  such that for all  $k \in [K + 1]$
$$
\mathbf {Q} _ {k, 0} \mathbf {h} _ {i} ^ {(2)} = \left[ \begin{array}{c} (f _ {k} (\mathbf {x} _ {i}) + R) \cdot \mathbf {1} _ {k} \\ \mathbf {0} \end{array} \right], \quad \mathbf {Q} _ {k, 1} \mathbf {h} _ {i} ^ {(2)} = \left[ \begin{array}{c} (f _ {k - 1} (\mathbf {x} _ {i}) + R) \cdot \mathbf {1} _ {k} \\ \mathbf {0} \end{array} \right],
$$
$$
\mathbf {K} _ {k, 0} \mathbf {h} _ {j} ^ {(2)} = \mathbf {K} _ {k, 1} \mathbf {h} _ {j} ^ {(2)} = \left[ \begin{array}{c} 1 \\ - u _ {1} \\ \vdots \\ - u _ {k - 1} \\ \mathbf {0} \end{array} \right], \qquad \mathbf {V} _ {k, 0} \mathbf {h} _ {j} ^ {(2)} = \mathbf {e} _ {D - 2} = - \mathbf {V} _ {k, 1} \mathbf {h} _ {j} ^ {(2)},
$$
for all  $i,j\in [N + 1]$ , where we understand  $f_{0} = f_{K + 1} = 0$  and  $\mathbf{1}_k$  is the  $k$ -dimensional vector with all entries being 1. By the structure of  $\mathbf{h}_i^{(2)}$ , these matrices indeed exist, and further it is straightforward to check that they have norm bounds
$$
\max  _ {k \in [ K + 1 ], w \in \{0, 1 \}} \| \mathbf {Q} _ {k, w} \| _ {\mathrm {o p}} \leq K R, \quad \max  _ {k \in [ K + 1 ], w \in \{0, 1 \}} \| \mathbf {K} _ {k} \| _ {\mathrm {o p}} \leq 1, \quad \sum_ {k \in [ K + 1 ], w \in \{0, 1 \}} \| \mathbf {V} _ {k, w} \| _ {\mathrm {o p}} \leq 2 K + 2.
$$
Now, for every  $i,j\in [N + 1],k\in [K + 1],w\in \{0,1\}$  , we have
$$
\begin{array}{l} \sigma \Big (\left\langle \mathbf {Q} _ {k, w} \mathbf {h} _ {i} ^ {(2)}, \mathbf {K} _ {k, w} \mathbf {h} _ {j} ^ {(2)} \right\rangle \Big) = \sigma ((1 - u _ {1} - \dots - u _ {k - 1}) (f _ {k - w} (\mathbf {x} _ {i}) + R)) \\ = \sigma \left(1 - u _ {1} - \dots - u _ {k - 1}\right) \cdot \left(f _ {k - w} (\mathbf {x} _ {i}) + R\right), \\ \end{array}
$$
where the last equality follows from  $f_{k}(\mathbf{x}_{i}) + R \geq 0 \forall k \in [K]$ . Therefore,
$$
\begin{array}{l} \sum_ {k \in [ K + 1 ], w \in \{0, 1 \}} \sigma \left(\left\langle \mathbf {Q} _ {m, k} \mathbf {h} _ {i} ^ {(2)}, \mathbf {K} _ {m, k} \mathbf {h} _ {j} ^ {(2)} \right\rangle\right) \mathbf {V} _ {m, k} \mathbf {h} _ {j} ^ {(2)} \\ = \sum_ {k = 1} ^ {K} \left[ \sigma \left(1 - u _ {1} - \dots - u _ {k - 1}\right) \cdot \left(f _ {k} (\mathbf {x} _ {i}) + R\right) - \sigma \left(1 - u _ {1} - \dots - u _ {k - 1}\right) \cdot \left(f _ {k - 1} (\mathbf {x} _ {i}) + R\right) \right] \cdot \mathbf {e} _ {D - 2} \\ \end{array}
$$
$$
= \widehat {f} (\mathbf {x} _ {i}) \cdot \mathbf {e} _ {D - 2},
$$
where the last equality is due to (41). Thus letting the attention layer  $\theta_{\mathrm{attn}}^{(3)} = \left\{(\mathbf{V}_{k,w},\mathbf{Q}_{k,w},\mathbf{K}_{k,w})\right\}_{(k,w)\in [K + 1]\times \{0,1\}}$  we have
$$
\begin{array}{l} \mathbf {h} _ {i} ^ {(3)} = \left[ \mathrm {A t t n} _ {\pmb {\theta}} (\mathbf {H} ^ {(2)}) \right] _ {i} = \mathbf {h} _ {i} + \frac {1}{N + 1} \sum_ {j = 1} ^ {N + 1} \sum_ {k, w} \sigma \Big (\left\langle \mathbf {Q} _ {k, w} \mathbf {h} _ {i} ^ {(2)}, \mathbf {K} _ {k, w} \mathbf {h} _ {j} ^ {(2)} \right\rangle \Big) \mathbf {V} _ {k, w} \mathbf {h} _ {j} ^ {(2)} \\ = \mathbf {h} _ {i} ^ {(2)} + \widehat {f} (\mathbf {x} _ {i}) \cdot \mathbf {e} _ {D - 2} \\ = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); u _ {1}; \dots ; u _ {K}; \widehat {f} (\mathbf {x} _ {i}); 1; t _ {i} ]. \\ \end{array}
$$
Now, we are ready to prove Theorem I.1.
Proof of Theorem I.1 As  $\ell(\cdot, \cdot)$  is  $(\gamma/3, R, M, C)$ -approximable by sum of relus, we can invoke Proposition I.1 to show that there exists a single attention layer  $\theta_{\mathrm{attn}}^{(1)}$  so that  $\mathrm{Att} \mathbf{\eta}_{\theta_{\mathrm{attn}}^{(1)}}$  maps
$$
\mathbf {h} _ {i} \quad \rightarrow \quad \mathbf {h} _ {i} ^ {\prime} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \widetilde {L} _ {\mathrm {v a l}} (f _ {1}); \dots ; \widetilde {L} _ {\mathrm {v a l}} (f _ {K}); 0; 1; t _ {i} ], \qquad i \in [ N + 1 ],
$$
for any input  $\mathbf{H} = [\mathbf{h}_i]_i$  of the form described in Theorem I.1, and  $\widetilde{L}_{\mathrm{val}}(\cdot)$  is a functional such that  $\max_k\left|\widetilde{L}_{\mathrm{val}}(f_k) - \widehat{L}_{\mathrm{val}}(f_k)\right|\leq \gamma /3.$
Next, by the proof of Proposition I.2, there exists  $(\pmb{\theta}_{\mathrm{mlp}}^{(1)},\pmb{\theta}_{\mathrm{mlp}}^{(2)},\pmb{\theta}_{\mathrm{attn}}^{(3)})$  that maps
$$
\mathbf {h} _ {i} ^ {\prime} \quad \rightarrow \quad \mathbf {h} _ {i} ^ {(3)} = \left[ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); *; \sum_ {k = 1} ^ {K} \lambda_ {k} f _ {k} (\mathbf {x} _ {i}); 1; t _ {i} \right], \qquad i \in [ N + 1 ],
$$
where  $\lambda = (\lambda_1, \dots, \lambda_K) \in \Delta([K])$  and  $\lambda_k > 0$  only when  $\widetilde{L}_{\mathrm{val}}(f_k) \leq \min_{k^{\star}} \widetilde{L}_{\mathrm{val}}(f_{k^{\star}}) + \gamma / 3$ . Using the fact that  $\max_k |\widetilde{L}_{\mathrm{val}}(f_k) - \widehat{L}_{\mathrm{val}}(f_k)| \leq \gamma / 3$ , we deduce that  $\lambda$  is supported on  $\{k : \widehat{L}_{\mathrm{val}}(f_k) \leq \min_{k^{\star} \in [K]} \widehat{L}_{\mathrm{val}}(f_{k^{\star}}) + \gamma\}$ .
Therefore,  $\pmb{\theta} = (\pmb{\theta}_{\mathrm{attn}}^{(1)}, \pmb{\theta}_{\mathrm{mlp}}^{(1)}, \pmb{\theta}_{\mathrm{mlp}}^{(2)}, \pmb{\theta}_{\mathrm{attn}}^{(3)})$  is the desired transformer, with
$$
\max _ {\ell \in [ 3 ]} M ^ {(\ell)} \leq (M + 3) K, \quad \max _ {\ell \in [ 3 ]} D ^ {(\ell)} \leq K ^ {2} + K + 1,
$$
$$
\begin{array}{l} \left\| \boldsymbol {\theta} \right\| \leq \max  \left\{3 R + \frac {2 N K C}{\left| \mathcal {D} _ {\text {v a l}} \right|} + 3 K + 1, K + 3 + \gamma^ {- 1}, K R + 2 K + 2 \right\} \\ \leq 7 K R + \frac {2 N K C}{| \mathcal {D} _ {\text {v a l}} |} + \gamma^ {- 1}. \\ \end{array}
$$
I.2 Proof of Theorem 11
We first restate Theorem 11 into the following version which provides additional size bounds for  $\pmb{\theta}$ . For the simplicity of presentation, throughout this subsection and Appendix J, we denote  $\mathcal{I}_t = \{i : (\mathbf{x}_i, y_i) \in \mathcal{D}_{\text{train}}\}$ ,  $\mathcal{I}_v = \{i : (\mathbf{x}_i, y_i) \in \mathcal{D}_{\text{val}}\}$ ,  $\mathbf{X}_{\text{train}} = [\mathbf{x}_i]_{i \in \mathcal{I}_t}$  to be the input matrix corresponding to the training split only, and  $N_{\text{train}} = |\mathcal{D}_{\text{train}}|$ ,  $N_{\text{val}} = |\mathcal{D}_{\text{val}}|$ .
Theorem I.2. For any sequence of regularizations  $\{\lambda_k\}_{k\in [K]}$ ,  $0\leq \alpha \leq \beta$  with  $\kappa \coloneqq \max_{k}\frac{\beta + \lambda_{k}}{\alpha + \lambda_{k}}$ ,  $B_w > 0$ ,  $\gamma >0$ , and  $\varepsilon < B_w / 2$ , suppose in input format (3) we have  $D\geq \Theta (Kd)$ . Then there exists an  $L$ -layer transformer  $\mathrm{TF}_{\theta}$  with
$$
L = \lceil 2 \kappa \log (B _ {w} / (2 \varepsilon)) \rceil + 4, \max _ {\ell \in [ L ]} M ^ {(\ell)} \leq 3 K + 1, \max _ {\ell \in [ L ]} D ^ {(\ell)} \leq K ^ {2} + K + 1,
$$
$$
\| \boldsymbol {\theta} \| \leq \mathcal {O} \left(K R + (\beta + \lambda) ^ {- 1} + \frac {N}{N _ {\mathrm {v a l}}} + \gamma^ {- 1}\right), \quad R := \max  \left\{B _ {x} B _ {w}, B _ {y}, 1 \right\},
$$
such that the following holds. On any input data  $(\mathcal{D},\mathbf{x}_{N + 1})$  such that the problem (ICRidge) is well-conditioned and has a bounded solution:
$$
\alpha \leq \lambda_ {\min } \left(\mathbf {X} _ {\text {t r a i n}} ^ {\top} \mathbf {X} _ {\text {t r a i n}} / N _ {\text {t r a i n}}\right) \leq \lambda_ {\max } \left(\mathbf {X} _ {\text {t r a i n}} ^ {\top} \mathbf {X} _ {\text {t r a i n}} / N _ {\text {t r a i n}}\right) \leq \beta , \quad \max  _ {k \in [ K ]} \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda_ {k}} \left(\mathcal {D} _ {\text {t r a i n}}\right) \right\| _ {2} \leq B _ {w} / 2, \tag {42}
$$
$\mathrm{TF}_{\theta}^{0}$  approximately implements ridge selection: its prediction
$$
\widehat {y} _ {N + 1} = \mathrm {r e a d} _ {\mathbf {y}} (\mathrm {T F} _ {\pmb {\theta}} ^ {0} (\mathbf {H})) = \langle \widehat {\mathbf {w}}, \mathbf {x} _ {N + 1} \rangle , \qquad \widehat {\mathbf {w}} = \sum_ {k = 1} ^ {K} \lambda_ {k} \widehat {\mathbf {w}} _ {k}
$$
satisfies the following.