Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

Introduction

Large neural sequence models have demonstrated remarkable in-context learning (ICL) capabilities , where models can make accurate predictions on new tasks when prompted with training examples from the same task, in a zero-shot fashion without any parameter update to the model. A prevalent example is large language models based on the transformer architecture , which can perform a diverse range of tasks in context when trained on enormous text . Recent models in this paradigm such as GPT-4 achieve surprisingly impressive ICL performance that makes them akin to a general-purpose agent in many aspects . Such strong capabilities call for better understandings, which a recent line of work tackles from various aspects .

Recent pioneering work of Garg et al. proposes an interpretable and theoretically amenable setting for understanding ICL in transformers. They perform ICL experiments where input tokens are real-valued (input, label) pairs generated from standard statistical models such as linear models (and the sparse version), neural networks, and decision trees. Garg et al. find that transformers can learn to perform ICL with prediction power (and fitted functions) matching standard machine learning algorithms for these settings, such as least squares for linear models, and Lasso for sparse linear models. Subsequent work further studies the internal mechanisms , expressive power , and generalization of transformers in this setting. However, these works only showcase simple mechanisms such as regularized regression or gradient descent , which are arguably only a small subset of what transformers are capable of in practice; or expressing universal function classes not specific to ICL . This motivates the following question:

How do transformers learn in context beyond implementing simple algorithms?

This paper makes steps on this question by making two main contributions: (1) We unveil a general mechanism—in-context algorithm selection—by which a single transformer can adaptively select different “base” ICL algorithms to use on different ICL instances, without any explicit prompting of the right algorithm to use in the input sequence. For example, a transformer may choose to perform ridge regression with regularization $\lambda_{1}$ on ICL instance 1, and $\lambda_{2}$ on ICL instance 2 (Fig. 2); or perform regression on ICL instance 1 and classification on ICL instance 2 (Fig. 5). This adaptivity allows transformers to achieve much stronger ICL performance than the base ICL algorithms. We both prove this in theory, and demonstrate this phenomenon empirically on standard transformer architectures. (2) Along the way, equally importantly, we present a first comprehensive theory for ICL in transformers by establishing end-to-end quantitative guarantees for the expressive power, in-context prediction performance, and sample complexity of pretraining. These results add upon the recent line of work on the statistical learning theory of transformers , and lay out a foundation for the intriguing special case where the learning targets are themselves ICL algorithms.

We prove that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares and ridge regression (Section 3.1), convex risk minimization for learning generalized linear models (such as logistic regression; Section 3.2), Lasso (Section 3.3), and gradient descent for two-layer neural networks (Section 3.4 & Appendix G). Our constructions admit mild bounds on the number of layers, heads, and weight norms, and achieve near-optimal prediction power on many in-context data distributions.

Technically, the above transformer constructions build on a new efficient implementation of in-context gradient descent (Section 3.5), which could be broaderly applicable. For a broad class of smooth convex empirical risks over the in-context training data, we construct an $(L+1)$-layer transformer that approximates $L$ steps of gradient descent. Notably, the approximation error accumulates only linearly in $L$, utilizing a stability-like property of smooth convex optimization.

We prove that transformers can perform in-context algorithm selection (Section 4). We construct two algorithm selection mechanisms: Post-ICL validation (Section 4.1), and Pre-ICL testing (Section 4.2). For both mechanisms, we provide general constructions as well as concrete examples. Fig. 1 provides a pictorial illustration of the two mechanisms.

As a concrete application, using the post-ICL validation mechanism, we construct a transformer that can perform nearly Bayes-optimal ICL on noisy linear models with mixed noise levels (Section 4.1.1), a more complex task than those considered in existing work.

We provide the first line of results for pretraining transformers to perform the various ICL tasks above, from polynomially many training sequences (Section 5).

Experimentally, we find that learned transformers indeed exhibit strong in-context algorithm selection capabilities in the settings considered in our theory (Section 6). For example, Fig. 2 shows that a single transformer can approach the individual Bayes risks (the optimal risk among all possible algorithms) simultaneously on two noisy linear models with different noise levels.

We humbly remark that the typical toolkit of a statistician contains much more beyond those covered in this work, including and not limited to inference, uncertainty quantification, and theoretical analysis. This work merely aims to show the algorithm selection capability of transformers, akin to what a statistician can do.

1 Related work

The in-context learning (ICL) capability of large language models (LLMs) has gained significant attention since demonstrated on GPT-3 Brown et al. . A number of subsequent empirical studies have contributed to a better understanding of the capabilities and limitations of ICL in LLM systems, which include but are not limited to . For an overview of ICL, see the survey by Dong et al. which highlights some key findings and advancements in this direction.

A line of recent work investigates why and how LLMs perform ICL . In particular, Xie et al. propose a Bayesian inference framework explaining how ICL works despite formatting differences between training and inference distributions. Garg et al. show empirically that transformers could be trained from scratch to perform ICL of linear models, sparse linear models, two-layer neural networks, and decision trees. Li et al. analyze the generalization error of trained ICL transformers from a stability viewpoint. They also experimentally show that transformers could perform “in-context model selection” (conceptually similar to in-context algorithm selection considered in this work) in specific tasks and presented related theoretical hypotheses. However, they do not provide concrete mechanisms or constructions for in-context model selection. A recent work shows that pretrained transformers can perform Bayesian inference in latent variable models, which may also be interpreted as a mechanism for ICL. Our experimental findings extend these results by unveiling and demonstrating the in-context algorithm selection capabilities of transformers.

Closely related to our theoretical results are , which show (among many things) that transformers can perform ICL by simulating gradient descent. However, these results do not provide quantitative error bounds for simulating multi-step gradient descent, and only handle linear regression models or their simple variants. Among these works, Akyürek et al. showed that transformers can implement learning algorithms for linear models based on gradient descent and closed-form ridge regression; it also presented preliminary evidence that learned transformers perform ICL similar to Bayes-optimal ridge regression. Our work builds upon and substantially extends this line of work by (1) providing a more efficient construction for in-context gradient descent; (2) providing an end-to-end theory with additional results for pretraining and statistical power; (3) analyzing a broader spectrum of ICL algorithms, including least squares, ridge regression, Lasso, convex risk minimization for generalized linear models, and gradient descent on two-layer neural networks; and (4) constructing more complex ICL procedures using in-context algorithm selection.

When in-context data are generated from a prior, the Bayes risk is a theoretical lower bound for the risk of any possible ICL algorithm, including transformers. Xie et al. , Akyürek et al. observe that learned transformers behave closely to the Bayes predictor on a variety of tasks such as hidden Markov models and noisy linear regression with a fixed noise level . Using the in-context algorithm selection mechanism (more precisely the post-ICL validation mechanism), we show that transformers can perform nearly-Bayes optimal ICL in noisy linear models with mixed noise levels (a strictly more challenging task than considered in ), with both concrete theoretical guarantees (Section 4.1.1) and empirical evidence (Fig. 2 & 4).

The transformer architecture, introduced by , has revolutionized natural language processing and been adopted in most of the recently developed large language models such as BERT and GPT . Broaderly, transformers have demonstrated remarkable performance in many other fields of artificial intelligence such as computer vision, speech, graph processing, reinforcement learning, and biological applications . Towards a better theoretical understanding, recent work has studied the capabilities , limitations , and internal workings of transformers.

