Can In-context Learning Really Generalize to Out-of-distribution Tasks?
Authors: Qixun Wang
Yifei Wang
Xianghua Ying
Yisen Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Large language models
In-context Learning
Out-of-distribution Generalization
Primary area: foundation or frontier models, including LLMs

Abstract: In this work, we explore the mechanism of in-context learning (ICL) on out-of-distribution (OOD) tasks that were not encountered during training. To achieve this, we conduct synthetic experiments where the objective is to learn OOD mathematical functions through ICL using a GPT-2 model. We reveal that Transformers may struggle to learn OOD task functions through ICL. Specifically, ICL performance resembles implementing a function within the pretraining hypothesis space and optimizing it with gradient descent based on the in-context examples. Additionally, we investigate ICL's well-documented ability to learn unseen abstract labels in context. We demonstrate that such ability only manifests in the scenarios without distributional shifts and, therefore, may not serve as evidence of new-task-learning ability. Furthermore, we assess ICL's performance on OOD tasks when the model is pretrained on multiple tasks. Both empirical and theoretical analyses demonstrate the existence of the \textbf{low-test-error preference} of ICL, where it tends to implement the pretraining function that yields low test error in the testing context. We validate this through numerical experiments. This new theoretical result, combined with our empirical findings, elucidates the mechanism of ICL in addressing OOD tasks.

Published as a conference paper at ICLR 2025
CAN IN-CONTEXT LEARNING REALLY GENERALIZE TO OUT-OF-DISTRIBUTION TASKS?
Qixun Wang' Yifei Wang'
Xianghua Ying!
Yisen Wang!
State Key Lab of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University
2 MIT CSAIL
3 Institute for Artificial Intelligence, Peking University
ABSTRACT
In this work, we investigate the mechanism of in-context learning (ICL) on out-of- distribution (OOD) tasks that were encountered during training. To this end, we conduct synthetic experiments using a GPT-2 model to learn OOD mathemat- ical functions through ICL. Our findings reveal that Transformers may struggle to learn OOD tasks via ICL. Specifically, ICL operates by selecting a function within the pretraining hypothesis space and optimizing via gradient descent us- ing in-context examples, rather than learning truly novel functions. Addition- ally, we examine ICL's well-documented ability to infer unseen abstract labels in context. We demonstrate that this ability only holds in scenarios without dis- tributional shifts, suggesting that does not constitute genuine new-task learn- ing. Furthermore, we analyze ICL's OOD performance when pretrained on mul- tiple tasks. Both empirical and theoretical results reveal a low-test-error pref- erence, where ICL tends select the pretraining function that minimizes test error rather than adapting to entirely new tasks, We validate this phenomenon through numerical experiments. Our theoretical insights, combined with em- pirical findings, provide a deeper understanding of ICL's limitations and its un- derlying mechanism when tackling OOD tasks. Code is available at https: //qithub.
INTRODUCTION
Pretrained large language models (LLMs) can perform in-context learning (ICL) (Brown 2020 where providing few examples of input-output pairs and query example improves the model's ability to generate the desired output, compared to zero-shot predictions. Understanding ICL's abil- ity to learn out-of-distribution (OOD) input-output relationships, which are unseen during training. has recently gained significant attention.
Recent studies have demonstrated that ICL can tackle seemingly new tasks. For instance, Garg et al. (2022); Raventós et al. (2023); Zhang et al. (2023a); Akkürek et al. (2023) found that ICL can learn new linear regression weights alter pretraining on large set of weight vectors. Moreover, Pan
2023); Kossen et al. (2024): Vacareanu et al. (2024) revealed that real-world LLMs like Llama-2
2023) are capable of solving artificially constructed lasks likely unseen in their pretraining data, such classification task with abstract labels.
Touvron et al. 2023) and GPT-4 (Achiam
However, another line of research (Yadlowsky et al., 2023 Ahuja & Lopez-Pat 2023) has raised
contrasting view, showing that ICI. struggles to generalize to OOD tasks where there are distri-
butional shifts either the input distribution P(X) or the input-label mapping (X). These findings raise several important questions:
Can ICL really learn new input-outpul mappings from the context? What underlying mechanism of ICL determines its performance on OOD tasks?
Corrrsponning Authors.

---

Published as a conference paper at ICLR 2025
This work aims to consolidate previous findings by addressing these questions. First, empirically show that when pretrained on a specific function class, the OOD performance of ICL approaches that of a model from the same function class optimized via gradient descent. This suggests that ICL tends to implement the function class encountered during pretraining. Additionally, we inves- tigate the widely observed phenomenon where ICL. can classify abstract unseen labels. We find that the ability to solve such task is a form of retrieval capability, which disappears when faced with OOD classification rules, indicating that success in these tasks does not necessarily reflect an inherent ability to learn new tasks. Finally, we explore scenarios in which the model is pretrained on multiple tasks, empirically uncovering the algorithm selection mechanism for OOD tasks. We also theoretically reveal the algorithm-selection mechanism for ICL. Our main contributions are:
1. We empirically show that ICL tends to implement the pretraining function based on the downstream task context, highlighting its limitation in solving OOD tasks (Section 3.1).
2. We further investigate ICL's ability to classify unseen abstract labels. We find that such tasks can be solved by retrieving similar examples from the context. This retrieval ability can arise from training on tasks with more diverse abstract labels (Section [4.1) and only emerges when the test function is in distribution (Section 4.2) We further validate our findings by showing that pretrained Llama-3-SB (Dubey et al. 2024) and Llama-2-7B fail to learn OOD functions through ICL in a synthetic vector classification task (Section 4.3
3. We explore the ICL's behavior when trained on multiple tasks, and observe that the algo- rithm selection mechanism broadly occurs in OOD scenarios. We theoretically prove the low-test-error preference of ICL prediction, i.e., the ICL prediction prefers implement the pretraining function with low test error (Section 5.1 numerical experiments (.2))
We also validate our theory with
2 EXISTING THEORETICAL PREDICTIONS OF ICL
Previous literature has provided some theoretical insights into the behavior of ICL. Here we briefly review some of the representative findings. 1) ICL makes Bayesian predictions. Xie et al. (2022); Wies et al. 20024); Zhang et al. (2023b) theoretically demonstrated that ICL behaves like a Bayesian- optimal predictor, i.e., it will infer task concept based on the given test context, and then predict using the inferred task and the input prompt. However, these Bayesian frameworks do not depict the concrete process of how the task concept inferred, especially for OOD scenarios. 2) ICL implements gradient descent (GD). Von Oswald et al.] (2023); Zheng et al. (2024) construct spe- cific Transformer weights on which the ICL. prediction is equivalent to a linear regression predictor or alignment objective (Wang et al. 2024a) optimized by gradient descent. 3) ICL implements
algorithm selection. Bai et al. (2023); ; Wang et al. (024b) demonstrate the existence of Transform- ers that can realize algorithm selection between linear classification and regression by constructing specific Transformer weights. 4) ICL performs retrieval. [i et al. (2024) proves that a trained non- linear Transformer will concentrate its attention of the query on the in-context examples possessing similar features to that of the query.
These theoretical findings may seem disparate, they describe different aspects of ICL under vary- ing assumptions and settings. Furthermore, many of them rely on oversimplified model architectures or deliberately constructed model weights to reach their conclusions, limiting their practical appli- cability. In the following sections, we aim to provide a unified perspective on ICL by conducting experiments with deep nonlinear Transformers on both synthetic and real-world OOD tasks.
3 EXPLORING THE PERFORMANCE OF ICL ON OOD TASKS
3.1 GPT-2 IMPLEMENTS THE FUNCTIONS CLASS SEEN DURING ICL PRETRAINING
Evaluating GPT-2 on unseen mathematical function classes. To investigate the ICI. performance on new tasks that are unseen during training, following Garg et al. (2022), we train a GPT-2 (Radford ctt al., 2019) from scratch on some simple functions and evaluate on functions different from the
training ones. Denote the Transformer model parameterized by 0 as Mg. The pretraining objective is
ming + ET Ey~> (S.0011))) - ((xx-))]13) where S,
(a y1 4 4D y2 @ (Br,0yb
2

---

Published as a conference paper at ICLR 2025
Rdx2 is the context, denotes concatenation. a, I are sampled from standard Gaussian distribution (0, 1) with dimension d 20. Let yi = f(x;) E R represent the labels, with F denoting the hypothesis class which f belongs. We train a GPT-2 model function classes F: linear regression (L.R), quadratic regression (QR), and a 2-layer ReLU neural network (ReLU NN, detailed descriptions in Appendix I.... We then evaluate its ICL performance on all three function classes. For comparison, we also train models of the corresponding F with gradient descent (GD) using the test in-context examples (details in Appendix C.I
of the three
Observations. We plot the test eтror on the three tasks in Figure | and observe that: 1) when evaluated on the same task FF pretraining, ICL can reach near-zero test error, which consistent with the findings in Garg et al. (2022). 2) when evaluated on new task, ICL performs similarly to the corresponding model the pretraining function class optimized by GD given enough in-context examples. This indicates that the ICI. prediction implements the training function classes. 3) The models trained on linear and quadratic regression exhibit a double descent error curve (Nakkiran) 2019). characterized by a high error when given exact d examples and evaluated on new task. This further demonstrates that ICL implements in-distribution (ID) predictions, as the double descent curve is distinctive phenomenon unique to linear regression models.
Dow rel
[Du, lnsar regraadoor E-yyer tnew my mygessine,
3.0
regravioon [CL nagroaa 24aye AN.
1.3
2.5
1.0
0.5
cantect ciamplss
In-content examoles
In-contcnt cxamples
(a) Evaluated on LR
(b) Evaluated on QR
(c) Evaluated on ReLU NN
Figure 1: The ICL test enror of Transformers trained different function classes (solid lines) and the performance of models from the corresponding pretraining functions classes trained by GD using the test in-context examples (dashed lines). Y-axis: test square error. X-axis: context length. We observe that as the test context length increases, the ICL performance pretrained on a particular function class closely approaches that of the model from this function class trained by GD.
3.2 REAL-WORLD LLMS TEND TO MAKE IN-DISTRIBUTION PREDICTIONS DURING ICL
In this section, we will demonstrate how the tendency of ICL to perform ID predictions manifests in real-world L.Mss. We designed task involving label prediction with reversed letters (e.g., in sentimental classification,
We found that in this task, pretrained Ilama- 3-8B (Dubey al.) 2024) prefers to output the inversion of the query word rather than predict the reversed correct label, as shown in Figure 2 Although both reversal tasks are rare, directly outputting the reversed version of a word is more common than first reasoning and then outputting the reversed prediction. This result suggests that LLMs, when performing ICL, are more inclined to implement ID tasks. For more details, refer Appendix C22] Inspired by |aaentós al. (2023). We also explore whether increasing the diversity of training tasks, while keeping them ID, can activate the OOD generalization ability. The results Appendix B.4 also suggest negative conclusion.
He camy
nottreeree de
090
Care hegh
Cenee lergh
Seeedlogh
tereclags
(a) Antonym
(b) Country-capital
(c) English-French
(d) English-German
Figure 2: The top-1 accuracy of predicting the reversed query word (bluc) and predicting the versed target label word (orange). The accuracy of predicting the reversed query word is higher than outputting the reversed target, indicating ICL makes ID predictions.
3

---

Published as a conference paper at ICLR 2025
Summary of the Empirical Results & Connections with the Existing Theories (I)
Given an OOD context, ICL. finds a near-optimal solution within its pretraining task space. Particularly, when learning OOD mathematical functions, ICL behaves as predictor of its pretraining function class optimized by GD using the in-context examples. This validates and extends existing results by Zhang al. (2023a) which theoretically shows that linear attention models trained on linear regression data still implement linear regression given arbitrary downstream context (see Appendix D
4 LEARNING ABSTRACT LABELS MAY NOT BE A REAL OOD CAPABILITY
4.1 CLASSIFYING ABSTRACT LABELS IS A PREDICT-THEN REHEI-EERRIEE PROCESS THAT CAN EMERGE FROM TRAINING
Kossen et al. 2024) have shown that LLMs can successfully perform classification tasks in which the labels are "abstract symbols" with no semantic meaning (e.g.. in sentimental classification, "positive" and "negative" are replaced with "foo" and "bar", respectively). These tasks are likely not encountered during pretraining. Pan (203) refer to this ability of ICL. to perform such classification as "task learning" (TL). In this section, we explore whether this really reflects a new-task-learning capability of ICL, or if it is something else.
Recent works Pan 2023
Training GPT-2 to perform a retrieval task through ICL. The classification of abstract labels can be approached by retrieving an example with semanties similar to the query and then outputting the label of that example. Therefore, the retrieval ability is crucial prerequisite for performing abstract- label classification. We design a retrieval task to investigate whether ICL's retrieval capability can emerge from training. Specifically, we generate predefined word embedding E R.N xd and randomly sample æ; E IRd from the first 5 rows of E. Suppose vector æ; is the 1z,-th row of E, i.., I, Ei.,- To generate the labels y. E Rdd, we first map the index Iæ: to Iy E [N] using the mapping rule In, = Ia. + where E N is randomly sampled. Next, we set y: Eig," All in-context examples in sequence share the same mapping rule defined by s. To succeed in this task, the model must retrieve the same token as the query example from the context and output its subsequent token. All models are trained with 200,000064 sequences, where 200,000 is the number of training steps and 64 is the batch size. We train three models with three different ranges of 8: 8 (50, 150), 8 (50, 250), and 8 ~ (50, 450) and evaluate on s (50, 150). (11,, 20), and (500, 600), where U denotes the uniform distribution.
Eval C200. 155)
Eval
3.0
taske yy [[500, 505)
c150,
2.5
2.5
e,000 251)
2.0
2.0
650 453)
2.5
3.5
2.0
0.
0.5
0.0
Invconkent
D-corfiast eccemples
in-comiaed ecemples
(a) Eval s ~ (500, 150)
(b) Eval ~ U(10, 20)
(c) Eval ~ U5500, 600)
Figure 3: The ICL test error of Transformers trained on the retrieval task with different numbers of label tokens. "Eval" denotes "evaluated on".. The ability to retrieve OOD labels emerges from increasing the number of training mappings (larger range of s).
Observations. In Figure 3 all three models perform well when labels are ID (a). When the labels are OOD, the ICL performance improves with the number of label vectors (random mappings) en- countered during training (b, c). This demonstrates that the ability to retrieve arbitrary labels from the context can emerge from training on diverse retrieval tasks. These findings may also offer new insights into how real-world LLMs develop in-context retrieval capabilities: when autoregressive pretraining includes numerous instances requiring the model to retrieve tokens from previous con- texts, such abilities can emerge. We further validate this finding by observing the emergence of induction heads in Appendix B.
4

