SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1. For each  $k \in [K]$ ,  $\widehat{\mathbf{w}}_k = \widehat{\mathbf{w}}_k(\mathcal{D}_{\mathrm{train}})$  approximates the ridge estimator  $\mathbf{w}_{\mathrm{ridge}}^{\lambda_k}(\mathcal{D}_{\mathrm{train}})$ , i.e.  $\left\| \widehat{\mathbf{w}}_k - \mathbf{w}_{\mathrm{ridge}}^{\lambda_k}(\mathcal{D}_{\mathrm{train}}) \right\|_2 \leq \varepsilon$ .
2.  $\lambda = (\lambda_1,\dots ,\lambda_K)\in \Delta ([K])$  so that
$$
\lambda_ {k} > 0 \text {o n l y i f} \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {k}) \leq \min  _ {k ^ {\star} \in [ K ]} \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {k ^ {\star}}) + \gamma .
$$
In particular, if we set  $\gamma' = 2(B_xB_w + B_y)B_x\varepsilon + \gamma$ , then it holds that
$$
\mathrm {d i s t} \Bigg (\widehat {\mathbf {w}}, \mathrm {c o n v} \{\widehat {\mathbf {w}} _ {\mathrm {r i d g e}, \mathrm {t r a i n}} ^ {\lambda_ {k}}: \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {\mathrm {r i d g e}, \mathrm {t r a i n}} ^ {\lambda_ {k}}) \leq \min  _ {k _ {\star} \in [ K ]} \widehat {L} _ {\mathrm {v a l}} (\widehat {\mathbf {w}} _ {\mathrm {r i d g e}, \mathrm {t r a i n}} ^ {\lambda_ {k _ {\star}}} + \gamma^ {\prime} \} \Bigg) \leq \varepsilon ,
$$
where we denote  $\widehat{\mathbf{w}}_{\mathrm{ridge,train}}^{\lambda_k} \coloneqq \mathbf{w}_{\mathrm{ridge}}^{\lambda_k}(\mathcal{D}_{\mathrm{train}})$ .
To prove Theorem I.2, we first show that, for the squared validation loss, there exists a 3-layer transformer that performs predictor selection based on the exactly evaluated  $\widehat{L}_{\mathrm{val}}(f_k)$  for each  $k \in [K]$ . (Proof in Appendix I.2.1.)
Theorem I.3 (Square-loss version of Theorem I.1). Consider the squared validation loss
$$
\widehat {L} _ {\mathrm {v a l}} (f) := \frac {1}{2 | \mathcal {D} _ {\mathrm {v a l}} |} \sum_ {(x _ {i}, y _ {i}) \in \mathcal {D} _ {\mathrm {v a l}}} \left(f (\mathbf {x} _ {i}) - y _ {i}\right) ^ {2}.
$$
Then 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 7 K R + \frac {2 N}{| \mathcal {D} _ {\mathrm {v a l}} |} + \gamma^ {- 1},
$$
such that for any input  $\mathbf{H}$  that takes form
$$
\mathbf {h} _ {i} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \mathbf {0} _ {K}; *; 1; t _ {i} ],
$$
where  $\mathrm{TF}_{\theta}$  outputs  $\mathbf{h}_{N + 1} = [\mathbf{x}_{N + 1};\widehat{f} (\mathbf{x}_{N + 1});*;1;0]$ , where the predictor  $\widehat{f}:\mathbb{R}^d\to \mathbb{R}$  is a convex combination of  $\{f_k:\widehat{L}_{\mathrm{val}}(f_k)\leq \min_{k_\star \in [K]}\widehat{L}_{\mathrm{val}}(f_{k_\star}) + \gamma \}$ . As a corollary, for any convex risk  $L:(\mathbb{R}^d\to \mathbb{R})\to \mathbb{R},$ $\widehat{f}$  satisfies
$$
L (\widehat {f}) \leq \min  _ {k _ {\star} \in [ K ]} L \left(f _ {k _ {\star}}\right) + \max  _ {k \in [ K ]} \left| \widehat {L} _ {\mathsf {v a l}} \left(f _ {k}\right) - L \left(f _ {k}\right) \right| + \gamma .
$$
Proof of Theorem 1.2 First, by the proof of Theorem 4 and Proposition B.6, for each  $k \in [K]$ , there exists a  $T = L - 3$  layer transformer  $\theta^{(1:T)}$  such that  $\mathrm{TF}_{\theta^{(1:T)}}$  maps
$$
\mathbf {h} _ {i} \quad \rightarrow \quad \mathbf {h} _ {i} ^ {(T)} = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; \langle \widehat {\mathbf {w}} _ {1}, \mathbf {x} _ {i} \rangle ; \dots ; \langle \widehat {\mathbf {w}} _ {K}, \mathbf {x} _ {i} \rangle ; \mathbf {0} _ {K}; 1; t _ {i} ],
$$
This is because  $\widehat{L}_{\mathrm{val}}(\mathbf{w})$  is  $(B_{x}B_{w} + B_{y})B_{x}$ -Lipschitz w.r.t.  $\mathbf{w} \in \mathsf{B}_2(B_w)$ .
Technically, an adapted version where the underlying ICGD mechanism operates on the training split (with  $t_i = 1$ ) with size  $N_{\mathrm{train}}$  instead of on all  $N$  training examples, which only changes  $\| \pmb{\theta} \|$  by at most a constant factor, and does not change the number of layers and heads.
so that if (42) holds, we have  $\left\| \widehat{\mathbf{w}}_k - \mathbf{w}_{\mathrm{ridge}}^{\lambda_k}\right\| _2\leq \varepsilon$  and  $\widehat{\mathbf{w}}_k\in \mathsf{B}_2(B_w)$
Next, by Theorem I.3, there exists a 3-layer transformer  $\pmb{\theta}^{(T + 1:T + 3)}$  that outputs
$$
\mathbf {h} _ {N + 1} ^ {(T + 3)} = [ \mathbf {x} _ {N + 1}; \langle \widehat {\mathbf {w}}, \mathbf {x} _ {N + 1} \rangle ; *; 1; t _ {i} ],
$$
where  $\widehat{\mathbf{w}} = \sum_{k=1}^{K} \lambda_k \widehat{\mathbf{w}}_k$ ,  $\lambda = (\lambda_1, \dots, \lambda_K) \in \Delta([K])$  so that
$$
\lambda_ {k} > 0 \text {o n l y i f} \widehat {L} _ {\mathbf {v a l}} (\widehat {\mathbf {w}} _ {k}) \leq \min  _ {k ^ {\star} \in [ K ]} \widehat {L} _ {\mathbf {v a l}} (\widehat {\mathbf {w}} _ {k ^ {\star}}) + \gamma .
$$
I.2.1 Proof of Theorem I.3
Similar to the proof of Proposition 10, Theorem I.3 is a direct corollary by combining Proposition I.3 with Proposition I.2.
Proposition I.3 (Evaluation layer for the squared loss). There exists an attention layer  $\mathrm{TF}_{\theta}$  with  $2K$  heads and  $\| \pmb{\theta} \| \leq 3R + 2NK / |\mathcal{D}_{\mathrm{val}}|$  such that  $\mathrm{TF}_{\theta}$  maps
$$
\begin{array}{l} \mathbf {h} _ {i} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \mathbf {0} _ {K}; *; 1; t _ {i} ] \\ \rightarrow \mathbf {h} _ {i} ^ {\prime} = [ *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \widehat {L} _ {\mathrm {v a l}} (f _ {1}); \dots ; \widehat {L} _ {\mathrm {v a l}} (f _ {K}); *; 1; t _ {i} ], \qquad i \in [ N + 1 ]. \\ \end{array}
$$
Proof of Proposition I.3. For every  $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]$ ,
$$
\begin{array}{l} \mathbf {Q} _ {k, 0} \mathbf {h} _ {i} = \left[ \begin{array}{c} 1 \\ - 1 \\ - 2 \\ \mathbf {0} \end{array} \right], \quad \mathbf {Q} _ {k, 1} \mathbf {h} _ {i} = \left[ \begin{array}{c} - 1 \\ 1 \\ - 2 \\ \mathbf {0} \end{array} \right], \quad \mathbf {K} _ {k, 0} \mathbf {h} _ {j} = \mathbf {K} _ {k, 1} \mathbf {h} _ {j} = \left[ \begin{array}{c} f _ {k} (\mathbf {x} _ {j}) \\ y _ {j} \\ R (1 + t _ {j}) \\ \mathbf {0} \end{array} \right], \\ \mathbf {V} _ {k, 0} \mathbf {h} _ {j} = - \mathbf {V} _ {k, 1} \mathbf {h} _ {j} = \frac {(N + 1)}{2 | \mathcal {D} _ {\mathrm {v a l}} |} \cdot (f _ {k} (\mathbf {x} _ {j}) - y _ {j}) \mathbf {e} _ {D - (K - k) - 3}. \\ \end{array}
$$
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 _ {k \in [ K ], w \in \{0, 1 \}} \| \mathbf {Q} _ {k, w} \| _ {\mathrm {o p}} \leq 3, \quad \max _ {k \in [ K ], w \in \{0, 1 \}} \| \mathbf {K} _ {k, w} \| _ {\mathrm {o p}} \leq 1 + R, \quad \sum_ {k \in [ K ], w \in \{0, 1 \}} \| \mathbf {V} _ {k, w} \| _ {\mathrm {o p}} \leq \frac {K (N + 1)}{| \mathcal {D} _ {\mathrm {v a l}} |}.
$$
$$
\begin{array}{l} \sum_ {w \in \{0, 1 \}} \sigma (\langle \mathbf {Q} _ {k, w} \mathbf {h} _ {i}, \mathbf {K} _ {k, w} \mathbf {h} _ {j} \rangle) \mathbf {V} _ {k, w} \mathbf {h} _ {j} \\ = [ \sigma (f _ {k} (\mathbf {x} _ {j}) - y _ {j} - 2 R (1 + t _ {j})) - \sigma (y _ {j} - f _ {k} (\mathbf {x} _ {j}) - 2 R (1 + t _ {j})) ] \cdot \frac {(N + 1)}{2 | \mathcal {D} _ {\mathrm {v a l}} |} (f _ {k} (\mathbf {x} _ {j}) - y _ {j}) \mathbf {e} _ {D - (K - k) - 3} \\ = 1 \{t _ {j} = - 1 \} \cdot \left[ \sigma \left(f _ {k} (\mathbf {x} _ {j}) - y _ {j}\right) - \sigma \left(y _ {j} - f _ {k} (\mathbf {x} _ {j})\right) \right] \cdot \frac {(N + 1)}{2 | \mathcal {D} _ {\mathrm {v a l}} |} \left(f _ {k} (\mathbf {x} _ {j}) - y _ {j}\right) \mathbf {e} _ {D - (K - k) - 3} \\ = 1 \{t _ {j} = - 1 \} \cdot \frac {(N + 1)}{2 | \mathcal {D} _ {\mathrm {v a l}} |} (f _ {k} (\mathbf {x} _ {j}) - y _ {j}) ^ {2} \mathbf {e} _ {D - (K - k) - 3}, \\ \end{array}
$$
where the second equality follows from the bound  $|f_k(\mathbf{x}_j) - y_j| \leq 2R$ , so that the relus equals 0 if  $t_j \leq 0$ . Thus letting the attention layer  $\pmb{\theta} = \{(\mathbf{V}_{k,w},\mathbf{Q}_{k,w},\mathbf{K}_{k,w})\}_{(k,w) \in [K] \times \{0,1\}}$ , we have
$$
\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_ {k, w} \sigma \left(\left\langle \mathbf {Q} _ {k, w} \mathbf {h} _ {i}, \mathbf {K} _ {k, w} \mathbf {h} _ {j} \right\rangle\right) \mathbf {V} _ {k, w} \mathbf {h} _ {j}
$$
$$
\begin{array}{l} = \mathbf {h} _ {i} + \frac {1}{2 \left| \mathcal {D} _ {\mathrm {v a l}} \right|} \sum_ {j = 1} ^ {N + 1} \sum_ {k = 1} ^ {K} \left(f _ {k} \left(\mathbf {x} _ {j}\right) - y _ {j}\right) ^ {2} \cdot 1 \{t _ {j} = - 1 \} \mathbf {e} _ {D - (K - k) - 3} \\ = \mathbf {h} _ {i} + \sum_ {k = 1} ^ {K} \left(\frac {1}{2 | \mathcal {D} _ {\mathrm {v a l}} |} \sum_ {(\mathbf {x} _ {j}, y _ {j}) \in \mathcal {D} _ {\mathrm {v a l}}} \left(f _ {k} (\mathbf {x} _ {j}) - y _ {j}\right) ^ {2}\right) \mathbf {e} _ {D - (K - k) - 3} \\ = \mathbf {h} _ {i} + \sum_ {k = 1} ^ {K} \widehat {L} _ {\text {v a l}} (f _ {k}) \cdot \mathbf {e} _ {D - (K - k) - 3} \\ = [ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \mathbf {0} _ {K + 1}; 1; t _ {i} ] + [ \mathbf {0} _ {D - K - 3}; \widehat {L} _ {\text {v a l}} (f _ {1}); \dots ; \widehat {L} _ {\text {v a l}} (f _ {K}); 0; 0; 0 ] \\ = \left[ \begin{array}{l} \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; f _ {1} (\mathbf {x} _ {i}); \dots ; f _ {K} (\mathbf {x} _ {i}); \widehat {L} _ {\text {v a l}} (f _ {1}); \dots ; \widehat {L} _ {\text {v a l}} (f _ {K}); 0; 1; t _ {i} \end{array} \right], \quad i \in [ N + 1 ]. \\ \end{array}
$$
I.3 Proofs for Section 4.2
I.3.1 Proof of Lemma 13
It is straightforward to check that the binary type check  $\psi : \mathbb{R} \to \mathbb{R}$  can be expressed as a linear combination of 6 relu's (recalling  $\sigma(\cdot) = \mathrm{ReLU}(\cdot)$ ):
$$
\begin{array}{l} \psi (y) = \sigma \left(\frac {y + \varepsilon}{\varepsilon}\right) - 2 \sigma \left(\frac {y}{\varepsilon}\right) + \sigma \left(\frac {y - \varepsilon}{\varepsilon}\right) + \sigma \left(\frac {y - (1 - \varepsilon)}{\varepsilon}\right) - 2 \sigma \left(\frac {y - 1}{\varepsilon}\right) + \sigma \left(\frac {y - (1 + \varepsilon)}{\varepsilon}\right) \\ =: \sum_ {m = 1} ^ {6} a _ {m} \sigma \left(b _ {m} y + c _ {m}\right), \\ \end{array}
$$
with  $\sum_{m}|a_{m}| = 8 / \varepsilon$ ,  $\max_{m}\max \{|b_{m}|,|c_{m}|\} \leq 2$ . We can thus construct an attention layer  $\pmb{\theta} = \{(\mathbf{Q}_m,\mathbf{K}_m,\mathbf{V}_m)\}_{m = 1}^6$  with 6 heads such that
$$
\mathbf {Q} _ {m} \mathbf {h} _ {i} = [ b _ {m}; c _ {m}; \mathbf {0} _ {D - 2} ], \quad \mathbf {K} _ {m} \mathbf {h} _ {j} = [ y _ {j}; 1; \mathbf {0} _ {D - 2} ], \quad \mathbf {V} _ {m} \mathbf {h} _ {j} = \left[ \frac {N + 1}{N} a _ {m} \cdot t _ {j}; \mathbf {0} _ {D - 1} \right],
$$
which gives that for every  $i \in [N + 1]$ ,
$$
\begin{array}{l} \sum_ {m = 1} ^ {6} \frac {1}{N + 1} \sum_ {j \in [ N + 1 ]} \sigma (\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \rangle) [ \mathbf {V} _ {m} \mathbf {h} _ {j} ] _ {1} \\ = \sum_ {m = 1} ^ {6} \frac {1}{N} \sum_ {j = 1} ^ {N} \sigma \left(b _ {m} y _ {j} + c _ {m}\right) a _ {m} = \frac {1}{N} \sum_ {j = 1} ^ {N} \psi \left(y _ {j}\right) = \Psi^ {\text {b i n a r y}} (\mathcal {D}). \\ \end{array}
$$
Further, we have  $\| \pmb {\theta}\| \leq 18 / \varepsilon = \mathcal{O}(1 / \varepsilon)$ . This is the desired result.
By composing the above attention layer with one additional layer (with 2 heads) that implement the following function
$$
\sigma (2 (t - 1 / 2)) - \sigma (2 (t - 1)),
$$
on the output  $\Psi^{\mathrm{binary}}(\mathcal{D})$ , we directly obtain the following corollary.
Corollary I.1 (Thresholded binary test). There exists a two-layer attention-only transformer with  $\max_{\ell \in [2]} M^{(\ell)} \leq 6$  and  $\|\pmb{\theta}\| \leq \mathcal{O}(1 / \varepsilon)$  that exactly implements the thresholded binary test
$$
\Psi_ {\text {t h r e s}} ^ {\text {b i n a r y}} (\mathcal {D}) := \left\{ \begin{array}{l l} 1, & \text {i f} \Psi^ {\text {b i n a r y}} (\mathcal {D}) \geq 1, \\ 0, & \text {i f} \Psi^ {\text {b i n a r y}} (\mathcal {D}) \leq \frac {1}{2}, \\ \text {l i n e a r i n t e r p o l a t i o n ,} & \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad o. w. \end{array} \right. \tag {43}
$$
at every token  $i\in [N + 1]$ , where we recall the definition of  $\Psi^{\mathrm{binary}}$  in Lemma 13.
I.3.2 Formal statement and proof of Proposition 14
We say a distribution  $\mathsf{P}_y$  on  $\mathbb{R}$  is  $(C,\varepsilon_0)$ -not-concentrated around  $\{0,1\}$  if
$$
\mathrm {P} _ {y} ([ - \varepsilon , \varepsilon ] \cup [ 1 - \varepsilon , 1 + \varepsilon ]) \leq C \varepsilon
$$
for all  $\varepsilon \in (0,\varepsilon_0]$ . A sufficient condition is that the density  $\mathfrak{p}_y$  is upper bounded by  $C$  within  $[- \varepsilon_0, \varepsilon_0] \cup [1 - \varepsilon_0, 1 + \varepsilon_0]$ .
Throughout this section, let  $\sigma_{\log}(t) \coloneqq (1 + e^{-t})^{-1}$  denote the sigmoid activation, and let  $\widehat{\mathbf{w}}_{\mathrm{log}}$  denote the solution to the in-context logistic regression problem, i.e. (ICGLM) with  $g(\cdot) = \sigma_{\log}(\cdot)$ .
Proposition I.4 (Adaptive regression or classification; Formal version of Proposition 14). For any  $B_w > 0$ ,  $\varepsilon \leq B_xB_w / 10$ ,  $0 < \alpha \leq \beta$  with  $\kappa \coloneqq \beta / \alpha$ , and any  $(C, \varepsilon_0)$ , there exists a L-layer attention-only transformer with
$$
L \leq \mathcal {O} \left(\kappa \log \frac {B _ {x} B _ {w}}{\varepsilon}\right), \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq \mathcal {O} \left(\left(1 + \frac {B _ {x} ^ {4}}{\alpha^ {2}}\right) \varepsilon^ {- 2}\right), \quad \| \boldsymbol {\theta} \| \leq \mathcal {O} \left(R + \frac {1}{\beta} + \frac {1}{\varepsilon}\right)
$$
(with  $R\coloneqq \max \{B_xB_w,B_y,1\}$  , and  $\varepsilon$  depending only on  $(C,\varepsilon_0))$  such that the following holds. Suppose the input format is (3) with dimension  $D\geq 3d + 4$
On any classification instance  $(\mathcal{D},\mathbf{x}_{N + 1})$  (such that  $\{y_{i}\}_{i\in [N]}\subset \{0,1\}$ ) that is well-conditioned for logistic regression in the sense of (35), it outputs  $\widehat{y}_{N + 1}$  that  $\varepsilon$ -approximates the prediction of in-context logistic regression:
$$
\left| \widehat {y} _ {N + 1} - \sigma_ {\log} \left(\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} _ {\log} \rangle\right) \right| \leq \varepsilon .
$$
On the contrary, for regression problems, i.e. any in-context distribution  $\mathsf{P}$  whose marginal  $\mathsf{P}_y$  is  $(C,\varepsilon_0)$ -not-concentrated around  $\{0,1\}$ , with probability at least  $1 - \exp (-cN)$  over  $\mathcal{D}$  (where  $c > 0$  depends only on  $(C,\varepsilon_0))$ ,  $\widehat{y}_{N + 1}$ $\varepsilon$ -approximates the prediction of in-context least squares if the data is well-conditioned:
$$
\left| \hat {y} _ {N + 1} - \langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} _ {\mathrm {L S}} \rangle \right| \leq \varepsilon \quad \text {w h e n e v e r} \mathcal {D} \text {s a t i s f i e s (5) w i t h} \lambda = 0,
$$
where  $\widehat{\mathbf{w}}_{\mathrm{LS}}$  denotes the in-context least squares estimator, i.e. (ICRidge) with  $\lambda = 0$ .
Proof. The result follows by combining the binary test in Corollary I.1 with Theorem 4 and Theorem G.1. By those results, there exists three attention-only transformers  $\theta_{\mathrm{LS}}, \theta_{\mathrm{log}}, \theta_{\mathrm{bin}}$ , with (below  $L_g, C_g = \Theta(1)$  for  $g = \sigma_{\log}(\cdot)$ )
$$
L _ {\mathrm {L S}} \leq \mathcal {O} \left(\kappa \log \frac {B _ {x} B _ {w}}{\varepsilon}\right), \quad \max  _ {\ell \in \left[ L _ {\mathrm {L S}} \right]} M _ {\mathrm {L S}} ^ {(\ell)} \leq 3, \quad \| \boldsymbol {\theta} _ {\mathrm {L S}} \| \leq \mathcal {O} \left(R + \frac {1}{\beta}\right),
$$
$$
L _ {\log} \leq \mathcal {O} \left(\kappa \log \frac {L _ {g} B _ {x} B _ {w}}{\varepsilon}\right), \quad \max  _ {\ell \in [ L _ {\log} ]} M _ {\log} ^ {(\ell)} \leq \mathcal {O} \left(C _ {g} ^ {2} \left(1 + \frac {L _ {g} ^ {2} B _ {x} ^ {4}}{\alpha^ {2}}\right) \varepsilon^ {- 2}\right), \quad \| \boldsymbol {\theta} _ {\log} \| \leq \mathcal {O} \left(R + \frac {C _ {g}}{\beta}\right),
$$
$$
L _ {\mathrm {b i n}} = 2, \quad \max  _ {\ell \in [ 2 ]} M _ {\mathrm {b i n}} ^ {(\ell)} \leq 6, \quad \| \boldsymbol {\theta} _ {\mathrm {b i n}} \| \leq \mathcal {O} (1 / \varepsilon),
$$
that outputs prediction  $\widehat{y}_{N + 1}^{\mathrm{LS}}$ ,  $\widehat{y}_{N + 1}^{\log}$  (at the  $(N + 1)$ -th token) and  $\Psi_{\mathrm{thres}}^{\mathrm{binary}}(\mathcal{D})$  (at every token) respectively, which satisfy
$$
\left| \widehat {y} _ {N + 1} ^ {\log} - \sigma_ {\log} (\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} _ {\log} \rangle) \right| \leq \varepsilon ,
$$
$$
\left| \widehat {y} _ {N + 1} ^ {\mathrm {L S}} - \left\langle \mathbf {x} _ {N + 1}, \widehat {\mathbf {w}} _ {\mathrm {L S}} \right\rangle \right| \leq \varepsilon .
$$
when the corresponding well-conditionednesses are satisfied. In particular, we can make  $\widehat{\mathbf{w}}_{\mathrm{log}}$  well-defined on non-binary data, by multiplying  $\Psi_{\mathrm{thres}}^{\mathrm{binary}}(\mathcal{D})$  onto the  $\mathbf{x}_i$ 's (which can be implemented by slightly modifying  $\theta_{\mathrm{log}}$  without changing the order of the number of layers, heads, and norms) so that  $\widehat{\mathbf{w}}_{\mathrm{log}} = \mathbf{0}$  on any data where  $\Psi_{\mathrm{thres}}^{\mathrm{binary}}(\mathcal{D}) = 0$ .
By joining  $\theta_{\mathrm{LS}}$  and  $\theta_{\mathrm{log}}$  using Proposition B.6, concatenating with  $\theta_{\mathrm{bin}}$  before, and concatenating with one additional attention layer with 2 heads after to implement
$$
\Psi_ {\text {t h r e s}} ^ {\text {b i n a r y}} (\mathcal {D}) \hat {y} _ {N + 1} ^ {\log} + \left(1 - \Psi_ {\text {t h r e s}} ^ {\text {b i n a r y}} (\mathcal {D})\right) \hat {y} _ {N + 1} ^ {\mathrm {L S}}, \tag {44}
$$
we obtain a single transformer  $\theta$  with
$$
L \leq \mathcal {O} \bigg (\kappa \log \frac {B _ {x} B _ {w}}{\varepsilon} \bigg), \max _ {\ell \in [ L ]} M ^ {(\ell)} \leq \mathcal {O} \bigg (\bigg (1 + \frac {B _ {x} ^ {4}}{\alpha^ {2}} \bigg) \varepsilon^ {- 2} \bigg), \| \pmb {\theta} _ {\mathrm {L S}} \| \leq \mathcal {O} \bigg (R + \frac {1}{\beta} + \frac {1}{\varepsilon} \bigg),
$$
which outputs (44) as its prediction (at the location for  $\widehat{y}_{N + 1}$ ).
It remains to show that (44) reduces to either one of  $\widehat{y}_{N+1}^{\log}$  or  $\widehat{y}_{N+1}^{\mathrm{LS}}$ . When the data are binary  $(y_i \in \{0,1\})$ , we have  $\Psi^{\mathrm{binary}}(\mathcal{D}) = 1$  and  $\Psi_{\mathrm{thres}}^{\mathrm{binary}}(\mathcal{D}) = 1$ , in which case (44) becomes exactly  $\widehat{y}_{N+1}^{\log}$ . By contrast, when data is sampled from a distribution that is  $(C, \varepsilon_0)$ -not-concentrated around  $\{0,1\}$ , we have for any fixed  $\varepsilon \leq \varepsilon_0 \wedge \frac{1}{4C}$  that, letting  $B_{\varepsilon} := [-\varepsilon, \varepsilon] \cup [1 - \varepsilon, 1 + \varepsilon]$  and  $\mathsf{p}_{\varepsilon} := \mathsf{P}_y(B_{\varepsilon}) \leq C\varepsilon \leq \frac{1}{4}$ , by Hoeffding's inequality,
$$
\begin{array}{l} \mathsf {P} \left(\Psi_ {\text {t h r e s}} ^ {\text {b i n a r y}} (\mathcal {D}) \neq 0\right) = \mathsf {P} \left(\Psi_ {\text {t h r e s}} ^ {\text {b i n a r y}} (\mathcal {D}) \geq \frac {1}{2}\right) = \mathsf {P} \left(\frac {1}{N} \sum_ {i = 1} ^ {N} 1 \{y _ {j} \in B _ {\varepsilon} \} \geq \frac {1}{2}\right) \\ \leq \exp \left(- c (1 / 2 - \mathsf {p} _ {\varepsilon}) ^ {2} N\right) \leq \exp (- c ^ {\prime} N), \\ \end{array}
$$
where  $c' > 0$  is an absolute constant. On the event  $\Psi_{\text{thres}}^{\text{binary}}(\mathcal{D}) = 0$  (which happens with probability at least  $1 - \exp(-c'N)$ ), (44) becomes exactly  $\widehat{y}_{N+1}^{\text{LS}}$ . This finishes the proof.
I.4 Linear correlation test and application
In this section, we give another instantiation of the pre-ICL testing mechanism by showing that the transformer can implement a linear correlation test that tests whether the correlation vector  $\mathbb{E}[\mathbf{xy}]$  has a large norm. We then use this test to construct a transformer to perform "confident linear regression", i.e. output a prediction from linear regression only when the signal-to-noise ratio is high.
For any fixed parameters  $\lambda_{\mathrm{min}}, B_w^\star > 0$ , consider the linear correlation test over data  $\mathcal{D}$  defined as
$$
\begin{array}{l} \Psi^ {\mathrm {l i n}} (\mathcal {D}) := \frac {1}{\lambda_ {\mathrm {m i n}} ^ {2} (B _ {w} ^ {\star}) ^ {2} / 2} \cdot \left[ \sigma \Big (\| \widehat {\mathbf {t}} \| _ {2} ^ {2} - (\lambda_ {\mathrm {m i n}} B _ {w} ^ {\star} / 4) ^ {2} \Big) - \sigma \Big (\| \widehat {\mathbf {t}} \| _ {2} ^ {2} - (3 \lambda_ {\mathrm {m i n}} B _ {w} ^ {\star} / 4) ^ {2} \Big) \right] \\ = \left\{ \begin{array}{l l} 0, & \| \widehat {\mathbf {t}} \| _ {2} ^ {2} \leq \left(\lambda_ {\min } B _ {w} ^ {\star} / 4\right) ^ {2}, \\ 1, & \| \widehat {\mathbf {t}} \| _ {2} ^ {2} \geq \left(3 \lambda_ {\min } B _ {w} ^ {\star} / 4\right) ^ {2}, \\ \text {l i n e a r i n t e r p o l a t i o n ,} & \text {o . w . ,} \end{array} \right. \tag {45} \\ \end{array}
$$
$$
\mathrm {w h e r e} \widehat {\mathbf {t}} = \mathbf {T} (\mathcal {D}) := \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbf {x} _ {i} y _ {i}.
$$
Recall that  $\sigma (\cdot) = \mathrm{ReLU}(\cdot)$  above denotes the relu activation.
We show that  $\Psi^{\mathrm{lin}}$  can be exactly implemented by a 3-layer transformer.
Lemma I.1 (Expressing  $\Psi^{\mathrm{lin}}$  by transformer). There exists a 3-layer attention-only transformer  $\mathrm{TF}_{\theta}$  with at most 2 heads per layer and  $\| \pmb {\theta}\| \leq \mathcal{O}(1 + \lambda_{\min}^2 (B_w^\star)^2)$  such that on input sequence  $\mathbf{H}$  of the form (3) with  $D\geq 2d + 4$  , the transformer exactly implements  $\Psi^{\mathrm{lin}}$  : it outputs  $\widetilde{\mathbf{H}}$  such that  $\widetilde{\mathbf{h}}_i = [\mathbf{x}_i;y_it_i;*\Psi^{\mathrm{lin}}(\mathcal{D});1]$  for all  $i\in [N + 1]$
Proof. We begin by noting the following basic facts:
- Identity function can be implemented exactly by two ReLUs:  $t = \sigma(t) - \sigma(-t)$ .
- Squared  $\ell_2$  norm can be implemented exactly by a single attention head (assuming every input  $\mathbf{h}_i$  contains the same vector  $\mathbf{g}$ ):  $\| \mathbf{g} \|_2^2 = \sigma (\langle \mathbf{g}, \mathbf{g} \rangle)$ .
We construct the transformer  $\theta$  as follows.
Layer 1: Use 2 heads to implement  $\widehat{\mathbf{t}} = \frac{1}{N}\sum_{i=1}^{N}\mathbf{x}_{i}y_{i}$ , where  $\mathbf{V}_{\{1,2\}}^{(1)}\mathbf{h}_{j} = [\pm \mathbf{x}_{j};\mathbf{0}_{D - d}]$ ,  $\mathbf{Q}_{\{1,2\}}^{(1)}\mathbf{h}_{i} = \left[\frac{N + 1}{N};\mathbf{0}_{D - 1}\right]$ , and  $\mathbf{K}_{\{1,2\}}^{(1)}\mathbf{h}_{j} = [\pm y_{j}t_{j};\mathbf{0}_{D - 1}] = [\pm y_{j}1\{j < N + 1\};\mathbf{0}_{D - 1}]$  (where we recall  $t_{j} = 1\{j < N + 1\}$  and note that  $y_{j}t_{j}$  corresponds exactly to the location for  $y_{j}$  in  $\mathbf{H}$ , cf. (3)).
By manipulating the output dimension in  $\mathbf{V}^{(1)}$ , write the result  $\hat{\mathbf{t}}$  into blank memory space with dimension  $d$  at every token  $i \in [N + 1]$ .