We remark that the transformer architecture used in our theoretical constructions differs from the standard one by replacing the softmax activation (in the attention layers) with a (normalized) ReLU function. Transformers with ReLU activations is experimentally studied in the recent work of Shen et al. , who find that they perform as well as the standard softmax activation in many NLP tasks.

Training models (such as transformers) to perform ICL can be viewed as an approach for the broader problem of learning-to-learn or meta-learning . A number of other approaches has been studied extensively for this problem, including (and not limited to) training a meta-learner on how to update the parameters of a downstream learner , learning parameter initializations that quickly adapt to downstream tasks , learning latent embeddings that allow for effective similarity search . Most relevant to the ICL setting are approaches that directly take as input examples from a downstream task and a query input and produce the corresponding output . For a comprehensive overview, see the survey .

Theoretical aspects of meta-learning have received significant recent interest . In particular, analyzed the benefit of multi-task learning through a representation learning perspective, and studied the statistical properties of learning the parameter initialization for downstream tasks.

We build on various existing techniques from the statistics and learning theory literature to establish our approximation and generalization guarantees for transformers. For the approximation component, we rely on a technical result of Bach on the approximation power of ReLU networks. We use this result to show that transformers can approximate gradient descent (GD) on a broad range of loss functions, substantially extending the results of who primarily consider the square loss. The recent work of Giannou et al. also approximates GD with general loss functions by transformers, though using a different technique of forcing the softmax activations to act as sigmoids. Our analyses of Lasso and generalized linear models build on . Our generalization bound for transformers (used in our pretraining results) build on a chaining argument .

Preliminaries

Above, Eq. 1 uses a normalized ReLU activation $t\mapsto\sigma(t)/N$ in place of the standard softmax activation, which is for technical convenience and does not affect the essence of our studyFor each query index $i$, the attention weights ${\{\sigma(\left\langle{\mathbf{Q}}_{m}{\mathbf{h}}_{i},{\mathbf{K}}_{m}{\mathbf{h}}_{j}\right\rangle)/N\}}_{j\in[N]}$ is also a set of non-negative weights that sum to $O(1)$ (similar as a softmax probability distribution) in typical scenarios..

We consider a transformer architecture with $L\geq 1$ transformer layers, each consisting of a self-attention layer followed by an MLP layer.

We additionally define the following norm of a transformer ${\rm TF}_{\bm{\theta}}$:

In (2), the choices of the operator norm and max/sums are for convenience only and not essential, as our results (e.g. for pretraining) depend only logarithmically on ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}$.

2 In-context learning

We emphasize that the setting above considers predicting only at the last token ${\mathbf{x}}_{N+1}$, which is without much loss of generality. Our constructions may be generalized to predicting at every token, by using a decoder architecture and a corresponding input format (cf. Appendix B). Our theory focuses on predicting at the last token only, which simplifies the setting. Our experiments test both settings.

Basic in-context learning algorithms

We begin by constructing transformers that approximately implement a variety of standard machine learning algorithms in context, with mild size bounds and near-optimal prediction power on many standard in-context data distributions.

Consider the standard ridge regression estimator over the in-context training examples $\mathcal{D}$ with regularization $\lambda\geq 0$ (reducing to least squares at $\lambda=0$ and $N\geq d$):

We show that transformers can approximately implement Eq. ICRidge (proof in Section D.1).

For any $\lambda\geq 0$, $0\leq\alpha\leq\beta$ with $\kappa\mathrel{\mathop{:}}=\frac{\beta+\lambda}{\alpha+\lambda}$, $B_{w}>0$, and $\varepsilon<B_{x}B_{w}/2$, there exists an $L$-layer attention-only transformer ${\rm TF}^{0}_{\bm{\theta}}$ with

(with $R\mathrel{\mathop{:}}=\max{\{B_{x}B_{w},B_{y},1\}}$) such that the following holds. On any input data $(\mathcal{D},{\mathbf{x}}_{N+1})$ such that the problem Eq. ICRidge is well-conditioned and has a bounded solution:

${\rm TF}^{0}_{\bm{\theta}}$ approximately implements Eq. ICRidge: The prediction $\widehat{y}_{N+1}={\sf read}_{\sf y}({\rm TF}^{0}_{\bm{\theta}}({\mathbf{H}}))$ satisfies

Further, the second-to-last layer approximates ${\mathbf{w}}_{{\rm ridge}}^{\lambda}$: we have $\|{{\sf read}_{\sf w}({\mathbf{h}}^{(L-1)}_{i})-{\mathbf{w}}_{{\rm ridge}}^{\lambda}}\|_{2}\leq\varepsilon/B_{x}$ for all $i\in[N+1]$ (see Section D.1 for the definition of ${\sf read}_{\sf w}$).

Theorem 4 presents the first quantitative construction for end-to-end in-context ridge regression up to arbitrary precision, and improves upon Akyürek et al. whose construction does not give (or directly imply) an explicit error bound like Eq. 6. Further, the bounds on the number of layers and heads in Eq. 4 are mild (constant heads and logarithmically many layers).

Combining Theorem 4 with standard analyses of linear regression yields the following corollaries (proofs in Section D.3 & D.4).

For any $N\geq\widetilde{\mathcal{O}}(d)$, there exists an $\mathcal{O}(\kappa\log(\kappa N/\sigma))$-layer transformer ${\bm{\theta}}$, such that on any $\mathsf{P}$ satisfying standard statistical assumptions for least squares (Assumption), its ICL prediction $\widehat{y}_{N+1}$ achieves

Assumption requires only generic tail properties such as sub-Gaussianity, and not realizability (i.e., $\mathsf{P}$ follows a true linear model); $\kappa,\sigma$ above denote the covariance condition number and the noise level therein. The $\widetilde{\mathcal{O}}(d\sigma^{2}/N)$ excess risk is known to be rate-optimal for linear regression , and Corollary 5 achieves this in context with a transformer with only logarithmically many layers.

Next, consider Bayesian linear models where each in-context data distribution $\mathsf{P}=\mathsf{P}^{\sf lin}_{{\mathbf{w}}_{\star}}$ is drawn from a Gaussian prior $\pi:{\mathbf{w}}_{\star}\sim\mathsf{N}(0,{\mathbf{I}}_{d}/d)$, and $({\mathbf{x}},y)\sim\mathsf{P}^{\sf lin}_{{\mathbf{w}}_{\star}}$ is sampled as ${\mathbf{x}}\sim\mathsf{N}({\mathbf{0}},{\mathbf{I}}_{d})$, $y=\left\langle{\mathbf{w}}_{\star},{\mathbf{x}}\right\rangle+\mathsf{N}(0,\sigma^{2})$. It is a standard result that the Bayes estimator of $y_{N+1}$ given $(\mathcal{D},{\mathbf{x}}_{N+1})$ is given by ridge regression Eq. ICRidge: $\widehat{y}^{\sf Bayes}_{N+1}\mathrel{\mathop{:}}=\langle{\mathbf{w}}_{\rm ridge}^{\lambda},{\mathbf{x}}_{N+1}\rangle$ with $\lambda=d\sigma^{2}/N$. We show that transformers achieve nearly-Bayes risk for this problem, and we use

to denote the Bayes risk of this problem under prior $\pi$.

2 In-context learning of generalized linear models

As a natural generalization of linear regression, we now show that transformers can recover learn generalized linear models (GLMs) (which includes logistic regression for linear classification as an important special case), by implementing the corresponding convex risk minimization algorithm in context, and achieve near-optimal excess risk under standard statistical assumptions.

