# When can in-context learning generalize out of task distribution?

arXiv: https://arxiv.org/abs/2506.05574 · HTML: https://mlsanity.com/paper/2506.05574/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 26 references that are in the map.
- First posted on arXiv 5 Jun 2025.

## Builds on (most cited first, 22 of 22)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Towards Out-Of-Distribution Generalization: A Survey](https://mlsanity.com/paper/2108.13624/) (2021)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [The Transient Nature of Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2311.08360/) (2023)
- [Out of Distribution Generalization in Machine Learning](https://mlsanity.com/paper/2103.02667/) (2021)
- [Asymptotic theory of in-context learning by linear attention](https://mlsanity.com/paper/2405.11751/) (2024)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [Pretraining Data Mixtures Enable Narrow Model Selection Capabilities in Transformer Models](https://mlsanity.com/paper/2311.00871/) (2023)
- [Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks](https://mlsanity.com/paper/2406.02550/) (2024)
- [A Closer Look at In-Context Learning under Distribution Shifts](https://mlsanity.com/paper/2305.16704/) (2023)
- [Differential learning kinetics govern the transition from memorization to generalization during in-context learning](https://mlsanity.com/paper/2412.00104/) (2024)
- [Breaking through the learning plateaus of in-context learning in Transformer](https://mlsanity.com/paper/2309.06054/) (2023)

## Built on by (most cited first, 16 of 16)

- [Pretrain-Test Task Alignment Governs Generalization in In-Context Learning](https://mlsanity.com/paper/2509.26551/) (2025)
- [Distinct mechanisms underlying in-context learning in transformers](https://mlsanity.com/paper/2604.12151/) (2026)
- [Bifrost: Steering Strategic Trajectories to Bridge Contextual Gaps for Self-Improving Agents](https://mlsanity.com/paper/2602.05810/) (2026)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces](https://mlsanity.com/paper/2606.06814/) (2026)
- [Sequential Correlations Change In-Context Learning: Effective Context Length and Architectural Mismatch](https://mlsanity.com/paper/2607.03660/) (2026)
- [In-context Learning of Single-index Targets: Comparing Kernel and Feature Learners](https://mlsanity.com/paper/2610.01712/) (2026)
- [In-Context Curiosity: Distilling Exploration for Decision-Pretrained Transformers on Bandit Tasks](https://mlsanity.com/paper/2510.00347/) (2025)
- [When Looks Do Not Lie: Discourse Structure Guided In-Context Learning for Faithful Diagram Generation](https://mlsanity.com/paper/2601.20476/) (2026)
- [The Generalization Spectrum: A Chromatographic Approach to Evaluating Learning Algorithms](https://mlsanity.com/paper/2606.25450/) (2026)
- [A Unified Framework for In-Context Learning with Causal and Masked Language Models](https://mlsanity.com/paper/2607.04081/) (2026)
- [Shortcut Before Circuit: Document Statistics Time In-Context Conflict Resolution](https://mlsanity.com/paper/2608.24460/) (2026)
- [The Illusion of Replacement: Rethinking Specialized Machine Learning Models in the Foundation Model Era](https://mlsanity.com/paper/2608.28980/) (2026)
- [Theory on Attention Dynamics for Out-of-Distribution In-Context Learning](https://mlsanity.com/paper/2609.36448/) (2026)
- [How Well Do Prior-Data Fitted Networks Amortize Bayesian Inference?](https://mlsanity.com/paper/2610.13764/) (2026)
- [Fisher information and the geometry of memorization in neural networks](https://mlsanity.com/paper/2610.16161/) (2026)

## Abstract (arXiv)

In-context learning (ICL) is a remarkable capability of pretrained transformers that allows models to generalize to unseen tasks after seeing only a few examples. We investigate empirically the conditions necessary on the pretraining distribution for ICL to emerge and generalize out-of-distribution. Previous work has focused on the number of distinct tasks necessary in the pretraining dataset. Here, we use a different notion of task diversity to study the emergence of ICL in transformers trained on linear functions. We find that as task diversity increases, transformers undergo a transition from a specialized solution, which exhibits ICL only within the pretraining task distribution, to a solution which generalizes out of distribution to the entire task space. We also investigate the nature of the solutions learned by the transformer on both sides of the transition, and observe similar transitions in nonlinear regression problems. We construct a phase diagram to characterize how our concept of task diversity interacts with the number of pretraining tasks. In addition, we explore how factors such as the depth of the model and the dimensionality of the regression problem influence the transition.