Research map / Architectures and efficiency / Attention and transformers
When can in-context learning generalize out of task distribution?
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.
Trains GPT-2-style transformers for in-context linear regression with task vectors drawn from hyperspherical caps of varying angle, finding a sharp transition from specialized to out-of-task-distribution generalizing ICL as task diversity grows.
Problem
When can in-context learning generalize to tasks outside the pretraining task distribution, and what pretraining task diversity is required for this out-of-distribution generalization?
Approach
Defines a new notion of task diversity based on the angular size of the pretraining task subset on the hypersphere (cap half-angle), distinct from the number of pretraining tasks, and shows a sharp specialization-to-generalization transition as this diversity increases; a phase diagram combines both diversity axes.
Key results
- Transformers pretrained on caps with phi >= ~120 degrees achieve low test error across all test angles, while smaller caps yield specialized solutions that fail out of distribution (transition at ~135 degrees with label noise sigma^2=0.25).
- Transformers pretrained on large-enough caps outperform the Bayes-optimal in-distribution estimator restricted to the cap on out-of-distribution test tasks (delta=175 deg).
- The transition threshold along phi stays near 120 degrees across dimensions d=3 to higher d and is unaffected by model depth; a similar transition appears in logistic classification (~135 degrees) and nonlinear ReLU regression.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase the pretraining task-cap half-angle beyond about 120° | Transformer trained on linear regression ICL · Linear regression, tasks outside the training cap | test MSE | Helped |
| Increase the pretraining task-cap half-angle from the specialized regime to about 135° | Transformer trained on noisy linear regression ICL · Linear regression with σ² = 1/4, OOD tasks | excess test error | Helped |
| Increase task diversity | Transformer trained on linear regression ICL · Linear regression, short in-distribution contexts | test error | Mixed |
| Increase task diversity to φ ≳ 45° | Transformer pretrained on tasks with radius R=1 · Tasks with R<1 | test error | Helped |
| Increase the number of pretraining tasks | Transformer trained on linear regression ICL · In-distribution tasks interpolated between training tasks | loss relative to dMMSE | Helped |
| Increase regression dimension (ablation) | Transformer trained on linear regression ICL · Linear regression, OOD tasks | transition angle | No clear effect |
| Vary transformer depth (ablation) | Transformer trained with N=2¹¹ tasks · Linear regression, OOD tasks | transition angle | No clear effect |
| Replace linear labels with logistic-regression classification labels (ablation) | Transformer trained for ICL · Logistic regression, OOD tasks | test error | Helped |
| Replace linear labels with nonlinear ReLU regression labels (ablation) | Transformer trained for ICL · Nonlinear regression, OOD tasks | test loss | Helped |
| Restrict input vectors to a hyperspherical cap while keeping task vectors uniform over the full sphere (ablation) | Transformer trained on linear regression ICL · Data distribution shift, OOD tasks | test error | Helped |
| Remove input components perpendicular to the cap-center vector during early training (ablation) | Transformer trained with φ=45° · Linear regression training | training loss | No clear effect |
Surprising findings
- Pretrained transformers outperform the Bayes-optimal estimator for out-of-task-distribution tasks, which is restricted to the pretraining cap.
- Increasing task diversity drives generalization beyond the unit sphere to tasks with smaller norm R<1, despite training only on R=1.
Limitations
- Experiments limited to simple synthetic function classes (linear, logistic, small nonlinear regression) rather than LLM-scale ICL.
- No analytic theory explaining the transition; the mechanism is only characterized empirically.
- Transition point depends on label noise and the specific geometry of the task-diversity measure.
Experimental setup
Datasets: synthetic linear regression tasks sampled from hyperspherical caps
Benchmarks: synthetic linear regression test bands on hypersphere (delta=175 deg OOD, delta near 0 ID)
Metrics: mean squared error normalized NSR transition metric
Baselines: ordinary least squares Bayes-optimal estimator dMMSE estimator
Builds on: Raventós et al. pretraining task diversity AdamW GPT-2
Closest papers by meaning
- Can In-context Learning Really Generalize to Out-of-distribution Tasks? (ICLR 2025 · similarity 0.93)
- Differential learning kinetics govern the transition from memorization to generalization during in-context learning (ICLR 2025 · similarity 0.93)
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.92)
- Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context (NeurIPS 2025 · similarity 0.91)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.90)
- Position: Do pretrained Transformers Learn In-Context by Gradient Descent? (ICML 2024 · similarity 0.90)
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.