---

Published as a conference paper at ICLR 2025
Training GPT-2 to perform predict-then-retriee task through ICL. To further explore the emergence of the ability to classify abstract labels, we design a predict-then-retriew task that em- ulates the natural language classification with abstract labels. In this task. y: Eiz,, where I. = f0. * (w T æx)) + 8, with E being the predefined word embedding and & E N+ shared in ICL sequence. Here, Z;. w NN(0, 1) R In this task, estimating and calculating simulates predicting the original label ("positive" and "negative")) by in-context learning the natural language task, while retrieving y; from æ; such that floor(0.4 (w _ xx)) floor(0.4 (0.4+( (w Equery)) resembles identifying the abstract labels ("foo" and "bar" where 2 Equery is the query example. Again, we train three models on different ranges of mappings: $ ~ (1000, 200). s ~ (1100, 1000), and (1000, 2000), and evaluate on (1100, 200), s (500,, 600), and 4(3000, 3100).
æ;
Observations. In Figure 4 the ability to predict and retrieve unseen labels also improves as the number of labels encountered during training increases. This suggests that as long the LLM has been exposed to sufficiently many similar tasks during training, it can effectively classify arbitrary OOD labels retrievable from context through ICL.
Eval task: a E [100.200]
Evall tast: E (500,600)
Dun teake (3000,3103)
31100 2000
10
01100 1000)
11100.
0.0
02
0.2
0.0
0.0
0.0
40
60
80
100
40
60
100
in-context examples
Dn D-confent esamples
in-cortiat examples
(a) Eval 8 (4100, 200)
(b) Eval (500,, 600)
(c) Eval U(3000, 3100)
Figure 4: The ICL test error of Transformers trained and tested on the linear regression + retrieval task with different numbers of label tokens. "Eval" denotes "evaluated on". The performance unseen labels increases with the number of labels seen during training. Note that in this experiment, the underlying test function class is still ID.
4.2 ABSTRACT LABEL CLASSIFCATION CAN ONLY BE ACHIEVED ID TASKS
A predict-then-retrice task with OOD testing functions. Given the above observations, one might question whether, once the target labels appear in the context, ICL can generalize beyond the training function class by retrieving the target label from the context. To investigate this, duct the same predict-then-retries task as in Figure 4 but replace the test functions with quadratic regression while preserving linear regression as the pretraining task.
Observations. The results in FigureSjshow that the generalization does not improve with training on more ID functions, even when the target label vectors appear in the context. Combining observations from Figure 4 we conclude that ICL can only solve classification with unseen labels over ID test function classes. This finding highlights limitation in improving an LLM's performance through in-context examples. While providing examples with shared labels may seem helpful, this approach may fail if the underlying prediction rule is too OOD for the LLM learn.
Summary of the Empirical Results & Connections with the Existing Theories (II)
To handle classification with abstract labels, the model infers an input-label mapping to implicitly establish a similarity metric. then retrieves in-context examples similar to the query to deduce the OOD labels using this metric. However, this process succeeds only when the underlying function class is ID, thus it does not represent a true OOD general- ization capability. This observation aligns with Bayesian frameworks for ICI.-the implicit similarity metric here corresponds to the task concept inferred by the model. We leave an intuitive Bayesian interpretation of the findings in Section 4.1 and 4.2 to Appendix(A.A A.4
In our experimental setup, given a sufficiently long context (Rs 50), the label of the query is highly likely to appear in the context, as the number of the possible classes far less than the number of in-context examples.
5

---

Published as a conference paper at ICLR 2025
tadk: 11103,
Eusl U1520
Eval tapis 10300000000000
10
15
3.0
20000
0.5
2.0
20
2.0
0.5
655
90
100
66
22
102
Invonneeat eccmples les
n-cooeed not ecamples
Ih In-conizat csampies
(a) Eval s (100 200)
(b) ~ u5500, 600)
(c) Eval s~ Z((000, 3100)
Figure 5: The ICL test error of Transformers evaluated on a quadratic regression + retrieval task. Different colors denote models trained on the linear regression retrieval task with different num- bers of label tokens. "Eval" denotes "evaluation". The model trained on s ~ U(100, 2000) doesn't generalize better than the other two models.
4.3 REAL-WORLD LLMs MAY NOT NECESSARILY IN-CONTEXT LEARN NEW TASKS
Evaluating Llama-3 on an OOD synthetic vector classification task. Now we assess whether real-world LLMs can tackle OOD tasks through ICL. We design synthetic vector classification task for a pretrained Llama-3-8B. Specifically, we randomly sample 2; E _d from the word embedding of Llama-3-8B (denoted as Eilama) and generate random linear mappings W Rdxx as task functions (where C = 10). The label vectors are created by mapping 2; to one of the ten label vectors in Estaa using W. Experimental details are in Appendix C.4 To complete this task, the model must learn W in context, which is unlikely to have been seen during the pretraining of Llama.
For comparison, we also evaluate the ICL performance of Llama-3-8B on a retrieval version of this task. Concretely, we first randomly sample different vectors from Etoaa as æ; and compute y, in the same way as the above classification task to get repeat S for 20 times to construct the input sequence S @ S + @ S]. where @ denotes concatenation operation. The goal is to predict the next token given a prefix of S'. To succeed in
[ 221,311. 2/, DC, Yc] (C = 10). Then we
this task, the model has to retrieve the same token as the query token from the context and output its subsequent token y/-
Observations. The results of these two tasks pre- sented in Fi We observe that the ICL performance on the synthetic classification task close random guessing (10% accuracy), while performance on the re- trieval task is significantly better (similar results also hold for Llama-2-7B in Appendix B.3). This suggests that pre-
dhessiiattoon paess casuiicater
accuracy
0.6
0.4
trained real-world LLMs may also struggle to learn new input-output mappings from context; instead, ICI. ap- pears to be more adept at retrieval tasks. To show that the failure in the synthetic vector classification task is mainly due to its OOD nature instead of some other factors that make it difficult to learn, train GPT-2 to perform the same task in Appendix [B.2] and find that the task can be well addressed after training.
0.2
0.0
200
400
600
800
1000
context examples
Figure 6: The ICL accuracy of Llama- 3-8B on the synthetic tasks.
5 THE ALGORITHM SELECTION MECHANISM EXISTS BROADLY WHEN EVALUATED ON OOD TASKS
Real-world LLMs are pretrained on a huge corpus that could contain massive tasks. Bai et al. 2023); Yadlowsky et al. (2023) have empirically found that the ICL performance of Transformers trained on multiple tasks approaches the optimal pretraining function when evaluated on one of the training tasks. In this section, will show that this algorithm-selection phenomenon of ICL broadly persists when evaluated on OOD tasks, regardless of the test distribution, and provide a comprehensive theoretical characterization of the algorithm-selectonn mechanism.

---

Published as a conference paper at ICLR 2025
The Model pretrained on single task the model pretrained on multiple tasks. In Figure we compare the performance of GPT-2 models trained on a single lisk-innar regression (LR), quadratic regression (QR), 2-layer ReLU neural network (ReLU NN) regression -gganst the model trained on all three tasks when encountering four kinds of OOD tasks. We also plot the error of a 2-layer ReLU NN trained by GD (dashed blue line).
Ove NOT.NN
Easl CENCHamTiie
Deal repression
Uvall tarc oignaic
150
200
(a) Evaluated on SRR (b) Evaluated on CR (c) Evaluated on LQR (d) Evaluated on Sigmoid NN
Figure 7: The ICL performance of models trained on the individual task: linear regression (LR), quadratic regression (QR). 2-layer ReLU network (ReLU NN) regression, and the model trained on the mixture of the three tasks (LR+OR+ReLU NN). The evaluation functions are (a) square root regression (SRR), (b) cubic regression (CR), (c) linear+quadratic regression (LQR), and (d) 2-layer Sigmoid network (Sigmoid NN). The details of these evaluation tasks are in Appendix C.I). The performance of the model trained on the mixed tasks is close to that of the model trained on the single task that performs the best the evaluation task.
Observations. 1) The ICL performance of the model trained on mixed tasks (LR+QR+ReLU NN) close the performance of the model trained on a single task with the lowest test error on the evaluation task. This suggests that ICL can automatically select the best pretraining functions ac- cording to the downstream context. 2) The ICL performance of training on the ReLU NN function class aligns well with the ReLU NN model trained by GD, demonstrating that our findings in Section 3.1 still hold when the Transformer is trained on a mixture of multiple tasks.
5.1 THEORETICALLY REVEALING THE MECHANISM OF ALGORITHM SELECTION
In this section, we will provide theoretical insights into the underlying mechanism of the algorithm selection of ICL. We find there simultaneously exist two parallel mechanisms: the Low-test-error preference and the Similar-input-diittibttton preference.
A mixed Gaussian pretraining dataset of multiple tasks. We adopt the theoretical framework by Lin & Lee (2024). Consider noisy linear regression pretraining dataset with the inputs and task weights following the mixed Gaussian distribution:
Assumption 5.1. (Mixed Gaussian pretraining data) "A pretraining task" corresponds to a com- ponent of the mixed Gaussian distribution containing mean of the input mean pm and task weight mean wm- The input means ja and task weights w are sampled from a mixed Gaussian Mm, of (a
distribution: P(μ, wD = EM Am 1, Vm. The process of sampling a training sequence ST [I1 O y1 @ ... a øT @ yr) D; and yt is follows: 1) Sample the input mean μ and the task weight ur according to P( w). 2) Sample a; from P( = N/ (eel,, e2T). 3) Sample
/ (w Wm, I.T), where ΣM Tm = 1,
0 < Tm < 1 and || ll = [|wwll
a
Y/ from Plylær w) = V (ul (ar, w) og)- Define 5, = and 8, c Denote this pretraining distribution as Per
oy
According to the Corollary 2 of Lin & Lee (2024) (see Lemma E.1 in Appendix E). the closed- form prediction of the model trained on the pretraining data of Assumption [5.]] given the testing context and the query I1+1, remains a Gaussian mixture of the reweighted pretraining task weights: F ( ΣM #m w. where 7m and w.,. are the posterior variables of mm and w.. given the downstream context. Hence, to analyze how ICL selects pretraining priors, the key lies in uncovering how सm evolves after seeing the test context. First, we introduce Lemma 5.2 from Lin & Lee (2024) that characterizes the ratio of the reweighted weight of two pretraining tasks:
Lemma 5.2. (Appendix H.l Li & Lee (2024)) Consider any two different pretraining component and B, given test context ST 31 11 and the well-pretrainee model M the ratio between the
7

---

Published as a conference paper at ICLR 2025
weights of the two task priors #a/#t in ' ICL prediction can be decomposed into two terms:
Ra/ A = na exp ( ( (a, 3) + (a, 3)), where
T+1
T+1 Ellma mill -II (20 (1 (T+,))),))
V pp(a,
pB 2.11² Σ
(I)
(ICL. favors the pretraining function with similar input distribution to the test data) Further, assum- ing the test in-context examples xi NN"",, T).1), ) if u1 1 1=11== Ma..
=1=11 11 2 0 holds, then (a. B3 + (1l 1 p.a - 1e 20) 2 O. as the context length T + oc, the first term V (a
However, Lin & Lee 2024) did not analyze how the second term Vw 3) would evolve given any downstream task, which we will demonstrate to play an important role in the algorithm selection mechanism. In the following theorem, we prove that L w (a, 33 converges to a non-negative value when the pretraining function class a performs better on the downstream context than B.
Theorem 5.3. (ICL. favors the pretraining function with low error on the context, proof is in Ap- pendix E.3) Given the context ST, if the empirical risk of implementing a function of the pretraining task a is less than that of B, i.e., 1 E.= I|w - -ysll² 1//7-111n, W - y/ll? 2 0. then, under some mild Assumptions E.2
on the distribution of ST, we have V (a. 2 O.
Combining Lemma 5.2 if the downstream inputs D. D; ~ NN'H; 1) and |A8 - -111 Pl 1=1l² 2 hold, then as T + 00, we have na/Ag 2 Ta/TT.
Summary of the algorithm-selection mechanism. Lemma 5.2 and Theorem 5.3 together elucidate the algorithm-selection mechanism of ICL. According to Lemma model pretrained on the mixed Gaussian data will be reweighted combination of the pretraining task vectors w; Whether the ratio between the weights of two pretraining tasks, / ft.. given downstream context, exceeds the original ratio Ta/ TT3 depends on two factors: 1) whether the pretraining input distribution of a is closer to the downstream input distribution than that of (3; 2) whether the task function of induces lower test error in downstream context than that of B. When both conditions are met, we have / / 2 Ta/Tg, indicating that ICL prefers over 3 in its predictions. We leave the discussions of the advantage of our theory result in Appendix A.5 and offer an intuitive Bayesian interpretation of the algorithm selection in Appendix A.4
the ICL prediction of the
5.2 EMPIRICAL VALIDATION OF THE CECTORRTH--SELECTTIIN MECHANISM OF ICL
Now we validate our theoretical findings regarding ICL's algorithm-selection mechanism in OOD tasks by conducting numerical experiments following [in & Lee (2024). In Figure Sal and 8b the training data is a linear regression Gaussian mixture with four components (see Assumption 5.T). while the test function is a two-layer ReLU network (Appendix C.T.. Both the training and the test data are in ICL format. We plot the test error of using each pretraining task function to predict the downstream function (the first row of Figure 8a and (Bb), the weights for each pretraining function during ICL inference (the second row), and the test error of the pretrained ICL model with the closed form prediction derived in Lemma T. (the third row). We evaluate five different noise levels (6 = سw E {1/81, 1/9, 1, 9, 81), greater value means larger noise) and consider two settings described below that respectively validate the two mechanism in Section 5.1
Low-test-error preference of ICL. To validate Theorem 5.3 we ensure that the distributional dis- tances between the inputs of each training task and the test data remain consistent. Specifically, all x; in both training and test data are sampled from N O. 0.0] o I.. The task weights for different pretraining tasks vary, as detailed in the top half of Table [ In this setup, only the test error of the pretraining functions influences algorithm selection. From the top two rows of Figure Bal we can observe a clear negative correlation between the task weight and the test error of the correspond- ing pretraining task. This result supports Theorem 5.3 confirming that ICL prefers the pretraining functions with low test error. Also, it's consistent with the observations in Figure 7
preference of ICL. We also empirically validate Lemma 5.2 in Figure Sb In this case, the distributional distances between the input of different pretraining tasks and that of the test context vary: the distances of different tasks ordered from largest to smallest as 1 >>> 2 >> 3 > 4, while the test errors of different pretraining functions are almost the same (detailed setup is in the bottom half of Table T). As shown in the bottom two rows in Figure 8b the task
8

---

Published as a conference paper at ICLR 2025
weight n; is positively correlated with the similarity between the training and test input distribution. This is consistent with Lemma 5.2 which demonstrates that ICL prefers to select the pretraining function whose input distribution is close to the downstream one.
6, 6 1/81
6, = 62 = 1,9
6, = 6u
6. = 6μ
6. 6 81
400
Text or of PT lask
Tedt en of PT tast
200
Test or of PT tack
Test em of PT task
11
12
13
14
12
13
11
12
13
12
13
14
tack weight
PT task weighs
weight
0.0
400
200
27
tontorttt Eamplt -
Numerical verification of the low-test-error preference
5. = 6,
1/81
6 = 6 1/9
6, = 6
6 = 6
6, 6, 61
10000
Tesk er of PT lask
Teat en of PT
5000
Test em of PT task
Tist or of PT task
11
12
13
14
12
15
14
11
12
13
13
14
11
12
13
PT task wwip
tank
weight
task weight
27
2'
200
07
at at "usiiar ar G.nanfod Ehampih 0 2F
(b) Numerical verification of the similar-ing similar-input-disee distr-iistrbuution preference
Figure 8: Empirical validation of the algorithm-selection mechanism of ICL. The first rows: the test of the four pretraining functions. The mid rows: the weights of each pretraining function in the closed-form downstream ICL prediction (given by Lemma E. The last rows: the test error of the pretrained ICL model with the closed form prediction derived In the first two rows of Figure Sa the value of the task weight म, is negatively correlated with the test error of pretraining task 2) In the first two rows of Figure the task weights are negatively correlated with the distance between the training and testing input distribution.
Lemma|E.] Observations. 1)
Table 1: Experiment setting of Figure 8a and Figure 8b "PT" and "DS" are short for "pretraining" and "downstream" respectively.
Experiment
inputs
lask
input distribution
Lask functions
PT-DS input distance
Figure Ra
110000.0TTo2.1
No e
N[[[55,533
.
N[[-5,5, _
M[[00 .0.r)
Mt|-5,- .0
Wi[L.1.11T.o..1
Figure
N - 021)
M[[
M[ (
1, [ 0
I,] .0.
12.77
9.11
- b, 55, .022)
1.
1.73
5.3 VERIFYING THE
MECHANISM ON REAL-WORLD LLMS
In this section, we investigate whether real-world LLMs can perform algorithm selection through ICI. To achieve this, design ambiguous sentence classification task, in which each sentence

---

Published as a conference paper at ICLR 2025
be classified based on one of three aspects: "sentiment", "type", or "location". For each ICL sequence, we select one of the aspects as the classification criterion and map the label words to meaningless strings. For instance, if choose to classify each sentence according to its sentiment, then "positive", "neutral", and "negative" are mapped to "RqF", "Iw7." and "SdK". respectively. Detailed experimental setups in Appendix Cc5] We compute the top-5 accuracy of different classification criteria. The results in Figure 9
show that as the context length increases, the LLM finds the most appropriate criterion, exhibiting the low-test-error preference.
Ope
surtamort
65-
cnntet engh
mitee wich
csnae leath
(a) Sentiment
(b) Type
(c) Location
Figure 9: The top-5 accuracy of using (a)" sentiment", (b)"type", or (c) (coclion"""" as the classification criterion for in-context examples in test prompt. The accuracy of using the true underlying criterion to predict significantly higher than the other two. This suggests that L.MMs can perform algorithm selection in natural language tasks.
Summary of the Empirical Results & Connections with the Existing Theories (III)
Despite the impressive empirical performance of real-world LLMs in solving seemingly novel tasks through ICL, we observe that when faced with an OOD task, ICL operates by identifying the most suitable pretraining meta-distribution based on test error and input distribution discrepancies, and then attempts to find an optimal solution within that meta- distribution. Notably, this process occurs consistently, independent of downstream test distribution. See Appendix [.4 for the potential connection between such empirical obser- vations and the Bayesian framework work.
6 CONCLUSION
In this work, we empirically find that Transformers struggle to generalize beyond the pretraining function classes when given downstream in-context examples of OOD tasks. Instead, ICL tries to seek near-optimal solution within the pretraining function classes. We further investigate the widely observed capability of ICL. to perform classification. We reveal that of II) prediction and retrieval rather than an OOD generalization ability. We also examine ICL's performance on OOD tasks after pretraining on multiple tasks. Our theoretical and empirical anal- ysis reveals ICL's preference for low-test-error functions, i.c., ICL tends to implement pretraining function classes with low test error in the test context. This finding highlights two key factors that determine how ICL will implement the prediction function based the testing context and pre- training tasks: the distance between the training and testing input distributions, and the ability of pretraining function to solve the test task.
composition
ACKNOWLEDGMENT
Yisen Wang was supported by National Key R&D Program of China (2022ZD0160300), Na- tional Natural Science Foundation of China (92370129, 62376010), and Beijing Nova Program (20230484344, 20240484642). Xianghua Ying was supported by the National Natural Science Foundation of China (NSFC) under Grant No. 62371009, and Beijing Natural Science Founda- tion under Grant No. L247029. Yifei Wang was supported in part by the NSF AI Institute TILOS, and Alexander von Humboldt Professorship.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, llge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint wwxxi:2203.087770 2023.
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning. In NewrIPS, 2023.
Kartik Ahuja and David Lopez-Paz. closer look at in-context learning under distribution shifts. arXiv preprint 2.ti::2335/6604 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In ICLR, 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In NeurlPS, 2023.
Tom B Brown. Language models are few-shot learners. NeurlPS, 2020.
Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts function lipschitzness. In NeurIPS, 2024.
Joy Crosbic and Ekaterina Shutova. Induction heads as an essential mechanism for pattern matching in in-context learning. arXiv preprint vrXix:2407.07011, 2024
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Ezra Edelman, Nikolaos Tsilivis. Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. In NeurIPS, 2024.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. In ICLR, 2022.
Yu Huang. Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Jannik Kossen, Yarin Gal, and Tom Rainforth. In-context learning learns label relationships but is not conventional learning. In ICLR, 2024.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. How do nonlinear trans- formers learn and generalize in-context learning? In ICMI, Workshop on Theoretical Founda- tions of Foundation Models, 2024.
Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. ICMI,, 2024.
Preetum Nakkiran. More data hurt for linear regression: Sample-wise double descent. arXiv preprint arXiv:19/2.07244 2019.
Jane Pan. What in-context learning "learss" in-context: Disentangling task recognition and task learning. Findings of ACI,, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario AmodeiW, Ilya Sutskever, et al. Lan- guage models are unsupervised multitask learners. OpenAl blog, 2019.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. In NewrIPS, 2023.
Jie Ren, Qipeng Guo, Hang Yan, Dongrui Liu, Quanshi Zhang, Xipeng Qiu, and Dahua Lin. Identifying semantic induction heads to understand in-context learning. arXiv preprint arXiv:2402. 13055, 2024.
11

---

Published as a conference paper at ICLR 2025
Aaditya K Singh, Ted Moskovitz, Felix Hill, Stephanie CY Chan, and Andrew M Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. arXiv preprint arXtv:2404.07129, 2024a.
Aaditya K Singh, Ted Moskovitz, Felix Hill, Stephanie CY Chan, and Andrew M Saxe. What needs to go right for an induction head? a mechanistic study of in-context learning circuits and their formation. ICML, 2024b.
Jiajun Song, Zhuoyan Xu, and Yiqiao Zhong. Outtof-distribution generalization via composition: lens through induction heads in transformers. PNAS, 2024.
Eric Todd, Millicent L Li, Arab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau. Function vectors in large language models. ICLR, 2024.
Hugo Touvron, Louis Martin, Kevin Stone. Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykow, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. JMLR, 2023.
Robert Vacareanu, Vlad-Andrei Negru, Vasile Suciu, and Mihai Surdeanu. From words to numbers: Your large language model secretly capable regressor when given in-context examples. In COLM, 2024.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In ICMI., 2023.
Yifei Wang, Yuyang Wu, Zeming Wei, Stefanie Jegelka, and Yisen Wang. A theoretical understand- ing of self-correction through in-context alignment. In NeurlPS, 2024a.
Zhijie Wang, Bo Jiang, and Shuai Li. In-context learning function classes unveiled for transform- ers. ICML, 2024b.
Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context leannnng. NeurlPS, 2024.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In ICLR, 2022.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities transformer models. arXiv preprint (rxXix2311.000877 2023.
Ruiqi Zhang. Spencer Frei, and Peter L Bartlett, Trained transformers learn linear models in-context. JMLR, 2023a.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. 19420, 2023b.
Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, and Chongxuan Li. On mesa- optimization in autoregressively trained transformers: Emergence and capability. In NewrIPS. 2024.
12

---

Published as a conference paper at ICLR 2025
A COMPARISON WITH RELATED WORKS AND ADDITIONAL DISCUSSIONS
A.1 THE CAPABILITY OF ICL TO LEARN NEW TASKS
Besides studies indicating that ICL can learn new weights of linear regression (Garg et al. 2022 Raventós al. 2023 Zhang al. 2023a Akyürek al., )023), other research has found that LLMs can tackle tasks that are unlikely to have been encountered during pretraining. For exam- Pan 2023) showed that LLMs perform better than random guessing on classification tasks with meaningless labels. Kossen et al. (024) demonstrate that ICL can identify authorship based on writing style in private communication messages not included in the pretraining corpus. Addition- ally, Vacareanu et al. (2024) found that large-scale LLMs can learn various linear and non-linear functions from context. We argue that these findings do not contradict our work. While the LLMs may not have seen exactly the same tasks, there is no guarantee that they haven't 't encountered tasks from a similar distribution in their pretraining corpus. For instance, the LLMs could have been pretrained on corpus containing authorship identification tasks or on statistical data encompass- ing different functions. Our work does not claim that ICL cannot generalize to new task instances of a seen distribution; rather, it highlights the limitation in generalizing to unseen input-label distribution. Additionally, Yadlowsky et al. (202)) finds that ICI. struggles to generalize to testing function classes that are unseen during training (e.g., convex combinations extreme versions of the pretraining functions). They didn't 't delve into how ICL behaves on OOD data, while we reveal that it implements the pretraining functions.
A.2 THE ALGORITHM--SLLECTIO MECHANISM OF ICL
Recent works by Bai et al. (2023): : Wang et al. (2024b) have uncovered the algorithm selection phe- nomenon, demonstrating that Transformers pretrained on both linear regression and classification tasks perform well when presented with the context of either task during ICL inference. Theo- retically, they show that a Transformer with specific parameters can achieve algorithm selection. Yadlowsky et al. (2023) empirically found that ICL selects the optimal pretraining function class after observing in-context examples from function class present in the pretraining data mixture. However, the algorithm selection experiments in these studies are limited to scenarios where the test functions are among the training functions. In this work, empirically and theoretically demon- strate that the algorithm selection phenomenon broadly occurs when given downstream context from arbitrary function classes. To the best of our knowledge, we are the first to reveal the factors that determine the selection process.
A.3 THE BAYESIAN-OPTIMAL PERSPECTIVE FOR UNDERSTANDING ICL
Many studies have found that ICL makes Bayes-optimal predictions (Xie et al. 2022 Wies et al., 2024 Zhang et al., 2023b; Lin & Lee 2024). However, these works have certain limitttions that may reduce their practical applicability in predieting ICL behavior in general scenarios. 1) Limited empirical verification. |Wies et al. (2024) and Zhang et al.] (2023b) lack empirical verification of their theory on real deep Transformer models; 2) Limited theoretical settings: in-distribution tasks. Wies et al. (2024) assumes the downstream tasks are components of the pretraining distribution: Xie al. (2022) assumes that the latent concept of the test task 0* is within the pretraining task set O; In Lin & Lee (2024), the training and testing tasks are all linear regression with weights sampled from Gaussian distribution. 3) Limited implications of the theoretical results: although Xie et al. (2022); Zhang et al. (202bb) prove that ICL can infer a task concept # based on the downstream test context Steet- they don't reveal how Steet concretely affects the posterior distribution P(@ Steet) of the latent task concept 0 inferred by the model that determines the downstream ICL prediction, especially when the true downstream task 0* is OOD. Our work verifies and extends previous findings to more general setting by using real deep Transformers and evaluating ICL OOD tasks that significantly differ from the training tasks. For the first time, we also reveal how the interaction between the downstream distribution and the pretraining distribution affects ICL predictions (see Section 5)
In contrast, Raventós et al. 2023) claim that ICL can exhibit non- Bon-Bayeiian properties. They empir- ically demonstrate that when given sufficiently diverse pretraining tasks (linear regression vectors), ICL can outperform the Bayesian estimator on a new test distribution. However, the distributional shift in their setup might not be substantial enough to show that ICL can truly adapt to new down-
13

---

Published as a conference paper at ICLR 2025
stream distribution, which is considered to be "non-Bayesian" by Raventós al. 2023) In their setting, both the test and training vectors are sampled from the standard Gaussian distribution, and the only source of "distributional shift" comes from the finite size of the training set, which can only partially reflect the test distribution. Our work refutes their findings by showing that when the test distribution is significantly shifted, increasing the number of ID tasks may not help ICL generalize to it.
A.4 THE BAYESIAN INTERPRETATION FOR OUR EMPIRICAL FINDINGS
Although current Bayesian theories for ICL are too vacuous to predict the performance of deep Transformers on real OOD tasks (see Section A.3 the Bayesian framework shows promise potential lens for interpreting our empirical findings. Here we provide some intuitive interpretations for the findings in Section 3 and from a Bayesian perspective.
Consider the predicted distribution pa (2:1) given by a pretrained model A. If we assume that ICL makes Bayesian-like predictions over the test context as Xie al.] (2022); Wies al. (2024);
Zhhang et al. (2023b): Lin & Lee (2024) suggested, then the model will first infer task concept based on the given context 21:T-1 and predict yr using @ and F1.T, i.e.,
Po (1))) = po F1.T Po
To explain the results in Section 3 and Section 5 since the true downstream task o' unseen during pretraining, the inferred posterior distribution pe (0 assigns probability mass only to tasks é within the pretraining distribution that maximize po (1:r) ). This accounts for why ICL can only make in- distribution predictions, as shown in Figure in Section and why ICL prefers pretraining priors with low test error and input distributions similar to those in the test context (Section (5) Once a task seen during pretraining is identified as best fitting the test context 21.--1 the model refines its predictions based on this context (in Figure [ and Figure 7 the test error decreases as the number of in-context examples increase). This refinement corresponds to the factor pa ( Φ), explaining how ICL optimizes predictions within its pretraining distribution.
In Section|4 the underlying task concept acts similarity metric that allows the model to retrieve examples from the context that align with the query. Training on more abstract labels improves the model's ability to estimate more accurate é, which explains the results in 3 and Figure 4 When the test task is ID, even with OOD labels, ICL. can succeed by leveraging the learned & to predict the true label. It accomplishes this by retrieving an example a; from the context that is similar to the query under the ó metric. However, when the underlying task o* is OOD, the model fails because the learned similarity metric no longer applies effectively (Figure 5pp
A.5 DISCUSSION OF THE SETUP OF OUR THEORY
Notably, our theoretical result in Section 5.] applies to any model architecture, while previous the- oretical works of understanding ICL often adopt Transformers with oversimplified assumptions on their parameters or structures (Ahn et al. 2023 Zhang et al. 2023a Huang et al. 2023 Collins et al. 2024). Additionally, our analysis shows that models pretrained on the ICL tasks can imple- ment algorithm selection during ICL inference following Lin & Lee (2024). In contrast, prior work on algorithm selection (Bai et al. (203)) only shows that specific set of parameters in a simplified ReLU Transformer can enable algorithm selection. However, the parameter construction is com- plex and somewhat tricky, and there is no theoretical or experimental guarantee that Transformers exhibiting algorithm selection will necessarily implement these parameters.
B ADDITIONAL EXPERIMENTAL RESULTS
B. UNDERSTANDING THE EFFECT OF TRAINING ON MORE DIVERSE RETRIEVAL TASKS FROM THE ATTENTION SCORES
To further validate that the retrieval ability is evoked after trained on more random mappings, follow- ing Crosbie Shutova (2024), we construct another retrieval task and visualize the prefix matching score of all attention heads of the three pretrained models in Figure TO The prefix matching score
14

---

Published as a conference paper at ICLR 2025
is calculated by averaging the attention values from each token &; to the tokens after the same token as ty in earlier positions in the sequence, which correlates positively with the retrieval performance (Singh et al. 2024b). In Figure 10, we observe that the model best the retrieval task in Figure
3 exhibits more heads with high matching scores (also known as "induction heads" (Song et al,,
2024 |Snghh et al. 2024a] Ren et al., 2024, Crosbie & Shutova 2024 [ddelman et al. 2024 further demonstrating it gains the retrieval ability by training on more retrieval sequences.
Matchiig score
Motching score
Matching score
2:14
5120
COM
layers
AMT
AM
PMN
PMM
S4
AM
441
640
3:M
3:01
PPM
CHT)
CAN
PM
COM
CCN
CM
60
CET
heeds
hesds
heads
(a) PT S (500 150)
(b) PT 8 (50, 250)
(c) PT (450 450)
Figure 10: The matching score of all attention heads of models trained on the retrieval task. "PT" denotes "pretrained on". Each subfigure corresponds to different pretrained model. The model of (c) exhibits more heads with high matching scores, which is also the most performant model in the retrieval task in Figure
B.2 THE SYNTHETIC VECTOR CLASSIFICATION IS NOT THAT HARD TO SOLVE IF IT'S IN DISTRIBUTION
To show the failure in the synthetic vector classification in Section 4.3 is mainly due to its OOD
nature rather than it's intrinsically too hard to learn, we train a GPT-2 to perform the same task in Section 4.3 In this task, the xi and y; are generated in the
way Section 4.3 The
only modification is that we use a smaller predefined vector embedding E" E 12 100000020 (Eooaa E 332000 4096 in the experiment Section 4.3) The results in Figure show that when W has been encountered during pretraining, ICL can well address this task.
B.3 EVALUATING THE SYNTHETIC OOD CLASSIFICATION TASK ON LLAMA-2-7BB
We also evaluate Llama-2-78 on the same OOD vector classification task and the retricval task as in Section 4.3 Figure T2 shows the same observations in Figure that the LLM well address the retrieval task but fails to learn the OOD function W. In this experiment, we set the length of the repeating sequence be 10. We can observe that the accuracy of retrieval rapidly increases after seeing 10 in-context examples. This demonstrates that learning novel functions from the context is challenging for real-world pretrained LLMs, but the LLMs are good at retrieving.
synthetic word classification
1.0
Random
0.6
0.4
0.0
80
100
in-context examples
Figure 1I: Test error of the GPT-2 trained and evaluated on the same synthetic OOD vector classifi- cation task in Section 4
15

---

Published as a conference paper at ICLR 2025
Accyracy (CL on the word classiffication task, classes
Aurrcyy of ICL on the context retrieval task, 10 classes
Accuracy
Accuracy
9.6
3.4
600
1003
length
length
(b)
Figure 12: The ICL accuracy of Llama-2-7B the synthetic tasks. (a) the synthetic vector classifi cation task. (b) the synthetic word retrieval task.
B.4 WILL GENERALIZATION CAPABILITIES EMERGE FROM INCREASING THE NUMBER OF TRAINING TASKS?
Recent work by Raventós al. 2023) empirically demonstrates that when both the training and test tasks are linear regression, and the number of training vectors exceeds certain "task diversity threshold" (approximately 2 2 2 ICL can generalize from finite training set sampled biasedly from W(0, 1) to the test distribution Pies NO, 1) (see Appendix A..3 for details). We investigate whether similar phenomena persist for test tasks with larger distributional shifts. We train models using varying numbers of linear regression vectors and evaluate them on quadratic and ReL.U neural network regression tasks. Figure [3 we find that training on a vast number of ID functions does not yield any improvements, providing further evidence that ICL may struggle to achieve OOD generalization.
Evail Lank: quadratic regression
regression
1:25
1000 PT vectors
10000 vectors
onor
100000 vectors
0.75
0.75
1000000 vectors
0.50
0.00
20
40
60
50
100
20
4.0
50
In-conteut excemples
Incoentedd excemples
(a) quadratic regression
(b) ReLU NN regression
Figure 13: The ICI test error of models trained on different numbers of linear regression vectors. Even if the number of training vectors (up to 1, 000, 000 220) surpasses the threshold (214 216) reported by Raventós et al. (202)), model exhibits generalization to OOD function classes.
C EXPERIMENTAL DETAILS
C.I EXPERIMENTAL DETAILS IN SECTION 3.1 AND SECTION 5
Definitions of the function classes. The function classes in Figure and Figure 7 are:
Linear regression: Y/ w x;, where E Rd and w, 2; (0.
Quadratic regression: = w (x;)² where w, E Rd and w, a, N(O, 1), (æx)) denotes the element-wise square of æ;.
2-layer ReLU network regression: and 2;; Rdd wT, W2, a;; ~ (0, 1).
w Re2LU( where wi E Rdd w2 E Rd'xd
Square root linear regression: y/- w T Vw., where w. Rd and w, NoO. 1),
(zx)) denotes the element-wise square root of a;,
Cubic linear regression: y; = w w (xx)) where w, æ; E Rd and w, x, ~ N(O 1). (z;)² denotes the element-wise cuhe of z;;
16

---

Published as a conference paper at ICLR 2025
Linear+quadratic regression: y4 "w (xx)² w. 2;, where w1, W2, 3, IR and wi, W2, 2; ~ (0,
2-layer Sigmoid network: yi = wi (igmoi((u2a:)) where wi E Rdd, W2 E Rd'xd xd and Z; E R. w W2, I, (0, 1).
Baseline models in Figure 1 The models of each pretraining hypothesis class implemented by training neural network that yields functions of that hypothesis class. For example, linear regression weight can be implemented by single linear layer. The models are optimized using SGD with learning rate 1e-3 for 1000 steps.
C.2 EXPERIMENTAL DETAILS FOR SECTION
For the reversed-labed experiment, we choose four tasks: Antonym, Capital-country, English- French, and Saglish-German. The original datasets adopted from Todd et al. (2024). The top-1 accuracy is computed as follows: compute the top-1 accuracy for each token predicted by the model, based on the token length of the ground-truth label word. For each context length, compute the average accuracy over 128 test examples.
C.3 EXPERIMENTAL DETAILS FOR SECTION
We now provide additional details regarding the experiments of Figure T0 Following Crosbic & Shutova (202)), we generated a dataset consisting of 100 sequences of random tokens, each containing repeated sub-sequences. The task predict the next token that follows the last token in each sequence. This task can only be completed by retrieving the last token from the context and predicting its subsequent token.
C.4 EXPERIMENTAL DETAILS FOR SECTION
We uniformly sample 1000 word vectors 2; E Rdd from the word embedding E E Rxd of the pretrained Llama-3-8B, where = 128256 and d = 4096. Then we sample a task weight W E Rd'xC' from standard Gaussian distribution that only takes the first d' dimensions of æ; (denoted a d')) to compute probability distribution we set the label vectors y Earg max; p4 +s]+s E Fd where s = 10000 is offset. We set d' = 30 < 4096 reduce the complexity of the task. Hence, æ, are classified into C labels vectors E(s s C). The predicted token of æ; is computed as: arg max, p=]j), {8, 8 1,.., where P the output of the last linear layer of Llama-3-8B given It.
C classes: P: æt|: d' W E RC Next,
C.5 EXPERIMENTAL DETAILS FOR SECTION 5.3
In this section, present some details about the setups for the ambiguous classification task. The label mapping rule is presented in Table For each context length, we compute the average accuracy 128 test examples.
Table 2: Experiment setting of Figure Sa and Figure 8b "PT" and "DS" are short for "pretraining" and "downstream", respectively.
Classification criterion
Original labels
Labels presented in the context
"positive"
"RqF" "IwZ"
sentiment
"neutral"
"negative"
"SdK"
"science"
"RqF""
type
"sports"
"IwZ"
"arts"
"SdK"
"Asia"
"RqF"
location
"Furope"
"IwZ"
"Africa"
"Sdk"
17

---

Published as a conference paper at ICLR 2025
Prompt examples. Here present some in-context examples of the input prompt of using different classification criteria.
Using "sentiment" as the classification criterion.
Q: The groundbreaking discovery made by Japanese scientists has revolutionized renewable energy.
A: RqF gray# Original label: "positive"
Q: A chess championship occurred Russia, featuring players from around the continent. "neutral"
A: IwZ gray# Original label:
Using "type" as the classification criterion.
Q: A regional basketball league was formed in Kenya to promote the sport locally.
A: IwZ gray# Original label: "sports"
Q: The breathtaking architectural exhibition in Dubai left visitors absolutely awestruck. "arts"
A: SdK gray# Original label:
Using "location" as the classification criterion.
Q: A scientific paper from Finland explores new methodologies data analysis.
A: IwZ gray# Original label: "Europe"
Q: An astronomy workshop was conducted in Ethiopia for students interested in space.
A: SdK gray# Original label: "Africa"
Accuracy computation. For given label, the method to calculate top-5 accuracy is follows: compute the top-5 accuracy for each token predicted by the model, based on the token length of the ground-truth label word. For classification criterion other than the selected in the current sequence, to verify whether the model's prediction distribution across all test samples approaches the label distribution under that criterion, we select the permutation among all possible mappings between original labels and meaningless strings that yields the highest model prediction accuracy to compute the accuracy.
D EXISTING THEORETICAL EVIDENCE SUPPORTING THAT ICL MAKES ID PREDICTIONS
One recent work (Zhang al. 2023a) theoretically proved that one-layer linear self-attention model (LSA, defined in Appendix D pretrained on a lincar regression task will still implement the linear predictor given downstream in-context examples of arbitrary function classes, under some assumptions on the initialization of the Transformer weight matrices. We restate the Theorem 4.2 Zhang et al. (2023a) as Lemma D.I below:
Lemma D.1. (Theorem 4.2 of Zhang al. (023aa), informal) Let 'D be distribution over (æ, y) Rd x R, whose marginal distribution on x is D2 -N, A). Assume the test prompt is of the form
P (x), V1...., FT, YT, Equery), ), where (æx, v:) (Iqueryy Youerr)
Lid. D.
The prediction risk on
the test query Vquery of an arbitrary task satisfies:
E (Uquery - Byuery ) min E ((w, Degnery> Yquery )2 +const,
Ervor best linear predictor
where const is a constant depending on the downstream context, and the expectation is over (xs; Wi) ( (Fequery Miueery ) lid.
D.
18

---

Published as a conference paper at ICLR 2025
Lemma D.1
serves shred of theoretical evidence that ICL just implement the pretraining
function class, while the role of the context examples is to optimize the model within the pretraining hypothesis space.
Below, we provide the necessary details of the theoretical setting of Zhang et al. (023a).
The linear self-attention (LSA) model considered in the Theorem 4.2 of Zhang al. (2023a)
(Lemma D.I is defined as follows:
SLSA (esA(B (E;0) E+ W E. E E
(3)
N
where the input embedding defined as follows:
= E(P)
22
EN
query
E RR(-I)((N+))
(4)
y/
1/2
UN
WPV is obtained by merging the projection and value matrices into a single matrix, and WKQ is attained by merging the query and key matrices into a single matrix. is the context length.
Now we present the assumption on the attention weights of the linear-attention model in Lemma D.I
Assumption D.2, (Assumption 3.3 in Zhang et al. and θ Rdddd be any matrix satisfying (00 F and 6A * Odxd- We assume
(223a), initialization). Let > 0 be parameter,
PVV (0) =
Oaxd 0
Od
1 )).
W = a T0T 0g
0
The training objective is to minimize the population risk of the linear regression task:
I(0)= lim T(0)
Eog aaa 20 1 -0,, qay (Ur, query
(Wr, Lr query ))²
(5)
B-000
2
where W, ~ NO0 Ia), Trii Br.query (0, Dr. query is the prediction of the LSA model.
E THE LEMMAS, ASSUMPTION, AND PROOF FOR THEOREM 5.3
In this section, we will first present the lemmas and assumption that Theorem 5.3 depends on and the provide its proof.
E. LEMMAS FOR THEOREM 5.3
The lemma below states that the closed-form prediction of the model trained on the pretraining data under Assumption given the testing context, remains Gaussian mixture of the reweighted pretraining task weights:
Lemma E.I. (Corollary of Lee (2024), closed-form ICL prediction of the pretrained model) Denote the model M that minimizes the risk on the pretraining data of Assumption 5.1 i.e., M E
[IM (S, @ x4+1) -yt+ 112]. then the prediction on any sequence S, @ 1i++ 1 by M as follows: F M (S. 1:+11) = =0 (a..1 Em T. Tm.im)) where ims and w. depending on both the pretraining task and the downstream context example are given in Lemma Lin & Lee 2024).
arg min + TT1 E SimPtr
E.2 THE ASSUMPTION FOR THEOREM 5.3
The assumption below impose some mild requirements on the distribution of the downstream con- text:
Assumption E.2. (Assumption on the distribution of the downstream context examples.) As- sume that: the minimum eigenvalue of the covariance matrix of any in-context example æ;
satisfies Amin (xx,x,
2 1; (I TII + 8m 2, = I; ET= 2(w,
wa) 0.28 T E/=
2 x. Mi 2Txa2 20
19

---

Published as a conference paper at ICLR 2025
E.3 PROOF FOR THEOREM 5.3
Now we restate Theorem as the Theorem below
Theorem E.3. (ICL prediction favors the pretraining function with low error the context) Given
the context Sk, if the empirical risk of implementing a function of the pretraining task a is less than
that of B, i.e. 1
EE= 1= [uga;;
W. yel 2 O, then, under some mild Assumptions E.2
T
we have V wo 3)
O.
Combining Lemma
5.2 if the downstream inputs Da , r221) and |uB 14 1=11 1o
μ² 1²1 2 0 hold, then as T + oo, we have Ra/#g 2 Ta/ TB.
Proof. According to Lemma
of & Lee 2024).
r(a, 8) Fa
To Coct
exp (V, + (a, B))
(6)
3
In the Appendix H.1 of Lin & Lee (2024), they have proved that when the context length T + o0, under the first condition in Assumption
limy V (,, =2 O.
Now prove that when the empirical risk the in-context example of pretraining task function no more than that of (3, the second term (a, B) 2 0.
V w (),
I|,, +T w
(I+TA, x2 -1
exp
log
||00 [2
11 + 10/
(1 13 E10
exp
|//p|]
100
+ To,, wi
(I+T6.E.))
||
l| + T8
(1 Ew)
20
20
|,00||2
10 + 6w Σn 1i= Ditli
||
W. + 8 Σi=
Dzlli
(1+T6, E.."
(I+T5,, Ew,)
20
20
| P
[ww 2 a.
(I TI
Bilk 2
(I+T5 Ew
2o w
|wa/ 11 2 (wa
ΣT
(I TISw)
ET a.yi
(I+T6, E.
202
T
T
(0)
[
Filli
||2
||w0
Dilli
T
(I+T8 E.w.
T
(I+TS.
Σr
(b)
1
Billi
||20a
1=
Ciili
1|2
T
T
oi Yor
114m
11
8-1
(7)
where equation (a) is due to the third condition in Assumption E.2 equation (b) is by applying the Morrion formula. Since
2 0. to prove that Va B) 2 O, we only need to
I+6., Σ.
show that
133
Σ Dcli
12
ΣT
u.Φ." - Wa
ΣT
Filli
2
R.Φ." 2 0.
(8)
T
T
Eis
20

---

Published as a conference paper at ICLR 2025
Further, derive that the term Σ negative by using the condition 2 in Assumption E.2
1= [wg- 2.yalll 20.7 1|wa 2.y/ill 2,
below is non-
T
F.y4 la
| -i:3i m.
T
Σ
[wg
xayi) "0,00 (wo x.yli - (w., x.91)) "a, 20.9i)
T
1=1
1
[ + 20,11,) ax (wg wa)
(9)
T
i=1
T
2xxy)) (wps wa)
w + wa
T
T Jun 0c - yill? 1m x:
Yill' 20 0 (d)
where the incquality (c) holds since according to the condition 2 in Assumption E.2. a . is positive semi-definite, and the inequality (d) holds since the population downstream risk of a is lower than B. Therefore, to prove inequality (8) we just need to prove that the Lh..,. of inequality
80 multiplying + is not less than Ei= || a.yli 11
in Equation 99
2,007
1 T
Σ -i=1-
1
Billi T
10 x. a. 1/
xzyhi (1:1 l.
T
DE
ET aa.t
(10)
First, let's simplify the l.h.s of inequality
TO
1 T
[w0
Σ 1=
||wa
Li=l
T
T
E1 an
1
Σ
Σ j=1
)T T (10g
j= T3
(wo
Σj Egbi )T "m (wr.
Σ T3
T
T
T
1
+
Σ
|10
0;; T
Σ
T Villi
-11u 0s.
T
j=1
j=1
7
Σ (w) az,))²
(w) (w,) + 2(w.
w0s))
T
ar
T
1
(11)
Then we simplify the r.h.s. of incquality (10):
T
wa
Z:lli l, 1110 20.yill2
(12)
T [w] xx)²
(w x;
2(w. a ways
T
i=1
Subtracting Equation
from Equation
we get
T
2 2(w, - wa) x
x T
2(w, - wa) T xx x..9
1=1
j=1
(13)
W wp) x.
x x
V1
x3 x3
1=1
21

---

Published as a conference paper at ICLR 2025
applying the condition 4 in Assumption E.2 get the final conclusion.
F LIMITATIONS
1) Most experimental results are based on a GPT-2 model pretrained on a limited set of mathematical functions. It is challenging to assess whether modern large-scale language models like GPT-4 and Claude Opus face similar difficulties in generalizing beyond their pretraining corpus, given the vast range of tasks and content in their pretraining data. Nevertheless, our findings highlight the limitations of ICL in solving challenging tasks for smaller models like Llama-2-7B and Llama- 3-8B. 2) The models trained on ICL data, while real-world LLMs are trained autoregressively. However, the ICL pretraining objective is also next-token prediction, so there is no clear gap between these two pretraining objectives.
G REPRODUCIBILITY
In the main text and Appendix C we've 'e stated all setups for reproducing our experimental results. For the theoretical part, we" ve included the assumptions (Assumption E.2) and proofs in Appendix
22


Reviews:
[{"review_id": "0OfL46c2dW", "reviewer": "Reviewer_KJfV", "summary": "This paper studied the limitations of in-context learning (ICL) on out-of-distribution (OOD) tasks through empirical experiments and theoretical analysis. They pointed out that ICL mainly implements functions from its pretraining distribution rather than learning new functions. Therefore, ICL cannot guarantee ideal performance when ICL data are from OOD distributions unseen during training. While ICL appears capable of handling tasks with abstract labels, the authors show this ability stems from learning retrieval mechanisms during training rather than true OOD generalization. They also provide theoretical insights into how ICL selects which pretraining function to implement when faced with OOD tasks.", "questions": "See the comments in the weakness section.\n\n\nUpdate:\n\nSorry for the late response. Most of my concerns were properly addressed. I have raised my score to 6.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Novel theoretical insights: The paper provides theoretical analysis of the \"low-test-error preference\" mechanism in ICL, explaining how models select which pretraining function to implement when handling OOD tasks. \n2. I found the experiments on the abstract label classification interesting. The paper provides a fresh perspective on how models handle abstract label classification, demonstrating it's more about retrieval capabilities than true OOD learning.\n3. Besides GPT-2, they use Llama2-7b, a pre-trained language model for evaluation.", "weaknesses": "1. Presentation issues. \n\n(a) In Eq.1, why the prediction is supervised by f(x_i)? Shouldn’t it be f(x_{I+1}) \n\n(b) In Fig.2, what is y_i id? Is it same as I_{y_i}? If yes, why not keep using the same notation? If no, please clarify. \n\n(c)Please correct the way how you are using the quote. See examples in the caption of Fig.3.\n\n2. Questions on the training set. The retrieval design in Sec.3.1 is interesting. When the training range is larger (say yi Id in [50, 455]), the test-time performance is better regardless that y_i in [50, 155] aligns better with the evaluation tasks. I wonder if you have same number of training tasks for different range.  If the answer is no, then is it  a fair comparison?  Is it possible that better performance simply comes more training data? Please clarify the number of training tasks for each range. Second, the setup is hard to understand, please improve the writing. \n\n\n3. Questions about the experiment design. The first and last experiments in Sec.3 are following the same idea, which maps the index of an embedding to new indices based on some pre-specified rules.  I find it hard to understand the motivation here. Are previous works using similar design? If yes, please cite. If no, could you clarify why this setup?\n\n\n4. Questions about the assumptions when evaluating pretrained models.In Sec.3.3, how is it possible to know an evaluation task is in-distribution or out-of-distribution for a pretrained llama-2-7b? In line 264, the authors said `To ensure the task is far from the pretraining distribution’, are there evidence supporting the assumption?\n\n\n5. Limited evaluation setups. All conclusions are derived based on either curve fitting or rule-based mapping. The setups are too simple. They are not sufficient to derive convincing conclusions especially when evaluating the behaviors of models with high parameter size.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "xqaRPQNFhP", "reviewer": "Reviewer_MhX7", "summary": "This paper investigates the ability of transformer models to generalize to tasks that are OOD compared to its pre-training distribution through in-context learning. They do this through empirical analysis with controlled synthetic tasks and on a pretrained Llama-2-7B model as well as provide some theoretical justification. The core contributions are:\n1. The performance of transformers on OOD tasks is similar to the performance of an estimator that learns the ICL examples through GD within the pretrianing function class.\n2. The transformers ability to learn and classify abstract concepts in-context can be attributed to its ability to do retrieval and can be learned by doing retrieval on a pre-training distribution of many diverse labels. This fails if the testing function is OOD.\n3. Provides theoretical and empirical evidence that when the pretraining distribution contains a mixture of function classes, when tested OOD, models select the pretraining function class that has the lowest test error and match between the train and test input distributions.", "questions": "1. In section 4, the authors demonstrate model selection within the pre-training task distribution. Can you connect this to your results in section 2, i.e. does the transformer get the same error as Gradient descent with the function class that gets the lowest error in Figure 7?\n2. What causes the sudden increase in performance after context length size 10 in Figure 6B\n3. What do you precisely mean by the claim that \"the ability of ICL to perform abstract label classification may not serve as an evidence of learning new tasks.\"? In this abstract label classification task, learning new tasks is the same as being able to classify OOD labels. This is something you show that models are capable of doing if pre-trained on enough diversity of abstract labels tasks (green line in Figure 4c). It seems like this presents evidence that (in a limited context) transformers can learn new OOD tasks.\n4. Can you please explain the motivation behind the experiment in Figure 5 better.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "**Originality**: This paper adds to the literature on understanding ICL in transformers on synthetic tasks by empirically and theoretically characterizing their behavior on OOD tasks. To my knowledge, this behavior on OOD tasks has not been carefully demonstrated yet.\n\n**Quality**: The paper does thorough, well thought out experiments and provides new theoretical insight.\n\n**Clarity**: The paper is clearly written, well structured and easy to understand.\n\n**Significance**:  The question asked by the paper are significant as it further clarifies the mystery of ICL especially regarding limitations to OOD tasks.", "weaknesses": "- In section 2, it would be helpful to further clarify what is existing knowledge and what is a new contribution. In fact, I think the presentation would be cleaner if Section 2 and Section 4 were combined into a unified presentation.\n- The explanation of the retrieval version of the Llama-2-7B task was hard to understand, it would be helpful to clarify that section\n- The conclusion that LLMs cannot learn OOD tasks because its ability to learn abstract concepts from context can be attributed to retrieval felt like a jump in conclusion: see questions).\n- I did not understand what hypothesis the experiment in Figure 5 was trying to test and a more clear explanation of why this experiment is useful and what the hypothesis is as an analogy to real world use case would be helpful. It seems like a really whacky OOD task that I am not sure is representative.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "yv8MzgSlp5", "reviewer": "Reviewer_7eJd", "summary": "The paper conducts a series of empirical studies on synthetic tasks aimed to characterize the ability of transformers to generalize to unseen function families using in-context learning (ICL). It finds that while ICL allows generalization to new tasks from known function families, it fails to generalize to tasks from unseen function families instead predicting solutions in terms of the known function families encountered during training that achieve the lowest test-error. It further demonstrates that the strong in-context retrieval ability of transformers can sometimes give a misleading impression of out-of-distribution generalization.", "questions": "- It is not clear to me whether the Llama2 OOD synthetic word classification task in section 3.3 fails due to a lack in OOD generalization or because of other factors that make this task difficult. One control would be to train a model explicitly on this ICL task to see if it is solvable when it is in-distribution. Would it make sense to run such a control?\n- I found Section 4.2 and Figure 8 and Figure 9 hard to read and understand. What exactly is shown in the last row of both of these figures? Why are the losses so high here throughout? I am probably misunderstanding this plot but I would have expected to see that the pretrained ICL model achieves the same loss as the corresponding PT task with the lowest test error.\n- Would it be possible to show the loss plots on a log scale (especially Figure 7) to get a better understanding how close the different methods are to each other? Many of the plots contain relatively large values (e.g. due to double descent) which make it difficult to assess how close the best models perform at small loss values. Maybe they could be added to the appendix.\n- Are the y-ticklabels of Figure 7b+c correct, i.e. is the loss 10x higher for these tasks? If so, why is that the case?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Characterizing the out-of-distribution generalization ability of in-context learning is an important problem with potential impact\n- The experiments are well-designed to investigate the respective hypotheses\n- The paper does a good job at referencing and contextualizing existing related work", "weaknesses": "- The presentation of the paper could be improved, in particular I found section 4.2 hard to follow (see questions) and with the many different experiments often following prior work it was difficult to identify what findings were novel and specific to this paper.\n- The theoretical results heavily borrow from [1] (but the paper makes this very transparent)\n\n[1] Lin, Ziqian, and Kangwook Lee. \"Dual operating modes of in-context learning.\" arXiv preprint arXiv:2402.18819 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "t0ZtidDyPk", "reviewer": "Reviewer_gxdx", "summary": "This paper studied OoD performance with in-context learning in Transformers. The notion of OoD was both with respect to the ground-truth data-generating process at inference time and the distribution of inputs. The main finding is that Transformers do not generalize in these OoD settings, and instead make predictions in accordance with the in-distribution task from the pre-training distribution that achieves the lowest test error on the novel OoD task. This core finding holds when considering simple regression tasks (e.g., linear and quadratic regression), retrieval tasks, and mixtures of tasks.", "questions": "1. Which findings in particular do the authors believe are genuinely surprising to most deep learning researchers?\n2. What findings from the paper can’t be explained through the theoretical Bayes-optimal prediction framework for ICL?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "1. The experimental setups are quite easily to follow and well-explained.\n2. The tasks introduced in Section 3.3 are quite interesting, as they constitute a simple and meaningful task that we would want a model to generalize to, but we are confident that no such task could have accidentally appeared in the training data.", "weaknesses": "1. The core results are to a large extent unsurprising, and don’t clearly tell us anything new. Transformers are essentially just deep neural network (DNN) sequence models, and we know that DNNs don’t generalize out-of-distribution in the absence of the right kinds of inductive biases. Why, then, would we expect a Transformer to generalize in the OoD settings studied in this paper? For instance, would anyone have ever genuinely expected a Transformer trained on (x,y) sequences with a linear relationship to then generalize to (x,y) sequences with quadratic relationships? I expect that the answer is no, and that almost everyone familiar with the deep learning literature will be unsurprised by the empirical findings in Sections 2, 3, and 4. Lack of unexpected results is my primary critique of this paper.\n2. This work did not engage with the dominant theory of ICL: that it does Bayes-optimal prediction. This is particularly surprising, since the Bayes-optimal prediction framework seems to account for all of the results. To account for Section 2: the Transformer learns a prior over the pre-training tasks and so when confronted with a new one, the best it can do is implement the function with non-zero prior that best explains the novel OoD data. To account for Section 4: the learned task prior includes a prior over task classes, and given some context the posterior is higher over the tasks that are consistent with that context.\n3. Several typos and ungrammatical sentences throughout the paper that make it more difficult to read. For example, there is even a typo in a subsection title “3.3 REAL-WORLD LLMS MAT NOT NECESSARILY IN-CONTEXT LEARN NEW TASKS”. The paper requires a thorough proofread.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)