Asymptotic theory of in-context learning by linear attention

Significance

Attention-based architectures are a powerful force in modern AI. In particular, the emergence of in-context learning abilities enables task generalization far beyond the original next-token prediction objective. There is still significant room for progress in a rigorous understanding of how and when such capacity emerges. Our work provides a sharp characterization of ICL in an analytically tractable model, offering insights into the sample complexity and data quality requirements. Furthermore, we demonstrate that these insights can be applied to more complex, realistic architectures.

Abstract

Transformers have a remarkable ability to learn and execute tasks based on examples provided within the input itself, without explicit prior training. It has been argued that this capability, known as in-context learning (ICL), is a cornerstone of Transformers’ success, yet questions about the necessary sample complexity, pretraining task diversity, and context length for successful ICL remain unresolved. Here, we provide a precise answer to these questions in an exactly solvable model of ICL of a linear regression task by linear attention. We derive sharp asymptotics for the learning curve in a phenomenologically rich scaling regime where the token dimension is taken to infinity; the context length and pretraining task diversity scale proportionally with the token dimension; and the number of pretraining examples scales quadratically. We demonstrate a double-descent learning curve with increasing pretraining examples, and uncover a phase transition in the model’s behavior between low and high task diversity regimes: in the low diversity regime, the model tends toward memorization of training tasks, whereas in the high diversity regime, it achieves genuine ICL and generalization beyond the scope of pretrained tasks. These theoretical insights are empirically validated through experiments with both linear attention and full nonlinear Transformer architectures.

Since their introduction by Vaswani et al. in (1), Transformers have become a cornerstone of modern AI. Originally designed for sequence modeling tasks, such as language modeling and machine translation, Transformers achieve state-of-the art performance across many domains, even those that are not inherently sequential (2). Most strikingly, they underpin the breakthroughs achieved by large language models such as BERT (3), LLaMA (4), the GPT series (5–8), the Claude model family (9), and DeepSeek R1 (10).

The technological advancements enabled by Transformers have inspired a substantial body of research aimed at understanding their working principles. One key observation is that language models gain new behaviors and skills as their number of parameters and the size of their training datasets grow (7, 11–13). A particularly important emergent skill is in-context learning (ICL), which describes the model’s ability to learn and execute tasks based on the context provided within the input itself, without the need for explicit prior training on those specific tasks. To give an example from natural language processing, a pretrained large language model might be able to successfully translate English to Italian after being prompted with a few example translations, even if it has not been specifically pretrained on that translation task (7). ICL enables language models to perform new, specialized tasks without retraining, which is arguably a key reason for their general-purpose abilities.

Despite recent progress in understanding ICL, fundamental questions about when and how ICL emerges in large language models remain unresolved. These models are typically trained or pretrained using a next-token prediction objective, but the impact of various algorithmic and hyperparameter choices during pretraining on ICL performance is still not well understood. What mechanisms do Transformers implement to facilitate ICL? How many pretraining examples are necessary for ICL capabilities to arise? Furthermore, how many examples must be provided within the input context for the model to successfully perform an in-context task? Another important question is the degree of task diversity in the training data: how diverse must the training tasks be to enable ICL for truly novel tasks that the model has not encountered before?

In this paper, we address these questions by investigating the ICL capabilities of a linear attention module for linear regression tasks. This simplified model setting, which exhibits the minimal architectural features to perform ICL for linear regression, enables us to derive an asymptotically precise theory of ICL performance, elucidating its exact dependence on various hyperparameters. In the remainder of this section, we first provide an overview of related work on ICL. We then summarize our main contributions.

1.1. Related Work.

1.1.1. ICL in transformer architectures.

The striking ICL abilities of Transformers were thrust to the fore by Brown et al.’s (7) work on GPT-3. Focusing on natural language processing (NLP) tasks, they showed that ICL performance dramatically improves with an increase in the number of model parameters, with an increase in the number of examples in the model’s context, and with the addition of a natural language task description. In subsequent work, Wei et al. (13) proposed that the emergence of ICL with increasing scale is an abrupt, unpredictable transition. This perspective has substantially influenced proposed theories for the emergence of ICL (14). However, Schaeffer et al. (15) have challenged the idea that the emergence of ICL is unpredictable; they suggest that appropriately chosen measures of otherwise hidden progress (16) reveal that ICL gradually develops with scale.

1.1.2. Empirical studies of synthetic ICL tasks.

Though ICL in NLP is both impressive and useful, these natural data do not allow precise experimental control and study. Toward a fine-grained understanding of the conditions required for ICL, many recent works have explored ICL of parametrically controllable synthetic tasks, notably linear regression and classification. These works have identified various features of pretraining data distributions that contribute to the emergence of ICL (17–21). Closely related to our work is a study of ICL of linear regression by Raventós et al. (20). Their work identified a task diversity threshold for the emergence of ICL, below which a pretrained Transformer behaves as a Bayesian estimator with a prior determined by the limited set of pretraining tasks. Above this threshold, the model’s performance approaches that of within-context optimal Bayesian ridge regression, corresponding to a Gaussian prior over all tasks, including those not seen during pretraining. A motivating objective of our work is to provide a theoretical account of the empirical findings made by Raventós et al. (20), which underscore the roles of task diversity, regularization, model capacity, and data structure in the emergence of ICL.

1.1.3. Theoretical studies of ICL.

Many theoretical studies of ICL have centered on the idea that Transformers learn a particular algorithm during pretraining, which is then flexibly deployed to solve in-context tasks. In broad strokes, papers from this program of research often consider a particular algorithm for solving an in-context task, prove that Transformers can approximately implement this algorithm, and then empirically compare the ICL performance of a pretrained Transformer to the performance of that algorithm (22–29). A clear consensus on which algorithm underlies ICL of linear regression in full transformers has yet to emerge (22–29). Within this line of research, closest to our work are a series of papers that consider ICL of linear regression by simplified Transformers using linear, rather than softmax, attention modules (25, 27–32). Zhang et al. (29) studied these models in the limit of infinite pretraining dataset size (i.e., the population risk limit), and show that their performance on in-context linear regression nearly matches that of the Bayes-optimal estimator for the ICL task. However, they found that linear Transformers are not robust to shifts in the within-context covariate distribution. Zhang et al. (28) then showed that any optimizer of the within-context risk for a linear Transformer solves the ICL task with an approximation to one step of gradient descent from a learnable initialization, and that the resulting estimator can saturate the Bayes error for tasks with a Gaussian prior and nonzero mean. As we will discuss in Section 2, our reduction of the linear attention module is inspired in part by these works. In very recent work, Duraisamy (32) has studied the fine-sample risk of in-context linear regression with a single step of gradient descent, without directly analyzing Transformers. An et al. (25) and Wu et al. (30) investigated how linear Transformers adapt to limited pretraining data and context length, again showing that in certain cases nearly optimal error is achievable. Like these studies, our work considers linear attention; but our analysis, with its asymptotically sharp predictions, allows us to elucidate the exact dependence of ICL performance on various hyperparameters and to pinpoint when and how the transition from memorization to ICL of linear regression occurs. In closing, we highlight the work of Reddy (21), who analyzed the emergence of in-context classification through a phenomenological model.

1.2. Summary of Contributions.