The following result (proof in Section E.1) shows that, as long as the empirical risk $\widehat{L}_{N}$ satisfies strong convexity and bounded solution conditions (similar as in Theorem 4), transformers can approximately implement the ERM predictor $g(\left\langle{\mathbf{x}}_{N+1},{\mathbf{w}}_{\rm GLM}\right\rangle)$, with ${\mathbf{w}}_{\rm GLM}$ given by Eq. ICGLM.

For any $0<\alpha<\beta$ with $\kappa\mathrel{\mathop{:}}=\frac{\beta}{\alpha}$, $B_{w}>0,B_{x}>0$, $\kappa_{w}:=L_{g}B_{x}^{2}/\alpha+1$ and $\varepsilon<B_{w}/2$, there exists an attention-only transformer ${\rm TF}^{0}_{\bm{\theta}}$ with

(where $L_{g}\mathrel{\mathop{:}}=\sup_{t}\left|g^{\prime}(t)\right|$, $R\mathrel{\mathop{:}}=\max{\{B_{x}B_{w},B_{y},1\}}$, and $C_{g}>0$ is a constant that depends only on $R$ and the $C^{2}$-smoothness of $g$ within $[-R,R]$), such that the following holds. On any input data $(\mathcal{D},{\mathbf{x}}_{N+1})$ such that

${\rm TF}^{0}_{\bm{\theta}}({\mathbf{H}}^{(0)})$ approximately implements Eq. ICGLM: We have ${\mathbf{h}}^{(L+1)}_{N+1}\mathrel{\mathop{:}}=[{\mathbf{x}}_{N+1};\widehat{y}_{N+1};\widehat{{\mathbf{w}}};1;1]$, where

In Theorem 7, the number of heads scales as $\widetilde{\mathcal{O}}(1/\varepsilon^{2})$ as opposed to $\Theta(1)$ as in ridge regression (Theorem 4), due to the fact that the gradient of the loss is in general a smooth function that can be only approximately expressed as a sum-of-relus (cf. Definition 12 & Lemma A.0) rather than exactly expressed as in the case for the square loss.

$\widehat{{\mathbf{w}}}$ achieves small excess risk under the population loss, i.e. for the linear prediction $\widehat{y}^{\sf lin}_{N+1}\mathrel{\mathop{:}}=\left\langle{\mathbf{x}}_{N+1},\widehat{{\mathbf{w}}}\right\rangle$,

Above, $\mathcal{O}\left(\cdot\right)$ hides constants that depend polynomially on the parameters in Assumption. Similar as in Corollary 5, the $\mathcal{O}(d/N)$ excess risk obtained here matches the optimal (fast) rate for typical learning problems with $d$ parameters and $N$ samples .

Applying Theorem 8 to logistic regression, we have the following result as a direct corollary. Below, the Gaussian input assumption is for convenience only and can be genearalized to e.g. sub-Gaussian input.

Consider any in-context data distribution $\mathsf{P}$ satisfying

For the link function $g=\sigma_{{\rm log}}$ and $B_{w}^{\star}=\mathcal{O}\left(1\right)$, we can choose $B_{w},B_{\mu},\mu_{g},L_{g},\mu_{x},K_{x},K_{y}=\Theta\left(1\right)$ so that Assumption holds. In that case, when $N\geq\mathcal{O}\left(d\right)$, there exists a transformer ${\bm{\theta}}$ with $L=\mathcal{O}\left(\log(N)\right)$ layers, such that for any $\mathsf{P}$ considered above,

The estimation $\widehat{{\mathbf{w}}}=\widetilde{\sf read}_{\sf w}({\rm TF}_{\bm{\theta}}({\mathbf{H}}))$ outputted by ${\bm{\theta}}$ achieves excess risk bound Eq. 8.

(Realizable setting) Consider the logistic in-context data distribution

Then, for any distribution $\mathsf{P}=\mathsf{P}^{\sf log}_{{\bm{\beta}}}$ with $\left\|{{\bm{\beta}}}\right\|_{2}\leq B_{w}^{\star}$, the prediction $\widehat{y}_{N+1}=\widetilde{\sf read}_{\sf y}({\rm TF}_{{\bm{\theta}}}({\mathbf{H}}))$ of ${\bm{\theta}}$ additionally achieves the square loss excess risk Eq. 9.

3 In-context Lasso

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 F).

For any $\lambda_{N}\geq 0$, $\beta>0$, $B_{w}>0$, and $\varepsilon>0$, there exists a $L$-layer transformer ${\rm TF}_{\bm{\theta}}$ with

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 ${\bm{\theta}}$ such that the following holds: For any $s$ and $N\geq\mathcal{O}\left(s\log(d/\delta)\right)$, suppose that $\mathsf{P}$ is an $s$-sparse linear model: ${\mathbf{x}}_{i}\sim\mathsf{N}(0,{\mathbf{I}}_{d})$, $y_{i}=\left\langle{\mathbf{w}}_{\star},{\mathbf{x}}_{i}\right\rangle+\mathsf{N}(0,\sigma^{2})$ for any $\left\|{{\mathbf{w}}_{\star}}\right\|_{2}\leq B_{w}^{\star}$ and $\left\|{{\mathbf{w}}_{\star}}\right\|_{0}\leq s$, then with probability at least $1-\delta$ (over the randomness of $\mathcal{D}$), the transformer output $\widehat{y}_{N+1}$ achieves

The $\widetilde{\mathcal{O}}(s\log d/N)$ excess risk obtained in Theorem 11 is optimal up to log factors . We remark that Theorem 11 is not a direct corollary of Theorem 10; Rather, the bound on the number of layers in Theorem 11 requires a sharper convergence analysis of the Eq. ICLasso problem under sparse linear models (Section F.2), similar to .

4 Gradient descent on two-layer neural networks

Thus far, we have focused on convex risks with (generalized) linear predictors of the form ${\mathbf{x}}\mapsto\left\langle{\mathbf{w}},{\mathbf{x}}\right\rangle$. To move beyond both restrictions, as a primary example, we show that transformers can approximate in-context gradient descent on two-layer neural networks (NNs).

5 Mechanism: In-context gradient descent

Technically, the constructions in Section 3.1-3.3 rely on a new efficient construction for transformers to implement in-context gradient descent and its variants, which we present as follows. We begin by presenting the result for implementing (vanilla) gradient descent on convex empirical risks.

such that $\sup_{{\mathbf{z}}\in[-R,R]^{k}}|g({\mathbf{z}})-f_{M,C}({\mathbf{z}})|\leq\varepsilon_{\rm approx}$.

Definition 12 is known to contain broad class of functions. For example, any mildly smooth $k$-variate function is approximable by a sum of relus for any $(\varepsilon_{\rm approx},R)$, with mild bounds on $(M,C)$ (Proposition A.0, building on results of Bach ). Also, any function that is a $(M,C)$-sum of relus itself (which includes all piecewise linear functions) is by definition $(0,\infty,M,C)$-approximable by sum of relus.

We show that $L$ steps of Eq. ICGD can be approximately implemented by an $(L+1)$-layer transformer.

Fix any $B_{w}>0$, $L>1$, $\eta>0$, and $\varepsilon\leq B_{w}/(2L)$. Suppose that

${\rm TF}^{0}_{\bm{\theta}}({\mathbf{H}}^{(0)})$ approximately implements Eq. ICGD with initialization ${\mathbf{w}}^{0}_{{\rm GD}}={\mathbf{0}}$:

(Prediction space) The final output ${\mathbf{H}}^{(L+1)}={\rm TF}^{0}_{\bm{\theta}}({\mathbf{H}}^{(0)})$ approximates the prediction of $L$ steps of Eq. ICGD: We have ${\mathbf{h}}^{(L+1)}_{N+1}=[{\mathbf{x}}_{N+1};\widehat{y}_{N+1};\widehat{{\mathbf{w}}}^{L};\mathbf{0}_{D-2d-3};1;t_{i}]$, where $\widehat{y}_{N+1}=\left\langle\widehat{{\mathbf{w}}}^{L},{\mathbf{x}}_{N+1}\right\rangle$ so that

Further, the transformer admits norm bound ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}\leq 2+R+2\eta C$.

The proof can be found in Section C.3. Theorem 13 substantially generalizes that of von Oswald et al. (which only does GD on square losses with a linear self-attention), and is simpler than the ones in Akyürek et al. and Giannou et al. . See Fig. 3 for a pictorial illustration of the basic component of the construction, which implements a single step of gradient descent using a single attention layer (Proposition C.0).

Technically, we utilize the stability of convex gradient descent as in the following lemma (proof in Section C.4) to obtain the linear error accumulation in Theorem 13; the error accumulation will become exponential in $L$ in the non-convex case in general; see Lemma G.0(b).

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 Fig. 1 for a pictorial illustration.

denote the validation loss of any predictor $f$.

We show that (proof in Section H.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}_{{\sf val}}$ concentrates around the expected loss $L$. Below, the input sequence ${\mathbf{H}}$ uses a generalized positional encoding ${\mathbf{p}}_{i}\mathrel{\mathop{:}}=[{\mathbf{0}}_{D-(d+3)};1;t_{i}]$ in Eq. 3, where $t_{i}\mathrel{\mathop{:}}=1$ for $i\in\mathcal{D}_{{\sf train}}$, $t_{i}\mathrel{\mathop{:}}=-1$ for $i\in\mathcal{D}_{{\sf val}}$, and $t_{N+1}\mathrel{\mathop{:}}=0$.

As an example, we use Proposition 15 to construct a transformer to perform in-context ridge regression with regularization selection according to the unregularized validation loss $\widehat{L}_{{\sf val}}({\mathbf{w}})\mathrel{\mathop{:}}=\frac{1}{2|\mathcal{D}_{{\sf val}}|}\sum_{(x_{i},y_{i})\in\mathcal{D}_{\sf val}}{\left(\left\langle{\mathbf{w}},{\mathbf{x}}_{i}\right\rangle-y_{i}\right)}^{2}$ (proof in Section H.2). Let $\lambda_{1},\dots,\lambda_{K}\geq 0$ be $K$ fixed regularization strengths.

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. Eq. 5) for all ${\{\lambda_{k}\}}_{k\in[K]}$, it outputs $\widehat{y}_{N+1}=\left\langle\widehat{{\mathbf{w}}},{\mathbf{x}}_{N+1}\right\rangle$, where

Above, $\widehat{{\mathbf{w}}}_{{\rm ridge},{\sf train}}^{\lambda}$ denotes the solution to Eq. ICRidge on the training split $\mathcal{D}_{{\sf train}}$, and $\gamma^{\prime}\mathrel{\mathop{:}}=2(B_{x}B_{w}+B_{y})B_{x}\varepsilon+\gamma$, where $B_{x},B_{w},B_{y}$ are the bounds in the well-conditioned assumption Eq. 5.

1.1 Nearly Bayes-optimal ICL on noisy linear models with mixed noise levels

We build on Theorem 16 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},\dots,\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 ${\{({\mathbf{x}}_{i},y_{i})\}}_{i\in[N+1]}\stackrel{{\scriptstyle\rm iid}}{{\sim}}\mathsf{P}_{{\mathbf{w}}_{\star},\sigma_{k}}$ as

For any fixed $(N,d)$, consider the Bayes risk for predicting $y_{N+1}$ under this model:

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 Section I.2 for a derivation). By using the post-ICL validation mechanism in Theorem 16, we construct a transformer that achieves nearly the Bayes risk.

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:

In particular, Theorem 17 applies in the proportional setting where $N,d$ are large and $N/d=\Theta(1)$ , in which case ${\sf 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 Akyürek et al. , who empirically find that transformers can achieve nearly Bayes risk under any fixed noise level. By contrast, Theorem 17 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 17). We demonstrate this empirically in Section 6 (cf. Fig. 4 & 2).

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 I.0.

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]}$.

