SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Corollary 6 (Nearly-Bayes linear regression with transformers by approximating ridge regression). Under the Bayesian linear model above with  $N \geq \max \{d / 10, \mathcal{O}(\log(1 / \varepsilon))\}$ , there exists a  $L = \mathcal{O}(\log(1 / \varepsilon))$ -layer transformer such that  $\mathbb{E}_{\mathbf{w}_{\star}, (\mathcal{D}, \mathbf{x}_{N+1}, y_{N+1})} \left[ \frac{1}{2} (\widehat{y}_{N+1} - y_{N+1})^2 \right] \leq \text{BayesRisk}_{\pi} + \varepsilon$ .
Generalized linear models In Appendix G, we extend the above results to generalized linear models [53] and show that transformers can approximate the corresponding convex risk minimization algorithm in context (which includes logistic regression for linear classification as an important special case), and achieve near-optimal excess risk under standard statistical assumptions.
3.2 In-context Lasso
Consider the standard Lasso estimator [82] which minimizes an  $\ell_1$ -regularized linear regression loss  $\widehat{L}_{\mathrm{lasso}}$  over the in-context training examples  $\mathcal{D}$ :
$$
\mathbf {w} _ {\mathrm {l a s s o}} := \arg \min _ {\mathbf {w} \in \mathbb {R} ^ {d}} \widehat {L} _ {\mathrm {l a s s o}} (\mathbf {w}) = \frac {1}{2 N} \sum_ {i = 1} ^ {N} \left(\langle \mathbf {w}, \mathbf {x} _ {i} \rangle - y _ {i}\right) ^ {2} + \lambda_ {N} \| \mathbf {w} \| _ {1}. \qquad (\mathrm {I C L a s s o})
$$
We show that transformers can also approximate in-context Lasso with a mild number of layers, and can perform sparse linear regression in standard sparse linear models (proofs in Appendix H).
Theorem 7 (Implementing in-context Lasso). For any  $\lambda_N \geq 0$ ,  $\beta > 0$ ,  $B_w > 0$ , and  $\varepsilon > 0$ , there exists a  $L$ -layer transformer  $\mathrm{TF}_{\theta}$  with
$$
L = \left\lceil \beta B _ {w} ^ {2} / \varepsilon \right\rceil + 1, \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq 2, \quad \max  _ {\ell \in [ L ]} D ^ {(\ell)} \leq 2 d, \quad \| \boldsymbol {\theta} \| \leq \mathcal {O} \big (R + (1 + \lambda_ {N}) \beta^ {- 1} \big)
$$
(where  $R\coloneqq \max \left\{B_xB_w,B_y,1\right\}$  ) such that the following holds. On any input data  $(\mathcal{D},\mathbf{x}_{N + 1})$  such that  $\lambda_{\mathrm{max}}(\mathbf{X}^{\top}\mathbf{X} / N)\leq \beta$  and  $\| \mathbf{w}_{\mathrm{lasso}}\| _2\leq B_w / 2$ $\mathrm{TF}_{\theta}(\mathbf{H}^{(0)})$  approximately implements (ICLasso), in that it outputs  $\widehat{y}_{N + 1} = \langle \mathbf{x}_{N + 1},\widehat{\mathbf{w}}\rangle$  with  $\widehat{L}_{\mathrm{lasso}}(\widehat{\mathbf{w}}) - \widehat{L}_{\mathrm{lasso}}(\mathbf{w}_{\mathrm{lasso}})\leq \varepsilon$
Theorem 8 (Near-optimal sparse linear regression with transformers by approximating Lasso). For any  $d, N \geq 1, \delta > 0, B_w^\star, \sigma > 0$ , there exists a  $\widetilde{\mathcal{O}}((B_w^\star)^2 / \sigma^2 \times (1 + (d / N)))$ -layer transformer  $\pmb{\theta}$  such that the following holds: For any  $s$  and  $N \geq \mathcal{O}$  ( $s \log(d / \delta)$ ), suppose that  $\mathsf{P}$  is an  $s$ -sparse linear model:  $\mathbf{x}_i \sim \mathsf{N}(0, \mathbf{I}_d)$ ,  $y_i = \langle \mathbf{w}_\star, \mathbf{x}_i \rangle + \mathsf{N}(0, \sigma^2)$  for any  $\| \mathbf{w}_\star \|_2 \leq B_w^\star$  and  $\| \mathbf{w}_\star \|_0 \leq s$ , then with probability at least  $1 - \delta$  (over the randomness of  $\mathcal{D}$ ), the transformer output  $\widehat{y}_{N+1}$  achieves
$$
\mathbb {E} _ {\left(\mathbf {x} _ {N + 1}, y _ {N + 1}\right) \sim \mathsf {P}} \left[ \left(\widehat {y} _ {N + 1} - y _ {N + 1}\right) ^ {2} \right] \leq \sigma^ {2} \left[ 1 + \mathcal {O} \left(s \log (d / \delta) / N\right) \right].
$$
The  $\widetilde{\mathcal{O}}(s\log d / N)$  excess risk obtained in Theorem 8 is optimal up to log factors [62, 87]. We remark that Theorem 8 is not a direct corollary of Theorem 7; Rather, the bound on the number of layers in Theorem 8 requires a sharper convergence analysis of the (ICLasso) problem under sparse linear models (Appendix H.2), similar to [1].
3.3 Proof technique: In-context gradient descent
The constructions in Section 3.1 and 3.2 is built on the following result for approximating in-context (proximal) gradient descent on (regularized) convex losses.
Theorem 9 (ICGD; Informal version of Theorem D.1 & D.2). For a broad class of convex losses of form  $\mathbf{w} \mapsto \frac{1}{N} \sum_{i=1}^{N} \ell(\mathbf{w}^{\top} \mathbf{x}_i, y_i) + R(\mathbf{w})$ , there exists an  $L$ -layer transformer that takes in any  $(\mathcal{D}, \mathbf{w}^0)$  and outputs  $\widehat{\mathbf{w}}^L$  such that  $\| \widehat{\mathbf{w}}^L - \mathbf{w}_{\{\mathrm{GD}, \mathrm{PGD}\}}^L \|_2 \leq \mathcal{O}(L\varepsilon)$ , by composing  $L$  identical layers each  $\mathcal{O}(\varepsilon)$ -approximating a single step of GD (so that  $\mathcal{O}(L\varepsilon)$  is a linear error accumulation).
Theorem 9 is established in two main steps:
- Approximating one-step of ICGD using one attention layer (Proposition E.1), which substantially generalizes that of von Oswald et al. [86] (which only does GD on square losses with a linear self-attention), and is simpler than the ones in Akyurek et al. [2] and Giannou et al. [32].
- Stacking  $L$  of the above layer to approximate  $L$  steps of ICGD. Done naively, the error accumulation of this stacking operation is exponential in  $L$  in the worst case. We utilize the stability of convex gradient descent (Lemma D.1) to obtain the linear in  $L$  error accumulation in Theorem 9.
In Appendix D.3, we also give results for non-convex GD on two-layer neural nets, though with a worse (exponential in  $L$ ) error accumulation as expected.
4 In-context algorithm selection
We now show that transformers can perform various kinds of in-context algorithm selection, which allows them to implement more complex ICL procedures by adaptively selecting different "base" algorithms on different input sequences. We construct two general mechanisms: Post-ICL validation, and Pre-ICL testing; See Figure 1 for a pictorial illustration.
4.1 Post-ICL validation mechanism
In our first mechanism, post-ICL validation, the transformer begins by implementing a train-validation split  $\mathcal{D} = (\mathcal{D}_{\mathrm{train}},\mathcal{D}_{\mathrm{val}})$ , and running  $K$  base ICL algorithms on  $\mathcal{D}_{\mathrm{train}}$ . Let  $\{f_k\}_{k\in [K]}\subset (\mathbb{R}^d\to \mathbb{R})$  denote the  $K$  learned predictors, and
$$
\widehat {L} _ {\text {v a l}} (f) := \frac {1}{| \mathcal {D} _ {\text {v a l}} |} \sum_ {\left(\mathbf {x} _ {i}, y _ {i}\right) \in \mathcal {D} _ {\text {v a l}}} \ell \left(f \left(\mathbf {x} _ {i}\right), y _ {i}\right) \tag {7}
$$
denote the validation loss of any predictor  $f$ .
We show that (proof in Appendix I.1) a 3-layer transformer can output a predictor  $\widehat{f}$  that achieves nearly the smallest validation loss, and thus nearly optimal expected loss if  $\widehat{L}_{\mathrm{val}}$  concentrates around the expected loss  $L$ . Below, the input sequence  $\mathbf{H}$  uses a generalized positional encoding  $\mathbf{p}_i \coloneqq [0_{D - (d + 3)}; 1; t_i]$  in (3), where  $t_i \coloneqq 1$  for  $i \in \mathcal{D}_{\mathrm{train}}$ ,  $t_i \coloneqq -1$  for  $i \in \mathcal{D}_{\mathrm{val}}$ , and  $t_{N + 1} \coloneqq 0$ .
Proposition 10 (In-context algorithm selection via train-validation split). Suppose that  $\ell (\cdot ,\cdot)$  in (7) is approximable by sum of relus (Definition D.1, which includes all  $C^3$  -smooth bivariate functions). Then there exists a 3-layer transformer  $\mathrm{TF}_{\theta}$  that maps (defining  $y_{i}^{\prime} = y_{i}1\{i < N + 1\}$ )
$$
\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] \quad \rightarrow \quad \mathbf {h} _ {i} ^ {\prime} = \left[ \mathbf {x} _ {i}; y _ {i} ^ {\prime}; *; \widehat {f} (\mathbf {x} _ {i}); 1; t _ {i} \right], i \in [ N + 1 ],
$$
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 (f _ {k _ {\star}}) + \max  _ {k \in [ K ]} \left| \widehat {L} _ {\text {v a l}} (f _ {k}) - L (f _ {k}) \right| + \gamma .
$$
Ridge regression with in-context regularization selection As an example, we use Proposition 10 to construct a transformer to perform in-context ridge regression with regularization selection according to the unregularized validation loss  $\widehat{L}_{\mathrm{val}}(\mathbf{w})\coloneqq \frac{1}{2|\mathcal{D}_{\mathrm{val}}|}\sum_{(x_i,y_i)\in \mathcal{D}_{\mathrm{val}}}\left(\langle \mathbf{w},\mathbf{x}_i\rangle -y_i\right)^2$  (proof in Appendix I.2). Let  $\lambda_1,\ldots ,\lambda_K\geq 0$  be  $K$  fixed regularization strengths.
Theorem 11 (Ridge regression with in-context regularization selection). There exists a transformer with  $\mathcal{O}(\log(1/\varepsilon))$  layers and  $\mathcal{O}(K)$  heads such that the following holds: On any  $(\mathcal{D}, \mathbf{x}_{N+1})$  well-conditioned (cf. (5)) for all  $\{\lambda_k\}_{k \in [K]}$ , it outputs  $\widehat{y}_{N+1} = \langle \widehat{\mathbf{w}}, \mathbf{x}_{N+1} \rangle$ , where
$$
\left. \right. \operatorname {d i s t} \left( \right.\widehat {\mathbf {w}}, \operatorname {c o n v} \left\{ \right.\widehat {\mathbf {w}} _ {\text {r i d g e}, \text {t r a i n}} ^ {\lambda_ {k}}: \widehat {L} _ {\text {v a l}} \left(\widehat {\mathbf {w}} _ {\text {r i d g e}, \text {t r a i n}} ^ {\lambda_ {k}}\right) \leq \min  _ {k _ {\star} \in [ K ]} \widehat {L} _ {\text {v a l}} \left( \right.\widehat {\mathbf {w}} _ {\text {r i d g e}, \text {t r a i n}} ^ {\lambda_ {k _ {\star}}} + \gamma \left. \right\}\left. \right) \leq \varepsilon .
$$
Above,  $\widehat{\mathbf{w}}_{\mathrm{ridge,train}}^{\lambda}$  denotes the solution to (ICRidge) on the training split  $\mathcal{D}_{\mathrm{train}}$
4.1.1 Nearly Bayes-optimal ICL on noisy linear models with mixed noise levels
We build on Theorem 11 to show that transformers can perform nearly Bayes-optimal ICL when data come from noisy linear models with a mixture of  $K$  different noise levels  $\sigma_{1},\ldots ,\sigma_{K} > 0$ .
Concretely, consider the following data generating model, where we first sample  $\mathsf{P} = \mathsf{P}_{\mathbf{w}_{\star},\sigma_k}\sim \pi$  from  $k\sim \Lambda \in \Delta ([K])$ $\mathbf{w}_{\star}\sim \mathsf{N}(\mathbf{0},\mathbf{I}_d / d)$  , and then sample data  $\{(x_i,y_i)\}_{i\in [N + 1]}\stackrel {\mathrm{iid}}{\sim}\mathsf{P}_{\mathbf{w}_{\star},\sigma_k}$  as
$$
\mathsf {P} _ {\mathbf {w} _ {\star}, \sigma_ {k}}: \mathbf {x} _ {i} \sim \mathsf {N} (\mathbf {0}, \mathbf {I} _ {d}), y _ {i} = \langle \mathbf {x} _ {i}, \mathbf {w} _ {\star} \rangle + \mathsf {N} (0, \sigma_ {k} ^ {2}).
$$
For any fixed  $(N,d)$ , consider the Bayes risk for predicting  $y_{N + 1}$  under this model:
$$
\operatorname {B a y e s R i s k} _ {\pi} := \inf  _ {\mathcal {A}} \mathbb {E} _ {\pi} \left[ \frac {1}{2} (\mathcal {A} (\mathcal {D}) (\mathbf {x} _ {N + 1}) - y _ {N + 1}) ^ {2} \right].
$$
By standard Bayesian calculations, the above Bayes risk is attained when  $\mathcal{A}$  is a certain mixture of  $K$  ridge regressions with regularization  $\lambda_{k} = d\sigma_{k}^{2} / N$ ; however, the mixing weights depend on  $\mathcal{D}$  in a highly non-trivial fashion (see Appendix J.2 for a derivation). By using the post-ICL validation mechanism in Theorem 11, we construct a transformer that achieves nearly the Bayes risk.
Theorem 12 (Nearly Bayes-optimal ICL; Informal version of Theorem J.1). For sufficiently large  $N$ ,  $d$ , there exists a transformer with  $\mathcal{O}(\log N)$  layers and  $\mathcal{O}(K)$  heads such that on the above model, it outputs a prediction  $\widehat{y}_{N+1}$  that is nearly Bayes-optimal:
$$
\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} + \mathcal {O} \left(\left(\log K / N\right) ^ {1 / 3}\right). \tag {8}
$$
In particular, Theorem 12 applies in the proportional setting where  $N, d$  are large and  $N / d = \Theta(1)$  [22], in which case BayesRisk $_{\pi} = \Theta(1)$ , and thus the transformer achieves vanishing excess risk relative to the Bayes risk at large  $N$ .
This substantially strengthens the results of Akyurek et al. [2], who empirically find that transformers can achieve nearly Bayes risk under any fixed noise level. By contrast, Theorem 12 shows that a single transformer can achieve nearly Bayes risk even under a mixture of  $K$  noise levels, with quantitative guarantees. Also, our proof in fact gives a stronger guarantee: The transformer approaches the individual Bayes risks on all  $K$  noise levels simultaneously (in addition to the overall Bayes risk for  $k \sim \Lambda$  as in Theorem 12). We demonstrate this empirically in Section 6 (cf. Figure 3b & 2).
Exact Bayes predictor vs. Post-ICL validation mechanism As BayesRisk $_{\pi}$  is the theoretical lower bound for the risk of any possible ICL algorithm, Theorem 12 implies that our transformer performs similarly as the exact Bayes estimator<sup>5</sup>. Notice that our construction builds on the (generic) post-ICL validation mechanism, rather than a direct attempt of approximating the exact Bayes predictor, whose structure may vary significantly case-by-case. This highlights post-ICL validation as a promising mechanism for approximating the Bayes predictor on broader classes of problems beyond noisy linear models, which we leave as future work.
Generalized linear models with adaptive link function selection As another example of the post-ICL validation mechanism, we construct a transformer that can learn a generalized linear model with adaptively chosen link function for the particular ICL instance; see Theorem J.2.
4.2 Pre-ICL testing mechanism
In our second mechanism, pre-ICL testing, the transformer runs a distribution testing procedure on the input sequence to determine the right ICL algorithm to use. While the test (and thus the mechanism itself) could in principle be general, we focus on cases where the test amounts to computing some simple summary statistics of the input sequence.
To showcase pre-ICL testing, we consider the toy problem of selecting between in-context regression and in-context classification, by running the following binary type check on the input labels  $\{y_{i}\}_{i\in [N]}$ .
$$
\Psi^{\text{binary}}(\mathcal{D}) = \frac{1}{N}\sum_{i = 1}^{N}\psi (y_{i}),\quad \psi (y):= \left\{ \begin{array}{ll}1, & y\in \{0,1\} ,\\ 0, & y\notin [-\varepsilon ,\varepsilon ]\cup [1 - \varepsilon ,1 + \varepsilon ],\\ \text{linear interpolation}, & \text{otherwise}. \end{array} \right.
$$
<sup>5</sup>By the Bayes risk decomposition for square loss, (8) implies that  $\mathbb{E}[(\widehat{y}_{N + 1} - \widehat{y}_{N + 1}^{\mathrm{Bayes}})^2 ]\leq \mathcal{O}((\log K / N)^{1 / 3})$
Lemma 13. There exists a single attention layer with 6 heads that implements  $\Psi^{\mathrm{binary}}$  exactly.
Using this test, we construct a transformer that performs logistic regression when labels are binary, and linear regression with high probability if the label admits a continuous distribution.
Proposition 14 (Adaptive regression or classification; Informal version of Proposition I.4). There exists a transformer with  $\mathcal{O}(\log(1/\varepsilon))$  layers such that the following holds: On any  $\mathcal{D}$  such that  $y_i \in \{0,1\}$ , it outputs  $\widehat{y}_{N+1}$  that  $\varepsilon$ -approximates the prediction of in-context logistic regression.
By contrast, for any distribution  $\mathsf{P}$  whose marginal distribution of  $y$  is not concentrated around  $\{0,1\}$ , with high probability (over  $\mathcal{D}$ ),  $\widehat{y}_{N + 1}\varepsilon$ -approximates the prediction of in-context least squares.
The proofs can be found in Appendix I.3. We additionally show that transformers can implement more complex tests such as a linear correlation test, which can be useful in certain scenarios such as "confident linear regression" (predict only when the signal-to-noise ratio is high); see Appendix I.4.
5 Analysis of pretraining
Building on the expressivity results in Section 3 & 4, we provide the first line of polynomial sample complexity results for pretraining transformers to perform ICL (including with in-context algorithm selection). We begin by providing a generic generalization guarantee for pretraining transformers.
Consider the pretraining ERM problem (TF-ERM), which minimizes the pretraining risk  $\widehat{L}_{\mathrm{icl}}(\cdot)$  over  $n$  pretraining sequences. Let  $L_{\mathrm{icl}}(\cdot)$  denote the corresponding population risk.
Theorem 15 (Generalization of transformers; Informal version of Theorem K.1). The solution  $\widehat{\pmb{\theta}}$  to (TF-ERM) over transformers with  $L$  layers,  $M$  heads per layer, and hidden dimension  $D^{\prime}$  satisfies
$$
L _ {\mathrm {i c l}} (\widehat {\boldsymbol {\theta}}) \leq \inf  _ {\boldsymbol {\theta}} L _ {\mathrm {i c l}} (\boldsymbol {\theta}) + \widetilde {\mathcal {O}} \left(\sqrt {\frac {L ^ {2} (M D ^ {2} + D D ^ {\prime})}{n}}\right).
$$
Theorem 15 builds on standard uniform concentration analysis via chaining (Proposition B.4). Combining Theorem 15 with the in-context linear regression construction in Theorem 4 gives the following end-to-end result on the excess in-context prediction risk of trained transformers.
Theorem 16 (Pretraining transformers for in-context linear regression; Informal version of Theorem K.2). Under Assumption A and  $N \geq \widetilde{\mathcal{O}}(d)$ , the solution  $\widehat{\pmb{\theta}}$  to (TF-ERM) with  $L = \mathcal{O}(\kappa \log(\kappa N / \sigma))$  layers,  $M = 3$  heads,  $D' = 0$  (attention-only as in Theorem 4) achieves small excess ICL risk over the best linear predictor  $\mathbf{w}_{\mathsf{P}}^{\star} \coloneqq \mathbb{E}_{\mathsf{P}}[\mathbf{x}\mathbf{x}^{\top}]^{-1}\mathbb{E}_{\mathsf{P}}[\mathbf{x}y]$  for each  $\mathsf{P}$ :
$$
L _ {\mathrm {i c l}} (\widehat {\pmb {\theta}}) - \mathbb {E} _ {\mathsf {P} \sim \pi} \mathbb {E} _ {(\mathbf {x}, y) \sim \mathsf {P}} \left[ \frac {1}{2} (y - \langle \mathbf {w} _ {\mathsf {P}} ^ {\star}, \mathbf {x} \rangle) ^ {2} \right] \leq \widetilde {\mathcal {O}} \left(\sqrt {\frac {\kappa^ {2} d ^ {2}}{n}} + \frac {d \sigma^ {2}}{N}\right),
$$
See Appendix K.2 for similar results in several additional settings.
We test our theory by studying the ICL and in-context algorithm selection capabilities of transformers, using the encoder-based architecture in our theoretical constructions (Definition 3). Due to limited space, additional experimental details can be found in Appendix M.1. Results with a decoder architecture as in [31, 47] (including the setup of Figure 2) can be found in Appendix M.2.
Training data distributions and evaluation We train a 12-layer transformer, with two modes for the training sequence (instance) distribution  $\pi$ . In the "base" mode, similar to [31, 2, 86, 47], we sample the training instances from one of the following base distributions (tasks), where we first sample  $\mathsf{P} = \mathsf{P}_{\mathbf{w}_{\star}} \sim \pi$  by sampling  $\mathbf{w}_{\star} \sim \mathsf{N}(\mathbf{0}, \mathbf{I}_d / d)$ , and then sample  $\{(x_i, y_i)\}_{i \in [N + 1]}$ $\stackrel{\mathrm{iid}}{\sim} \mathsf{P}_{\mathbf{w}_{\star}}$  as  $\mathbf{x}_i \stackrel{\mathrm{iid}}{\sim} \mathsf{N}(\mathbf{0}, \mathbf{I}_d)$ , and  $y_i$  from one of the following models studied in Section 3:
1. Linear model:  $y_{i} = \langle \mathbf{w}_{\star}, \mathbf{x}_{i} \rangle$ ;
2. Noisy linear model:  $y_{i} = \langle \mathbf{w}_{\star},\mathbf{x}_{i}\rangle +\sigma z_{i}$ , where  $\sigma >0$  is a fixed noise level, and  $z_{i}\sim \mathsf{N}(0,1)$ .
(a) Base ICL capabilities
(b) Noisy reg with two noises
Figure 3: ICL capabilities of the transformer architecture used in our theoretical constructions. (a) On five representative base tasks, transformers approximately match the best baseline algorithm for each task, when pretrained on the corresponding task.  $(b,c)$  A single transformer TF_alg_select simultaneously approaches the performance of the strongest baseline algorithm on two separate tasks:  $(b)$  noisy linear regression with two different noise levels  $\sigma \in \{0.1,0.5\}$ , and  $(c)$  adaptively selecting between regression and classification.
(c) Reg + Classification
3. Sparse linear model:  $y_{i} = \langle \mathbf{w}_{\star}, \mathbf{x}_{i} \rangle$  with  $\| \mathbf{w}_{\star} \|_0 \leq s$ , where  $s < d$  is a fixed sparsity level, and in this case we sample  $\mathbf{w}_{\star}$  from a special prior supported on  $s$ -sparse vectors;
4. Linear classification model:  $y_{i} = \mathrm{sign}(\langle \mathbf{w}_{\star},\mathbf{x}_{i}\rangle)$ .
These base tasks have been empirically investigated by Garg et al. [31], though we remark that our architecture (used in our theory) differs from theirs in several aspects, such as encoder-based architecture instead of decoder-based, and ReLU activation instead of softmax. All experiments use  $d = 20$ . We choose  $\sigma \in \{\sigma_1, \sigma_2\} = \{0.1, 0.5\}$  and  $N = 20$  for noisy linear regression,  $s = 3$  and  $N = 10$  for sparse linear regression, and  $N = 40$  for linear regression and linear classification.
In the "mixture" mode,  $\pi$  is the uniform mixture of two or more base distributions. We consider two representative mixture modes studied in Section 4:
- Linear model + linear classification model;
- Noisy linear model with four noise levels  $\sigma \in \{0.1, 0.25, 0.5, 1\}$ .
Transformers trained with the mixture mode will be evaluated on multiple base distributions simultaneously. When the base distributions are sufficiently diverse, a transformer performing well on all of them will likely be performing some level of in-context algorithm selection. We evaluate transformers against standard machine learning algorithms in context (for each task respectively) as baselines.
Results Figure 3a shows the ICL performance of transformers on five base tasks, within each the transformer is trained on the same task. Transformers match the best baseline algorithm in four out of the five cases, except for the sparse regression task where the Transformer still outperforms least squares and matches Lasso with some choices of  $\lambda$  (thus utilizing sparsity to some extent). This demonstrates the strong ICL capability of the transformer architecture considered in our theory.
Figure 3b & 3c examine the in-context algorithm selection capability of transformers, on noisy linear regression with two different noise levels (Figure 3b), and regression + classification (Figure 3c). In both figures, the transformer trained in the mixture mode (TF_alg_select) approaches the best baseline algorithm on both tasks simultaneously. By contrast, transformers trained in the base mode for one of the tasks perform well on that task but behave suboptimally on the other task as expected. The existence of TF_alg_select showcases a single transformer that performs well on multiple tasks simultaneously (and thus has to perform in-context algorithm selection to some extent), supporting our theoretical results in Section 4.