We now summarize the main contributions of our paper relative to the prior art reviewed above. Building on recent literature, we focus on a simplified model of a Transformer that captures its key architectural motif: the linear self-attention module (25, 27–31). Linear attention includes the quadratic pairwise interactions between inputs that lie at the heart of softmax attention, but it omits the normalization steps and fully connected layers. This simplification makes the model more amenable to theoretical analysis. Our main result is a sharp asymptotic analysis of ICL for linear regression using linear attention, leading to a more precisely predictive theory than previous population risk analyses or finite-sample bounds (28, 29). The main contributions of our paper are structured as follows:

1. We begin in Section 2 by developing a simplified parameterization of linear self-attention that allows pretraining on the ICL linear regression task to be performed using ridge regression.

2. Within this simplified model, we identify a phenomenologically rich scaling limit in which the ICL performance can be analyzed (Section 3). As the token dimension tends to infinity, we allow the number of pretraining examples scale quadratically with the token dimension, while the context length and pretraining task diversity scale linearly. In this joint limit, we compute sharp asymptotics for ICL performance using random matrix theory.

3. The asymptotically precise theory curves we derive reveal several interesting phenomena (Section 3). First, we observe double-descent in the model’s ICL generalization performance as a function of pretraining dataset size, reflecting our assumption that it is pretrained to interpolation. Second, we study the nonmonotonic dependence of ICL performance on context length. Last, we uncover a transition from memorization to ICL as the pretraining task diversity increases. This transition recapitulates the empirical findings of Raventós et al. (20) in full Transformer models.

4. In Section 4, we demonstrate through numerical experiments that the insights from our theory, derived using the simplified linear attention model, transfer to full Transformer models with softmax self-attention. In particular, the scaling of pretraining sample complexity and task diversity with token dimension required for successful ICL is consistent.

2. Problem Formulation

We begin by describing the setting of our study.

2.1. ICL of Linear Regression.

In an ICL task, the model takes as input a sequence of tokens $\{x_1,y_1,x_2,y_2,\dots ,x_{\ell },y_{\ell },x_{\ell +1}\}$, and outputs a prediction of $y_{\ell +1}$. We will often refer to an input sequence as a “context.” The pairs $\{x_i,y_i{\}}_{i=1}^{\ell +1}$ are i.i.d. samples from a context-dependent joint distribution $P(x,y)$. Hence, the model needs to gather information about $P(x,y)$ from the first ℓ examples and use this information to predict $y_{\ell +1}$ from $x_{\ell +1}$. We will refer to ℓ as the “context length.”

In this work, we focus on an approximately linear mapping between $x_i\in {\mathbb{R}}^d$ and $y_i\in \mathbb{R}$:

$$\begin{array}{c}\hfill y_i=⟨x_i,w⟩+{ϵ}_i,\end{array}$$ [1]

where ${ϵ}_i$ is a Gaussian noise and $w\in {\mathbb{R}}^d$ is referred to as a task vector. We note that the task vector w is fixed within a context but can change between different contexts. The model has to learn w from the ℓ pairs presented within the context, and use it to predict $y_{\ell +1}$ from $x_{\ell +1}$.

2.2. Linear Self-Attention.

The model that we will analytically study is the linear self-attention block (33). Linear self-attention takes as input an embedding matrix Z, whose columns hold the sequence tokens. The mapping of sequences to matrices is not unique. Here, following the convention in refs. 29, 30, and 33, we will embed the input sequence $\{x_1,y_1,x_2,y_2,\dots ,x_{\ell },y_{\ell },x_{\ell +1}\}$ as

$$\begin{array}{c}\hfill Z=\left[\begin{array}{ccccc}{x}_1& x_2& \dots & x_{\ell }& x_{\ell +1}\\ y_1& y_2& \dots & y_{\ell }& 0\end{array}\right]\in {\mathbb{R}}^{(d+1)\times (\ell +1)},\end{array}$$ [2]

where 0 in the lower-right corner is a token that prompts the missing value $y_{\ell +1}$ to be predicted.

For value matrix $V\in {\mathbb{R}}^{(d+1)\times (d+1)}$ and key and query matrices $K,Q$ such that $K^{\top }Q\in {\mathbb{R}}^{(d+1)\times (d+1)}$, the output of a linear-attention block (33–35) is given by

$$\begin{array}{c}\hfill A:=Z+\frac{1}{\ell }VZ{(KZ)}^{\top }(QZ).\end{array}$$ [3]

The output A is a matrix while our goal is to predict a scalar, $y_{\ell +1}$. Following the choice of positional encoding in Eq. 2, we will take $A_{d+1,\ell +1}$, the element of A corresponding to the 0 prompt, as the prediction for $y_{\ell +1}$:

$$\begin{array}{c}\hfill \widehat{y}:=A_{d+1,\ell +1}.\end{array}$$ [4]

2.3. Pretraining Data.

The model is pretrained on n sample sequences, where the μth sample is a collection of $\ell +1$ vector–scalar pairs ${\{x_i^{\mu }\in {\mathbb{R}}^d,y_i^{\mu }\in \mathbb{R}\}}_{i=1}^{\ell +1}$ related by the approximate linear mapping in Eq. 1: $y_i^{\mu }=⟨x_i^{\mu },w^{\mu }⟩+{ϵ}_i^{\mu }$. Here, $w^{\mu }$ denotes the task vector associated with the μth sample. We make the following statistical assumptions:

1. $x_i^{\mu }$ are d-dimensional random vectors, sampled i.i.d. over both i and μ from an isotropic Gaussian distribution $\mathcal{N}(0,I_d/d)$.

2. At the beginning of training, construct a finite set with k task vectors, denoted by

   $$\begin{array}{c}\hfill {\omega }_k=\{w_1,w_2,\dots ,w_k\}.\end{array}$$

   The elements of this set are independently drawn once from

   $$\begin{array}{c}\hfill w_i{\sim }_{\text{i.i.d.}}\mathcal{N}(0,I_d).\end{array}$$

   For $1\le \mu \le n$, the task vector $w^{\mu }$ associated with the μth sample context is uniformly sampled from ${\omega }_k$. Note that the variable k controls the task diversity in the pretraining dataset. Importantly, k can be less than n, in which case the same task vector from ${\omega }_k$ will be repeated multiple times.

3. The noise terms ${ϵ}_i^{\mu }$ are i.i.d. over both i and μ, and drawn from a normal distribution $\mathcal{N}(0,\rho )$.

We denote a sample from this distribution by $(Z,y_{\ell +1})\sim {\mathcal{P}}_{\mathrm{train}}$.

2.4. Parameter Reduction.

Before specifying the training procedure, it is insightful to first examine the prediction mechanism of the linear attention module for the ICL task. This proves to be a fruitful exercise, shedding light on key questions: can linear self-attention learn linear regression in-context? If so, what do the model parameters learn from the data to solve this ICL problem? By closely analyzing these aspects, we can also formulate a simplified problem that lends itself to analytical study.

We start by rewriting the output of the linear attention module, Eq. 4, in an alternative form. Following ref. 29, we define