There exists a single attention layer with 6 heads that implements $\Psi^{\sf 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.

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 Section H.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 Section H.4.

Analysis of pretraining

Thus far, we have established the existence of transformers for performing various ICL tasks with good in-context statistical performance. We now analyze the sample complexity of pretraining these transformers from a finite number of training ICL instances.

Above, $\mathsf{clip}_{B_{y}}(t)\mathrel{\mathop{:}}=\max{\{\min{\{t,B_{y}\}},-B_{y}\}}$ is the standard clipping operator onto $[-B_{y},B_{y}]$, and ${\rm TF}^{R}_{{\bm{\theta}}}$ the transformer architecture as in Definition 3 with clipping operators after each layer: let ${\mathbf{H}}^{(0)}=\mathsf{clip}_{\sf R}({\mathbf{H}})$,

The clipping operator is used to control the Lipschitz constant of ${\rm TF}_{\bm{\theta}}$ with respect to ${\bm{\theta}}$, and we typically choose a sufficiently large clipping radius ${\sf R}$ so that it does not modify the behavior of the transformer on any input sequence of our concern.

We draw ICL instances ${\mathbf{Z}}\mathrel{\mathop{:}}=({\mathbf{H}},y_{N+1})=(\mathcal{D},({\mathbf{x}}_{N+1},y_{N+1}))$ from a (meta-)distribution denoted as $\pi$, which first sample an in-context data distribution $\mathsf{P}\sim\pi$, then sample iid examples $({\mathbf{x}}_{i},y_{i})_{i=1}^{N+1}\stackrel{{\scriptstyle\rm iid}}{{\sim}}\mathsf{P}^{\otimes(N+1)}$ and form $\mathcal{D}={\{({\mathbf{x}}_{i},y_{i})\}}_{i\in[N]}$. Our pretraining loss is the average ICL loss on $n$ pretraining instances ${\mathbf{Z}}^{(1:n)}\stackrel{{\scriptstyle\rm iid}}{{\sim}}\pi$, and we consider the corresponding test ICL loss on a new test instance:

Our pretraining algorithm is to solve a standard constrained empirical risk minimization (ERM) problem over transformers with $L$ layers, $M$ heads, and norm bound $B$ (recall the definition of the ${\left|\!\left|\!\left|{\cdot}\right|\!\right|\!\right|}$ norm in Eq. 2):

By standard uniform concentration analysis via chaining arguments (Proposition A.0; see also [84, Chapter 5] for similar arguments), we have the following excess loss guarantee for Eq. TF-ERM. The proof can be found in Section J.2.

With probability at least $1-\xi$ (over the pretraining instances ${\{{\mathbf{Z}}^{j}\}}_{j\in[n]}$), the solution $\widehat{{\bm{\theta}}}$ to Eq. TF-ERM satisfies

where $\iota=\log(2+\max{\left\{B,{\sf R},B_{y}\right\}})$ is a log factor.

2 Examples of pretraining for in-context regression problems

In Theorem 20, the comparator $\inf_{{\bm{\theta}}\in\Theta_{L,M,D^{\prime},B}}L_{{\sf icl}}({\bm{\theta}})$ is simply the smallest expected ICL loss for ICL instances drawn from $\pi$, among all transformers within the norm ball $\Theta_{L,M,D^{\prime},B}$. Using our constructions in Section 3 & 4, we show that this comparator loss is small on various (meta-)distribution $\pi$’s, by which we obtain end-to-end guarantees for pretraining transformers with small ICL loss at test time. Here we showcase this argument on several representative regression problems.

Suppose $\mathsf{P}\sim\pi$ is almost surely well-posed for in-context linear regression (Assumption) with the canonical parameters. Then, for $N\geq\widetilde{\mathcal{O}}{\left(d\right)}$, with probability at least $1-\xi$ (over the training instances ${\mathbf{Z}}^{(1:n)}$), the solution $\widehat{{\bm{\theta}}}$ of Eq. TF-ERM with $L=\mathcal{O}(\kappa\log(\kappa N/\sigma))$ layers, $M=3$ heads, $D^{\prime}=0$ (attention-only), and $B=\mathcal{O}(\sqrt{\kappa d})$ achieves small excess ICL risk over ${\mathbf{w}}^{\star}_{\mathsf{P}}$:

where $\widetilde{\mathcal{O}}(\cdot)$ only hides polylogarithmic factors in $\kappa,N,1/\sigma$.

To our best knowledge, Theorem 21 offers the first end-to-end result for pretraining a transformer to perform in-context linear regression with explicit excess loss bounds. The $\widetilde{\mathcal{O}}(\sqrt{\kappa^{2}d^{2}/n})$ term originates from the generalization of pretraining (Theorem 20), where as the $\widetilde{\mathcal{O}}(d\sigma^{2}/N)$ term agrees with the standard fast rate for the excess loss of linear regression . Further, as long as $n\geq\widetilde{\mathcal{O}}(\kappa^{2}N/\sigma^{2})$, the excess risk achieves the optimal rate $\widetilde{\mathcal{O}}(d\sigma^{2}/N)$ (up to log factors).

By similar arguments as in the proof of Theorem 21, we can directly turn most of our other expressivity results into results on the pretrained transformers. Here we present three such additional examples (proofs in Section J.4-J.6). The first example is for the sparse linear regression problem considered in Theorem 11.

Suppose each $\mathsf{P}\sim\pi$ is almost surely an instance of the sparse linear model specified in Theorem 11 with parameters $B_{w}^{\star}$ and $\sigma$. Suppose $N\geq\widetilde{\mathcal{O}}{\left(s\log((d\vee N)/\sigma)\right)}$ and let $\kappa\mathrel{\mathop{:}}=B_{w}^{\star}/\sigma$.

Then with probability at least $1-\xi$ (over the training instances ${\mathbf{Z}}^{(1:n)}$), the solution $\widehat{{\bm{\theta}}}$ of Eq. TF-ERM with $L=\widetilde{\mathcal{O}}(\kappa^{2}(1+d/N))$ layers, $M=2$ heads, $D^{\prime}=2d$, and $B=\widetilde{\mathcal{O}}{\left({\rm poly}(d,B_{w}^{\star},\sigma)\right)}$ achieves small excess ICL risk:

where $\widetilde{\mathcal{O}}(\cdot)$ only hides polylogarithmic factors in $d,N,1/\sigma$.

Our next example is for the problem of noisy linear regression with mixed noise levels considered in Theorem 17 and Theorem I.0. There, the constructed transformer uses the post-ICL validation mechanism to perform ridge regression with an adaptive regulariation strength depending on the particular input sequence.

Suppose $\pi$ is the data generating model (noisy linear model with mixed noise levels) considered in Theorem I.0, with $\sigma_{\max}\leq\mathcal{O}(1)$. Let $N\geq d/10$.

Then, with probability at least $1-\xi$ (over the training instances ${\mathbf{Z}}^{(1:n)}$), the solution $\widehat{{\bm{\theta}}}$ of Eq. TF-ERM with input dimension $D=\Theta(dK)$, $L=\mathcal{O}(\sigma_{\min}^{-2}\log(N/\sigma_{\min}))$ layers, $M=\mathcal{O}(K)$ heads, $D^{\prime}=\mathcal{O}(K^{2})$, and $B=\mathcal{O}({\rm poly}(K,\sigma_{\min}^{-1},d,N))$ achieves small excess ICL risk:

where $\widetilde{\mathcal{O}}(\cdot)$ only hides polylogarithmic factors in $d,N,K,1/\sigma_{\min}$.

Our final example is for in-context logistic regression. For simplicity we consider the realizable case.

Suppose for $\mathsf{P}\sim\pi$, $\mathsf{P}$ is almost surely a realizable logistic model (i.e. $\mathsf{P}=\mathsf{P}^{\sf log}_{{\bm{\beta}}}$ with $\left\|{{\bm{\beta}}}\right\|_{2}\leq B_{w}^{\star}$ as in Corollary 9). Suppose that $B_{w}^{\star}=\mathcal{O}\left(1\right)$ and $N\geq\mathcal{O}\left(d\right)$.

Then, with probability at least $1-\xi$ (over the training instances ${\mathbf{Z}}^{(1:n)}$), the solution $\widehat{{\bm{\theta}}}$ of Eq. TF-ERM with $L=\mathcal{O}(\log(N))$ layers, $M=\widetilde{\mathcal{O}}\left(d^{3}N\right)$ heads, $D^{\prime}=0$, and $B=\mathcal{O}({\rm poly}(d,N))$ achieves small excess ICL risk:

where $\widetilde{\mathcal{O}}(\cdot)$ only hides polylogarithmic factors in $d,N$.

All results above are established by the expressivity results in Section 3 & 4 for transformers to implement various ICL procedures (such as least squares, Lasso, GLM, and ridge regression with in-context algorithm selection), combined with the generalization bound (Theorem 20). However, the transformer itself was not specified to encode any actual structure about the problem at hand in any result above, other than having sufficiently large number of layers, number of heads, and weight norms, which illustrates the flexibility of the transformer architecture.

Experiments

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). Additional experimental details can be found in Section K.1.

We train a 12-layer transformer, with two modes for the training sequence (instance) distribution $\pi$. In the “base” mode, similar to , 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 ${\{({\mathbf{x}}_{i},y_{i})\}}_{i\in[N+1]}\stackrel{{\scriptstyle\rm iid}}{{\sim}}\mathsf{P}_{{\mathbf{w}}_{\star}}$ as ${\mathbf{x}}_{i}\stackrel{{\scriptstyle\rm iid}}{{\sim}}\mathsf{N}({\mathbf{0}},{\mathbf{I}}_{d})$, and $y_{i}$ from one of the following models studied in Section 3:

Linear model: $y_{i}=\left\langle{\mathbf{w}}_{\star},{\mathbf{x}}_{i}\right\rangle$;

Noisy linear model: $y_{i}=\left\langle{\mathbf{w}}_{\star},{\mathbf{x}}_{i}\right\rangle+\sigma z_{i}$, where $\sigma>0$ is a fixed noise level, and $z_{i}\sim\mathsf{N}(0,1)$.

