When can in-context learning generalize out of task distribution?
Authors: Chase Goddard
Lindsay M. Smith
Vudtiwat Ngampruetikorn
David J. Schwab
Venue: ICML 2025 poster (2025)
Keywords: In-context learning
Generalization
OOD
out-of-distribution
Machine Learning
ICL
Primary area: theory->domain_adaptation_and_transfer_learning
TL;DR: We observe a transition between transformers that learn a specialized solution and a generalizing solution when trained to do in-context learning of linear functions with varying task diversity.

Abstract: 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 \emph{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.

When can in-context learning generalize out of task distribution?
Chase Goddard Lindsay M. Smith
Vudtiwat Ngampruetikorn *2 David J. Schwab
Abstract
without re-training. Of particular interest how the ability for transformers to perform ICL arises from pretraining and what the limits of ICL generalization are: What conditions must be met in order for ICL to emerge and generalize outside of the pretraining distribution?
In-context learning (ICL) is remarkable capabil- ity of pretrained transformers that allows models to generalize to unseen tasks after seeing only few examples. We investigate empirically the conditions necessary the pretraining distribu- tion for ICL 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 transition from specialized solution, which exhibits ICI. 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 Icarned by the transformer on both sides of the transition, and observe similar transitions in nonlinear re- gression problems. We construct 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.
Prior work (Raventos et al., 2023; Lu et al., 2024; He et al., 2024) has focused on understanding how the number of pretraining tasks affects the ability of the model gener- alize to unseen tasks (generated from the same distribution as the pretraining tasks). Here, we ask a related but dis- tinct question: If model is pretrained only on tasks from a subset of the full task space, what conditions are necessary for it to generalize to the rest of the space? We think of this question as asking about the out-of-distribution (OOD) performance of models trained do in-context learning. This question prompts to consider more general notion of task diversity that depends not only on task enumeration but also on how different or similar they are. Sampling a distribution with many similar tasks has the potential to induce the model towards a more specialized ICL solution that performs well only on novel tasks within its pretraining distribution. However, we observe that transformers trained to do ICL of linear functions undergo transition from specialized solution to one that generalizes over the full task space as we increase the degree of task diversity. This phe- nomenon of -ottt--distribttion task generalization sheds new light on in-context learning behavior
1. Introduction
1.1. Related work
The ability of transformers (Vaswani et al., 2017) to do few-shot learning from examples seen in their context is 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
Our investigation into the effects of task diversity on the emergence of ICL was motivated by the results of Raventos et al. (2023) on pretraining task diversity for linear regres- sion tasks in-context, where they find that the number of pre- training tasks impacts the emergence of ICL in-distribution.
Here, we investigate in-context learning of linear functions in the out-of -iit--iiiitiiuttionn setting, and investigate the domain generalization performance of transformers. Our experimen- tal setups are intentionally similar to those in Raventos et al. (2023) in order to aid comparison between in- and out-of- distribution ICL behavior. See e.g., Gulrajani & Lopez-Paz (2020); Arjovsky (2021); Liu et al. (2023); Yang et al. (2023)
'Equal contribution 'Joseph Henry Laboratories of Physics, Princeton University, Princeton. NJ, USA "School of Physies, Uni- versity of Sydney. Sydney, Australia "Iitiaative for the Theoretical Sciences, The Graduate Center, CUNY, New York, NY, USA. Correspondence to: Chase Goddard <cgoddard@ Segoddard@princctoo.dde pdi..
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).
Code available at http:///gthhb.ccm/ooooddd//////

---

