SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yu Bai*
Salesforce Research
yu.bai@salesforce.com
Fan Chen*
fanchen@mit.edu
Huan Wang
huan.wang@salesforce.com
Caiming Xiong
cxiong@salesforce.com
Song Mei*
songmei@berkeley.edu
Neural sequence models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) abilities, where they can perform new tasks when prompted with training and test examples, without any parameter update to the model. This work first provides a comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, learning generalized linear models, and gradient descent on two-layer neural networks, with near-optimal predictive power on various in-context data distributions. Using an efficient implementation of in-context gradient descent as the underlying mechanism, our transformer constructions admit mild size bounds, and can be learned with polynomially many pretraining sequences.
Building on these "base" ICL algorithms, intriguingly, we show that transformers can implement more complex ICL procedures involving in-context algorithm selection, akin to what a statistician can do in real life—A single transformer can adaptively select different base ICL algorithms—or even perform qualitatively different tasks—on different input sequences, without any explicit prompting of the right algorithm or task. We both establish this in theory by explicit constructions, and also observe this phenomenon experimentally. In theory, we construct two general mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL validation. As an example, we use the post-ICL validation mechanism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task—noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.
Large neural sequence models have demonstrated remarkable in-context learning (ICL) capabilities [12], 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 [84], which can perform a diverse range of tasks in context when trained on enormous text [12, 90]. Recent models
*Equal technical and directional contributions.
Code is available at https://github.com/allenbai01/transformers-as-statisticians.
Figure 1: Illustration of in-context algorithm selection, and two mechanisms constructed in our theory. Left, middle-left: A single transformer can perform ridge regression with different  $\lambda$ 's on input sequences with different observation noise; we prove this by the post-ICL validation mechanism (Section 4.1). Middle-right, right: A single transformer can perform linear regression on regression data and logistic regression on classification data; we prove this via the pre-ICL testing mechanism (Section 4.2).
in this paradigm such as GPT-4 achieve surprisingly impressive ICL performance that makes them akin to a general-purpose agent in many aspects [65, 14]. Such strong capabilities call for better understandings, which a recent line of work tackles from various aspects [49, 94, 28, 72, 15, 57, 64].
Recent pioneering work of Garg et al. [31] 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. [31] 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 [2, 86, 18], expressive power [2, 32], and generalization [47] of transformers in this setting. However, these works only showcase simple mechanisms such as regularized regression [31, 2, 47] or gradient descent [2, 86, 18], which are arguably only a small subset of what transformers are capable of in practice; or expressing universal function classes not specific to ICL [89, 32]. 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 (Figure 2); or perform regression on ICL instance 1 and classification on ICL instance 2 (Figure 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 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 [97, 89, 27, 39], and lay out a foundation for the intriguing special case where the learning targets are themselves ICL algorithms.
A detailed summary of our contributions is as follows.
- We prove that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, convex risk minimization for learning generalized linear models (such as logistic regression), and gradient descent for two-layer neural networks (Section 3). 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 (Appendix D), which could be broadly 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
(a) Noisy linear reg with noise  $\sigma_{1}$
Figure 2: In-context algorithm selection on two separate noisy linear regression tasks with noise  $(\sigma_1, \sigma_2) = (0.1, 0.5)$ .  $(a, b)$  A single transformer TF_alg_select simultaneously approaches the performance of the two individual Bayes predictors ridge_lam_1 on task 1 and ridge_lam_2 on task 2.  $(c)$  At token 20 (using example  $\{0, \dots, 19\}$  for training), TF_alg_select approaches the Bayes error on two tasks simultaneously, and outperforms ridge regression with any fixed  $\lambda$ .  $(a, b, c)$  Note that transformers pretrained on a single task (TF-noise_1, TF-noise_2) perform near-optimally on that task but suboptimally on the other task. More details about the setup and training method can be found in Appendix M.2.
(b) Noisy linear reg with noise  $\sigma_{2}$
(c) Task 1 vs. task 2 at token 20
(Section 4.2). For both mechanisms, we provide general constructions as well as concrete examples. Figure 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 & Appendix K).
- 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, Figure 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.
Transformers as statisticians 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.
Related work Our work is intimately related to the lines of work on in-context learning, theoretical understandings of transformers, as well as other formulations for learning-to-learn such as meta-learning. Due to limited space, we discuss these related work in Appendix A.
We consider a sequence of  $N$  input vectors  $\{\mathbf{h}_i\}_{i=1}^N \subset \mathbb{R}^D$ , written compactly as an input matrix  $\mathbf{H} = [\mathbf{h}_1, \ldots, \mathbf{h}_N] \in \mathbb{R}^{D \times N}$ , where each  $\mathbf{h}_i$  is a column of  $\mathbf{H}$  (also a token). Throughout this paper, we let  $\sigma(t) \coloneqq \mathrm{ReLU}(t) = \max \{t, 0\}$  denote the standard relu activation.
We consider transformer architectures that process any input sequence  $\mathbf{H} \in \mathbb{R}^{D \times N}$  by applying (encoder-mode²) attention layers and MLP layers formally defined as follows.
Definition 1 (Attention layer). A (self-)attention layer with  $M$  heads is denoted as  $\mathrm{Att}\mathbf{n}_{\theta}(\cdot)$  with parameters  $\pmb{\theta} = \{(\mathbf{V}_m,\mathbf{Q}_m,\mathbf{K}_m)\}_{m\in [M]}\subset \mathbb{R}^{D\times D}$ . On any input sequence  $\mathbf{H}\in \mathbb{R}^{D\times N}$ ,
$$
\widetilde {\mathbf {H}} = \operatorname {A t t n} _ {\boldsymbol {\theta}} (\mathbf {H}) := \mathbf {H} + \frac {1}{N} \sum_ {m = 1} ^ {M} \left(\mathbf {V} _ {m} \mathbf {H}\right) \times \sigma \left(\left(\mathbf {Q} _ {m} \mathbf {H}\right) ^ {\top} \left(\mathbf {K} _ {m} \mathbf {H}\right)\right) \in \mathbb {R} ^ {D \times N}, \tag {1}
$$
where  $\sigma : \mathbb{R} \to \mathbb{R}$  is the ReLU function. In vector form,
$$
\widetilde {\mathbf {h}} _ {i} = [ \mathrm {A t t n} _ {\pmb {\theta}} (\mathbf {H}) ] _ {i} = \mathbf {h} _ {i} + \sum_ {m = 1} ^ {M} \frac {1}{N} \sum_ {j = 1} ^ {N} \sigma (\langle \mathbf {Q} _ {m} \mathbf {h} _ {i}, \mathbf {K} _ {m} \mathbf {h} _ {j} \rangle) \cdot \mathbf {V} _ {m} \mathbf {h} _ {j}.
$$
2Many of our results can be generalized to decoder-based architectures; see Appendix C for a discussion.
Above, (1) uses a normalized  $\mathrm{ReLU}^3$  activation  $t\mapsto \sigma (t) / N$  in place of the standard softmax activation; we remark this activation is also found to work well empirically in recent studies [78, 93].
Definition 2 (MLP layer). A (token-wise) MLP layer with hidden dimension  $D'$  is denoted as  $\mathrm{MLP}_{\theta}(\cdot)$  with parameters  $\pmb{\theta} = (\mathbf{W}_1, \mathbf{W}_2) \in \mathbb{R}^{D' \times D} \times \mathbb{R}^{D \times D'}$ . On any input sequence  $\mathbf{H} \in \mathbb{R}^{D \times N}$ ,
$$
\tilde {\mathbf {H}} = \operatorname {M L P} _ {\boldsymbol {\theta}} (\mathbf {H}) := \mathbf {H} + \mathbf {W} _ {2} \sigma \left(\mathbf {W} _ {1} \mathbf {H}\right),
$$
where  $\sigma : \mathbb{R} \to \mathbb{R}$  is the ReLU function. In vector form, we have  $\tilde{\mathbf{h}}_i = \mathbf{h}_i + \mathbf{W}_2\sigma (\mathbf{W}_1\mathbf{h}_i)$ .
We consider a transformer architecture with  $L \geq 1$  transformer layers, each consisting of a self-attention layer followed by an MLP layer.
Definition 3 (Transformer). An  $L$ -layer transformer, denoted as  $\mathrm{TF}_{\pmb{\theta}}(\cdot)$ , is a composition of  $L$  self-attention layers each followed by an MLP layer:  $\mathbf{H}^{(L)} = \mathrm{TF}_{\pmb{\theta}}(\mathbf{H}^{(0)})$ , where  $\mathbf{H}^{(0)} \in \mathbb{R}^{D \times N}$  is the input sequence, and
$$
\mathbf {H} ^ {(\ell)} = \mathrm {M L P} _ {\boldsymbol {\theta} _ {\mathrm {m l p}} ^ {(\ell)}} \Big (\mathrm {A t t n} _ {\boldsymbol {\theta} _ {\mathrm {a t t n}} ^ {(\ell)}} \big (\mathbf {H} ^ {(\ell - 1)} \big) \Big), \quad \ell \in \{1, \ldots , L \}.
$$
Above, the parameter  $\pmb{\theta} = (\pmb{\theta}_{\mathrm{attn}}^{(1:L)}, \pmb{\theta}_{\mathrm{mlp}}^{(1:L)})$  consists of the attention layers  $\pmb{\theta}_{\mathrm{attn}}^{(\ell)} = \{(\mathbf{V}_m^{(\ell)}, \mathbf{Q}_m^{(\ell)}, \mathbf{K}_m^{(\ell)})\}_{m \in [M^{(\ell)}]} \subset \mathbb{R}^{D \times D}$  and the MLP layers  $\pmb{\theta}_{\mathrm{mlp}}^{(\ell)} = (\mathbf{W}_1^{(\ell)}, \mathbf{W}_2^{(\ell)}) \in \mathbb{R}^{D^{(\ell)} \times D} \times \mathbb{R}^{D \times D^{(\ell)}}$ . We will frequently consider "attention-only" transformers with  $\mathbf{W}_1^{(\ell)}, \mathbf{W}_2^{(\ell)} = \mathbf{0}$ , which we denote as  $\mathrm{TF}_{\pmb{\theta}}^0(\cdot)$  for shorthand, with  $\pmb{\theta} = \pmb{\theta}^{(1:L)} := \pmb{\theta}_{\mathrm{attn}}^{(1:L)}$ .
We additionally define the following norm of a transformer  $\mathrm{TF}_{\theta}$ :
$$
\| \boldsymbol {\theta} \| := \max  _ {\ell \in [ L ]} \left\{\max  _ {m \in [ M ]} \left\{\| \mathbf {Q} _ {m} ^ {(\ell)} \| _ {\mathrm {o p}}, \| \mathbf {K} _ {m} ^ {(\ell)} \| _ {\mathrm {o p}} \right\} + \sum_ {m = 1} ^ {M} \| \mathbf {V} _ {m} ^ {(\ell)} \| _ {\mathrm {o p}} + \| \mathbf {W} _ {1} ^ {(\ell)} \| _ {\mathrm {o p}} + \| \mathbf {W} _ {2} ^ {(\ell)} \| _ {\mathrm {o p}} \right\}. \tag {2}
$$
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  $\| \pmb{\theta}\|$ .
2.2 In-context learning
In an in-context learning (ICL) instance, the model is given a dataset  $\mathcal{D} = \{(\mathbf{x}_i,y_i)\}_{i\in [N]}\stackrel {\mathrm{~iid}}{\sim}\mathsf{P}$  and a new test input  $\mathbf{x}_{N + 1}\sim \mathsf{P}_{\mathbf{x}}$  for some data distribution P, where  $\{\mathbf{x}_i\}_{i\in [N]}\subseteq \mathbb{R}^d$  are the input vectors,  $\{y_{i}\}_{i\in [N]}\subseteq \mathbb{R}$  are the corresponding labels (e.g. real-valued for regression, or  $\{0,1\}$ -valued for binary classification), and  $\mathbf{x}_{N + 1}$  is the test input on which the model is required to make a prediction. Different from standard supervised learning, in ICL, each instance  $(\mathcal{D},\mathbf{x}_{N + 1})$  is in general drawn from a different distribution  $\mathsf{P}_j$ , such as a linear model with a new ground truth coefficient  $\mathbf{w}_{\star ,j}\in \mathbb{R}^{d}$ . Our goal is to construct fixed transformer to perform ICL on a large set of  $\mathsf{P}_j$ 's.
We consider using transformers to perform ICL, in which we encode  $(\mathcal{D},\mathbf{x}_{N + 1})$  into an input sequence  $\mathbf{H}\in \mathbb{R}^{D\times (N + 1)}$ . In our theory, we use the following format, where the first two rows contain  $(\mathcal{D},\mathbf{x}_{N + 1})$  (zero at the location for  $y_{N + 1}$ ), and the third row contains fixed vectors  $\{\mathbf{p}_i\}_{i\in [N + 1]}$  with ones, zeros, and indicator for being the train token (similar to a positional encoding vector):
$$
\mathbf {H} = \left[ \begin{array}{l l l l l} \mathbf {x} _ {1} & \mathbf {x} _ {2} & \dots & \mathbf {x} _ {N} & \mathbf {x} _ {N + 1} \\ y _ {1} & y _ {2} & \dots & y _ {N} & 0 \\ \mathbf {p} _ {1} & \mathbf {p} _ {2} & \dots & \mathbf {p} _ {N} & \mathbf {p} _ {N + 1} \end{array} \right] \in \mathbb {R} ^ {D \times (N + 1)}, \quad \mathbf {p} _ {i} := \left[ \begin{array}{c} \mathbf {0} _ {D - (d + 3)} \\ 1 \\ 1 \{i <   N + 1 \} \end{array} \right] \in \mathbb {R} ^ {D - (d + 1)}. \tag {3}
$$
We will choose  $D = \Theta(d)$ , so that the hidden dimension of  $\mathbf{H}$  is at most a constant multiple of  $d$ . We then feed  $\mathbf{H}$  into a transformer to obtain the output  $\widetilde{\mathbf{H}} = \mathrm{TF}_{\boldsymbol{\theta}}(\mathbf{H}) \in \mathbb{R}^{D \times (N + 1)}$  with the same shape, and read out the prediction  $\widehat{y}_{N + 1}$  from the  $(d + 1, N + 1)$ -th entry of  $\widetilde{\mathbf{H}} = [\widetilde{\mathbf{h}}_i]_{i \in [N + 1]}$  (the entry corresponding to the missing test label):  $\widehat{y}_{N + 1} = \mathrm{read}_{\mathbf{y}}(\widetilde{\mathbf{H}}) \coloneqq (\widetilde{\mathbf{h}}_{N + 1})_{d + 1}$ . The goal is to predict
For each query index  $i$ , the attention weights  $\{\sigma (\langle \mathbf{Q}_m\mathbf{h}_i,\mathbf{K}_m\mathbf{h}_j\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. Also, our approximation results can potentially be generalized to softmax attention e.g. using the technique of [32].
$\widehat{y}_{N + 1}$  that is close to  $y_{N + 1} \sim \mathsf{P}_{y|\mathbf{x}_{N + 1}}$  measured by proper losses. We emphasize that we consider predicting only at the last token  $\mathbf{x}_{N + 1}$ , which is without much loss of generality.
Miscellaneous setups We assume bounded features and labels throughout the paper (unless otherwise specified, e.g. when  $\mathbf{x}_i$  is Gaussian):  $\| \mathbf{x}_i\| _2\leq B_x$  and  $|y_{i}|\leq B_{y}$  with probability one. We use the standard notation  $\mathbf{X} = [\mathbf{x}_1^\top ;\ldots ;\mathbf{x}_N^\top ]\in \mathbb{R}^{N\times d}$  and  $\mathbf{y} = [y_1;\dots;y_N]\in \mathbb{R}^N$  to denote the matrix of inputs and vector of labels, respectively. To prevent the transformer from blowing up on tail events, in all our results concerning (statistical) in-context prediction powers, we consider a clipped prediction  $\widehat{y}_{N + 1} = \widetilde{\mathrm{read}}_{\mathrm{y}}(\widetilde{\mathbf{H}})\coloneqq \mathrm{clip}_{R}((\widetilde{\mathbf{h}}_{N + 1})_{d + 1})$ , where  $\mathrm{clip}_R(t)\coloneqq \mathrm{Proj}_{[-R,R]}(t)$  is the standard clipping operator with (a suitably large) radius  $R\geq 0$  that varies in different problems.
3 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.
3.1 In-context ridge regression and least squares
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$ ):
$$
\mathbf {w} _ {\text {r i d g e}} ^ {\lambda} := \arg \min  _ {\mathbf {w} \in \mathbb {R} ^ {d}} \frac {1}{2 N} \sum_ {i = 1} ^ {N} \left(\langle \mathbf {w}, \mathbf {x} _ {i} \rangle - y _ {i}\right) ^ {2} + \frac {\lambda}{2} \| \mathbf {w} \| _ {2} ^ {2}. \tag {ICRidge}
$$
We show that transformers can approximately implement (ICRidge) (proof in Appendix F.1).
Theorem 4 (Implementing in-context ridge regression). For any  $\lambda \geq 0$ ,  $0 \leq \alpha \leq \beta$  with  $\kappa := \frac{\beta + \lambda}{\alpha + \lambda}$ ,  $B_w > 0$ , and  $\varepsilon < B_x B_w / 2$ , there exists an  $L$ -layer attention-only transformer  $\mathrm{TF}_{\theta}^0$  with
$$
L = \left\lceil 2 \kappa \log \left(B _ {x} B _ {w} / (2 \varepsilon)\right)\right\rceil + 1, \quad \max  _ {\ell \in [ L ]} M ^ {(\ell)} \leq 3, \quad \| \boldsymbol {\theta} \| \leq 4 R + 8 (\beta + \lambda) ^ {- 1}. \tag {4}
$$
(with  $R \coloneqq \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 (ICRidge) is well-conditioned and has a bounded solution:
$$
\alpha \leq \lambda_ {\min } \left(\mathbf {X} ^ {\top} \mathbf {X} / N\right) \leq \lambda_ {\max } \left(\mathbf {X} ^ {\top} \mathbf {X} / N\right) \leq \beta , \quad \left\| \mathbf {w} _ {\text {r i d g e}} ^ {\lambda} \right\| _ {2} \leq B _ {w} / 2, \tag {5}
$$
$\mathrm{TF}_{\theta}^{0}$  approximately implements (ICRidge): The prediction  $\widehat{y}_{N + 1} = \mathrm{read}_{\mathrm{y}}(\mathrm{TF}_{\theta}^{0}(\mathbf{H}))$  satisfies
$$
\left| \widehat {y} _ {N + 1} - \left\langle \mathbf {w} _ {\text {r i d g e}} ^ {\lambda}, \mathbf {x} _ {N + 1} \right\rangle \right| \leq \varepsilon . \tag {6}
$$
Theorem 4 presents the first quantitative construction for end-to-end in-context ridge regression up to arbitrary precision, and improves upon Akyurek et al. [2] whose construction does not give (or directly imply) an explicit error bound like (6). Further, the bounds on the number of layers and heads in (4) are mild (constant heads and logarithmically many layers).
Near-optimal in-context prediction power for linear problems Combining Theorem 4 with standard analyses of linear regression yields the following corollaries (proofs in Appendix F.3 & F.4).
Corollary 5 (Near-optimal linear regression with transformers by approximating least squares). For any  $N \geq \widetilde{\mathcal{O}}(d)$ , there exists an  $\mathcal{O}(\kappa \log(\kappa N / \sigma))$ -layer transformer  $\theta$ , such that on any  $\mathsf{P}$  satisfying standard statistical assumptions for least squares (Assumption A), its ICL prediction  $\widehat{y}_{N+1}$  achieves
$$
\mathbb {E} _ {(\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P}} [ (\widehat {y} _ {N + 1} - y _ {N + 1}) ^ {2} ] \leq \inf  _ {\mathbf {w}} \mathbb {E} _ {(\mathbf {x}, y) \sim \mathsf {P}} \bigl [ (y - \langle \mathbf {w}, \mathbf {x} \rangle) ^ {2} \bigr ] + \widetilde {\mathcal {O}} (d \sigma^ {2} / N).
$$
Assumption A requires only generic tail properties such as sub-Gaussianity, and not realizability (i.e.,  $\mathbb{P}$  follows a true linear model);  $\kappa, \sigma$  above denote the covariance condition number and the
<sup>4</sup>Our constructions may be generalized to predicting at every token, by using a decoder architecture and potentially different input formats correspondingly (cf. Appendix C). Our theory focuses on predicting at the last token only, which simplifies the setting. Our experiments test both settings.
noise level therein. The  $\tilde{\mathcal{O}}(d\sigma^2/N)$  excess risk is known to be rate-optimal for linear regression [38], 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}_{\mathbf{w}_{\star}}^{\mathrm{lin}}$  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}_{\mathbf{w}_{\star}}^{\mathrm{lin}}$  is sampled as  $\mathbf{x} \sim \mathsf{N}(\mathbf{0}, \mathbf{I}_d)$ ,  $y = \langle \mathbf{w}_{\star}, \mathbf{x} \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 (ICRidge):  $\hat{y}_{N+1}^{\mathrm{Bayes}} := \langle \mathbf{w}_{\mathrm{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
$$
\text {B a y e s R i s k} _ {\pi} := \mathbb {E} _ {\mathbf {w} _ {\star} \sim \pi , (\mathcal {D}, \mathbf {x} _ {N + 1}, y _ {N + 1}) \sim \mathsf {P} _ {\mathbf {w} _ {\star}} ^ {\mathrm {l i n}}} \left[ \frac {1}{2} \left(\widehat {y} _ {N + 1} ^ {\text {B a y e s}} - y _ {N + 1}\right) ^ {2} \right]
$$
to denote the Bayes risk of this problem under prior  $\pi$ .