Sparse linear model: $y_{i}=\left\langle{\mathbf{w}}_{\star},{\mathbf{x}}_{i}\right\rangle$ with $\left\|{{\mathbf{w}}_{\star}}\right\|_{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;

Linear classification model: $y_{i}=\operatorname{sign}(\left\langle{\mathbf{w}}_{\star},{\mathbf{x}}_{i}\right\rangle)$.

These base tasks have been empirically investigated by Garg et al. , 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.

Fig. 4 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.

Fig. 4 & 4 examine the in-context algorithm selection capability of transformers, on noisy linear regression with two different noise levels (Fig. 4), and regression + classification (Fig. 4). 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.

2 Decoder-based architecture & details for Figure 2

ICL capabilities have also been demonstrated in the literature for decoder-based architectures . 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 Fig. 4), we refer the readers to Garg et al. .

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.1, except that the architecture is GPT-2 following Garg et al. , and the input format is changed to Eq. 15 (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 Section 6.1. The rest of the setup (optimization, training, and evaluation) is the same as in Section 6.1 & K.1. Note that we also train on the objective Eq. 48 for all tokens averaged, instead of for the last test token as in Section 6.1.

Fig. 2 shows the results for noisy linear models with two different noise levels, and Fig. 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 Fig. 2 & 2, TF_alg_select matches ridge regression with the optimal $\lambda$ on all tokens $i\in{\{1,\dots,N\}}$ ($N=40$). In Fig. 5 & 5, 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.

Conclusion

This work shows that transformers can perform complex in-context learning procedures with strong in-context algorithm selection capabilties, by both explicit theoretical constructions and experiments. We believe our work opens up many exciting directions, such as (1) more mechanisms for in-context algorithm selection; (2) Bayes-optimal ICL on other problems by either the post-ICL validation mechanism or new approaches; (3) understanding the internal workings of transformers performing in-context algorithm selection; (4) other mechanisms for implementing complex ICL procedures beyond in-context algorithm selection; (5) further statistical analyses, e.g. of pretraining.

Acknowledgment

The authors would like to thank Tengyu Ma and Jason D. Lee for the many insightful discussions. S. Mei is supported in part by NSF DMS-2210827 and NSF CCF-2315725.

References

Appendix A Technical tools

A.1 Concentration inequalities

Let ${\bm{\beta}}\sim\mathsf{N}({\mathbf{0}},{\mathbf{I}}_{d}/d)$. Then we have

The following lemma is a standard result of covariance concentration, see e.g. [82, Theorem 4.6.1].

Suppose that ${\mathbf{x}}_{1},\cdots,{\mathbf{x}}_{N}$ are independent $d$-dimensional $K$-sub-Gaussian random vectors. Then as long as $N\geq C_{0}d$, with probability at least $1-\exp(-N/C_{0})$ we have

A.2 Approximation theory

The following result for expressing smooth functions as a random feature model with relu activation is adapted from Bach [4, Proposition 5].

where $\phi(\cdot;\cdot):\mathcal{X}\times\mathcal{W}\to[-B,B]$ is $L$-Lipschitz (in $d_{\mathcal{X}}$) in the first argument, and $\mu$ is a signed measure over $\mathcal{W}$ with finite total measure $A=\operatorname{TV}(\mu)<\infty$. Then for any $\varepsilon>0$, there exists $\alpha_{1},\cdots,\alpha_{K}\in{\{\pm 1\}}$, ${\mathbf{w}}_{1},\cdots,{\mathbf{w}}_{K}\in\mathcal{W}$ with $K=\mathcal{O}(A^{2}B^{2}\log\mathcal{N}(\mathcal{X},d_{\mathcal{X}},\frac{\varepsilon}{3AL})/\varepsilon^{2})$, such that

where $\mathcal{N}(\mathcal{X},d_{\mathcal{X}},\frac{\varepsilon}{3AL})$ denotes the $(\frac{\varepsilon}{3AL})$-covering number of $\mathcal{X}$ in $d_{\mathcal{X}}$.

Let $\alpha({\mathbf{w}})\mathrel{\mathop{:}}=\operatorname{sign}(d\mu({\mathbf{w}}))\in{\{\pm 1\}}$ denote the sign of the density $d\mu({\mathbf{w}})$. We have

Note that $|d\mu({\mathbf{w}})|/A$ is the density of a probability distribution over $\mathcal{W}$. Thus for any ${\mathbf{x}}\in\mathcal{X}$, as long as $K\geq\mathcal{O}(A^{2}B^{2}\log(1/\delta)/\varepsilon^{2})$, we can sample ${\mathbf{w}}_{1},\dots,{\mathbf{w}}_{K}\stackrel{{\scriptstyle\rm iid}}{{\sim}}|d\mu(\cdot)|/A$, and obtain by Hoeffding’s inequality that with probability at least $1-\delta$,

Let $\mathcal{N}(\frac{\varepsilon}{3AL})\mathrel{\mathop{:}}=\mathcal{N}(\mathcal{X},d_{\mathcal{X}},\frac{\varepsilon}{3AL})$ for shorthand. By union bound, as long as $K\geq\mathcal{O}(A^{2}B^{2}\log(\mathcal{N}(\frac{\varepsilon}{3AL})/\delta)/\varepsilon^{2})$, we have with probability at least $1-\delta$ that for every $\widehat{{\mathbf{x}}}$ in the covering set corresponding to $\mathcal{N}(\frac{\varepsilon}{3AL})$,

Taking $\delta=1/2$ (for which $K=\mathcal{O}(A^{2}B^{2}\log\mathcal{N}(\frac{\varepsilon}{3AL})/\varepsilon^{2})$), by the probabilistic method, there exists a deterministic set ${\{{\mathbf{w}}_{i}\}}_{i\in[K]}\subset\mathcal{W}$ and ${\{\alpha_{i}\mathrel{\mathop{:}}=\alpha({\mathbf{w}}_{i})\}}_{i\in[K]}\in{\{\pm 1\}}$ such that the above holds.

Next, note that both $g$ (by Eq. 13) and the function ${\mathbf{x}}\mapsto\frac{A}{K}\sum_{i=1}^{K}\alpha({\mathbf{w}}_{i})\phi({\mathbf{x}};{\mathbf{w}}_{i})$ are $(AL)$-Lipschitz. Therefore, for any ${\mathbf{x}}\in\mathcal{X}$, taking $\widehat{{\mathbf{x}}}$ to be the point in the covereing set with $d_{\mathcal{X}}({\mathbf{x}},\widehat{{\mathbf{x}}})\leq\frac{\varepsilon}{3AL}$, we have

such that $\sup_{{\mathbf{z}}\in[-R,R]^{k}}\left|f({\mathbf{z}})-g({\mathbf{z}})\right|\leq\varepsilon_{\rm approx}$.

A.3 Optimization

The following convergence result for minimizing a smooth and strongly convex function is standard from the convex optimization literature, see e.g. Bubeck [13, Theorem 3.10].

The following convergence result of proximal gradient descent (PGD) on convex composite minimization problem is also standard, see e.g. .

A.4 Uniform convergence

The following result is shown in [84, Section 5.6].

Suppose that $\{X_{\theta}\}_{\theta}$ is a zero-mean random process indexed by $\theta\in\Theta$ such that $\left\|X_{\theta}-X_{\theta^{\prime}}\right\|_{\psi}\leq\rho(\theta,\theta^{\prime})$ for some metric $\rho$ on the space $\Theta$. Then it holds that

where $D$ is the diameter of the metric space $(\Theta,\rho)$, and the generalized Dudley entropy integral $J$ is given by

where $N(\delta;\Theta,\rho)$ is the $\delta$-covering number of $(\Theta,\rho)$.

As a corollary of Theorem A.0, we have the following result.

Suppose that $\{X_{\theta}\}_{\theta\in\Theta}$ is a zero-mean random process given by

The index set $\Theta$ is equipped with a distance $\rho$ and diameter $D$. Further, assume that for some constant $A$, for any ball $\Theta^{\prime}$ of radius $r$ in $\Theta$, the covering number admits upper bound $\log N(\delta;\Theta^{\prime},\rho)\leq d\log(2Ar/\delta)$ for all $0<\delta\leq 2r$.

Then with probability at least $1-\delta$, it holds that

where $C$ is a universal constant, and we denote $\kappa=1+B^{1}D/B^{0}$.

Furthermore, if we replace the ${\rm SG}$ in assumption (b) and (c) by ${\rm SE}$, then with probability at least $1-\delta$, it holds that

Fix a $D_{0}\in(0,D]$ to be specified later. We pick a $(D_{0}/2)$-covering $\Theta_{0}$ of $\Theta$ so that $\log\left|\Theta_{0}\right|\leq d\log(2AD/D_{0})$. Then, by the standard uniform covering of independent sub-Gaussian random variables, we have with probability at least $1-\delta/2$,

Assume that $\Theta_{0}=\{\theta_{1},\cdots,\theta_{n}\}$. For each $j\in[n]$, we consider $\Theta_{j}$ is the ball centered at $\theta_{j}$ of radius $D_{0}$ in $(\Theta,\rho)$. Then $\theta\in\Theta_{j}$ has diameter $D_{0}$ and admits covering number bound $\log\mathcal{N}(\Theta_{j},\delta)\leq d\log(AD_{0}/\delta)$. Hence, we can apply Theorem A.0 with the process $\{X_{\theta}\}_{\theta\in\Theta_{j}}$, then

Therefore, we can let $t\leq\sqrt{\log(2n/\delta)/N}$ in the above inequality and taking the union bound over $j\in[n]$, and hence with probability at least $1-\delta/2$, it holds that for all $j\in[n]$,

Notice that for each $\theta\in\Theta,$ there exists $j\in[n]$ such that $\theta\in\Theta_{j}$, and hence

Thus, with probability at least $1-\delta$, it holds

Taking $D_{0}=D/\kappa$ completes the proof of ${\rm SG}$ case.

We next consider the ${\rm SE}$ case. The idea is the same as the ${\rm SG}$ case, but in this case we need to consider the following Orlicz-norm:

Then Bernstein’s inequality of ${\rm SE}$ random variables yields

for some universal constant $C_{0}$. Therefore, we can repeat the argument above to deduce that with probability at least $1-\delta$, it holds

Taking $D_{0}=D/\kappa$ completes the proof. ∎

A.5 Useful properties of transformers

The following result can be obtained immediately by “joining” the attention heads and MLP layers of two single-layer transformers.

Further, ${\bm{\theta}}$ has at most $M\leq M_{1}+M_{2}$ heads, $D^{\prime}\leq D^{\prime}_{1}+D^{\prime}_{2}$ hidden dimension in its MLP layer, and norm bound ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}\leq{\left|\!\left|\!\left|{{\bm{\theta}}_{1}}\right|\!\right|\!\right|}+{\left|\!\left|\!\left|{{\bm{\theta}}_{2}}\right|\!\right|\!\right|}$.

