Research map / Learning theory and algorithms / Kernels and regression theory
How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows theoretically that a single-layer linear attention model (one-step GD with matrix stepsize) pretrained by SGD on Gaussian-prior linear regression needs only a dimension-independent number of tasks, and nearly matches Bayes-optimal ridge regression at inference when context lengths match.
Problem
How many independent pretraining tasks are needed for in-context learning of linear regression with a Gaussian prior, and how close does a pretrained single-layer linear attention model get to the Bayes-optimal (optimally tuned ridge) predictor?
Approach
Reparameterizes single-layer linear attention as one-step GD with a matrix stepsize, gives a sharp task-complexity bound for SGD pretraining via an operator method extended to 8th-order tensors (using diagonalization and operator polynomials), and shows the pretrained attention matches ridge regression risk when pretraining and inference context lengths are close.
Key results
- Pretrained one-step GD ICL risk approaches optimally tuned ridge regression as the number of pretraining tasks grows (experiments up to 1e5 tasks, d=100, N=M=200).
- Theory: for uniform spectrum the excess risk is bounded by N/s + T_eff/s^2 (or s^2/T_eff), so dimension-independent task counts suffice under fast-decaying spectra.
- Pretrained attention matches ridge risk when inference context length M is close to pretraining length N; the gap grows when M << N (e.g., exponential spectrum: log M/M vs log N/M).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase the number of independent pretraining tasks | One-step GD ICL · Exponential and polynomial spectra; M = N = 200 | MSE | Helped |
| Increase the number of independent pretraining tasks | Three-layer transformer · Polynomial spectrum λᵢ = i⁻⁴; M = N | ICL error | Helped |
| Use an inference context length close to the pretraining context length | Optimally pretrained single-layer linear attention · Linear regression with Gaussian prior | average regression risk | Helped |
| Use an inference context length much smaller than the pretraining context length | Optimally pretrained single-layer linear attention · Linear regression with Gaussian prior | average regression risk versus tuned ridge regression | Hurt |
| Increase ambient dimension d (ablation) | One-step GD ICL · Fast-decaying covariance spectra (λᵢ ∼ 2⁻ⁱ and λᵢ ∼ i⁻²) | ICL performance | No clear effect |
| Replace Gaussian additive noise with uniform noise of matched variance (ablation) | One-step GD ICL · Uniform-noise misspecification | ICL performance | No clear effect |
| Replace the linear mean with a sigmoid of βᵀx (ablation) | One-step GD ICL · Sigmoid misspecification | ICL performance | Hurt |
| Replace the linear mean with (βᵀx)² (ablation) | One-step GD ICL · Square-mean misspecification | ICL performance | Hurt |
Surprising findings
- Under a uniform spectrum with M<N the average risk of the attention model can be worse than ridge in the regime where the gap scales as N^(1/a)/M for polynomial spectra.
Limitations
- Analysis restricted to one-step GD (restricted linear attention) parameterization, not the original softmax-free attention parameterization which is non-convex.
- Fixed context length N during pretraining; suboptimal ICL when inference context length differs substantially from pretraining length.
- Gaussian prior and well-specified linear model assumptions; misspecification experiments degrade performance.
Experimental setup
Datasets: synthetic Gaussian linear regression tasks
Benchmarks: synthetic linear regression with Gaussian prior
Metrics: excess ICL risk regression MSE
Baselines: optimally tuned ridge regression ordinary least squares uniform convergence pretraining bound (Ahn et al.)
Builds on: Garg et al. in-context learning of linear regression Ahn et al. linear attention as one-step GD with matrix stepsize Bach and Moulines / Jain et al. operator method for SGD
Closest papers by meaning
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.91)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.90)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.90)
- Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond (NeurIPS 2024 · similarity 0.90)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.90)
- In-Context Learning with Representations: Contextual Generalization of Trained Transformers (NeurIPS 2024 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.