When can in-context learning generalize out of task distribution?
for more background on domain generalization.
1.3. Contributions
Similarly, He et al. (2024) look at the setting of modular arithmetic ICL tasks and analyze how the number of pre- training tasks affects generalization. Lu et al. (2024) provide an exactly solvable model of linear attention transformers for linear regression ICL tasks and empirically show agree- ment of their theory with traditional transformers. Other work has analyzed how transformers implement an inter- nal wradient-descent algorithm to learn tasks in-context, shedding light on the mechanistic properties behind ICL (Akyürek et al., 2023; Ahn et al., 2023; Von Oswald et al., 2023). Fu et al. (2024a) argue that transformers must be implementing second-orde optimization methods instead in order to solve ICI. linear regression tasks. Chan et al. (2022); Edelman al. (2024); Singh al. (2023); Nguyen & Reddy (2024) investigate how properties of the data and optimiza- tion dynamics impact the emergence of ICL throughout the course of training.
Our core contributions are as follows:
We train transformers to exhibit ICL of linear func- tions with weight vectors drawn from a subset of the unit hypersphere. As the size of this subset in- creases, we observe a transition from specialized mod- els, which perform well only on the training portion of the hypersphere, to models that generalize out-of-task distribution to the entire hypersphere.
We investigate the nature of the solutions found by transformers, and find that specialized solutions out- perform optimal Bayesian solutions to the regression problem on small numbers of examples. In contrast, transformers that generalize to the entire hypersphere exhibit performance similar to optimal solutions.
1.2. Distribution Shift& & Generalization
We examine how two notions of task diversity (number of tasks and task similarity") interact, and construct a phase diagram that reveals three distinct regimes of ICL generalization.
In-context learning is a powerful capability of language models, and in order to build trust in models, we should better understand how well such capabilities extend to novel tasks beyond those in the training data. In addition to generalizing to tasks that interpolate between those seen in training, we'd like models that generalize to genuinely novel tasks. Whether this is possible depends on the nature of the distribution shift, and this question is in the purview of out-of- di-srriibbttonn generalization but at the level of tasks instead of samples.
We show that specializaioonngneraalzzzzooo transitions generalization
also occur in nonlinear regression problems, suggesting that the phenomenon may be a general feature of ICL in transformers.
2. Training setup and task distribution geometry
In this work, we characterize when transformers develop in- context learning behaviors that are robust to task distribution shift. Since ICL can be thought of as generalization to novel tasks, here ask when models succeed at domain general- ization with respect to tasks. Indeed, common assumption in the literature on out-of-distributin generalization (Gulra- jani & Lopez-Paz, 2020; Arjovsky, 2021; Liu et al., 2023; Yang et al., 2023) is the so-called "covariate shift" assump- tion where the conditional distribution of the label given the input is held fixed, but the distribution of inputs changes at test time. In the context of tasks, we consider sampling tasks during training from subset of the task space and ask about generalization to tasks outside of that subset. All tasks share the common property of being linear relationships, and thus we are studying domain generalization of tasks.
ICL of linear functions: We investigate the ability of transformers perform in-context learning of linear func- tions when pretraining tasks are drawn from distributions with varying levels of task diversity. Specifically, each task is a linear map in d dimensions, w E Rd, and we con- trol task diversity by sampling tasks from hyperspherical caps of varying half-angles. The transformer takes as in- put a sequence of up to n pairs {I1, ,i...., In where Y/ = = r; + G, with 2; N'O0, Ia) and c ~ oO).
Pretraining task distribution: We define a family of task distributions parameterized by o E [0, IB). We take sa to be section of the surface of the hypersphere in dimensions, i.e. sd-1( = {w E $ angle(w, v) < Φ). with v E Id a fixed vector. We then define the task distribution as a uniform distribution on this spherical cap, i.c., pv(w) Unif( sdd For details how to effectively sample from such a distribution, section A.5 in the appendix.
T]] (See Fig
"We note that another perspective is to train on as large and diverse pretraining dataset as possible nothing is out of dis- tribution, but there will always be novel tasks that have not been encountered before.
For linear problems, is natural to choose the inner product between tasks wf uy as a similarity measure.
2

---

When can in-context learning generalize out of task distribution?
Input sequence
C2
X1
Y1
X2
32
X3
Y3
Xn
15
Transformer
Wuain
f1
82
(2
93
Bn=To(Cp)
Wtest
Output sequence
(A)
Figure 1: Testing ICL generalization via task similarity. A: The transformer takes as input sequence of pairs {2;, Ui J and is trained to predict yı from context Cs (Ii, ,1,..., ...) The elements and are related linearly by task via = w" x4 + Eq. B: The training tasks Wtrain are drawn from hyperspherical cap with half-angle á (with ø = 180 corresponding to the entire hypersphere). The test tasks Whest are drawn from a hyperspherical band of width Ad starting an angle å away from the "pole" the sphere.
During pretraining, the transformer To op- timized to minimize the mean squared error (MSE) between context of data C {I1, V1 ... ................................ F& and the target yk- During pretraining, the tasks are drawn i.i.d. for each context from (w) ). We use Adam W (Loshchilov & Hutter, 2017) to optimize the MSE,
Pretraining:
layers (except in Fig 9), and attention heads. We use learned linear embedding map v; and yı the hidden dimension d,, 128. The target values yi are padded with d - zeroes. For further training details, see Appendix section A.1.
During pretraining, we train 12 models over pretraining distributions (w(w) for ø E [15", 180%) in 15" increments. We observe that repeated runs, with different initializations and trained data generated from different sampled tasks w Pa. yield consistent results (see Fig 11). For small we also see the model go through two stages of specializa- tion over the course of training (see Appendix A.2), the first of which is before this plateau.
1 Σ k= (To(Ck) - 3nm)
Lirain (0) = Euwrps
(1)
Test task distribution: We evaluate the performance of the transformer over family of task distributions parame- terized by 6, A& E [0, m] (See Fig IB). We define the hyper- spherical band starting at angle 8 with width AO to be the set (8, A6) = (w E sd-1 8 I angle(w, v) S 8 + Ad), with some fixed vector. The uniform over this set: ps (w) Unif( (5, A6)). For more how to sample from these bands, section A.5.2
3.1. Modes of out-of-distribution generalization
task distribution then
Before we investigate experimental results, it instructive to discuss the ways in which models may (or may not) generalize out-of-task-disttiiuuion
in the appendix.
1. Models may fail to meaningfully generalize out-of- dask-disrribution. (For example, when prompted to solve a task Whe not in the support of the pretraining distribution 5d ( models may simply pick the task Welose E (d- that closest to Whett but fail to generalize beyond this level.)
Evaluation:
We evaluate models by computing the MSE between the full context Cn, and the final target yn- During test time, we draw w i.idd for each context from P5,48:
Laua (
Er (To(C.) ynn?]
(2)
2. Models may generalize out-of-task-distrribuue but only achieve maximum performance when pretrained on the entire task space (ó 180")
3. Experimental Results
Unless stated otherwise, we study d = 10 dimensional regression with n = 50 examples in each context. We use a GPT-2 style transformer (Radford et al., 2019) with learned positional embeddings, hidden dimension of d, 128,
"During training, the loss generically "plateaus," staying constant value for many steps before beginning decrease. This phenomenon been observed before, (Fu et al., 2024b).
3

---

When can in-context learning generalize out of task distribution?
x² = 0
8 (
4
Optimal Bayes
175
x² 0
160
Test MSE
145
Test MSE
0.1
130
115
2
100
10
85
70
10-3
55
40
0
25
10
0
45 & (degrees)
90
135
180
Figure 3: Pretrained transformers outperform Bayes- generaliza-
optimal solutions tion. For d 175 (shaded grey region), plot the excess test loss for models with varying pretraining distributions. The dashed line shows the test error for the optimal in-task- distribution Bayesian solution (see section 3.2).
Test MSE
0.1
45
90
135
180
& (degrees)
5 and examine range of (see legend). For models with pretraining task distribution with ø x 120", observe good test performance only within the portion of the hypersphere covered by the pretraining distribution, and performance degrades outside of this range. However, for models trained over pa (w) with Φ Z 120", we observe low test error for all 8. In fact, the model performs similarly for test angles suggesting that transformers learn a general- purpose solution in this regime. This occurs despite the fact that these models were trained using only data restricted to subxet of the full task space. Notice also that even before the transition, models trained on a cap with & 2 45" exhibit nontrivial out-of distrid-dttiittiion task generalization (Fig 3).
Figure 2: Task distribution diversity induces a transition from specialized general-purpose ICL. A: Test error in Ad 5 bands (see Fig 1) for transformers pretrained to do in-context learning of linear functions with pretraining task distributions Po (w). For distributions with 6 S 120%, the transformer learns specialized solution that performs well on unseen tasks drawn from the (w))), but fails for tasks outside this distribution. However, for pretraining distribu- tions with Φ Z 120%, the transformer learns solution that performs well for all test angles 8. Here, the label noise is x² 0. B: With x² 0.25, still observe a transition from specialized to generic solution, but the transition point has moved to 2 135". The vertical axis measures the excess test error above the noise floor set by o².
In Fig 2B, we see that the transition sensitive to the level of label noise o². For noisy regression with x² + we see that the transition now occurs around 135%. but that the qualitative behavior of the transition is unchanged.
3. Models may generalize out-of-task-distribunn in a sharp way with increasing d: there may be some oe < 180" such that models achieve maximum performance for all Φe.
We view the last option as perhaps the most striking: it implies that models can achieve optimal performance with incomplete data. In the following, we show that transform- ers can generalize
Optimal Bayes OOD generalization How well do transformers generalize til-offtask-distrbbutiee Here, we compare their out-of-task distriiiiinntionn performance that of the optimal Bayesian estimator for given Po(w). Fol- lowing (Raventos et al., 2023), we derive an expression for the optimal estimator of yn. under the prior p& (w). Observe that in order to minimize Less((), (f), the optimal estimator is given by the posterior mean of yn conditioned on the full
in this fashion.
3.2. Sperializationngeeerallaae
transition
We show the results from evaluating these models on the test task distributions ps,4s(w) in Fig 2. We pick A6

---

When can in-context learning generalize out of task distribution?
min-norm least squares
Φ (9)
A
0 ())
180
180
165
165
Test MSE
0.1
135
150
Test MSE
135
150
120
120
105
105
10a2
90
90
75
75
60
60
45
45
10-3
15
15
0
10
20
30
0.5
1.5
2
# in-context examples
Test radius, R
B
Figure 4: Specialized ICL outperforms OLS for small context length. We evaluate the models in-task-distribution for varying context lengths, and plot the performance of the transformer (solid) and ordinary least squares (dashed) for the same data. For low context length, the specialized solution learned by models with Φ $ 90" outperforms OLS. For = 15", the specialized solution worse than OLS for large context length.
0.3
Test MSE
0.4
0.1
0.5
0.6
0.7
0.84
10-2
0.9
0
45
90
135
180
context Cn.
(degrees)
/ dw dyn yor(w, ya Cn.)
E yr [Cn]
Figure 5: Transformers trained to do ICL on the sphere generalize beyond it. A: The test error for tasks drawn uniformly from subsets of hypersphere of radius R, when a model is pretrained on tasks taken only from subsets of the unit hypersphere. When 6 2 45,, the model generalizes to tasks with R (shaded), despite being pretrained with R 1. B: Increasing task diversity drives generalization beyond the sphere: With sufficient task diversity (ø Z 45"), transformers generalize not only to OOD tasks on the sphere (Fig 2), but also to OOD tasks within it.
/ dw dyn Buplynlen w)p( Cn 1n-1)
/ dw wa 2.p((ulCn11,
w Fn
with
I du pl w) plun loxs w)
(3)
dw p( I p(y lxk, w)
We now consider the choice pww) = (w)) Although the integrals in Eqn are intractable in this case, is clear that because (ww lacks support outside of (d-)(o), v must always be contained in (ő)) The optimal Bayesian estimator (WoR) for this problem therefore fails to meaning- fully generalize out-of-task wew-oof-aa--isrriiiie distribution, as the best one can hope for to assign ú to that vector woB E 5 (() that closest to the target task.
Comparison to ordinary least squares We next investi- gate the solutions learned by the transformer on both sides of the transition. In Fig 4, we compare the performance of the transformer and ordinary least squares (OLS), solid and dashed curves, respectively. For short contexts, the special- ized solution which is learned for Φ x 120" outperforms OLS within the task distribution. For Φ > 120%, the perfor- mance the transformer is similar to OLS, except that the models' test error is not identically zero after d 10 exam- ples, unlike the least-squares solution. This sheds light on the nature of the specialization occurring: by fitting strong prior to the pretraining data, models with low & sacrifice out-of out-of-distribution performance, but this bias enables them to outperform the general-purposs solution (OLS) for low context length. See also Appendix Figs 16 17 for models' performance for varying context length other settings.
In Fig 3, the test loss for woB given by the dashed line. Notice that the ability of pretrained transformers to out- perform this optimal estimator tal----taa--ttidutttiii therefore a direct consequence of the failure of these trans- formers to fit the optimal Bayesian solution. It remains an interesting avenue for further work to ask what form the prior p(w) should take in order to give rise to the out-of- dist-iistinuution performance we see here.
is
5

---

When can in-context learning generalize out of task distribution?
OuC 364 ditiouu onneraaation
190
150
160
bution ganosalaziion
0.6
120
120
120-
0.8
Training
80-
80
60
40-
42
40-
022
40
20
20
8.
0.0
20
Number of taoks
Number of taoks
Number of tasks
(A)
(B)
(C)
Figure 6: Phase diagrams reveal three phases of generalization. A: Phase diagram for in-task- tast-distribution test loss
(6 09). B: Phase diagram for out-of-task-distiiuttion test loss (6 175"). Diagonal structure reveals interplay between two the measures of task diversity (Training o, number of tasks N). C: Combining the earlier two phase diagrams reveals three phases of in-context learning. In-weights learning, teal: the model fits the training data but fails to generalize, both in- and waw-oofftas-dissribbuue pretraining distribution, but fails to generalize out of distribution. buu-oofttux--iistibbtioo generalization, purple: The model generalizes well both in- and disttrrutttdtttituuttinnn In constructing the phase diagram, we set the threshold between high and low generalization losses 10
generalization, yellow: The model generalizes within the support of the
Beyond the unit hypersphere What happens to the gen- eralization ability of the model as the radius of the task distribution changes? We train several models on data gen- erated from tasks on the surface of the unit hypersphere, and evaluate them on tasks drawn from spheres of varying radii. Each model is trained on tasks from Pa (o(w) and evalu- ated on the equivalent distribution (with the same d) on hypersphere with a different radius. In Fig 5, we observe that for Z 45 the model is able to generalize perfectly to tasks with R 1 (shaded region), despite being trained only on tasks with R = 1. Increasing task diversity therefore drives the model to generalize not only to new portions of the hypersphere, but also beyond the hypersphere entirely.
and large o perform poorly in-distribution, suggesting that the density of tasks may be important. For more detailed analysis of the in-distribution performance of these models, see Appendix A.4.
In Fig 6B, we plot the resulting tut-of-task-distribue loss averaged over the sets, corresponding to test angles 8 175". In order to compare effectively with the normal- ized t--task-distributtin results, we normalize these losses by dividing by 4, the maximum loss possible when w, w are on the sphere We see that models with small ø perform poorly, and observe a diagonal boundary dividing models that generalize well and those that do not, suggesting in- terplay between these two forms of task diversity. In 6C, we summarize these results as a phase diagram, depicting three distinct phases with threshold of 0.01 as a cut-off between low loss (good generalization) and high loss (poor generalization):
3.3. Interplay between the two forms of task diversity
In order to examine the effect of both forms of task di- versity (number of tasks and task similarity), we train 4 sets of 120 models (480 models total) with task sim- ilarity @ and number of tasks N in the set: N) E {150, 30 ..... 180%} x {22 23 23!.....!!! In Fig 6A, we plot the resulting 4 sets: the loss for test angle 5 between 0 and 5 (these test angles are always in the training task distribution). In order to more effectively compare across different Φ, we normalize the loss for each ø by the maximum possible loss achievable from predictor and target w* in (o). Without such a normalization scheme, models with small & would trivially outperform those with larger o, since even an arbitrary choice of estimator ú E gd-1 cannot be too far from the target w*. We see that models with low N
1. Good generalization both in- and out-of-task- distribution (top right. purple).
2. Good in-task-distribution generalization, poor out-of- task-distribution generalization (bottom right, yellow).
loss averaged over the
3. Poor generalization both in- and out-of-task- distribution (left, teal); the model exhibits only in-weights learning (IWL).
Scaling Dimension and Depth In Fig 8, we investigate how the transition changes with changes in d, the dimen-
"To see this, set v = ww.
6

---

When can in-context learning generalize out of task distribution?
sionality of the regression problem. Unlike what is observed in Raventos et al. (2023), where d increases, greater pre- training task diversity N is necessary to induce ICL, here we observe that as d increases the transition along the & axis does not seem to change location. This suggests that the transition along this axis is not merely a result of high- dimensional geometry - in low dimensions (e.g., 3), 120"
Test Loss (Cross Entropy)
10
())
175
160
145
130
-15
100
the transition is still around
85
70
Similarly, we vary the number of layers in the transformer in Fig 9 to study how the transition from specialization to generalization changes with respect to model depth. We find that depth does not affect the angle (ø ~ 120%) of the transition for out-of out-of-task- teff-ttsk--srriii distribution test loss (6 = 1759). These models were trained using N = 2 pretraining tasks. We further show in Fig 19 in the appendix that the transi- tion point remains the same across depth even for different testing angles 8. These results further corroborate the phase diagram results in Fig 6 and show that model depth does not affect the transition point.
55
40
0.1
10
0
45
90
135
180
& (degrees)
Figure 7: Task distribution diversity induces a transi- tion from specialized to general-purpose ICL in clas- sification tasks: We now consider logistic regression. Vi H (aww z;)). We see a similar transition to the one observed in the case of regression, but the transition occurs at 6 R 135%.
3.4. Classification
To investigate the generality of -eniraattiinneneerrzzaatiooo transitions, we investigate the possibility of seeing a tran- sition in a classification task: logistic regression. We now take the mapping between x and y to be:
non-monotonicity can be observed where the red and blue lines, corresponding to ó near the poles, have lower test loss than the yellow lines, corresponding
90%.
In contrast, in Fig 10B, pick from two separate hyper- spheres: vec(W1) E 5 and W2 E S². This choice leads to a qualitatively similar transition to those we see in the linear case, with a transition around & ~ 135". Observation of a specialization generaaizannaaaiitiion transition beyond the lin- ear regime suggests that such transitions may be general phenomenon ICL.
= H2 (1 x:))
(4)
where H H.).) is the Heaviside step function with threshold and (.) is the logistic function. In Fig 7, we observe generalization
ppecialzzioon-gnneaalizzziooono transition for this task. The transition now occurs at & x 135 The fact that we observe a -onaliaatiionnnnnrllzzzioo transition in a classification setting hints that such transitions may be a more universal phenomenon.
4. Discussion and future work
We propose another "axis" to training task diversity, distinct from the task diversity measure in (Raventos et al., 2023) (i.e., the number of pretraining tasks). This new axis of task diversity, based on the size of the subset of the task space, accounts for the similarity present between tasks. Depending the level of task diversity present during pretraining, we have shown that transformers learn either specialized solution that fails to generalize out-of-task. distribution, or a generic solution with good performance across the entire task space.
3.5. Nonlinear regression
We now change the mapping between input and label for the regression be a nonlinear function of the weights. Specifically, we consider y/ w ReLUUW).::). with I, W2 [d and W1 E [ddxd We choose d = 3 that the model has 12 parameters. In Fig 10, we see that specialization generaaiaannna transitions still occur, and in- vestigate two ways of choosing the parameters. In Fig 10A, we pick the full 12-dimensional parameter vector [Vec,W1), w2) from the surface of S This choice induces bias towards ||w2]| < 1 for angles ó near the 'poles' (v (=1, O) This bias is relaxed, however, when 90" near the equator of the sphere. This leads to nonmonotonic behavior as ô changes- the tasks near the poles are more similar to each other than to those near the equator. In Fig 10A, we only show 8 < 90º for this reason. In Fig 18 in the appendix, we show results for all 8, and the
While we focus on the case of learning linear tasks, the phenomena of apecialization-geeeraiizzzz likely more general. In particular, such transitions appear also in the presence of label noise (Fig 2) and in nonlinear regression problems (Fig 10), To extend our analysis to more complex tasks, it would be interesting to investigate ICL performance in richer settings as more general notions of "task similarity" are varied. In the context of linear
transitions are

---

When can in-context learning generalize out of task distribution?
Test MSE
d
Test MSH
0.1
25
30
0.1
layers
20
10
10
15
10
3
103
10
10-3
45
90
135
180
0
45
90
135
180
& (degrees)
& (degrees)
Figure 8: Distributional diversity threshold is unaffected by task dimension. siloof-tak-ddisributtoo test loss training spherical cap polar angle @ at various task dimen- sions d (see legend). The test tasks are drawn from the spherical cap opposite to the "pole" of the training task distribution, = 175 (shaded region). We see that the threshold for out-of-taskddisttibttton generalization stays close 120 (dashed line), regardless of the task di- mension.
Figure 9: Distributional diversity threshold is unaffected by model depth. ing spherical cap polar angle & from transformer models with two (blue), three (red), and ten (black) layers. The test tasks are drawn from the spherical cap opposite to the training task distribution, 5 = 175 (shaded region). We see that the threshold for out-of-task-distributte generalization stays close Φ 120 (dashed line), regardless of the model depth.
test loss vs train-
data generating process to:
problems the sphere, the similarity between tasks wn and W2 is naturally measured by their inner product wf w2, but for more general problem settings it is less clear how to measure the similarity between tasks. Even within the linear setting, it may be interesting to explore task geometries beyond the sphere.
p(Cn, Vn) p' (Cn)p(unlCc)
(7)
a
p'(w) Iplas) 8(yn w==x))
(8)
Notice that this is a special type of covariate shift: we do not allow the full distribution p(C,)) vary rather, we impose that 1) the sample generation distribution p(xr,) re- mains unchanged; and that 2) still draw w independently from the samples z;. It may be the case that this additional structure allows the network to more easily deal with the distribution shift.
4.1. Types of distribution shift
In this work, we considered how transformers' ICL perfor- mance responds to a domain shift in the task distribution. To make this precise, we can formulate ICL as a supervised learning problem. Let's just consider the linear regression setting for now. Our supervised learning dataset is then:
This is not the only type of distribution shift one may care about, however, and prior work has investigated ICL under different distribution shifts. Ahuja & Lopez-Paz (2023) investigate instead domain shift of the x, 's (i.e. p(r,) + p' (x;;) and find that transformers fail to generalize when exposed to such shifts. We suggest that this may be because the authors do not expose the model to sufficient data diversity during pretraining. In Appendix A.6 we in- vestigate the effect of data diversity the linear regression setting and (noisy) transition in the behavior of the model's OOD generalization. Yadlowsky et al. (2024; 2023) instead consider ICI. under concept shift:: the mapping of data to labels ((x) changes between train and test. The authors here find that transformers largely fail to generalize to such distribution shifts, even when exposed to a diver- sity of tasks during pretraining. Hill et al. (2025) consider
D (Cn.yrh
(5)
),), with C₂ = (x1, "i,...,,,. and y/s = wT xt. The problem of task generalization we consider is then a special case of the domain generalization problem. To see this, notice that during training we see a given context On according to distribution C. p(w) I p(x:). The label function p(y)(Cr) is deterministic given the context:
p' C' 8(yn w xn)
(6)
When we ask about contexts with w outside the support of the training distribution, what we have done change the
8

---

When can in-context learning generalize out of task distribution?
Additionally, it would be interesting study the dynamies of how sal-offtaskkdistribttie ICL. emerges during training, shedding light on the learning processes & implicit biases of the optimization algorithms that enable models to generalize beyond their training task distribution.
A
0.1
8 (9)
85
70
Test MSE
10a
55
40
25
Our experiments open a new direction for understanding how general -urposse models are able to solve unseen tasks using only few examples in their context: we show em- pirically that transformers can learn to do ICL over much more of the task space than they are trained on. Understand- ing the generality of this behavior may help explain why language models are able to perform well on ICL tasks not present in their pretraining distribution. Although our exper- iments here are limited by their focus on relatively simple functions the ICL task, we believe investigations into specialization ppecializtioonggeeeralzzzoooo transitions for more complex tasks are a promising direction for future study. Building trust in LLMs is an important challenge with positive so- cietal impacts, and understanding the degree and nature of task generalization via ICL takes a step towards this goal.
10-3
10-4
8 (
175
160
Test MSE
145
0.1
130
115
100
2
70
55
10-3
40
25
0
45
90
135
180
Impact statement
& (degrees)
We present work on linear regression ICL tasks with the goal of understanding more broadly the factors that lead to the development of ICL behavior. While we focus simplified experimental settings compared to the more complex ICL phenomena seen in LLMs, the impact understanding the underpinnings of specialization and generalization in ICL will provide insights into the learning dynamics of L.LMs, building trust in Al and producing positive consequences for society.
Figure 10: Task distribution diversity induces a transi- tion from specialized to general-purpose ICL in nonlin- ear regression tasks. A: All parameters in the nonlinear model (a small one-hidden-layee network) are drawn from the same hypersphere. The transition occurs at é 60%. Although the curves for each 8 do not collapse onto each other after the transition, they remain within a band of size ~10-3, B: The parameters in the nonlinear model are drawn separately from different hypersphere for each layer in the model. The transition occurs Φ 135º.
A Acknowledgments
CG acknowledges travel support from the National Science Foundation and by DoD OUSD (R&E) under Co-operative Agreement PHY-2229929 (The NSF AI Institute for Ar- tificial and Natural Intelligence). LMS is supported by the National Science Foundation Graduate Research Fel- lowship Program under Grant No. DGE-2039656. VN acknowledges research funds from the University of Syd- ney. DIS was partially supported by a Simons Fellowship in the MMLS, a Sloan Fellowship, and the National Science Foundation, through the Center for the Physics of Biological Function (PHY-1734030),
ICL of linear functions where task vectors at test time are orthogonal to those seen during training, and again see that transformers fail to generalize to this distribution shift. We expect that this is because the natural task similarity mea- sure for linear tasks, the inner product, is zero for orthogonal tasks. The ability our models display in successfully adapt- ing to distribution shift between train and test time may therefore depend on the type of distribution shift considered, in addition to task/data diversity thresholds.
4.2. Future directions
References
A natural next step would be to develop analytic the- transitions
ory for analyzing sesiiaiiattioongneralzzzzan in transformers, similar to the analysis Lu et al. (2024) per- formed for the setting studied in Raventos et al. (2023). Such a model could provide further insight into the condi- tions necessary for out-of-task mask-iistribution ICL to emerge.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers leam to implement preconditioned gradient descent for in-context learning. In Thirry-sevnnth Conference on Neural Information Processing Sys-

---

When can in-context learning generalize out of task distribution?
Tems, 2023. URL https: http:///ppneeeeeeeeeeeee Forum?id=t.ziiiAN
Fu, J., Yang, T., Wang, Y., Lu, Y., and Zheng, N. Break- ing through the learning plateaus of in-context lear- ing in Transformer. June 2024b. URL https://
Ahuja, K. and Lopez-Paz, A closer look at in-context learning under distribution shifts. In Workshop on Effi- cient Systems for Foundation Models @ ICML2023, 2023. URL hitps:/ http:///eeeeeeeeeeeeeeeeeeee 2NbpM2KPOB.
Garg, S., Tsipras, D., Liang P., and Valiant, G. What Can Transformers Learn In-Comtext? A Case Study of Simple Function Classes, August 2023. URL http: [cs].
arXiv:2208.01066
Akyürek, E., Schuurmans, D., Andreas, J., Ma. T., and Zhou, D. What learning algorithm in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https: http///pppeeeeetetttt.// Zorum?idd=0004HHBBB
Guirajani, I. and Lopez-Paz, D. In search of lost domain gen- eralization, 2020. URL https: https:/:///aaxiii http:///wwwwwwa/aaaa/e/aaaaoee/e 2007.01434
He, T., Doshi, D., Das, A., and Gromov, A. Learning to grok: Emergence of in-context learning and skill compo- sition in modular arithmetic tasks. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: http//ppeeeeeeeeeeeteeeeeeeeeeee forum aerum?id=ahhaaa 9KRZdRk.
Arjovsky, M. Out of distribution generalization in ma- chine learning, 2021. URL https://arxiv.org) abs/ 2103.0266..
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D.,
Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Hill, J., Eyre, B., and Creager, E. Transformers don't in-context learn least squares regression, 2025. URL https://arxiv.or http://rrg//ab//50.00g/0
Askell, A., Agarwal, S., Herbert- verbert-Voss.. A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL https::/pproceednnggeexxxxpp (BsToofbbccc39671118858883222211 Paper.pdf.
Liu, J., Shen, Z., He, Y., Zhang, X., Xu, R., Yu, H. and Cui, P. Towards out-of-distribution generalization: A survey, 2023. URL 2108.13624.
Loshchilov, I. and Hutter, Fixing weight decay regular- ization in adam. CoRR, abs/1711.05101, 2017. URL http:
Lu, Y. M., Letey, M. I., Zavatone- Veth, J. A., Maiti, A., and Pehlevan, Asymptotic theory of in-context learning by linear attention, 2024. URL https: https::///xxxv.rg/// abs/2405.11151.
Chan, S. C. Y., Santoro, A., Lampinen, A. K., Wang, X., Singh, A., Richemond, P. H., McClelland, J., and Hill, F. Data distributional properties drive emer- gent in-context learning in transformers, 2022. URL https: / /at.aaaww.rr/aa./a////./......aa http://///bb/22.000055555
Nguyen, A. and Reddy, G. Differential learning kinetics govern the transition from memorization to generaliza- tion during in-context learning, 2024. URL https: //arxiv.org/aaaaeeiet httpabb/////.000..
Edelman, E., Tsilivis. N., Edelman, B. L.,, eran malach, and Goel, S. The evolution of statistical induction heads: In-context learning markov chains. In The Thirty-eighth Neural Information Processing Annual Conference / pttp//eeie.neeet/ Systems, 2024. URL https: forum? d=t.dddaaaaarrTrTeo...
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T.. Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019, URL
Fu. D., Chen, T.-Q., Jia, R., and Sharan, V. Transform- ers learn to achieve second-order convergence rates for in-context linear regression. In The Thihty-eighth An- nual Conference on Neural Information Processing Sys- tems, 2024a. URL https:/ : http//pppeeeeeeeeeeeeeeee forum? d=L.h 6cozebn.
paper.fflee//ppeer///////iii babe:22Beer152222255557722222222 Abstract.htm).
10

---

When can in-context learning generalize out of task distribution?
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models unsupervised multitask learners. Technical report, OpenAl, 2019.
cessing, 2023. URL https: http:::/pperreevee..en ppenrevieeneeee/
Forum?id=ivSJdhhe
Raventos, A., Paul, M., Chen, and Ganguli, S. Pretraining task diversity and the emergence of non- bayesian in-context learning for regression. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL /tttp:://penneeeeew..eett Forum?1d=BtA24a5x00
Singh, A. K., Chan, S. C., Moskovitz, T., Grant, E., Saxe, A. M., and Hill, F. The transient nature of emer- gent in-context learning in transformers. seventh Conference on Newral Information Processing Systems, 2023. URL https:://pppeeeeeeee /ttpss////pezeeee.net... forum? d=0ffG0000000000
Thirty-
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, u., and Polosukhin, Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. OODT.URLLtttss//rrooceigggeeerre paper.filees/aaeer//0//////e 3ffo2223$766e1111103cccc33333333 Paper.pdf.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J.. Mordvintsev, A., Zhmoginov, A., and Vladymy- rov, M. Transformers learn in-context by gradient de- scent. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learn- ing, volume 202 of Proceedings of Machine Learning Research, pp. 35151-35174, PMLR, 23-29 Jul 2023. URL htw.:cppps/ppooceiinn.g..pppsssp von-oswa 1d2.aahml.......m.mm
Yadlowsky, S., Doshi, L., and Tripuraneni, N. Pretrain- ing data mixtures enable narrow model selection ca- pabilities in transformer models, 2023. URL https:
Yadlowsky, S., Doshi, L., and Tripuraneni, N. Can trans- former models generalize via in-context learning be- yond pretraining data? In NeurIPS 2023 Workshop Distribution Shifts: New Frontiers with Foundation Models, 2024. URL
forum? d=ow 7o@Has3m.
Yang, I.,, Song. Y., Ren, X., Lyu, C., Wang, Y., Zhuo, J. Liu, L., Wang, J., Foster, J., and Zhang, Y. Out- of-distribution generalization in natural language pro- cessing: Past, present, and future. In The 2023 Confer- ence on Empirical Methods in Natural Language Pro-

---

When can in-context learning generalize out of task distribution?
A. Appendix supplemmnnal material
A.2. Two stages of specialization during training
A.I. Further experimental details
In Fig 12, we compare the training loss of transformer trained normally on data with ø 45 to transformer trained modified data. To modify the data, we zero out all components of z; that are perpendicular to the vector e defining the center of the hyperspherical training cap. We see that during early stages of training, the transformer trained on unmodified data performs similarly to the trans- former trained on modified data, suggesting that early in training, transformers trained to do linear regression only take into account the component of z; parallel to v. Later in training, the unmodified transformer learns to take into account other directions in the training data. This suggests that there are two distinct specialized solutions learned by transformers when e is small, the first of which is transient and disappears after training long enough.
All code was written in Python using the PyTorch library (Paszke et al., 2019). All models were trained for 58,000 steps using batch size of 128 and constant learning rate of 3 x 10-1 All models were converged at the end of training. All models were trained MIG GPU with 10GB of memory or an A100 with 40GB of memory, and took ~v 3hrs to train.
single GPU, either
4.0
$=15
3.5
$=30
$=120
4:455
#=135
4-60
4-150
Test Error (MSE)
2.5
$=
$=100
2.0
A.3. Defining the transition point
1.5
1.0
To more precisely define the transition point, we quantify the degree to which a model trained with given á performs similarly across test angles 6. The intuition for this defini- tion is that for model with ó above the transition point, we should expect similar performance across all ó due to the rotational symmetry in the linear regression problem. To quantify this, we measure the standard deviation of the model's performance across 6, and normalize this by the mean performance across 8:
0.5
0.0
20
60
100
Starting angle of band (degrees)
Figure 11: A second run of Fig 2, with a different initial- ization and sampling of w ~ Po-
10
Vars [Eu ao ((To(Cn) - y/ )
45 f 45, modified x,
NSR
Es E EEwpps (To 3
0.8
This definition resembles an inverse signal-to-noise ratio, or "NSR". To identify the transition point, we then look for sharp drop in the NSR. In Fig 13, we plot the NSR against & for the models in Fig 2A. We see a sharp drop in the NSR on a logarithmic scale, and identify which phase of learning we are in by thresholding the NSR. The fact that the drop is sharp means that our phase identification is not very sensitive to our choice of threshold.
MSE on train
0.6
0.2
0.0
3000
6000
7000
Step
A.4. Comparison with dMMSE estimator
Figure 12: Transformers undergo two stages of special- ization during training: A: For small Φ, the transformer rapidly (within the first epoch) learns a solution that only takes into account the component of x; in the direction of the vector forming the center of the hyperspherical cap. Blue: A transformer trained normally on training data with s 45 Orange: A transformer trained on data with the components of z; perpendicular to e zeroed out. The train- ing loss is smoothed with an exponential moving average for clarity of visualization.
Plugging the uniform distribution over a finite pretraining set W {wi, u22 ... wvv) } into Eqn 3 yields the discrete minimum mean-squared error (dMMSE) estimator (for the case of zero label noise):
WaMMSE = arg min
Σ
(w=xx-y) - vi)²
(10)
WEW
In Fig 14A, we investigate the in-distribution performance of pretrained transformers by comparing their performance to a dMMSE estimator. To do this, we interpolate between
12

---

When can in-context learning generalize out of task distribution?
2 tasks
1.0
0.12
2 tasks
2 tasks
08
0.10
0.08
835
0.06
0.4
0.
0.02
10-1
0.0
0.00
0.0
0.2
0.4
0.6
0.8
Training angle
Interpolation Step
Figure 13: Blue: NSR (Eqn 9) vs á for the models in Fig 2A. Green: Phase identification by thresholding the NSR, with threshold of 0.5.
(A)
180
160
0.10
140
Excess Inter wewpppoliiionnoommm Lons
120
0.05
tasks in along great circle, and test both the transformer and dMMSE along this path. For low number of tasks, transformers perform about well as dMMSE, but trans- formers outperform dMMSE for larger numbers of tasks.
Training
100
80
In Fig 14B, we plot a heatmap of the quantity:
20
-0.30
D = max (wAMMSE* -ya)² (To(ra) -va) va)" (11)
23
24
26
2
210
Number tasks
where Taz Vo is the regression problem generated by the interpolant weight vector wo at interpolation step The quantity D measures the performance of the transformer relative to the dMMSE estimator: when D < O, the trans- former outperforms dMMSE, the heatmap, we see that increasing number tasks drives the transformer to solu- tion that outperforms dMMSE for all c. order to compare across different values of scheme described in section 3.3. Notice that this normaliza- tion aligns the transition points for the transition from IWL to ICL across increasing number of tasks.
(B)
Figure 14: In-task- generalization. A: Solid: The loss for transformer with Φ 30 along a great circle connecting two weight vectors in the pretraining set. Dashed: The loss along the same great circle for the dMMSE estimator (Eqn 10) corresponding to the pretrain- ing set. B: Transformers outperform dMMSE for all é with sufficiently many pretraining tasks. The color shows the excess loss the transformer over the dMMSE estimator. When it is negative, the transformer outperforms dMMSE (see Eqn 11).
normalize D using the
A.5. Sampling from a portion of the hypersphere
The curse of dimensionality precludes sampling from portion of the hypersphere via rejection sampling. Instead, we consider the problem of sampling uniformly from the intersection of the sphere with cone d+1 dimensions: i.e. sampling from the sphere Sd(R) subject the restriction:
writing w (Z y) with R. and y E Rd, we see that
RZ > R2 - 12/2
(13)
We will show that, up this restriction, 7 R((U 1), with U ~ BBea d. z). To generate draw a vector uniformly from the sphere (d-1 7 Proposition A.1. Z = R(2U - 1), with U ~ Beta d 2 2
first draw Z, then
w| 2 2 2 angle(w, v) < of
R
12/2
(12)
where is our sampled vector and is fixed vector that defines the hypercone. The opening half-angle of the cone 0 is related to via , = 2R² cos 8).
Proof. Observe that p(Z) dZ is proportional to the (d- dimensional) surface area of the conical frustum with height dZ and radii / R 22, / R (7 + dz)². Expanding
WLOG, we take Ré; = (R 0), where Ō E d Then,
13

---

When can in-context learning generalize out of task distribution?
first order in dZ, the slant height of the frustum is given by:
10"
R2
540
e (dZ)²
(14)
R22 2
Tea une (Dres Eneray))
120
10"1
300
so that:
89
dz x ((VV((P VR2 7.2
d-1
(15)
(R² (22 (112)222)22 dZ
(16)
Making the change of variables 7 R(2u - 1), we obtain:
puu) du x (u - u²² 22)(d-22)/2 du
Training Cur Hall Anglle
(17)
x u² 11 (1 -u) du
(18)
Data diversity transition: When instead drawing x, from a hyperspherical cap. we see (noisy) ppecializatoon-enneaaazazzzzzi transition at á 105".
Figure 15:
which has the form of Beta( 3: 4) random variable.
1. SAMPLING Z
Transformer
OLS
As we have seen, Z follows scaled & shifted Beta( 4. 4) distribution, except for the constraint that RZ 2 R² /2/2. To implement this constraint, observe first that it is equivalent to:
Error (458))
0.8
0.6
U 2
(19)
4R
0.4
Since the distribution of U ~ Beta d () is symmetric 2-
02-
about 1/2, this condition equivalent to:
Examples in context
U
S 4R
(20)
Figure 16: In-distribution test error vs context length for the noisy case (a² }).
It follows that to sample from 7, we can perform the fol- lowing sequence of transformations:
T ~ Unif 0, F
(21) A.6. Data diversity thresholds
4H²
U F (
(22) In this section we investigate the level of diversity necessary in the data: i.e. when we instead draw 2; ~ 5 (Ø), with w drawn uniformly over the entire unit sphere. In Fig 15, we show the results. We still observe transition, suggesting that sufficient diversity is also required in the data in order generalize out-of-distrit oisiittiibutiin. We notice that the transition is much noisier than that observed in Fig 2. This difference in behavior may be due to the fact that the model directly observes in the context, whereas the task vector w is a latent variable.
Z R(UU 1)
(23)
where F(-) is the cdf of the ( 2 distribution.
A.5.2. MINIMUM ANGLES
straightforward implement an additional constraint corresponding to minimum distance/angle away from the vector If o' is this minimum distance, then we perform:
(r')²
A.7. Different Context Lengths
Unif
F
(24)
4R²
2R
In Figures 16 and 17, we investigate models trained with (25) various & across many context lengths different settings
= F-((T)
R(2U 1)
(26) to the setting Fig 4.
14

---

When can in-context learning generalize out of task distribution?
9-15
layers 2
--20
120
150
165
Test Erro (MSS)
180
10%
10-1
40
Examples context
10
Figure 17: Out-of-distribation (6 1755) test error
layers 3
context length.
6
10"
10
A.8. Supplemental plots for section 3.3
25
40
In Fig 19, we show the data from Fig 9, but for all ô, high- lighting that the transition point does not change with in- creasing model depth.
55
10-1
70
85
100
In Fig 18, reproduce Fig 10, but show many values of 8, highlighting the non-monotonic behavior that arises as result of the rotational symmetry of the sphere being violated by the nonlinearity.
115
10-2.
130
145
160
175
layers 10
10º
Test MSE
10-1
10-2
10-3
0.1
(*)
180
15
45
105 & (degrees)
75
135 165
10.2
Test Test MSE
90
10:3
Distributional diversity threshold is unaf- fected by model depth, even across test angles. Test loss vs training spherical cap polar angle á from transformer models with two (blue), three (red). and (black) layers. We show results for test tasks drawn from A6 5" bands from d 15t to 175 where lighter shades indicate smaller 8. We see that the threshold for -at--aaittibbuttiin gen- eralization stays close to á ~ 120", regardless of the model depth.
Figure 19:
10-4
45
90
135
180
(degrees)
Figure Same as Fig 10A, but shows all values of 5. Notice how the behavior of the test loss is non-monotonic 8 increases.
15


Reviews:
[{"review_id": "oe5voLwu4x", "reviewer": "Reviewer_STGS", "summary": "This paper studies empirically the task diversity to the generalizability, focusing on the transformer trained to learn a linear regression problem. It proposes the new axis of task diversity, namely the task similarity, independent of the unique task numbers seen during pretraining.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful comments and suggestions!\n\n&nbsp;\n> why [do we] mention the \"loss plateaus\"? \n\nWe mention the loss plateaus because they are characteristic of ICL behavior (Fu et al, ICML 2024; Reddy, ICLR 2023), confirming that our training is consistent with ICL phenomenology. \n\nWe will move this sentence to a footnote.\n\n&nbsp;\n> Are … [these] runs that \"escape\" the plateaus?\n\nYes, all our experiments escape the plateau.\n\n&nbsp;\n> How to [understand that], the \"effective training degree\" is larger than the set training degree?\n\nWith non-zero label noise *σ*, the effective training degree does not change. The training degree, *ɸ*, controls the portion of the hypersphere that the tasks, *w*, are drawn from. The tasks do not change when noise is added as the noise level only affects the label, *y*. Specifically, the noise, *ε* ~ 𝒩(0,*σ*²), is added to *w·x*.\n\n&nbsp;\n> Could you explain why the noise level sigma affects the transition point? \n\nWe posit that because of the noisy information that the model sees during training, it takes a larger proportion of the hypersphere to learn a general solution to the linear regression task, moving the transition slightly.\n\n&nbsp;\n> In Figure 2, when \\delta is small, the test MSE is better for small \\phi than large \\phi. … [Does the model learn a shortcut solution?] \n\nYes, this “shortcut” solution is the specialized solution the transformer develops when only training on small *ɸ*. The transformer has not seen enough of the hypersphere in training to generalize to large test angles (*δ*) it has never seen before.\n\n&nbsp;\n> Is there any way to normalize the test MSE …? \n\nWe have made a plot with normalized test MSE to disentangle any effects from unnormalized test MSE and to instead focus on the transition from specialization at small *ɸ* to generalization at large *ɸ*. Here, the test MSE starts at the same value and avoids this behavior of the unnormalized test MSE being higher at small test angles *δ* when trained on large *ɸ*. However, this does not change the transition to generalization, illustrated in Fig 2.\n\n&nbsp;\n> How do you plot figure 6C? What is the threshold chosen, and why? \n\nCurrently, the information is only in the figure caption—Thank you for reminding us to add this important detail to the main text. The threshold is set to 0.01 where losses below this threshold for the in-distribution (ID) and out-of-distribution (OOD) losses are characterized as out of task distribution generalization (purple), losses below the threshold for the ID losses but above this threshold for OOD losses are characterized as in task distribution generalization (yellow), and losses above this threshold for both ID and OOD losses are characterized as IWL (teal). \n\n&nbsp;\n> How do you validate the IWL regime?\n\nWe validate the IWL regime in appendix section A.3. We note that this is the same methodology used by Raventos et al. \n\n&nbsp;\n> Could you properly define the \"transition\" point? \n\nWe have produced new plots clarifying how we define the transition point that we add to the text. Here, we define the transition point as the reduction in the standard deviation of the mean over *δ*, leading to the “collapse” in the MSE. In Figs 7 and 8, we only plot the results for *δ* = 175°. Results for all values of *δ* for Fig 8 are shown in Fig 15 in the appendix, which emphasizes that the transition is the same regardless of the value of *δ* plotted, where we see the collapse of the MSE (similar to the results in Fig 2).\n\n&nbsp;\n> In Figure 7, what is the context length? What training algorithm did you use …\n\nWe use the same context length (*n*=50 examples) and training algorithm (AdamW) across all runs, regardless of task dimension, *d*.\n\n&nbsp;\n> Are all the runs have saturated performance …?\n\nThe loss converges for all runs. The loss increases across task dimensions because the regression problem becomes more difficult for the transformer to solve. However, we still see the same transition point in the training angle *ɸ* where the transformer begins to learn a general solution.\n\n&nbsp;\n> [Explain] why the transition happens independent of the problem dimension and model depth? Could it have to do with model embedding size?\n\nThe model embedding size remains the same for all experiments. We believe that the transition occurs independently of the linear regression task dimension (as in Fig 7) and the number of layers of the transformer (as in Figs 8 and 15) because of the structure in the data itself. The similarity in the data, independent of the dimension of the hypersphere, is enough for the model to learn a general solution that extends to the entire hypersphere after seeing tasks drawn from only *ɸ*=120°. While previous work, such as Raventos et al (2023), focuses on the number of tasks needed to generalize, our experiments suggest that this task similarity measure is also crucial in moderating the learning of a generalized solution.", "claims_and_evidence": "See experimental section for questions.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "NA", "experimental_designs_or_analyses": "1. Could you explain why you mention the \"loss plateaus\" behavior in section 3? Are the results presented with the runs that \"escape\" the plateaus? \n2. Could you explain why the noise level sigma affects the transition point? How to filter out the effect that because of the non-zero noise level, the \"effective training degree\" is larger than the set training degree?\n3. In Figure 2, when \\delta is small, the test MSE is better for small \\phi than large \\phi. I assume this has to do with the fact that the test MSE for a shortcut solution trained for small \\phi is small, instead of transformer trained on small \\phi learns a better solver. Is there any way to normalize the test MSE such that we will not observe the confusing behavior where an increase in test MSE when \\phi increases?  Maybe use the way you normalize the loss in section 3.2?\n4. How do you plot figure 6C? What is the threshold chosen, and why? How do you validate the IWL regime?\n5. Could you properly define the \"transition\" point? Is it the \\phi that across \\delta the test MSE collapses? Or is it the \\phi that the test MSE plateaus when \\delta=175? In Figure 2 I assume the definition goes by the first, and in figure 7 the definition goes by the second. Are these two definitions always collapses?\n6. In Figure 7, what is the context length? What training algorithm did you use to train especially the task with large d? Are all the runs have saturated performance (that escape the loss plateaus)?\n7. Out of the whole results presented in the paper, I am mostly interested in the results in Figure 7 & 8. Could there be any explanation on why the transition happens independent of the problem dimension and model depth? Could it have to do with model embedding size?", "relation_to_broader_scientific_literature": "This paper expand on the previous work on the effect of task diversity in generalizability of the transformer.", "essential_references_not_discussed": "Not that I aware of.", "other_strengths_and_weaknesses": "I appreciate the paper’s effort in introducing a new perspective on task diversity measurement and providing extensive experimental evidence to empirically analyze model behavior. The work is well-motivated and contributes valuable insights into understanding task generalization.\n\nThat said, the transition from an exemplar-based solver to a more generalizable solver is a well-studied phenomenon, so some of the initial findings may not be particularly surprising. For example, Figure 2 illustrates the existence of a transition point, while Figures 3 and 4 highlight that transformers learn to solve tasks using a different prior than OLS. Figure 5 further shows generalization beyond the unit sphere, where models trained with small \\phi struggle with OOD tasks within the unit sphere. Lastly, Figure 6 provides a helpful visualization of the phase transition in terms of task numbers and similarity. While these results are interesting and support the overall argument, they largely align with existing expectations.", "other_comments_or_suggestions": "- radii in line 272", "questions_for_authors": "Please see the experimental section. Thanks!", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "biw3meTdRC", "reviewer": "Reviewer_u8n5", "summary": "The authors study a new notion of task diversity--task similarity--and investigate the condition on task diversity for out-of-distribution generalization to emerge. The authors find that there is a transition from specialized models to generalizable models with increasing task diversity. They also show that this specialization-generalization transitions also occur in nonlinear regression problems.\n\n\n## update after rebuttal\n\nI do not have further concerns. I will maintain my score and suggest acceptance of the paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the helpful comments and suggestions!\n\n&nbsp;\n\n> Specifically, the authors have not discussed how they determine the transition point. For example, in Fig. 8 and 15, it appears from the plot that the threshold slightly increases with number of layers.\n\nThank you for the opportunity to clarify our presentation here. We will introduce a more formal definition of the transition point, which captures our intuition. We say a transition has occurred when the spread of the test error across test angles *δ* = {5°,10°,…} decreases below a specified threshold as the training cone angle *ɸ* increases. We quantify this spread by the relative standard deviation (ie, standard deviation normalized by mean). When plotted against *ɸ*, we generically see a sharp drop in this spread measure as *ɸ* increases.\n\nWe include plots of this measure and the thresholded phase boundary in the revised appendix. We will update our presentation, including Figs 8 and 15, to use this more precise definition of the transition.\n\n&nbsp;\n\n> Specifically, there is also rescaling symmetry in this nonlinear parameterization. Considering some measure in the function space may make more sense than in the parameter space.\n\nWhile we agree that a function space measure would indeed be a more principled approach for nonlinear functions, developing such a measure is beyond the scope of this work. For the nonlinear regression experiments, our main aim is to show that, at least qualitatively, specialization-generalization transitions exist for more complex settings than linear regression. We view the fact that our simple parameter-space measure (which ignores many features important to the nonlinear problem) works at all in this setting as a sign that specialization-generalization transitions may be somewhat robust to a choice of similarity measure.\n\nWe will edit the text to clarify this point.\n\n&nbsp;\n\n> Previous works have shown that the emergence of in-context learning on out-of-distribution linear regression tasks…\n\nWe would like to emphasize that although previous work has shown the ability of transformers to generalize to linear tasks beyond those they were pretrained on, our work means something different (and, to the best of our knowledge, novel) by “out-of-distribution.” In particular, previous work (eg, Raventos et al) has focused on the ability of transformers to generalize to new tasks *within* the support of the pretraining distribution. Instead, we focus on the ability of transformers to generalize to new tasks which are completely disjoint from those in the pretraining distribution.\n\nIn particular, we show in our revised manuscript how “task generalization” can be seen as a particularly structured type of *covariate shift* or *domain generalization* by defining the ICL problem as a special case of supervised learning. Our analysis suggests that it may be the type of OOD problem considered, not only the scale of the model/data, that allows for nontrivial OOD generalization in LLMs.\n\n&nbsp;\n\n> In Fig. 4, it seems that the models with ϕ > 120 will deviate away from OLS for in-context examples more than 10. I suggest the authors adding discussions for this observation.\n\nThank you for the suggestion to clarify our presentation here. Indeed, in our experiments, models with *ɸ* > 120 do not achieve zero test error after 10 examples. (The test error is just small). While the log scale on the plot exacerbates the difference between these curves, it is possible that this discrepancy is part of what allows transformers to generalize nontrivially out-of-distribution while an optimal Bayesian algorithm (with a prior matching the pretraining distribution) does not. \n\nWe will update our manuscript to comment on this phenomenon.\n\n&nbsp;\n\n> For Fig. 5B, it may be good to plot normalized MSE loss instead of the test MSE.\n\nThank you for the suggestion. If we normalize the MSE accordingly in Fig 5B, the transition remains qualitatively similar. We have added this plot to the appendix (but we leave the current plot in the main text to ease presentation).\n\n&nbsp;\n\n> In Fig 5A, why does training on larger ϕ  have a worse generalization for the cases where test radius is larger than 1?\n\nWe believe that this effect is due to initialization/optimization noise in the training of the models. We update the plot with a version that averages over more models to clarify this.", "claims_and_evidence": "I feel that some of the claims may be a little problematic. Specifically, the authors have not discussed how they determine the transition point. For example, in Fig. 8 and 15, it appears from the plot that the threshold slightly increases with number of layers.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense.", "theoretical_claims": "N. A.", "experimental_designs_or_analyses": "I have checked the soundness and validity of experimental designs and analyses. For nonlinear regression experiments, concatenating the weight vectors may cause some unexpected behaviors as the authors have also discussed in section 3.3. Specifically, there is also rescaling symmetry in this nonlinear parameterization. Considering some measure in the function space may make more sense than in the parameter space.", "relation_to_broader_scientific_literature": "Understanding the emergence of in-context learning is a very important question. Previous works have shown that the emergence of in-context learning on out-of-distribution linear regression tasks will only occur if the pretraining task is diverse enough. The diversity discussed in previous literatures are about the number of pretraining tasks. Here, the authors take another perspective and view task similarity as another dimension of task diversity. The results of the paper contribute to this broader scientific understanding of in-context learning.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "The paper studies the effect of task diversity on the out-of-distribution performance of in-context learning linear regression from a new perspective. The finding of the paper is interesting, and the paper is overall well written. \n\nHowever, there are still a few points that the authors can improve.\n\n1.\tFirst, the authors should more clearly write down how they define and find the transition point.\n\n2.\tIn Fig. 4, it seems that the models with $\\phi>120$ will deviate away from OLS for in-context examples more than 10. I suggest the authors adding discussions for this observation.", "other_comments_or_suggestions": "1.\tFor Fig. 5B, it may be good to plot normalized MSE loss instead of the test MSE. One can normalize the loss by the variance of y so that the effect of R on the variance of y will be factored out for a clearer point.", "questions_for_authors": "1.\tIn Fig 5A, why does training on larger $\\phi$ have a worse generalization for the cases where test radius is larger than 1?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "77h7NszhuS", "reviewer": "Reviewer_h9P4", "summary": "The paper explores the effect of pretraining task diversity, i.e., instead of the number of pretraining tasks, the paper considers the diversity of the fixed number of pretraining tasks.\nThe same number $N$ of tasks could be more diverse than the other $N$ tasks.\nSpecifically, the paper draws samples from a subset of a unit hypersphere and tests on the entire hypersphere.\nThe paper measures the new concept of diversity and shows that there is a transition of the trained Transformer from a specialized solution to generalization over the full tasks when the new diversity of tasks is increased.\nThere are multiple ablation studies provided to make the experiment complete.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and suggestions! We are pleased that you find our work interesting and well supported.\n\n&nbsp;\n> The paper could do better via diversifying the experimental setting on the tasks.\n\n> …consider more either proposing another task, possibly on classification rather than regression, or consider a task distribution other than the sphere.\n\nThank you for suggesting an intuitive way to diversify our experimental setting. We have performed experiments on a classification task with binary class labels {0,1}, generated from $y=\\mathrm{Heaviside}(w^Tx)$. The tasks $w$ are drawn uniformly from a hypersphere cap similar to our regression setting. We measure model performance on unseen tasks in and out of the training distribution. In this classification setting, we also observe a specialization-generalization transition similar to Fig 2. We update our manuscript to include the new results.\n\nThese new results will add to existing variations we considered, such as noisy regression (Fig 2B), increasing task dimensions (Fig 7), and nonlinear regression (Fig 9). A transition between specialized and general purpose solutions was evident in all settings investigated, including the new classification experiments described above. These results suggest that this transition, driven by distributional task diversity, is a universal feature of out-of-task-distribution generalization in exemplar in-context learning.\n\n&nbsp;\n> …consider to discuss the relationship to: [Yadlowsky et al (2023)]\n\nThank you for pointing out this interesting and relevant work. Yadlowsky et al (2023) asked whether a transformer can in-context learn a novel task and highlight that the coverage of the training task distribution (ie, *distributional task diversity* in our work) is an important factor in ICL ability.\n\nWhile related, our work differs from theirs in several important aspects:\n\n* Yadlowsky et al’s setting can best be described as “concept shift” of the task distribution, in the context of standard OOD literature. Our work fits within the framework of “covariate shift” of the task distribution.\n* We develop an experimental framework that allows us to clearly define task similarity, which enables quantitative investigations of its effects.\n* We demonstrate that distributional task diversity drives a sharp transition between specialized and general-purpose ICL.\n\nWe include discussion of Yadlowsky et al and clarify the points above in our revised manuscript.\n\n&nbsp;\n> …experiments are good but not surprising, and the claim is very reasonable, well supported…\n\nWe emphasize that the existence of a specialization-generalization transition in the out-of-task-distribution performance of transformers is nontrivial. One could expect at least three potential outcomes a priori:\n\n1. No out-of-distribution generalization. This is what one would observe from an optimal Bayesian solution with a prior matching the pretraining distribution.\n2. Nontrivial out-of-distribution generalization with performance gradually improving with increasing coverage of training task distribution.\n3. Nontrivial out-of-distribution generalization with abrupt performance improvement as the coverage of training task distribution exceeds a threshold (without encompassing the entire task space).\n\nIn our view, the third option, which we found empirically, is perhaps the most interesting. It implies a sharp change in ICL solution patterns where the model learns a new concept that generalizes beyond both the training tasks and the training task distribution.\n\nWe discuss these potential options in the introduction of our revised manuscript.\n\n&nbsp;\n> If the finding could shed some light on how such generalization happens, the paper will score 4.\n\nWhile we don’t have a full theory of the transition to OOD generalization, we suggest that our “covariate shift” of the task distribution setting is critical and helps explain why Yadlowsky et al did not find such a transition. To clarify this, we will add a new section to the discussion that casts the ICL problem as a special case of supervised learning. In this framework, we show how “task generalization” can be seen as a particularly structured type of *covariate shift* or *domain generalization*. In contrast, we show that Yadlowsky et al instead study a type of *concept shift*. We argue that this distinction in the mode of OOD generalization considered is what allows our models to generalize OOD while the models in Yadlowsky et al do not. Our analysis suggests that it may be the type of OOD problem considered, not only the scale of the model/data, that allows for nontrivial OOD generalization in LLMs.\n\n&nbsp;\n> … should consider arranging some experiments from the appendix into the main paper…\n\nWe appreciate this helpful suggestion and revise our manuscript accordingly.\n\n&nbsp;\n> There are multiple misuses in the citation…\n\nThank you for spotting these typos and formatting errors. We fix these in the revision.", "claims_and_evidence": "Claims: The paper claims there should be another angle of tasks diversity, i.e., the scale of pretraining task distribution. Via special experimental design, the paper illustrates the claim is true.\n\nSupported? Yes. The claim itself is reasonable and straightforward, and supported by the experiments. The paper could do better via diversifying the experimental setting on the tasks. The paper concentrates on the pretraining task's parameters on the sphere, the paper could consider more either proposing another task, possibly on classification rather than regression, or consider a task distribution other than the sphere.", "methods_and_evaluation_criteria": "Yes, the experimental setup makes sense.", "theoretical_claims": "No theory claim.", "experimental_designs_or_analyses": "Yes, I read the full paper so I checked the experimental design. The experiments are straightforward to me.", "relation_to_broader_scientific_literature": "The key contribution could be regarded as a complement to the definition of pretraining task diversity. There are indeed two angles: the number of training tasks and the scale of pretraining task distribution. The paper highlights the second part. Prior findings are concentrated on the first part.", "essential_references_not_discussed": "The author could consider to discuss the relationship to:\n\nCan Transformer Models Generalize Via In-Context Learning Beyond Pretraining Data?", "other_strengths_and_weaknesses": "Strength:\n\n(i) Fig 2 illustrates an interesting phenomenon: the pretraining task does not need to cover all the space to achieve good performance on all the space. There is a sweet spoint on the pretraining task distribution which depends on the level of label noise.\n\nWeakness:\n\n(i) The paper should consider arranging some experiments from the appendix into the main paper, especially when the main paper is not full 8 pages and the main paper refers to some figure in the appendix. (P3 col1 line 142 Fig 10 is in the appendix)\n\nThe paper is a bit beyond borderline for me, but cannot reach 4, so I score 3. The experiments are good but not surprising, and the claim is very reasonable, well supported, but not surprising. The contribution is decent, well, but not huge. The writing is also good. If the finding could shed some light on how such generalization happens, the paper will score 4.", "other_comments_or_suggestions": "There are multiple misuses in the citation. For instance, P1 col2 line038 \"the results of (...)\". I think the author used \\citep, but here \\citet should be used since there should be no \"()\".\n\nP3 col1 line 142 \"Fig 10\" should be \"Fig. 10\".\n\nP4 col1 line 201-204 \"set by $\\sigma^2$. with\". should be \"With\"?\n\nThe Optimal Bayes* solution depends on the label noise as shown in equation (5), where $p(w|C,y))$ depends on the label noise over $y$. Therefore, the author may connect the Optimal Bayes* solution to the case of $\\sigma\\rightarrow 0$.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "skP98heucc", "reviewer": "Reviewer_BMPu", "summary": "This paper examines the generalization properties of in-context learning (ICL) in transformers. Specifically, it explores the conditions necessary for ICL to emerge and extend beyond the pretraining distribution. To investigate this, the authors conduct a series of experiments across various tasks and summarize their key findings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for highlighting the clarity and importance of our work! We appreciate your comments and suggestions, which help us improve the paper.\n\n&nbsp;\n\n> The experiments are relatively simple…\n\nWe agree that our experimental setups are relatively simple. However, this simplicity is precisely what enables controlled experiments that are essential for developing an understanding of complex learning phenomena. We view our work as establishing an important foundation upon which future research can build more sophisticated models of out-of-task-distribution generalization in increasingly complex settings. This minimal model approach has proved crucial in advancing our understanding of nontrivial learning phenomena, such as generalization and learning dynamics of deep networks (eg, Saxe et al, ICLR 2014; Lampinen & Ganguli, ICLR 2019; Ji & Telgarsky, ICLR 2019; Arora et al, ICLR 2019), double descent and benign overfitting (eg, Bartlett et al, PNAS 2020; Wu & Xu, NeurIPS 2020; Hastie et al, Ann Stat 2022; Richards et al, AISTATS 2021; Mel & Ganguli, ICML 2021), and emergence of in-context learning (eg, Chan & Santoro et al, NeurIPS 2022; Singh & Chan et al, NeurIPS 2023; Raventós et al, NeurIPS 2023; Reddy, ICLR 2024; ​​Nguyen & Reddy, ICLR 2025).\n\n&nbsp;\n\n> …and do not explore the effects of increasing task complexity.\n\nOur work considered several settings with varying task complexity, including noisy regression (Fig 2B), increasing task dimensions (Fig 7), and nonlinear regression (Fig 9). We have also recently performed experiments on a classification task with $y = \\mathrm{Heaviside}(w^Tx)$ and $w$ drawn from a hyperspherical cap similar to our linear regression setting. We will include the results in our revised manuscript. In all settings considered, we observe a specialization-generalization transition, similar to those observed in Fig 2A. Our results hint at universal features of out-of-task-distribution generalization in exemplar in-context learning. \n\nWe will clarify this point in the discussion section in our revised manuscript.\n\n&nbsp;\n\n> …incorporating vision-language tasks post-pretraining could provide deeper insights into ICL.\n\nWe would love to be able to do this, but doing so would require developing a more robust notion of task similarity for vision-language tasks, which is outside the scope of our current work.\n\n&nbsp;\n\n> The experimental design appears highly similar to [Raventós et al], yet the paper does not adequately clarify this resemblance.\n\n> Can the authors explicitly discuss the similarities between this paper and [Raventós et al] in terms of experimental design?\n\nThank you for pointing out an opportunity to clarify the differences between our work and Raventós et al.\n\nWe intentionally chose our experimental setup to begin from Raventós et al. This similarity enables a seamless transition from this earlier work, which considered the effects of the number of training tasks on in-task-distribution generalization, through to our work on out-of-task-distribution generalization. This approach allows for a more direct comparison between previous results and ours, which helps readers build intuitions on the effects of task diversity on the emergence of general purpose in-context learning, without the need to translate between two different experimental settings.\n\nWe will revise our manuscript to better explain the rationale behind our experimental design, where we describe our setup. \n\n&nbsp;\n\n> The paper does not clearly answer the key questions posed in the “Contributions” section. While the experimental results section attempts to address them, the discussion lacks clarity and direct conclusions.\n\n> Can the authors provide clearer conclusions drawn from the experimental results regarding the key questions outlined in the “Contributions” section?\n\nWe will include a new Conclusions section to the discussion to further clarify how our experiments support our claims in the Contributions section. In particular, we will add the following texts: \n\n* We investigate the ability of transformers trained on simple tasks to generalize out-of-*task*-distribution: where “task generalization” is a particular type of covariate shift.\n* We identify a novel transition in the task generalization performance of transformers trained to do ICL, and show that pretraining tasks need not cover the full task space in order for models to generalize\n* Our experiments identify specialization-generalization transitions in several different ICL problems, suggesting that these transitions may be a universal phenomenon.", "claims_and_evidence": "See strength and weaknesses", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "yes", "other_strengths_and_weaknesses": "## Strengths\n- Paper is written clearly and easy to follow\n- Paper ask important questions regarding how transformers generalize, a crucial topic for advancing the field’s understanding.\n\n## Weaknesses\n- The experiments are relatively simple and do not explore the effects of increasing task complexity. For instance, incorporating vision-language tasks post-pretraining could provide deeper insights into ICL.\n- The experimental design appears highly similar to [1], yet the paper does not adequately clarify this resemblance.\n- The paper does not clearly answer the key questions posed in the “Contributions” section. While the experimental results section attempts to address them, the discussion lacks clarity and direct conclusions.\n\n1. Raventós, A., Paul, M., Chen, F., & Ganguli, S. (2023). Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in neural information processing systems, 36, 14228-14246.", "other_comments_or_suggestions": "none found", "questions_for_authors": "## Questions to ask\n- Can the authors explicitly discuss the similarities between this paper and [1] in terms of experimental design?\n- Can the authors provide clearer conclusions drawn from the experimental results regarding the key questions outlined in the “Contributions” section?\n\n1. Raventós, A., Paul, M., Chen, F., & Ganguli, S. (2023). Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in neural information processing systems, 36, 14228-14246.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this work, the authors propose a framework for probing the emergence of in-context learning in transformers.  They focus on the linear regression setting where each task (regression vector) has angle at most $\phi$ with another fixed vector.  By changing this angle from pretraining and test time, this allows for a nice, quantitative measure to probe whether a good algorithm is learned based on how diverse your pretraining data is.  The authors find that there are sharp transitions in the ability for models to generalize to arbitrary data based on the diversity of the pretraining data, as well as a number of other interesting results.  The reviewers were largely in favor of the work, and I'm inclined to agree, and recommend publication.
Decision: Accept (poster)