$$\begin{array}{c}\hfill V=\left[\begin{array}{cc}{V}_{11}& v_{12}\\ v_{21}^{\top }& v_{22}\end{array}\right],M=\left[\begin{array}{cc}{M}_{11}& m_{12}\\ m_{21}^{\top }& m_{22}\end{array}\right]:=K^{\top }Q,\end{array}$$ [5]

where $V_{11}\in {\mathbb{R}}^{d\times d}$, $v_{12},v_{21}\in {\mathbb{R}}^d$, $v_{22}\in \mathbb{R}$, $M_{11}\in {\mathbb{R}}^{d\times d}$, $m_{12},m_{21}\in {\mathbb{R}}^d$, and $m_{22}\in \mathbb{R}$. We assume that the inner dimension of $K^{\top }Q$ is greater than equal to $d+1$ so that matrix M can achieve full rank. From Eqs. 3 and 4, one can check that

$$\begin{array}{cc}\hfill \widehat{y}& =\frac{1}{\ell }\langle x_{\ell +1},v_{22}{M}_{11}^{\top }{\displaystyle \sum _{i=1}^{\ell }}{y}_i{x}_i+v_{22}{m}_{21}{\displaystyle \sum _{i=1}^{\ell }}{y}_i^2\hfill \\ \hfill & +M_{11}^{\top }{\displaystyle \sum _{i=1}^{\ell +1}}{x}_i{x}_i^{\top }{v}_{21}+m_{21}{\displaystyle \sum _{i=1}^{\ell }}{y}_i{x}_i^{\top }{v}_{21}\rangle ,\hfill \end{array}$$

where $⟨·,·⟩$ stands for the standard inner product.

This expression reveals several interesting points, including how this model could express a solution to the ICL task. First, not all parameters in Eq. 5 contribute to the output: we can discard all the parameters except for the last row of V and the first d columns of M. Second, the first term

$$\begin{array}{c}\hfill \frac{1}{\ell }{v}_{22}{M}_{11}^{\top }{\displaystyle \sum _{i=1}^{\ell }}{y}_i{x}_i\end{array}$$

offers a hint about how the linear attention module might be solving the task. The sum $\frac{1}{\ell }\sum _{i\le \ell }{y}_i{x}_i$ is a noisy estimate of $\mathbb{E}[x{x}^{\top }]w$ for that context. Hence, if the parameters of the model are such that $v_{22}{M}_{11}^{\top }$ is approximately $\mathbb{E}{[x{x}^{\top }]}^{-1}$, this term alone makes a good prediction for the output. Third, the third term, $M_{11}^{\top }\sum _{i=1}^{\ell +1}{x}_i{x}_i^{\top }{v}_{21}$ does not depend on outputs y, and thus does not directly contribute to the ICL task that relies on the relationship between x and y. Finally, the fourth term, $m_{21}\sum _{i=1}^{\ell }{y}_i{x}_i^{\top }{v}_{21}$, only considers a one-dimensional projection of x onto $v_{21}$. Because the task vectors w and x are isotropic in the statistical models that we consider, there are no special directions in the problem. Consequently, we expect the optimal $v_{21}$ to be approximately zero by symmetry considerations.

Motivated by these observations, and for analytical tractability, we study the linear attention module with the constraint $v_{21}=0$. In this case, collecting the remaining parameters in a matrix

$$\begin{array}{c}\hfill \mathrm{\Gamma }:=v_{22}\left[\begin{array}{cc}{M}_{11}^{\top }/d& m_{21}\end{array}\right]\in {\mathbb{R}}^{d\times (d+1)}\end{array}$$ [6]

and the input sequence in another matrix $H_Z$, defined as

$$\begin{array}{c}\hfill H_Z:=x_{\ell +1}\left[\begin{array}{cc}\frac{d}{\ell }\sum _{i\le \ell }{y}_i{x}_i^{\top }& \frac{1}{\ell }\sum _{i\le \ell }{y}_i^2\end{array}\right]\in {\mathbb{R}}^{d\times (d+1)},\end{array}$$ [7]

we can rewrite the predicted label as

$$\begin{array}{c}\hfill \widehat{y}=⟨\mathrm{\Gamma },H_Z⟩.\end{array}$$ [8]

The $1/d$ scaling of $M_{11}$ in Γ is chosen so that the columns of $H_Z$ scale similarly; it does not affect the final predictor $\widehat{y}$.

We note that ref. 29 provides an analysis of the population risk (whereas we focus on empirical risk) for a related reduced model in which both $v_{21}$ and $m_{21}$ are set to 0. Consequently, the predictors considered in ref. 29 slightly differ from ours in Eqs. 6–8 by an additive term (due to $m_{21}$). The authors of ref. 29 justify this reduced model through an optimization argument: if these parameters $v_{21}$ and $m_{21}$ are initialized to zero, they remain zero under gradient descent optimization of the population risk.

In the remainder of this paper, we will examine the ICL performance of the reduced model given in Eqs. 7 and 8, except when making comparisons to a full, nonlinear Transformer architecture. Henceforth, unless explicitly stated otherwise, we will refer to this reduced model as the linear attention module.

2.5. Model Pretraining.

The parameters of the linear attention module are learned from n samples of input sequences,

$$\begin{array}{c}\hfill \{x_1^{\mu },y_1^{\mu },\dots ,x_{\ell +1}^{\mu },y_{\ell +1}^{\mu }\},\hspace{1em}\mu =1,\dots ,n.\end{array}$$

We estimate model parameters by minimizing MSE loss on next-output prediction with ridge regularization, giving

$$\begin{array}{cc}\hfill {\mathrm{\Gamma }}^{\ast }& ={\displaystyle \underset{\mathrm{\Gamma }}{argmin}}{\displaystyle \sum _{\mu =1}^n}{\left(y_{\ell +1}^{\mu }-⟨\mathrm{\Gamma },H_{Z^{\mu }}⟩\right)}^2+\frac{n}{d}\lambda {\mathrm{\Gamma }}_{\mathrm{F}}^2,\hfill \end{array}$$ [9]

where $\lambda >0$ is a regularization parameter, and $H_{Z^{\mu }}$ refers to the input matrix Eq. 7 populated with the μth sample sequence. The factor $n/d$ in front of λ makes sure that, when we take the $d\to \infty$ or $n\to \infty$ limits later, there is still a meaningful ridge regularization when $\lambda >0$. The solution to the optimization problem in Eq. 9 can be expressed explicitly as

$$\begin{array}{cc}\hfill \text{vec}({\mathrm{\Gamma }}^{\ast })& ={\left(\frac{n}{d}\lambda I+{\displaystyle \sum _{\mu =1}^n}\text{vec}(H_{Z^{\mu }})\text{vec}{(H_{Z^{\mu }})}^{\top }\right)}^{-1}\hfill \\ \hfill & \times {\displaystyle \sum _{\mu =1}^n}{y}_{\ell +1}^{\mu }\text{vec}(H_{Z^{\mu }}),\hfill \end{array}$$ [10]

where $\mathrm{vec}(·)$ denotes the vectorization operation. Throughout this paper, we adopt the row-major convention. Thus, for a $d_1\times d_2$ matrix A, $\mathrm{vec}(A)$ is a vector in ${\mathbb{R}}^{d_1{d}_2}$, formed by stacking the rows of A together.

2.6. Evaluation.

For a given set of parameters Γ, the model’s generalization error is defined as