When $L_{1}=L_{2}$ (${\bm{\theta}}_{1}$ and ${\bm{\theta}}_{2}$ have the same number of layers), the result follows directly by applying Proposition A.0 repeatedly for all $L_{1}$ layers and the definition of the norm Eq. 2.

If (without loss of generality) $L_{1}<L_{2}$, we can augment ${\bm{\theta}}_{1}$ to $L_{2}$ layers by adding $(L_{2}-L_{1})$ layers with zero attention heads, and zero MLP hidden dimension (note that this does not change $M_{1}$, $D_{1}^{\prime}$, and ${\left|\!\left|\!\left|{{\bm{\theta}}_{1}}\right|\!\right|\!\right|}$). Due to the residual structure, the transformer maintains the output $P_{1}({\mathbf{H}}_{1})$ throughout layer $L_{1}+1,\dots,L_{2}$, and it reduces to the case $L_{1}=L_{2}$. ∎

Appendix B Extension to decoder-based architecture

Decoder TFs are the same as encoder TFs, except that the attention layers are replaced by masked attention layers with a specific decoder-based (causal) attention mask.

Notice that standard masked attention definitions use the pre-activation additive masks (with mask value $-\infty$) . The post-activation multiplicative masks we use is equivalent to the pre-activation additive masks, and the modified presentation is for notational convenience. We also use a normalized ReLU activation $t\mapsto\sigma(t)/i$ in place of the standard softmax activation to be consistent with Definition 1. Note that the normalization $1/i$ is to ensure that the attention weights ${\{\sigma(\left\langle{\mathbf{Q}}_{m}{\mathbf{h}}_{i},{\mathbf{K}}_{m}{\mathbf{h}}_{j}\right\rangle)/i\}}_{j\in[i]}$ is a set of non-negative weights that sum to $O(1)$. The motivation of masked attention layer is to ensure that, when processing a sequence of tokens, the computations at any token do not see any later token.

We next define the decoder-based transformers with $L\geq 1$ transformer layers, each consisting of a masked attention layer (c.f. Definition B.0) followed by an MLP layer (c.f. Definition 2). This definition is similar to the definition of encoder-based transformers (c.f., Definition 3), except that we replace the attention layers by masked attention layers.

We also use Eq. 2 to define the norm of ${\rm DTF}_{\bm{\theta}}$.

B.2 In-context learning with decoder-based transformers

Eq. 15 is different from out input format Eq. 3 for encoder-based TFs. The main difference is that $({\mathbf{x}}_{i},y_{i})$ are in different tokens in (15), whereas $({\mathbf{x}}_{i},y_{i})$ are in the same token in (3). The reason for the former (i.e., different tokens in decoder) is that we want to avoid every $[{\mathbf{x}}_{i};0]$ token seeing the information of $y_{i}$, since we will evaluate the loss at every token. The reason for the latter (i.e., the same token in encoder) is for presentation convenience: since we only evaluate the loss at the last token, it is not necessary to alternate between $[{\mathbf{x}}_{i};0]$ and $[{\mathbf{0}};y_{i}]$ to avoid information leakage.

B.3 Results

We discuss how our theoretical results upon encoder TFs can be converted to those of the decoder TFs. Taking the implementation of (ICGD) (a key mechanism that enables most basic ICL algorithms such as ridge regression; cf. Section 3.5) as an example, this conversion is enabled by the following facts: (a) the input format (15) of decoders can be converted to the input format (3) of encoders by a 2-layer decoder TF; (b) the encoder TF that implements (ICGD) with input format (3), by a slight parameter modification, can be converted to a decoder TF that implements the (ICGD) algorithm with a converted input format.

Despite the difference between the input format (15) and (3), we show that there exists a 2-layer decoder TF that can convert the input format (15) to format (3). The proof can be found in Section B.4.

