SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where the second inequality follows from Lemma L.2 and Lemma L.1 and the fact that  $\| \mathrm{Attn}_{\theta_{\mathrm{attn}}}(\mathbf{H}) \|_{2,\infty} \leq \overline{\mathsf{R}} \coloneqq \mathsf{R} + B^3 \mathsf{R}^3$  for all  $\mathbf{H} \in \mathcal{H}_{\mathbb{R}}$ .
Furthermore, for  $\mathbf{H}' \in \mathcal{H}_{\mathbb{R}}$ , we have
$$
\begin{array}{l} \left\| \mathrm {T F} _ {\boldsymbol {\theta}} (\mathbf {H}) - \mathrm {T F} _ {\boldsymbol {\theta}} \left(\mathbf {H} ^ {\prime}\right) \right\| _ {2, \infty} \leq (1 + B ^ {2}) \left\| \operatorname {A t t n} _ {\boldsymbol {\theta} _ {\text {a t t n}}} (\mathbf {H}) - \operatorname {A t t n} _ {\boldsymbol {\theta} _ {\text {a t t n}}} \left(\mathbf {H} ^ {\prime}\right) \right\| _ {2, \infty} \\ \leq \left(1 + B ^ {2}\right) \left(1 + B ^ {3} \mathsf {R} ^ {2}\right) \| \mathbf {H} - \mathbf {H} ^ {\prime} \| _ {2, \infty}, \\ \end{array}
$$
which also follows from Lemma L.2 and Lemma L.1.
Proposition L.1 (Lipschitzness of transformers). For a fixed number of heads  $M$  and hidden dimension  $D'$ , we consider
$$
\Theta_ {\mathrm {T F}, L, B} = \left\{\boldsymbol {\theta} = \left(\boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {(1: L)}, \boldsymbol {\theta} _ {\mathrm {m l p}} ^ {(1: L)}\right): M ^ {(\ell)} = M, D ^ {(\ell)} = D ^ {\prime}, \| \boldsymbol {\theta} \| \leq B \right\}.
$$
Then the function  $\mathrm{TF}^{\mathsf{R}}$  is  $(LB_{H}^{L - 1}B_{\Theta})$ -Lipschitz w.r.t  $\pmb{\theta}\in \Theta_{\mathrm{TF},L,B}$  for any fixed  $\mathbf{H}$ .
Proof. For  $\pmb{\theta} = \pmb{\theta}^{(1:L)} \in \Theta_{\mathrm{TF},L,B}$ ,  $\widetilde{\pmb{\theta}} = \widetilde{\pmb{\theta}}^{(1:L)} \in \Theta_{\mathrm{TF},L,B}$ , we have
$$
\left\| \mathrm {T F} _ {\boldsymbol {\theta}} ^ {\mathrm {R}} (\mathbf {H}) - \mathrm {T F} _ {\boldsymbol {\theta}} ^ {\mathrm {R}} (\mathbf {H}) \right\| _ {2, \infty}
$$
$$
\begin{array}{l} \leq \sum_ {\ell = 1} ^ {L} \left\| \mathrm {T F} _ {\boldsymbol {\theta} ^ {(\ell + 1: L)}} ^ {\mathsf {R}} \left(\mathrm {T F} _ {\boldsymbol {\theta} ^ {(\ell)}} ^ {\mathsf {R}} \left(\mathrm {T F} _ {\hat {\boldsymbol {\theta}} ^ {(1: \ell - 1)}} ^ {\mathsf {R}} (\mathbf {H})\right)\right) - \mathrm {T F} _ {\boldsymbol {\theta} ^ {(\ell + 1: L)}} ^ {\mathsf {R}} \left(\mathrm {T F} _ {\hat {\boldsymbol {\theta}} ^ {(\ell)}} ^ {\mathsf {R}} \left(\mathrm {T F} _ {\hat {\boldsymbol {\theta}} ^ {(1: \ell - 1)}} ^ {\mathsf {R}} (\mathbf {H})\right)\right) \right\| _ {2, \infty} \\ \leq \sum_ {\ell = 1} ^ {L} B _ {\Theta} ^ {L - \ell} \left\| \mathrm {T F} _ {\boldsymbol {\theta} ^ {(\ell)}} ^ {\mathsf {R}} \left(\mathrm {T F} _ {\tilde {\boldsymbol {\theta}} ^ {(1: \ell - 1)}} ^ {\mathsf {R}} (\mathbf {H})\right) - \mathrm {T F} _ {\tilde {\boldsymbol {\theta}} ^ {(\ell)}} ^ {\mathsf {R}} \left(\mathrm {T F} _ {\tilde {\boldsymbol {\theta}} ^ {(1: \ell - 1)}} ^ {\mathsf {R}} (\mathbf {H})\right) \right\| _ {2, \infty} \\ \leq \sum_ {\ell = 1} ^ {L} B _ {H} ^ {L - \ell} B _ {\Theta} \cdot \left\| \boldsymbol {\theta} ^ {(\ell)} - \widetilde {\boldsymbol {\theta}} ^ {(\ell)} \right\| \leq L B _ {H} ^ {L - 1} B _ {\Theta} \cdot \left\| \boldsymbol {\theta} - \widetilde {\boldsymbol {\theta}} \right\|, \\ \end{array}
$$
where the second inequality follows from Corollary L.1, and the last inequality is because  $B_H \geq 1$ .
L.2 Proof of Theorem K.1
In this section, we prove a slightly more general result by considering the general ICL loss
$$
\ell_ {\mathrm {i c l}} (\boldsymbol {\theta}; \mathbf {Z}) := \ell (\widetilde {\operatorname {r e a d} _ {\mathcal {Y}}} (\operatorname {T F} _ {\boldsymbol {\theta}} ^ {\mathbf {R}} (\mathcal {H})), y _ {N + 1}).
$$
We assume that the loss function  $\ell$  satisfies  $\sup |\ell| \leq B_{\ell}^{0}$  and  $\sup |\partial_1\ell| \leq B_{\ell}^{1}$ . For the special case  $\ell(s,t) = \frac{1}{2}(s - t)^2$ , we can take  $B_{\ell}^{0} = 4B_{y}^{2}$ ,  $B_{\ell}^{1} = 2B_{y}$ .
We then consider
$$
X _ {\boldsymbol {\theta}} := \frac {1}{n} \sum_ {j = 1} ^ {n} \ell_ {\mathrm {i c l}} (\boldsymbol {\theta}; \mathbf {Z} ^ {j}) - \mathbb {E} _ {\mathbf {Z}} [ \ell_ {\mathrm {i c l}} (\boldsymbol {\theta}; \mathbf {Z}) ],
$$
where  $\mathbf{Z}^{(1:n)}$  are i.i.d copies of  $\mathbf{Z}\sim \mathsf{P},\mathsf{P}\sim \pi$ . It remains to apply Proposition B.4 to the random process  $\{X_{\theta}\}$ . We verify the preconditions:
(a) By [87, Example 5.8], it holds that  $\log N(\delta; \mathsf{B}_{\| \cdot \|}(r), \| \cdot \|) \leq L(3MD^2 + 2DD') \log (1 + 2r / \delta)$ , where  $\mathsf{B}_{\| \cdot \|}(r)$  is any ball of radius  $r$  under norm  $\| \cdot \|$ .
(b)  $|\ell_{\mathrm{icl}}(\pmb {\theta};\mathbf{Z})|\leq B_{\ell}^{0}$  and hence  $B_{\ell}^{0}$  -sub-Gaussian.
(c)  $\left| \ell_{\mathrm{icl}}(\pmb{\theta}; \mathbf{Z}) - \ell_{\mathrm{icl}}(\widetilde{\pmb{\theta}}; \mathbf{Z}) \right| \leq B_{\ell}^{1} \cdot (LB_{H}^{L-1}B_{\Theta}) \cdot \left\| \pmb{\theta} - \widetilde{\pmb{\theta}} \right\|$ , by Proposition L.1.
Therefore, we can apply the uniform concentration result in Proposition B.4 to obtain that, with probability at least  $1 - \xi$ ,
$$
\sup  _ {\boldsymbol {\theta}} | X _ {\boldsymbol {\theta}} | \leq C B _ {\ell} ^ {0} \sqrt {\frac {L (M D ^ {2} + D D ^ {\prime}) \iota + \log (1 / \xi)}{n}},
$$
where  $\iota = \log (2 + B\cdot LB_H^{L - 1}B_\Theta B_\ell^1 /B_\ell^0)\leq 20L\log (2 + \max \{B,\mathsf{R},B_\ell^1 /B_\ell^0\})$  . Recalling that
$$
L _ {i c l} (\widehat {\boldsymbol {\theta}}) \leq \inf  _ {\boldsymbol {\theta}} L _ {i c l} (\boldsymbol {\theta}) + 2 \sup  _ {\boldsymbol {\theta}} | X _ {\boldsymbol {\theta}} |
$$
L.3 Proof of Theorem K.2
By Corollary 5, there exists a transformer  $\mathrm{TF}_{\theta}$  such that for every  $\mathsf{P}$  satisfying Assumption A with canonical parameters (and thus in expectation over  $\mathsf{P} \sim \pi$ ) and every  $N \geq \widetilde{\mathcal{O}}(d)$ , it outputs prediction  $\widehat{y}_{N+1} = \widetilde{\mathrm{read}}_{\mathsf{y}}(\mathrm{TF}_{\theta}(\mathbf{H}))$  such that
$$
L _ {\mathrm {i c l}} (\boldsymbol {\theta}) = \mathbb {E} _ {\mathsf {P} \sim \pi , (\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1} \sim \mathsf {P})} \left[ \frac {1}{2} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] \leq \mathbb {E} _ {\mathsf {P} \sim \pi} \left[ L _ {\mathsf {P}} \left(\mathbf {w} _ {\mathsf {P}} ^ {\star}\right) \right] + \mathcal {O} \left(\frac {d \sigma^ {2}}{N}\right),
$$
where we recall that  $L_{\mathsf{P}}(\mathbf{w}_{\mathsf{P}}^{\star}) \coloneqq \frac{1}{2}\mathbb{E}_{(\mathbf{x},y)\sim \mathsf{P}}[(y - \langle \mathbf{w}_{\mathsf{P}}^{\star},\mathbf{x}\rangle)^{2}]$ . By inspecting the proof, the same result holds if we change  $\mathrm{TF}_{\theta}$  to the clipped version  $\mathrm{TF}_{\theta}^{\mathsf{R}}$  if we choose  $\mathsf{R}^2 = \mathcal{O}(B_x^2 +B_y^2 +B_w^2 +1) = \mathcal{O}(d + \kappa)$ , so that on the good event  $E_{\mathrm{cov}}\cap E_w$  considered therein, all intermediate outputs within
$\mathrm{TF}_{\theta}$  has  $\| \cdot \|_{2,\infty} \leq R$  and thus the clipping does not modify the transformer output on  $E_{\mathrm{cov}} \cap E_w$ . Further, recall by (33) that  $\theta$  has size bounds
$$
L \leq \mathcal {O} \left(\kappa \log \frac {N \kappa}{\sigma}\right), \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq 3, \quad \| \boldsymbol {\theta} \| \leq \mathcal {O} (\sqrt {\kappa d}).
$$
We can thus apply Theorem K.1 to obtain that the solution  $\widehat{\pmb{\theta}}$  to (TF-ERM) with the above choice of  $(L,M,B)$  and  $D^{\prime} = 0$  (attention-only) satisfies the following with probability at least  $1 - \xi$ :
$$
\begin{array}{l} L _ {i c l} (\widehat {\boldsymbol {\theta}}) \leq \inf  _ {\boldsymbol {\theta} ^ {\prime} \in \Theta_ {L, M, D ^ {\prime}, B}} L _ {i c l} \left(\boldsymbol {\theta} ^ {\prime}\right) + \mathcal {O} \left(\sqrt {\frac {L ^ {2} M D ^ {2} \iota + \log (1 / \xi)}{n}}\right) \\ \leq L _ {\mathrm {i c l}} (\boldsymbol {\theta}) + \widetilde {\mathcal {O}} \left(\sqrt {\frac {L ^ {2} M D ^ {2} + \log (1 / \xi)}{n}}\right) \leq \widetilde {\mathcal {O}} \left(\sqrt {\frac {\kappa^ {2} d ^ {2} + \log (1 / \xi)}{n}} + \frac {d \sigma^ {2}}{N}\right). \\ \end{array}
$$
Above,  $\iota = \mathcal{O}(\log (1 + \max \{B_y,\mathsf{R},B\})) = \widetilde{\mathcal{O}} (1)$ . This finishes the proof.
L.4 Proof of Theorem K.3
We invoke Theorem 8 (using the construction in Theorem 7 with a different choice of  $L$ ) with the following parameters:
$$
\begin{array}{l} L = \widetilde {\mathcal {O}} \big ((B _ {w} ^ {\star}) ^ {2} / \sigma^ {2} \times (1 + d / N) \big) = \widetilde {\mathcal {O}} \big (\kappa^ {2} (1 + d / N) \big), \quad M = \Theta (1), \quad D ^ {\prime} = 2 d, \\ B _ {x} = \widetilde {\mathcal {O}} (\sqrt {d}), \quad B _ {y} = \widetilde {\mathcal {O}} (B _ {w} ^ {\star} + \sigma), \quad \delta = \left(\sigma^ {2} \frac {1}{B _ {y} ^ {2} N}\right) ^ {2}, \\ \| \boldsymbol {\theta} \| \leq B = \mathcal {O} \left(R + (1 + \lambda_ {N}) \beta^ {- 1}\right) \leq \mathcal {O} \left(R + \sigma \sqrt {\log d}\right) \leq \widetilde {\mathcal {O}} (\operatorname {p o l y} (d, B _ {w} ^ {\star}, \sigma)), \\ \end{array}
$$
where  $\tilde{\mathcal{O}} (\cdot)$  hides polylogarithmic factors in  $d,N,B_w^\star ,\kappa$
Then, Theorem 8 shows that there exists a transformer  $\pmb{\theta}$  with  $L$  layers,  $\max_{\ell \in [L]}M^{(\ell)}\leq M$  heads,  $D^{\prime}$  hidden dimension for the MLP layers, and  $\| \pmb {\theta}\| \leq B$  such that, on almost surely every  $\mathsf{P}\sim \pi$ , it returns a prediction  $\widehat{y}_{N + 1}$  such that, on the good event  $\mathcal{E}_0$  considered therein (over  $\mathcal{D}\sim \mathsf{P}$ ) which satisfies  $\mathsf{P}(\mathcal{E}_0)\geq 1 - \delta$ ,
By inspecting the proof, the same result holds if we change  $\mathrm{TF}_{\theta}$  to the clipped version  $\mathrm{TF}_{\theta}^{\mathsf{R}}$  if we choose  $\mathsf{R}^2 = \mathcal{O}(B_x^2 + B_y^2 + (B_w^\star)^2 + 1) = \mathcal{O}(d + (B_w^\star)^2 + \sigma^2)$ , so that on the good event  $\mathcal{E}_0$  considered therein, all intermediate outputs within  $\mathrm{TF}_{\theta}$  has  $\| \cdot \|_{2,\infty} \leq \mathsf{R}$  and thus the clipping does not modify the transformer output on the good event. On the bad event  $\mathcal{E}_0^c$ , using the same argument as in the proof of Theorem 8, we have
$$
\mathbb {E} _ {\mathcal {D}, (\mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P}} \left[ 1 \{\mathcal {E} _ {0} ^ {c} \} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] \leq \sqrt {\mathsf {P} _ {\mathcal {D}} \left(\mathcal {E} _ {0} ^ {c}\right)} \cdot \left(8 \mathbb {E} _ {y _ {N + 1} \sim \mathsf {P}} \left[ B _ {y} ^ {4} + y _ {N + 1} ^ {4} \right]\right) ^ {1 / 2} \leq \widetilde {\mathcal {O}} \left(\frac {\sigma^ {2}}{N}\right).
$$
Combining the above two bounds and further taking expectation over  $\mathsf{P} \sim \pi$  gives
$$
L _ {\mathrm {i c l}} (\boldsymbol {\theta}) = \mathbb {E} _ {\mathsf {P} \sim \pi , (\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P}} \left[ \frac {1}{2} \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] \leq \sigma^ {2} + \widetilde {\mathcal {O}} \big (\sigma^ {2} s \log d / N \big).
$$
We can thus apply Theorem K.1 to obtain that the solution  $\widehat{\pmb{\theta}}$  to (TF-ERM) with the above choice of  $(L,M,B,D^{\prime})$  satisfies the following with probability at least  $1 - \xi$ :
$$
\begin{array}{l} L _ {i c l} (\widehat {\boldsymbol {\theta}}) \leq \inf  _ {\boldsymbol {\theta} ^ {\prime} \in \Theta_ {L, M, D ^ {\prime}, B}} L _ {i c l} \left(\boldsymbol {\theta} ^ {\prime}\right) + \mathcal {O} \left(\sqrt {\frac {L ^ {2} \left(M D ^ {2} + D D ^ {\prime}\right) \iota + \log (1 / \xi)}{n}}\right) \\ \leq L _ {\mathrm {i c l}} (\boldsymbol {\theta}) + \widetilde {\mathcal {O}} \left(\sqrt {\frac {L ^ {2} \left(M D ^ {2} + D D ^ {\prime}\right) + \log (1 / \xi)}{n}}\right) \\ \leq \sigma^ {2} + \widetilde {\mathcal {O}} \left(\sqrt {\frac {\kappa^ {4} d ^ {2} (1 + d / N) ^ {2} + \log (1 / \xi)}{n}} + \sigma^ {2} \frac {s \log d}{N}\right). \\ \end{array}
$$
L.5 Proof of Theorem K.4
We invoke Theorem 12 and Theorem J.1, which shows that (recalling the input dimension  $D = \Theta(Kd)$ ) there exists a transformer  $\theta$  with the following size bounds:
$$
L \leq \mathcal {O} \left(\sigma_ {\min } ^ {- 2} \log (N / \sigma_ {\min })\right), \quad \max _ {\ell \in [ L ]} M ^ {(\ell)} \leq M = \mathcal {O} \left(K\right), \quad \max _ {\ell \in [ L ]} D ^ {(\ell)} \leq D ^ {\prime} = \mathcal {O} (K ^ {2}),
$$
such that it outputs  $\widehat{y}_{N + 1}$  that satisfies
$$
\mathbb {E} _ {\pi} \left[ \frac {1}{2} \left(y _ {N + 1} - \widehat {y} _ {N + 1}\right) ^ {2} \right] \leq \text {B a y e s R i s k} _ {\pi} + \widetilde {\mathcal {O}} \left(\frac {\sigma_ {\max } ^ {2}}{\sigma_ {\min } ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right).
$$
By inspecting the proof, the same result holds if we change  $\mathrm{TF}_{\theta}$  to the clipped version  $\mathrm{TF}_{\theta}^{\mathsf{R}}$  if we choose  $\mathsf{R}^2 = \mathcal{O}(B_x^2 + B_y^2 + (B_w^\star)^2 + 1) = \mathcal{O}(d + \sigma_{\max}^2)$ , so that on the good event considered therein, all intermediate outputs within  $\mathrm{TF}_{\theta}$  has  $\| \cdot \|_{2,\infty} \leq \mathsf{R}$  and thus the clipping does not modify the transformer output on the good event. Using this clipping radius, we obtain
$$
L _ {\mathrm {i c l}} (\boldsymbol {\theta}) = \mathbb {E} _ {\mathsf {P} \sim \pi , (\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P}} \left[ \frac {1}{2} (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} \right] \leq \text {B a y e s R i s k} _ {\pi} + \widetilde {\mathcal {O}} \left(\frac {\sigma_ {\max } ^ {2}}{\sigma_ {\min } ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right).
$$
$$
\begin{array}{l} L _ {\mathrm {i c l}} (\widehat {\boldsymbol {\theta}}) \leq \inf  _ {\boldsymbol {\theta} ^ {\prime} \in \Theta_ {L, M, D ^ {\prime}, B}} L _ {\mathrm {i c l}} \left(\boldsymbol {\theta} ^ {\prime}\right) + \mathcal {O} \left(\sqrt {\frac {L ^ {2} \left(M D ^ {2} + D D ^ {\prime}\right) \iota + \log (1 / \xi)}{n}}\right) \\ \leq L _ {i c l} (\boldsymbol {\theta}) + \widetilde {\mathcal {O}} \left(\sqrt {\frac {L ^ {2} \left(M D ^ {2} + D D ^ {\prime}\right) + \log (1 / \xi)}{n}}\right) \\ \leq \text {B a y e s R i s k} _ {\pi} + \widetilde {\mathcal {O}} \left(\sqrt {\frac {\sigma_ {\operatorname* {m i n}} ^ {- 4} K ^ {3} d ^ {2} + \log (1 / \xi)}{n}} + \frac {\sigma_ {\operatorname* {m a x}} ^ {2}}{\sigma_ {\operatorname* {m i n}} ^ {2 / 3}} \left(\frac {\log K}{N}\right) ^ {1 / 3}\right). \\ \end{array}
$$
L.6 Proof of Theorem K.5
The proof follows from similar arguments as of Theorem K.3 and Theorem K.4, where we plug in the size bounds (number of layers, heads, and weight norms) from Theorem G.2 and Corollary G.1.  $\square$
M Experimental details and additional studies
M.1 Additional details for Section 6
Architecture and optimization We train a 12-layer encoder-only transformer, where each layer consists of an attention layer as in Definition 1 with  $M = 8$  heads, hidden dimension  $D = 64$ , and ReLU activation (normalized by the sequence length), as well as an MLP layer as in Definition 2. Hidden dimension  $D' = 64$ . We add Layer Normalization [3] after each attention and MLP layer to help optimization, as in standard implementations [84]. We append linear read-in layer and linear read-out layer before and after the transformer respectively, both applying a same affine transform to all tokens in the sequence and are trainable. The read-in layer maps any input vector to a  $D$ -dimensional hidden state, and the read-out layer maps a  $D$ -dimensional hidden state to a 1-dimensional scalar.
Each training sequence corresponds to a single ICL instance with  $N$  in-context training examples  $\{(\mathbf{x}_i, y_i)\}_{i=1}^N \subset \mathbb{R}^d \times \mathbb{R}$  and test input  $\mathbf{x}_{N+1} \in \mathbb{R}^d$ . The input to the transformer is formatted as
in (3) where each token has dimension  $d + 1$  (no zero-paddings). The transformer is trained by minimizing the following loss with fresh mini-batches:
$$
L (\boldsymbol {\theta}) = \mathbb {E} _ {\mathsf {P} \sim \pi , (\mathbf {H}, y _ {N + 1}) \sim \mathsf {P}} [ \ell_ {\mathsf {P}} (\operatorname {r e a d} _ {\mathcal {Y}} (\operatorname {T F} _ {\boldsymbol {\theta}} (\mathbf {H})), y _ {N + 1}) ], \tag {47}
$$
where the loss function  $\ell_{\mathsf{P}}:\mathbb{R}^2\to \mathbb{R}$  may depend on the training data distribution  $\mathsf{P}$  in general; we use the square loss when  $\mathsf{P}$  is regression data, and the logistic loss when  $\mathsf{P}$  is classification data. We use the Adam optimizer with a fixed learning rate  $10^{-4}$ , which we find works well for all our experiments. Throughout all our experiments except for the sparse linear regression experiment in Figure 3a, we train the model for 300K steps, where each step consists of a (fresh) minibatch with batch size 64 in the base mode, and  $K$  minibatches each with batch size 64 in the mixture mode.
For the sparse linear regression experiment, we find that minimizing the training objective (47) alone was not enough, e.g. for the learned transformer to achieve better loss than the least squares algorithm (which achieves much higher test loss than the Lasso; cf. Figure 3a). To help optimization, we augment (47) with another loss that encourages the second-to-last hidden states to recover the true (sparse) coefficient  $\mathbf{w}_{\star}$ :
$$
L _ {\text {f i t - w}} (\boldsymbol {\theta}) = \frac {1}{N _ {0}} \sum_ {j = 1} ^ {N _ {0}} \mathbb {E} _ {\mathrm {P} = \mathrm {P} _ {\mathbf {w} ^ {\star}} \sim \pi , (\mathbf {H}, y _ {N + 1}) \sim \mathrm {P}} \left[ \left\| \left[ \mathrm {T F} _ {\boldsymbol {\theta}} ^ {(1: L - 1)} (\mathbf {H}) \right] _ {j, (D - d + 1): D} - \mathbf {w} ^ {\star} \right\| _ {2} ^ {2} \right]. \tag {48}
$$
Specifically, the above loss encourages the first  $N_0 \leq N$  tokens within the second-to-last layer to be close to  $\mathbf{w}^{\star}$ . We choose  $N_0 = 5$  (recall that the total number of tokens is  $N = 10$  and sequence length is  $N + 1 = 11$  for this experiment). We minimize the loss  $L(\pmb{\theta}) + \lambda L_{\mathrm{fit - w}}(\pmb{\theta})$  with  $\lambda = 0.1$  for 2M steps for this task.
Evaluation All evaluations are done on the trained transformer with 6400 test instances. We use the square loss for regression tasks, and the classification error (1-accuracy) between the true label  $y_{N + 1} \in \{0,1\}$  and the predicted label  $1\{\widehat{y}_{N + 1} \geq 1/2\}$ . We report the means in all experiments, as well as their standard deviations (using one-standard error bars) in Figure 2a, 2b, 5a, 5b. In Figure 2c, 3b, 3c, 5c, all standard deviations are sufficiently small (not significantly exceeding the width of the markers), thus we did not show error bars in those plots.
Baseline algorithms We implement various baseline machine learning algorithms to compare with the learned transformers. A superset of the algorithms is shown in Figure 3a:
- Least squares, Logistic regression: Standard algorithms for linear regression and linear classification, respectively. Note that least squares is also a valid algorithm for classification.
- Averaging: The simple algorithm which computes the linear predictor  $\widehat{\mathbf{w}} = \frac{1}{N}\sum_{i=1}^{N}y_{i}\mathbf{x}_{i}$  and predicts  $\widehat{y}_{N+1} = \langle \widehat{\mathbf{w}},\mathbf{x}_{N+1}\rangle$ ;
- 3-NN: 3-Nearest Neighbors.
- Ridge: Standard ridge regression as in (ICRidge). We specifically consider two  $\lambda$ 's (denoted as lam_1 and lam_2):  $\lambda_1, \lambda_2 = (0.005, 0.125)$ . These are the Bayes-optimal regularization strengths for the noise levels  $(\sigma_1, \sigma_2) = (0.1, 0.5)$  respectively under the noisy linear model (cf. Corollary 6), using the formula  $\lambda^{\star} = d\sigma^{2}/N$ , with  $(d, N) = (20, 40)$ .
- Lasso: Standard Lasso as in (ICLasso) with  $\lambda \in \{1,0.1,0.01,0.001\}$ .
In Figure 2c, the ridge_analytical curve plots the expected risk of ridge regression under the noisy linear model over 20 geometrically spaced values of  $\lambda$ 's in between  $(\lambda_1, \lambda_2)$ , using analytical formulae (with Monte Carlo simulations). The Bayes_err_{\{1,2\}} indicate the expected risks of  $\lambda_1$  on task 1 (with noise  $\sigma_1$ ) and  $\lambda_2$  on task 2 (with noise  $\sigma_2$ ), respectively.
M.2 Decoder-based architecture
ICL capabilities have also been demonstrated in the literature for decoder-based architectures [31, 2, 47]. There, the transformer can do in-context predictions at every token  $\mathbf{x}_i$  using past tokens  $\{(\mathbf{x}_j,\mathbf{y}_j)\}_{j\leq i - 1}$  as training examples. Here we show that such architectures is also able to perform in-context algorithm selection at every token; For results for this architecture on "base" ICL tasks (such as those considered in Figure 3a), we refer the readers to Garg et al. [31].
(a) Linear regression
(b) Linear classification
Figure 5: In-context algorithm selection abilities of transformers between linear regression and linear classification.  $(a,b)$  On these two tasks, a single transformer TF_alg_select simultaneously approaches the performance of the strongest baseline algorithm Least Squares on linear regression and Logistic Regression on linear classification. (c) At token 40 (using example  $\{0,\dots ,39\}$  for training), TF_alg_select matches the performance of the best baseline algorithm for both tasks.  $(a,b,c)$  Note that transformers pretrained on a single task (TF_reg, TF_c1s) perform near-optimally on their pretraining task but suboptimally on the other task.
(c) Reg vs. cls at token 40
Setup Our setup is the same as the two "mixture" modes (linear model + linear classification model, and noisy linear models with two different noise levels) as in Section 6, except that the architecture is GPT-2 following Garg et al. [31], and the input format is changed to (11) (so that the input sequence has  $2N + 1$  tokens) without positional encodings. For every  $i \in [N + 1]$ , we extract the prediction  $\widehat{y}_i$  using a linear read-out function applied on output token  $2i - 1$ , and the (learnable) linear read-out function is the same across all tokens, similar as in Appendix M.1. The rest of the setup (optimization, training, and evaluation) is the same as in Section 6 & M.1. Note that we also train on the objective (47) for all tokens averaged, instead of for the last test token as in Section 6.
Result Figure 2 shows the results for noisy linear models with two different noise levels, and Figure 5 shows the results for linear model + linear classification model. We observe that at every token, in both cases, TF_alg_select nearly matches the strongest baseline for both tasks simultaneously, whereas transformers trained on a single task perform suboptimally on the other task. Further, this phenomenon consistently shows up at every token. For example, in Figure 2a & 2b, TF_alg_select matches ridge regression with the optimal  $\lambda$  on all tokens  $i \in \{1, \dots, N\}$  ( $N = 40$ ). In Figure 5a & 5b, TF_alg_select matches least squares on the regression task and logistic regression on the classification task on all tokens  $i \in [N]$ . This demonstrates the in-context algorithm selection capabilities of standard decoder-based transformer architectures.
M.3 Computational resource
All our experiments are performed on 8 Nvidia Tesla A100 GPUs (40GB memory). The total GPU time is approximately 5 days (on 8 GPUs), with the largest individual training run taking about a single day on a single GPU.