$$\begin{array}{c}\hfill e(\mathrm{\Gamma }):={\mathbb{E}}_{{\mathcal{P}}_{\mathrm{test}}}{\left(y_{\ell +1}-⟨\mathrm{\Gamma },H_Z⟩\right)}^2,\end{array}$$

where $(Z,y_{\ell +1})\sim {\mathcal{P}}_{\mathrm{test}}$ is a new sample drawn from the probability distribution of the test dataset. We consider two different test data distributions ${\mathcal{P}}_{\mathrm{test}}$:

1. ICL task: $x_i$ and ${ϵ}_i$ are i.i.d. Gaussians as in the pretraining case. However, each task vector $w^{\text{test}}$ associated with a test input sequence of length ℓ is drawn independently from $\mathcal{N}(0,I_d)$. We will denote the test error under this setting by $e^{\mathrm{ICL}}(\mathrm{\Gamma })$.

2. In-distribution generalization (IDG) task: The test data are generated in exactly the same manner as the training data, i.e., ${\mathcal{P}}_{\mathrm{test}}={\mathcal{P}}_{\mathrm{train}}$, hence the term in-distribution generalization. In particular, the set of unique task vectors $\{w_1,\dots ,w_k\}$ is identical to that used in the pretraining data. We will denote the test error under this setting by $e^{\text{IDG}}(\mathrm{\Gamma })$. This task can also be referred to as in-weight learning.

The ICL task evaluates the true ICL performance of the linear attention module. The task vectors in the test set differ from those seen in training, requiring the model to infer them from context. The IDG task assesses the model’s performance on task vectors it has previously encountered during pretraining. High performance on the IDG task but low performance on the ICL task indicates that the model memorizes the training task vectors. Conversely, high performance on the ICL task suggests that the model can learn genuinely new task vectors from the provided context.

To understand the performance of our model on both ICL and IDG tasks, we will need to evaluate these expressions for the pretrained attention matrix ${\mathrm{\Gamma }}^{\ast }$ given in Eq. 10. An asymptotically precise prediction of $e^{\mathrm{ICL}}({\mathrm{\Gamma }}^{\ast })$ and $e^{\mathrm{IDG}}({\mathrm{\Gamma }}^{\ast })$ will be a main result of this work.

2.7. Baselines: Bayes-Optimal Within-Context Estimators.

Following ref. 20, it is useful to compare the predictions made by the trained linear attention to optimal estimators that use only the current context information. These estimators rely solely on the data within the given context for their predictions but need oracle knowledge of the full statistical models underlying the data. Under the mean square loss, the optimal Bayesian estimator ${\widehat{y}}_{\text{Bayes}}={\mathbb{E}}_{{\mathcal{P}}_{\mathrm{test}}}[y_{\ell +1}|x_1,y_1,x_2,y_2,\dots ,x_{\ell },y_{\ell },x_{\ell +1}]$ in our setting has the form

$$\begin{array}{c}\hfill {\widehat{y}}_{\text{Bayes}}={(w_{\text{Bayes}})}^{\top }{x}_{\ell +1},\end{array}$$

where $w_{\text{Bayes}}$ is the Bayes estimator of the task vector w.

For the ICL task, the Bayes-optimal ridge regression estimator is given by

$$\begin{array}{c}\hfill w_{\text{Bayes}}^{\text{ridge}}:={\left({\displaystyle \sum _{i=1}^{\ell }}{x}_i{x}_i^{\top }+\rho I_d\right)}^{-1}\left({\displaystyle \sum _{i=1}^{\ell }}{y}_i{x}_i\right),\end{array}$$ [11]

where the ridge parameter is set to the noise variance ρ. We will refer to it as the ridge estimator. For the IDG task, the Bayes-optimal estimator is given by

$$\begin{array}{c}\hfill w_{\text{Bayes}}^{\text{dMMSE}}:=\frac{{\sum }_{j=1}^k{w}_j{e}^{-\frac{1}{2\rho }{\sum }_{i=1}^{\ell }{\left(y_i-w_j^{\top }{x}_i\right)}^2}}{{\sum }_{j=1}^k{e}^{-\frac{1}{2\rho }{\sum }_{i=1}^{\ell }{\left(y_i-w_j^{\top }{x}_i\right)}^2}}.\end{array}$$ [12]

Here, we assume that the training task vectors $\{w_1,\dots ,w_k\}$ are known to the estimator. Following ref. 20, we will refer to this estimator as the discrete minimum mean squared error (dMMSE) estimator.

The test performance of these estimators are calculated by

$$\begin{array}{c}\hfill e_{{\mathcal{P}}_{\mathrm{test}}}^{\text{Bayes}}={\mathbb{E}}_{{\mathcal{P}}_{\mathrm{test}}}{\left(y_{\ell +1}-{(w_{\text{Bayes}})}^{\top }{x}_{\ell +1}\right)}^2,\end{array}$$

where ${\mathcal{P}}_{\mathrm{test}}$ can be associated with either the ICL or the IDG task, and $w_{\text{Bayes}}$ can be the ridge or the dMMSE estimator. To avoid possible confusion, we emphasize that we will sometimes plot the performance of an estimator on a task for which it is not optimal. For example, we will test the dMMSE estimator, which is Bayes-optimal for the pretraining distribution, on the ICL task, where it is not optimal. This will be done for benchmarking purposes.

3. Theoretical Results

To answer the questions raised in the introduction, we provide a precise asymptotic analysis of the learning curves of the linear attention module for ICL of linear regression. We then focus on various implications of these equations, and verify through simulations that our insights gained from this theoretical analysis extend to more realistic nonlinear Transformers.

3.1. Joint Asymptotic Limit.

We have now defined both the structure of the training data as well as the parameters to be optimized. For our theoretical analysis, we consider a joint asymptotic limit in which the input dimension d, the pretraining dataset size n, the context length ℓ, and the number of task vectors in the training set k, go to infinity together such that

$$\begin{array}{c}\hfill \frac{\ell }{d}:=\alpha =\mathrm{\Theta }(1),\hspace{1em}\frac{k}{d}:=\kappa =\mathrm{\Theta }(1),\hspace{1em}\frac{n}{d^2}:=\tau =\mathrm{\Theta }(1).\end{array}$$ [13]

Identification of these scalings constitutes one of the main results of our paper. As we will see, the linear attention module exhibits rich learning phenomena in this limit.

The intuition for these scaling parameters can be seen as follows. Standard results in linear regression (36–38) show that to estimate a d-dimensional task vector w from the ℓ samples within a context, one needs at least $\ell =\mathrm{\Theta }(d)$. The number of unique task vectors that must be seen to estimate the covariance matrix of the true d-dimensional task distribution $\mathcal{N}(0,I_d)$ should also scale with d, i.e., $k=\mathrm{\Theta }(d)$. Finally, we see from Eq. 6 that the number of linear attention parameters to be learned is $\mathrm{\Theta }(d^2)$. This suggests that the number of individual contexts the model sees during pretraining should scale similarly, i.e., $n=\mathrm{\Theta }(d^2)$.

3.2. Learning Curves for ICL of Linear Regression by a Linear Attention Module.