There exists a 2-layer decoder TF ${\rm DTF}$ with $3$ heads per layer, hidden dimension $2$ and ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}\leq 12$ such that upon taking input ${\mathbf{H}}$ of format (15), it outputs $\widetilde{{\mathbf{H}}}={\rm DTF}({\mathbf{H}})$ with

In particular, format Eq. 16 contains format Eq. 3 as a submatrix, by restricting to the $\{1,2,\ldots,D-1,D-2,D\}$ rows and $\{2,4,\ldots,2N-2,2N,2N+1\}$ columns.

The construction in Theorem 13 can be generalized to using the input format Eq. 16 along with a decoder TF, by using the scratch pad within the last token to record the gradient descent iterates. Further, if we slightly change the normalization in ${\rm MSK}$ from $1/i$ to $1/((i-1)\vee 1)$, then the same construction performs Eq. ICGD (with training examples ${\{1,\dots,j\}}$) at every token $i=2j+1$ (corresponding to predicting at ${\mathbf{x}}_{j+1}$). Building on this extension, all our constructions in Section 3 and Section 4.2 can be generalized to decoder TFs.

B.4 Proof of Proposition B.0

In the following, we construct the desired ${\bm{\theta}}=({\bm{\theta}}^{(1)},{\bm{\theta}}^{(2)})$ as follows.

for all $i,j$. By the structure of ${\mathbf{h}}_{i}$, these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Notice that $t_{i}\neq 0$ only when $2\mid i$, we then compute for $i=2k$ that

for all $i,j$. By the structure of ${\mathbf{h}}_{i}^{(1)}$, these matrices indeed exist, and further it is straightforward to check that they have norm bounds

where the last equality follows from the fact that

Finally, we only need to take a MLP layer ${\bm{\theta}}_{\tt mlp}^{(2)}=({\mathbf{W}}_{1}^{(2)},{\mathbf{W}}_{2}^{(2)})$ with hidden dimension 2 that maps

Combining the two steps above, we complete the proof of Proposition B.0. ∎

Appendix C Proofs for Section 3.5 and additional results

The definition above captures the proximal operator $\mathbf{prox}_{\eta\mathcal{R}}$ for a broad class of regularizers, such as the (commonly-used) $L_{1}$ and $L_{2}$ regularizer listed in the following proposition, for all of which one can directly check that they can be exactly implemented by an MLP as stated below.

For $\mathcal{R}=\lambda\left\|{\cdot}\right\|_{1}$, $\mathbf{prox}_{\eta\mathcal{R}}$ is $(0,+\infty,4d,4+2\eta\lambda)$-approximable by MLP.

For $\mathcal{R}=\frac{\lambda}{2}\left\|{\cdot}\right\|_{2}^{2}$, $\mathbf{prox}_{\eta\mathcal{R}}$ is $(0,+\infty,2d,2+2\eta\lambda)$-approximable by MLP.

Fix any $B_{w}>0$, $L>1$, $\eta>0$, and $\varepsilon+\varepsilon^{\prime}\leq B_{w}/(2L)$. Suppose that

$\mathcal{R}$ convex, and the proximal operator $\mathbf{prox}_{\eta\mathcal{R}}({\mathbf{w}})$ is $(\eta\varepsilon^{\prime},R^{\prime},D^{\prime},C^{\prime})$-approximable by MLP with $R^{\prime}=\sup_{\left\|{{\mathbf{w}}}\right\|_{2}\leq B_{w}}\left\|{{\mathbf{w}}_{\eta}^{+}}\right\|_{2}+\eta\varepsilon$.

${\rm TF}_{\bm{\theta}}({\mathbf{H}}^{(0)})$ approximately implements Eq. ICGD:

(Prediction space) The final output ${\mathbf{H}}^{(L+1)}={\rm TF}_{\bm{\theta}}({\mathbf{H}}^{(0)})$ approximates the prediction of $L$ steps of Eq. ICGD: We have ${\mathbf{h}}^{(L+1)}_{N+1}=[{\mathbf{x}}_{N+1};\widehat{y}_{N+1};\widehat{{\mathbf{w}}}^{L};\mathbf{0}_{D-2d-3};1;t_{i}]$, where $\widehat{y}_{N+1}=\left\langle\widehat{{\mathbf{w}}}^{L},{\mathbf{x}}_{N+1}\right\rangle$ so that

Further, the weight matrices have norm bounds ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}\leq 3+R+2\eta C+C^{\prime}$.

The proof of Theorem C.0 is essentially similar to the proof of Theorem 13, using the following generalized version of Lemma 14.

The proof of the above lemma is done by utilizing the non-expansiveness of the PGD operator ${\mathbf{w}}\mapsto\mathbf{prox}_{\eta\mathcal{R}}({\mathbf{w}}-\eta\nabla f({\mathbf{w}}))$ and otherwise following the same arguments as for Lemma 14.

C.2 Approximating a single GD step

Then, for any $\varepsilon>0$, there exists an attention layer ${\bm{\theta}}={\{({\mathbf{Q}}_{m},{\mathbf{K}}_{m},{\mathbf{V}}_{m})\}}_{m\in[M]}$ with $M$ heads such that, for any input sequence that takes form ${\mathbf{h}}_{i}=[{\mathbf{x}}_{i};y_{i}^{\prime};{\mathbf{w}};\mathbf{0}_{D-2d-3};1;t_{i}]$ with $\left\|{{\mathbf{w}}}\right\|_{2}\leq B_{w}$, it gives output $\widetilde{{\mathbf{h}}}_{i}={\left[{\rm Attn}_{\bm{\theta}}({\mathbf{H}})\right]}_{i}=[{\mathbf{x}}_{i};y_{i}^{\prime};\widetilde{{\mathbf{w}}};\mathbf{0}_{D-2d-3};1;t_{i}]$ for all $i\in[N+1]$, where

Further, ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}\leq 2+R+2\eta C$.

for all $i,j\in[N+1]$. As the input has structure ${\mathbf{h}}_{i}=[{\mathbf{x}}_{i};y_{i}^{\prime};{\mathbf{w}};\mathbf{0}_{D-2d-3};1;t_{i}]$, these matrices indeed exist, and further it is straightforward to check that they have norm bounds

Consequently, ${\left|\!\left|\!\left|{{\bm{\theta}}}\right|\!\right|\!\right|}\leq 2+R+2\eta C$.

where the last equality follows from the bound

so that the above relu equals if $t_{j}\leq 0$. Therefore,

Thus letting the attention layer ${\bm{\theta}}={\{({\mathbf{V}}_{m},{\mathbf{Q}}_{m},{\mathbf{K}}_{m})\}}_{m\in[M]}$, we have

C.3 Proof of Theorem 13

We first prove part (a), which requires constructing the first $L$ layers of ${\bm{\theta}}$. Note that by our precondition $L\leq B_{w}/(2\varepsilon)$.

We now prove part (b), which requires constructing the last attention layer ${\bm{\theta}}^{(L+1)}$. Recall ${\mathbf{h}}^{(L)}_{i}=[{\mathbf{x}}_{i};y_{i}^{\prime};\widehat{{\mathbf{w}}}^{L};\mathbf{0}_{D-2d-3};1;t_{i}]$ for all $i\in[N+1]$. We construct a 2-head attention layer ${\bm{\theta}}^{(L+1)}={\{({\mathbf{Q}}^{(L+1)}_{m},{\mathbf{K}}^{(L+1)}_{m},{\mathbf{V}}^{(L+1)}_{m})\}}_{m=1,2}$ such that for every $i,j\in[N+1]$,