ICML 2025

When can in-context learning generalize out of task distribution?
Chase Goddard 1 Lindsay M. Smith 1 Vudtiwat Ngampruetikorn * 2 David J. Schwab * 3
Abstract
In-context learning (ICL) is a remarkable capabil-
ity 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 distribu-
tion for ICL to emerge and generalize out-of-
distribution. Previous work has focused on the
number of distinct tasks necessary in the pretrain-
ing dataset. Here, we use a different notion of task
diversity to study the emergence of ICL in trans-
formers 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 dis-
tribution, to a solution which generalizes out of
distribution to the entire task space. We also in-
vestigate the nature of the solutions learned by
the transformer on both sides of the transition,
and observe similar transitions in nonlinear re-
gression 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 regres-
sion problem influence the transition.
1. Introduction
The ability of transformers (Vaswani et al., 2017) to do
few-shot learning from examples seen in their context is a
striking phenomenon exhibited by modern large language
models called in-context learning (ICL) (Brown et al., 2020).
ICL has been extensively studied (Raventos et al., 2023; Lu
et al., 2024; Garg et al., 2023; Chan et al., 2022; Singh
et al., 2023) and enables models to solve certain new tasks
*Equal contribution
1Joseph Henry Laboratories of Physics,
Princeton University, Princeton, NJ, USA 2School of Physics, Uni-
versity of Sydney, Sydney, Australia 3Initiative for the Th...