Our theoretical analysis, explained in detail in SI Appendix, leads to asymptotically precise expressions for the generalization errors under the two test distributions being studied. Specifically, our theory predicts that, as $d,n,\ell ,k\to \infty$ in the joint limit given in Eq. 13,

$$\begin{array}{c}\hfill e^{\mathrm{ICL}}({\mathrm{\Gamma }}^{\ast })⟶e^{\mathrm{ICL}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )\hspace{1em}\text{almost surely},\end{array}$$

and

$$\begin{array}{c}\hfill e^{\mathrm{IDG}}({\mathrm{\Gamma }}^{\ast })⟶e^{\mathrm{IDG}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )\hspace{1em}\text{almost surely},\end{array}$$

where $e^{\mathrm{ICL}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )$ and $e^{\mathrm{IDG}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )$ are two deterministic functions of the parameters τ, α, κ, ρ and λ. The exact expressions of these two functions can be found in SI Appendix, section 5. For simplicity, we only present in what follows the ridgeless limit (i.e., $\lambda \to 0^{+}$) of the asymptotic generalization errors.

Result 1.

(ICL generalization error in the ridgeless limit). Let

$$\begin{array}{c}\hfill q^{\ast }:=\frac{1+\rho }{\alpha },\hspace{1em}{m}^{\ast }:={\mathcal{M}}_{\kappa }\left(q^{\ast }\right),\hspace{1em}{\mu }^{\ast }:=q^{\ast }{\mathcal{M}}_{\kappa /\tau }(q^{\ast }),\end{array}$$ [14]

where ${\mathcal{M}}_{\kappa }(·)$, defined in SI Appendix, section 8, is a function related to the Stieltjes transform of the Marchenko–Pastur law. Then

