How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?

TL;DR

<3-5 sentence high-level summary> The work analyzes in-context learning (ICL) in the simplest yet nontrivial setting by pretraining a restricted single-layer linear attention model for Gaussian-prior linear regression. It proves a dimension-free task-complexity bound showing that only a small number of independent tasks are needed for effective pretraining, and that the pretrained model can achieve nearly Bayes-optimal (ridge) risk on unseen tasks when the inference context length matches pretraining. To establish these results, the authors develop novel operator-theoretic tools—diagonalization and operator polynomials—to analyze 8th-order tensor recursions arising in SGD dynamics. They also provide a detailed comparison between the pretrained attention estimator and the Bayes-optimal ridge estimator, clarifying when ICL can be optimally or near-optimally realized, and complement prior empirical findings with a solid statistical foundation.

Abstract

Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities, enabling them to solve unseen tasks solely based on input contexts without adjusting model parameters. In this paper, we study ICL in one of its simplest setups: pretraining a linearly parameterized single-layer linear attention model for linear regression with a Gaussian prior. We establish a statistical task complexity bound for the attention model pretraining, showing that effective pretraining only requires a small number of independent tasks. Furthermore, we prove that the pretrained model closely matches the Bayes optimal algorithm, i.e., optimally tuned ridge regression, by achieving nearly Bayes optimal risk on unseen tasks under a fixed context length. These theoretical findings complement prior experimental research and shed light on the statistical foundations of ICL.

Figures (5)

• Figure 1: Task complexity of ICL (of the one-step GD model), ridge regression, and OLS. The context length is $M=N=200$. The ambient dimension is $d=100$. We observe that as the number of pretraining tasks increases, one-step GD achieves smaller MSE and becomes closer to the Bayes algorithm, ridge regression. This is consistent with our theory.
• Figure 2: The effect of the ambient dimension for ICL (of one-step GD), ridge regression, and OLS. The context length is $M=N=200$. The number of pretraining tasks is $10^5$ for ICL. We observe that when the spectrum of the data covariance $\mathbf{H}$ decays relatively fast, for example, $\lambda_i\sim 2^{-i}$ and $\lambda_i\sim i^{-2}$, the performances of the three considered algorithms are not sensitive to the ambient dimension. This is consistent with our theory.
• Figure 3: The effect of the number of context examples during inference for ICL (of one-step GD) and ridge regression. The number of context examples during pretraining is $N= 40$. The ambient dimension is $d = 20$. The MSE of OLS is significantly worse than ICL and ridge regression when $M\leq N=20$, so we ignore OLS in this plot for a better visualization. We observe that the ICL achieves a similar MSE to ridge regression when $M$ is close to $N$. However, the gap becomes larger when $M$ is much smaller than $N$. This is consistent with our theory.
• Figure 4: The effect of data misspecification for the ICL of one-step GD. The base setup, $y\sim\mathcal{N}(\bm{\beta}^\top \mathbf{x}, \sigma^2)$ with $\sigma^2=1$, is well-specified. We then consider three misspecification scenarios. Uniform: $y\sim \bm{\beta}^\top\mathbf{x}+\mathrm{uniform}[-\sqrt{3},\sqrt{3}]$. Sigmoid: $y \sim\mathcal{N}( \mathrm{sigmoid}(\bm{\beta}^{\top}\mathbf{x}), \sigma^2)$. Square: $y \sim \mathcal{N}((\bm{\beta}^{\top}\mathbf{x})^{2}, \sigma^2)$. We observe that the type of misspecification affects the ICL performance. In particular, the ICL performance declines less when the ground-truth model is closer to a linear model.
• Figure 5: ICL of a three-layer transformer. The linear regression tasks are generated according to Assumption \ref{['assump:data']}, with $d = 20$, $N = 2d$, standard deviation $\sigma = 0.5$, scaling factor $\psi = 1$, and a polynomial decay spectrum $\lambda_i = i^{-4}$. We fix the number of context examples during inference to $M = N$. We observe that the ICL error decreases as the number of pretraining tasks increases, approaching the performance of the Bayes optimal algorithm, ridge regression.

Theorems & Definitions (69)

• Theorem 3.1: ICL risk
• Theorem 4.1: Task complexity for pretraining
• Corollary 4.2: Large stepsize
• Proposition 5.1: Optimally tuned ridge regression
• Corollary 5.2: Average risk of ridge regression, corollary of tsigler2020benign
• Theorem 5.3: Average risk of the pretrained attention model
• Corollary 5.4: Examples
• Lemma C.1
• proof : Proof of Lemma \ref{['lemma:4-moment']}
• proof : Proof of Theorem \ref{['thm:population-risk']}