$$\begin{array}{cc}& e_{\mathrm{ridgeless}}^{\mathrm{ICL}}:={\displaystyle \underset{\lambda \to 0^{+}}{lim}}{e}^{\mathrm{ICL}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )=\hfill \\ \hfill & \left\{\begin{array}{cc}\frac{\tau (1+q^{\ast })}{1-\tau }\left[1-\tau {(1-{\mu }^{\ast })}^2+{\mu }^{\ast }(\rho /q^{\ast }-1)\right]\hfill & \\ -2\tau (1-{\mu }^{\ast })+(1+\rho )\hfill & \tau <1\hfill \\ \left(q^{\ast }+1\right)\left(1-2q^{\ast }{m}^{\ast }-{(q^{\ast })}^2{\mathcal{M}}_{\kappa }^{\prime }(q^{\ast })+\frac{(\rho +q^{\ast }-{(q^{\ast })}^2{m}^{\ast })m^{\ast }}{\tau -1}\right)\hfill & \\ -2(1-q^{\ast }{m}^{\ast })+(1+\rho )\hfill & \tau >1\hfill \end{array}\right.,\hfill \end{array}$$ [15]

where ${\mathcal{M}}_{\kappa }^{\prime }(·)$ means the derivative of ${\mathcal{M}}_{\kappa }(q^{\ast })$ with respect to $q^{\ast }$.

Result 2.

(IDG generalization error in the ridgeless limit). Let $q^{\ast }$, $m^{\ast }$, and ${\mu }^{\ast }$ be the scalars defined in Eq. 14. We have

$$\begin{array}{cc}& e_{\mathrm{ridgeless}}^{\mathrm{IDG}}:={\displaystyle \underset{\lambda \to 0^{+}}{lim}}{e}^{\mathrm{IDG}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )=\hfill \\ \hfill & \left\{\begin{array}{cc}\frac{\tau }{1-\tau }\left(\frac{\rho +q^{\ast }-2q^{\ast }(1-\tau )(q^{\ast }/{\xi }^{\ast }+1)}{1-p^{\ast }(1-\tau )}+\frac{\tau {\mu }^{\ast }{(q^{\ast }+{\xi }^{\ast })}^2}{q^{\ast }}\right)\hfill & \tau <1\hfill \\ \frac{\tau }{\tau -1}[\rho +q^{\ast }(1-q^{\ast }{m}^{\ast })]\hfill & \tau >1\hfill \end{array}\right.,\hfill \end{array}$$ [16]

where ${\xi }^{\ast }=\frac{(1-\tau )q^{\ast }}{\tau {\mu }^{\ast }}$ and $p^{\ast }=(1-\kappa (\frac{\kappa {\xi }^{\ast }}{1-\tau }+1{)}^{-2}{)}^{-1}$.

We will discuss various implications of these equations in the next sections.

We derived these results using techniques from random matrix theory. The full setup and technical details are presented in SI Appendix. The computations involve analysis of the properties of the finite-sample optimal parameter matrix ${\mathrm{\Gamma }}^{\ast }$ (Eq. 10). A key technical component of our analysis involves characterizing the spectral properties of the sample covariance matrix of $n=\mathrm{\Theta }(d^2)$ i.i.d. random vectors in dimension $\mathrm{\Theta }(d^2)$. Each of these vectors is constructed as the vectorized version of the matrix in Eq. 7. Related but simpler versions of this type of random matrices involving the tensor product of i.i.d. random vectors have been studied in recent work (39). Some of our derivations are based on nonrigorous yet technically plausible heuristics. We support these predictions with numerical simulations and discuss in SI Appendix the steps required to achieve a fully rigorous proof.

3.3. Sample-Wise Double-Descent.

How large should n, the pretraining dataset size, be for the linear attention to successfully learn the task in-context? In Fig. 1, we plot our theoretical predictions for the ICL and IDG error as a function of $\tau =n/d^2$ and verify them with numerical simulations. Our results demonstrate that the quadratic scaling of sample size with input dimensions is indeed an appropriate regime where nontrivial learning phenomena can be observed.

Fig. 1.

ICL performance as a function of τ: theory (solid lines) vs. simulations (dots). Plots of (A) $e_{\mathrm{ridgeless}}^{\mathrm{IDG}}(\tau ,\alpha ,\kappa ,\rho )$, (B) $e_{\mathrm{ridgeless}}^{\mathrm{ICL}}(\tau ,\alpha ,\kappa ,\rho )$, and (C) $e^{\mathrm{ICL}}(\tau ,\alpha ,\kappa ,\rho ,\lambda )$ against τ. Simulated errors calculated by evaluating the corresponding test error on the corresponding optimized ${\mathrm{\Gamma }}^{\ast }$. Parameters: $d=100$, $\rho =0.01$ for all; (A and B) $\kappa =0.5$, (C) $\alpha =10,\kappa =\infty$. Averages and SD are computed over 10 runs.

As apparent in Fig. 1, we find that the generalization error for both ICL and IDG tasks are not monotonic in the number of samples. In the ridgeless limit, both ICL and IDG errors diverge at $\tau =1$, with the leading order behavior in the $\tau \uparrow 1$ (respectively $\tau \downarrow 1$) limit given by $\frac{c_1}{1-\tau }$ (respectively $\frac{c_2}{\tau -1}$), where $c_1$ (respectively $c_2$) is a τ-independent constant. This leads to a “double-descent” behavior (38, 40) in the number of samples. As in other models exhibiting double-descent (38, 40, 41), the location of the divergence is at the interpolation threshold: the number of parameters of the model (elements of Γ) is, to leading order in d, equal to $d^2$, which matches the number of pretraining samples at $\tau =1$. Further, we can investigate the effect of ridge regularization on the steepness of the double descent, as illustrated in Fig. 1C for the ICL task. As we would expect from other models exhibiting double-descent (38, 40, 41), increasing the regularization strength suppresses the peak in error around the interpolation threshold.

Finally, we note that if we take the limit of $\kappa \to \infty$ and $\alpha \to \infty$ in Result 1 (in either order), the ICL generalization error in the ridgeless case reduces to the generalization error of simple ridgeless interpolation with isotropic Gaussian covariates in $d^2$ dimensions (38, 41):

$$\begin{array}{cc}\hfill {\displaystyle \underset{\alpha \to \infty }{lim}}{\displaystyle \underset{\kappa \to \infty }{lim}}{e}_{\mathrm{ridgeless}}^{\mathrm{ICL}}& ={\displaystyle \underset{\kappa \to \infty }{lim}}{\displaystyle \underset{\alpha \to \infty }{lim}}{e}_{\mathrm{ridgeless}}^{\mathrm{ICL}}\hfill \\ \hfill & =\left\{\begin{array}{cc}1-\tau +\frac{\rho }{1-\tau }\hfill & \tau <1,\hfill \\ \frac{\rho \tau }{\tau -1}\hfill & \tau >1.\hfill \end{array}\right.\hfill \end{array}$$

This result makes sense, given that in this limit the ICL generalization problem reduces to the generalization error of ridge regression in $d^2$ dimensions with covariates formed as the tensor product of i.i.d. Gaussian vectors, which by universality results in ref. 39 should in turn be asymptotically equal to that for isotropic Gaussian covariates (38). We note that taking the limit $\alpha \to \infty$ is not strictly necessary for this universality, but doing so simplifies the expressions and makes the double-descent obvious.

3.4. ICL and IDG Error Curves Can Have Nonmonotonic Dependence on Context Length.

How large should the context length be? In Fig. 2, we plot our theoretical results verified with experiments. We observe that we have correctly identified the regime where in-context and in-weight learning appear: context length indeed scales linearly with input dimensions, as numerical simulations computed using finite d fit the asymptotic error curves.

Fig. 2.

Error curves as functions of α: theory (solid lines) vs. simulations (dots). Plots of (A) $e_{\mathrm{ridgeless}}^{\mathrm{IDG}}$ and (B and C) $e_{\mathrm{ridgeless}}^{\mathrm{ICL}}$ against α. Hollowed markers, when plotted, indicate ${\alpha }^{\ast }$ value minimizing error if it exists. Parameters: $d=100$; (A) $\tau =0.5,\rho =0.5$; (B) $\tau =0.5,\rho =0.1$; (C) $\tau =20,\rho =0.5$. Averages and SD are computed over 20 to 100 runs.

An interesting observation is that the IDG and ICL errors do not always monotonically decrease with context length; we see that there are parameter configurations for which the error curves are minimized at some finite ${\alpha }^{\ast }$.

While the functional form of $e_{\mathrm{ridgeless}}^{\mathrm{IDG}}$ and $e_{\mathrm{ridgeless}}^{\mathrm{ICL}}$ is too complex to study their minimizers analytically, we can investigate the nonmonotonicity of IDG and ICL error in α numerically. This is done in Fig. 3, where the blue-green colormap corresponds to the value ${\alpha }^{\ast }$, if it exists, that minimizes $e_{\mathrm{ridgeless}}^{\mathrm{IDG}}(\tau ,\alpha ,\kappa ,\rho )$ or $e_{\mathrm{ridgeless}}^{\mathrm{ICL}}(\tau ,\alpha ,\kappa ,\rho )$ at fixed $\tau ,\kappa ,\rho$. The gray color indicates cases where the corresponding function is monotonic in α. The results reveal that the nonmonotonicity of the errors with respect to context length α depends in a nontrivial manner on both the task diversity κ and the noise-to-signal ratio ρ.

Fig. 3.

Phase diagram of nonmonotonicity in α of ridgeless IDG and ICL error, for fixed τ against task diversity κ and label noise ρ. The colormap corresponds to the value ${\alpha }^{\ast }$ that minimizes error in α, if it exists, at that particular $\kappa ,\rho$ pair; gray is plotted if the error curve is monotonic at that $\kappa ,\rho$ pair. Figures (A–C) in this plot correspond to the respective setup in Fig. 2, and points A–E correspond to the respective curve in Fig. 2. The dashed vertical lines are plotted at $\kappa =\text{min}(\tau ,1)$. We know from Eq. 17 and surrounding discussion that ICL error diverges in α for all $\kappa ,\rho$ to the left of this line, and IDG error diverges on this line for $\tau <1$.

In addition to the nonmonotonicity of error curves in α, we also notice a divergence in error as α increases for some $\kappa ,\tau$ values. To determine when this occurs, we compute the large α limit of the above error curves. For ICL error, we have

$$\begin{array}{cc}\hfill {\displaystyle \underset{\alpha \to \infty }{lim}}{e}_{\mathrm{ridgeless}}^{\mathrm{ICL}}& =\left\{\begin{array}{cc}\infty \hfill & \kappa \le \text{min}(\tau ,1),\hfill \\ 1-\tau +\rho +\frac{\rho \kappa \tau }{(\kappa -\tau )(1-\tau )}\hfill & \tau <1,\kappa >\tau ,\hfill \\ \rho +\frac{\rho \kappa }{(\kappa -1)(\tau -1)}\hfill & \tau >1,\kappa >1.\hfill \end{array}\right.\hfill \end{array}$$ [17]

Note the explicit divergence of ICL error for $\kappa \le \text{min}(\tau ,1)$ (independent of ρ). Similar investigation of Result 2 under the limit $\alpha \to \infty$ reveals an explicit divergence of IDG error for $\kappa =\tau$. The divergence in α of ICL and IDG errors is a peculiarity of the ridgeless limit of errors we have taken in Results 1 and 2 and is further discussed in SI Appendix, section 5.

3.5. Learning Transition with Increasing Pretraining Task Diversity.

It is important to quantify if and when a given model is actually learning in-context, i.e., solving a new regression problem by adapting to the specific structure of the task rather than relying solely on memorized training task vectors. We refer to this phenomenon as task generalization.

We expect pretraining task diversity (k) to play a key role in task generalization. To provide some motivation for this, suppose the model was pretrained on only a single task vector ($k=1$). It is reasonable to expect that the model will memorize this task vector and make predictions under the prior that the task vector in the test examples is identical to that particular pretraining task vector. On the other extreme, suppose the model was pretrained by drawing a fresh task vector for every training sequence. Here, the expectation would be that the model achieves task generalization. Indeed, results from previous computational studies on the linear regression task (20) suggest that transformers are better in-context learners if they have seen a greater task diversity in the training set, and further the models transition from a task memorization phase to a task generalization phase as task diversity increases. Hence, we study the task generalization behavior of the linear transformer as the task diversity increases and precisely quantify the task diversity scale at which such transition happens.

First, we need to discuss how to measure task generalization. To motivate our solution, we take a closer look at the ridge estimator ($w_{\text{Bayes}}^{\text{ridge}}$ as in Eq. 11), which is the Bayes-optimal estimator given the ICL task structure. We analytically characterize $e_{\text{ICL}}$ and $e_{\text{IDG}}$ errors for this estimator in SI Appendix, section 6, giving

$$e_{\text{ICL}}^{\text{ridge}}=\rho \left(1+\frac{1}{\alpha }{\mathcal{M}}_{\alpha }\left(\frac{\rho }{\alpha }\right)\right)=e_{\text{IDG}}^{\text{ridge}}.$$

Note that the performance of the ridge estimator is identical on the ICL and IDG test distributions, meaning that it does not rely at all on memorizing task vectors.

Motivated by these arguments, we propose to measure the task generalization capability of a model by studying the difference between its performance on the ICL test distribution and its performance on the IDG test distribution. Specifically, we consider the quantity $g_{\text{task}}=e_{\text{ICL}}-e_{\text{IDG}}$ for a given model or estimator. This difference being large implies the model, performing better on training tasks, has not learned the true task distribution and is not generalizing in task. Conversely, a small difference between ICL error and IDG error suggests that the model is leveraging the underlying structure of the task distribution rather than overfitting to and interpolating specific task instances seen in training.

Simulations of $g_{\text{task}}$ as a function of task diversity, $\kappa =k/d$, are shown in Fig. 4 for two inference models: dMMSE estimator, and the linear transformer. The dMMSE estimator is included as a benchmark quantifying the performance of a “pure memorizer,” as it is the Bayes-optimal estimator that assumes task vectors only seen in the training test. We see that, as the task diversity parameter κ increases, $g_{\text{task}}$ falls for both estimators but at very different rates.

Fig. 4.

Plots of $g_{\text{task}}$ against κ for linear transformer model vs. dMMSE estimator. Panel (A) has plots on log–log scale highlighting the substantial difference in their rates of convergence toward 0. Dashed and solid lines are theory predictions; dots and squares are numerical simulation. Panel (B) has plots on log–linear scale with the addition of the green $\alpha \to \infty$ curve given by Eq. 18 demonstrating the phase transition at $\kappa =1$. Parameters: Linear transformer: $d=100\mathrm{to}140$, $\tau =0.2\alpha$ throughout, simulations computed over 20 runs; dMMSE: $d=100$, simulations computed over 1,000 to 5,000 runs.

The $g_{\text{task}}$ for the dMMSE estimator falls very slowly as a function of κ. This is expected because the explicit form of $w_{\text{Bayes}}^{\text{dMMSE}}$, given in Eq. 12, is that of a kernel smoother employing a Gaussian kernel as the weighting function around each task vector. Known results about this class of estimators show the performance of the dMMSE estimator on the ICL task suffers from the “curse of dimensionality” (42–44), i.e., the required number of training samples $k=d\kappa$ would need to be exponential in d to outperform some given $e_{\text{ICL}}$ error tolerance. Intuitively, the IDG task distribution ${\mathcal{P}}_{\text{test}}={\mathcal{P}}_{\text{train}}=\text{Unif}\{w_1,\cdots ,w_k\},w_i\sim \mathcal{N}(0,I_d)$ and the ICL task distribution ${\mathcal{P}}_{\text{test}}=\mathcal{N}(0,I_d)$ are indistinguishable (by appropriate measures) only if k is exponentially large in d. We present more formal arguments in SI Appendix, section 6.

On the other hand, the $g_{\text{task}}$ for the linear transformer estimator first behaves similarly to the dMMSE estimator, but starts falling sharply after around $\kappa =1$. How quickly does $g_{\text{task}}^{\text{LT}}$ limit to 0 as $\kappa \to \infty$? By expanding $e_{\text{ICL}}-e_{\text{IDG}}$ from Eq. 15 and Eq. 16 in κ we have the leading asymptotic behavior given by

$$g_{\text{task}}^{\text{LT}}=0+\frac{1}{\kappa }\left\{\begin{array}{cc}{c}_1\hfill & \tau <1\hfill \\ c_2\hfill & \tau >1\hfill \end{array}\right.+\mathcal{O}\left(\frac{1}{{\kappa }^2}\right)$$

for constants $c_1,c_2$ dependent on $\alpha ,\tau ,\rho$, provided in SI Appendix, section 6. This sharp fall signals a transition of the model to the task generalization regime. To further understand the nature of this transition and the role of κ in the solution learned by the linear attention mechanism, consider the regime where $\tau ,\alpha \to \infty$ with $\alpha /\tau \equiv c^{\ast }$ kept fixed. Under this setting, we have

$$\begin{array}{c}\hfill {\displaystyle \underset{\begin{array}{c}\tau \to \infty \\ \alpha \to \infty \end{array}}{lim}}{g}_{\text{task}}^{\text{LT}}=\left\{\begin{array}{cc}(1-\kappa )\left(1+\frac{\rho }{1+\rho }{c}^{\ast }\right)\hfill & \kappa <1\hfill \\ 0\hfill & \kappa >1\hfill \end{array}\right..\end{array}$$ [18]

This change in analytical behavior indicates a phase transition at $\kappa =1$. We see this phenomenon in both Fig. 4 A and B where the $g_{\text{task}}$ curve plotted against κ becomes nondifferentiable at $\kappa =1$ for $\alpha \to \infty$.

We conclude that the linear transformer model is a much more efficient task generalizer than the dMMSE estimator which is the Bayes-optimal memorizer, as shown in Fig. 4. The $1/\kappa$ decay in $g_{\text{task}}^{\text{LT}}$ vs. the much slower decay of $g_{\text{task}}^{\text{dMMSE}}$ suggests that the linear transformer quickly learns an inference algorithm that generalizes in-context rather than interpolates between training tasks.

4. Experiments with Fully Parameterized Nonlinear Transformers

Given that our theoretical results are derived from a simplified linear attention setting, we aim to determine whether these insights can be extended to nonlinear and deep Transformer models with full $K,Q,V$ parameterization. Even though we would not expect the specific algorithm discussed in Section 2.4 to transfer to these nonlinear settings, we will test whether our proposed scalings and main qualitative predictions remain. Specifically, we will investigate the following: 1) whether we have identified the correct scaling regime for nontrivial learning in an ICL task; 2) whether the full Transformer exhibits a sample-wise double descent, where the location of the peak error scales quadratically with input dimensions as predicted by our theory; and 3) whether the transition from memorization to task generalization occurs.

Our experiments^* are done with a standard Transformer architecture, where each sample context initially takes the form given by Eq. 2. The fully parameterized linear transformer Fig. 5A and the softmax-only transformer Fig. 5B do not use a multilayer perceptron (MLP) in addition to attention. If MLPs are used (e.g., Fig. 5 C and D), the architecture consists of blocks with: 1) a single-head softmax self-attention with $K,Q,V\in {\mathbb{R}}^{(d+1)\times (d+1)}$ matrices, followed by 2) a two-layer dense MLP with GELU activation and hidden layer of size $d+1$ (1). Residual connections are used between the input tokens (padded from dimension d to $d+1$), the pre-MLP output, and the MLP output. We use a variable number of attention+MLP blocks before returning the final logit corresponding to the $(d+1,\ell +1)$th element in the original embedding structure given by Eq. 2. The loss function is the mean squared error (MSE) between the predicted label (the output of the model for a given sample Z) and the true value $y_{\ell +1}$. We train the model in an offline setting with n total samples $Z_1,\cdots ,Z_n$, divided into 10 batches, using the Adam optimizer (45) with a learning rate ${10}^{-4}$ until the training error converges, typically requiring 10,000 epochs.^† The structure of the pretraining and test distributions exactly follows the setup for the ICL and IDG tasks described in Section 2.

Fig. 5.

Experimental verification in full linear attention (A) (full $K,Q,V$ matrices) and nonlinear models (B) (one block softmax attention only), (C) (two blocks of softmax attention and 1-layer MLP), of both scaling definition for τ and double descent behavior in n. (A–C) show error curves against τ for various architectures, consistent across token dimension $d=20,40,80$. Some deviations arise near double descent in the linear model (A), possibly due to parameter inversion or instabilities, and for large τ in the deep nonlinear model (C), possibly due to the training procedure. Double-descent phenomena are confirmed: increasing n will increase error until an interpolation threshold is reached. Colored dashed lines indicate experimental interpolation threshold for that architecture and d configuration. (D) shows that the location of the interpolation threshold occurs for n proportional to $d^2$, as predicted by the linear theory. Dots are experimental interpolation thresholds for various architectures, and dashed lines are best fit curves correspond to fitting $log(n)=alog(d)+b$, each with $a\approx 2$ (explicitly, $a_{\text{full linear}}=1.87$, $a_{\text{softmax}}=1.66$, $a_{\text{2 blocks}}=2.13$, $a_{\text{3 blocks}}=2.08$). Interpolation threshold was computed empirically by searching for location in τ of sharp increase in value and variance of training error at a fixed number of gradient steps. Parameters:$\alpha =1,\kappa =\infty ,\rho =0.01$. For (A–C): Variance shown comes from model trained over different samples of pretraining data; lines show averages over 10 runs and shaded region shows SD.

In Fig. 5, we plot the average error for a range of fully parameterized linear (A) and nonlinear architectures (B and C) on the ICL task against the parameter $\tau =n/d^2$, which measures the pretraining sample complexity. Notice that the large-d full $K,Q,V$ parameterized linear architecture studied in (A) follows the general trend of the theory derived on a simplified parameterization. Further, in all figures, plots across different values of d roughly collapse, suggesting again that our scaling relations capture the observed behavior.

Across all architectures, we observe a double descent in τ. The peak of this curve occurs at the interpolation threshold, which we identify by tracking the value of τ at which the training loss becomes nonzero (see figure caption). Our theory predicts that the number of samples n at the peak, as well as the interpolation threshold, should scale with $d^2$. Indeed, for the fully parameterized linear transformer in (Fig. 5A), the double descent occurs very close to $\tau =1$ for larger d, as predicted by the reduced-parameter linear theory. This general quadratic scaling is summarized in (Fig. 5D) across a range of architectures. These findings suggest that nonlinear Transformers exhibit a similar interpolation transition, with the same quadratic scaling in sample size as seen in the simplified model.

Finally, we can investigate the effects of task diversity on $g_{\text{task}}$ in a nonlinear model, shown in Fig. 6. This provides empirical evidence toward the choice of κ scaling, as errors are consistent over a range of token dimension choices. Further, we recover the memorization-to-task-generalization transition in κ, as observed before in ref. 20. In Fig. 6, we see the value of κ for which $g_{\text{task}}$ appears to converge toward 0 is below 1 (as in Eq. 18), suggesting that the more powerful architecture used here may be more “task-efficient” than the linear architecture in achieving memorization-to-generalization transition.

Fig. 6.

Plot of $g_{\text{task}}$ against κ (log–linear scale) for the nonlinear architecture given in Fig. 5C. This demonstrates consistency of κ scaling across increasing dimension choices $d=20,40,80$, as well as a transition in task-generalization measure $g_{\text{task}}$ near $\kappa =1$, familiar from the linear theory. Parameters:$\tau =10,\alpha =1,\rho =0.01$. Variance shown comes from 10 models trained over different samples of pretraining data.

5. Discussion

In this paper, we computed sharp asymptotics for the ICL performance in a simplified model of ICL for linear regression using linear attention. This exactly solvable model demonstrates a transition from memorizing pretraining task vectors to generalizing beyond them as the diversity of pretraining tasks increases, echoing empirical findings in full Transformers (20). Additionally, we observe a sample-wise double descent as the amount of pretraining data increases. Our numerical experiments show that full, nonlinear Transformers exhibit qualitatively similar behavior in the scaling regime relevant to our solvable model. Understanding the mechanistic underpinnings of ICL of well-controlled synthetic tasks in solvable models is an important prerequisite to understanding how it emerges from pretraining on natural data (21).

Our paper falls within a broader program of research that seeks sharp asymptotic characterizations of the performance of machine learning algorithms. This program has a long history in statistical physics (41, 46, 47), and has in recent years attracted substantial attention in machine learning (38, 41, 48–54). For simplicity, we have assumed that the covariates in the in-context regression problem are drawn from an isotropic Gaussian distribution. However, our technical approach could be extended to anisotropic covariates, and, perhaps more interestingly, to featurized linear attention models in which the inputs are passed through some feature map before linear attention is applied (34, 35). This extension would be possible thanks to an appropriate form of Gaussian universality: for certain classes of regression problems, the asymptotic error coincides with that of a model where the true features are replaced with Gaussian features of matched mean and covariance (38, 39, 48–53, 55). This would allow for a theoretical characterization of ICL for realistic data structure in a closer approximation of full softmax attention, yielding more precise predictions of how performance scales in real Transformers.

In our analysis, we have assumed that the model is trained to interpolation on a fixed dataset. This allows us to cast our simplified form of linear attention pretraining as a ridge regression problem, which in turn enables our random matrix analysis. In contrast, Transformer-based large language models are usually trained in a nearly online setting, where each gradient update is estimated using fresh examples with no repeating data (56). Some of our findings, such as double-descent in the learning curve as a function of the number of pretraining examples, are unlikely to generalize to the fully online setting. It will be interesting to probe these potential differences in future work.

Finally, our results have some bearing on the broad question of what architectural features are required for ICL (7, 13, 21). Our work shows that a full Transformer—or indeed even full linear attention—is not required for ICL of linear regression. However, our simplified model retains the structured quadratic pairwise interaction between inputs that is at the heart of the attention mechanism. It is this quadratic interaction that allows the model to solve the ICL regression task, which it does essentially by reversing the data correlation (Section 2.4). One would therefore hypothesize that our model is minimal in that further simplifications within this model class would impair its ability to solve this ICL task. In the specific context of regression with isotropic data, a simple point of comparison would be to fix $\mathrm{\Gamma }=I_d$, which gives a pretraining-free model that should perform well when the context length is very long. However, this further-reduced model would perform poorly if the covariates of the in-context task are anisotropic. Generally, it would be interesting to investigate when models lacking this precisely engineered quadratic interaction can learn linear regression in-context (57), and if they are less sample-efficient than the attention-based models considered here.

Supplementary Material

Appendix 01 (PDF)

Data, Materials, and Software Availability

Code to reproduce experiments’ data have been deposited in Github (https://github.com/Pehlevan-Group/incontext-asymptotics-experiments) (58).