Can In-context Learning Really Generalize to Out-of-distribution Tasks?

Response to Weakness 3: We sincerely apologize for any typos or grammatical errors in the previous version of our paper and regret any inconvenience they may have caused during your reading. We have now thoroughly revised the manuscript.

---

Response to Question 1: We believe the most insightful finding of our work is as follows: 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 th·en attempts to find an optimal solution within that meta-distribution. Notably, this process occurs independently of the downstream test distribution. Our work provides grounded empirical and theoretical evidence to clarify both the capabilities and limitations of ICL, which we believe offers valuable insights for the community in understanding the boundaries of ICL's potential.

---

Response to Question 2: We acknowledge that the existing Bayes-optimal prediction framework can explain the findings in Sec.2 that ICL implements in-distribution predictions. However:

• The framework cannot account for the emergence of OOD retrieval ability in Sec. 3.1, where ICL retrieves OOD labels after training on a sufficient number of retrieval tasks. Bayes-optimal predictors are not inherently designed to generalize to novel input-label relationships.

• The framework lacks insight into algorithm selection for OOD contexts. While classic Bayesian theories, such as [10], have demonstrated similar selection mechanisms for mixed Gaussian models, they fall short in predicting the behavior of ICL performed by complex deep Transformers in OOD tasks. Moreover, existing Bayesian frameworks ([5-9]) have not identified the factors that influence ICL's algorithm selection in arbitrary OOD contexts.

  Our work addresses this gap by characterizing the specific factors that affect downstream ICL predictions (Sec. 4). Furthermore, our theory is notably general: unlike [11-12], we do not assume that test tasks are seen during training, nor does our proof rely on manually constructed model parameters as in [11-12].

We have added detailed discussions to Appendix A.3.

---

[1] Disentangling task recognition and task learning, Pan et al., Master’s thesis, Princeton University, 2023

[2] From words to numbers: Your large language model is secretly a capable regressor when given in-context examples, Vacareanu et al., Arxiv 2024

[3] Many-Shot In-Context Learning, Agarwal et al., Arxiv 2024.

[4] In-context learning learns label relationships but is not conventional learning, Kossen et al., ICLR 2024

[5] The learnability of in-context learning, NeurIPS 2023

[6] An explanation of in-context learning as implicit Bayesian inference, Xie et al., Arxiv 2021

[7] What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization., Zhang et al., Arxiv 2023

[8] Dual operating modes of in-context learning, Lin et al., ICML 2024

[9] Pretraining task diversity and the emergence of non-bayesian in-context learning for regression, Raventós, et al., NeurIPS 2023

[10] Pattern recognition and machine learning. Bishop, C. M. (2006). Springer. Chapter 9

[11] Transformers as statisticians: Provable in-context learning with in-context algorithm selection, Bai et al., NeurIPS 2023

[12] In-context learning on function classes unveiled for transformers, Wang et al., ICML 2024

We thank Reviewer gxdx for your careful reading and valuable comments! We'd like to address your concerns on the following points:

---

Response to Weakness 1: Although it might seem natural that Transformers cannot achieve OOD generalization without appropriate inductive biases, however:

• **ICL is not conventional learning, whether ICL can solve OOD tasks remains controversial. ** Wealth of empirical evidence has demonstrated the emergence of generalization capabilities in real-world LLMs with ICL across various tasks that are unlikely to have been seen during pretraining. These capabilities include performing classification tasks with meaningless labels [1], fitting arbitrary non-linear functions [2], translating very rare languages [3], and classifying writing styles on private datasets not included in the pretraining corpus [4]. Therefore, it remains unclear whether the ability of ICL to solve these seemingly novel tasks represents true OOD generalization or is simply because LLMs have been exposed to similar tasks during pretraining. This OOD issue in ICL is considered "an important problem with potential impact" by reviewer 7eJd and "has not been carefully demonstrated yet" by reviewer MhX7. Our work addresses this uncertainty by conducting experiments with rigorously controlled pretraining and test distributions, indicating that ICL may struggle to learn arbitrary new tasks, challenging previous literature [1-4].

Here we summarize our novel findings and insights:

• to the best of our knowledge, we are the first to demonstrate that the OOD behavior of ICL closely resembles finding a near-optimal solution within the pretraining distribution through experiments with deep Transformers (Sec. 2).
• we find that being able to classify OOD abstract labels is a capability of retrieval rather than learning OOD input-output relationships (Sec. 3.1), challenging the findings suggested by [4]. Further, such retrieval ability only emerge when the underlying task is ID (Sec. 3.2).
• we reveal that the algorithm selection mechanism consistently exist regardless of the downstream distribution and demonstrate how ICL selects its pretraining prior based on the test context (Sec. 4).

---

Response to Weakness 2: Thank you for highlighting the connection between our work and the existing Bayes-optimal theory for ICL. While several studies attempt to understand ICL from a Bayesian perspective ([5-9]), these works have certain limitations that may reduce their practical applicability:

• Limited empirical verification. [5] and [7] lack empirical verification of their theories on real deep transformer models. [9] claims that training on sufficiently many linear regression vectors can cause ICL to generalize to the test distribution, which is considered a non-Bayesian property. However, the distributional shifts in their experiments are not significant (see Appendix A.3 for details) enough to conclude that ICL can truly generalize.
• Assumption of in-distribution downstream tasks. [5] assumes that downstream tasks are components of the pretraining distribution. Similarly, [6] assumes that the latent concept of the test task $\theta^*$ is within the pretraining task set $\Theta$. In [8], both training and testing tasks are linear regression problems with weights sampled from Gaussian distributions.
• Limited implications of the theoretical results: Although [6] and [7] prove that ICL can infer a latent task concept $\theta$ based on the downstream test context $S_{\text{test}}$, they do not reveal how $S_{\text{test}}$ concretely affects the posterior distribution $P(\theta \mid S_{\text{test}})$ inferred by the model, which determines the downstream ICL prediction—especially when the true downstream task $\theta^*$ is OOD.

Our work addresses the aforementioned limitations through the following contributions:

• Comprehensive empirical evidence. Our findings are supported by extensive experiments, including synthetic tasks evaluated on deep GPT-2 and natural-language tasks tested on pretrained Llama-2 and Llama-3 models (newly added in Sec. 2.3 and 4.3). These experiments validate and extend the results of [5-7] to broader, more practical scenarios.
• More realistic settings reflective of real-world tasks. We explore OOD tasks where the test function classes differ significantly from the training ones, better mirroring real-world challenges. The results presented in Sec. 2.2 (newly added) suggest a revision of findings from [9], showing that ICL cannot generalize to OOD tasks even after extensive training on ID tasks.
• Novel insights into ICL's capabilities. Our analysis uncovers how the downstream and pretraining distributions influence ICL predictions (Sec. 4). Additionally, we provide a new perspective on the widely observed phenomenon of abstract-label learning, showing that it reflects a retrieval capability rather than an ability to learn unseen tasks.

Dear reviewer gxdx,

I hope this message finds you well. I wanted to kindly follow up regarding our previous response to your comments. If there are any further questions or clarifications needed, we would be more than happy to provide additional information or explanations to assist in the review process. Thank you for your time and consideration.

We thank Reviewer 7eJd for your careful reading and valuable comments! We'd like to address your concerns on the following points:

---

Q1: 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.

A1: Thank you for your feedback! We apologize for any lack of clarity. We have refined the presentation of Section 4.2 and highlighted the novel contributions of each experiment in the updated PDF. Here is a brief summary of the novel findings from all the experiments (all figures below use the indices in the updated PDF):

• Fig. 1: When evaluated on OOD tasks, ICL performs similarly to the corresponding model of the pretraining function class optimized by GD given enough in-context examples. This suggests that ICL may lack the capability to learn new functions.
• Fig. 2 (newly added): The OOD performance cannot improve with training on more diverse (up to $\sim 2^{20}$) ID vectors.
• Fig. 3: A pretrained Llama-3-8b tends to make in-distribution predictions in a natural-language task.
• Fig. 4: The retrieval ability can emerge from training on retrieval tasks with more diverse labels.
• Fig. 5: ICL can perform "predict-then-retrieve" tasks with OOD labels after training on similar tasks with diverse labels. This offers insights into how real-world LLMs develop the capability to perform abstract label classification.
• Fig. 6: In the "predict-then-retrieve" task, ICL succeeds only if the prediction rule is ID. Otherwise, the model fails even if the target label appears in the context.
• Fig. 7: A pretrained Llama-3-8b cannot solve the OOD synthetic word classification task, indicating the limits of real-world LLMs in solving OOD tasks through ICL.
• Fig. 8: Algorithm selection in ICL broadly happens on OOD tasks.
• Fig. 9: The selection of pretraining priors on OOD tasks depends on the test error of the pretraining functions and the distance between training and test inputs.

---

Q2: The theoretical results heavily borrow from [1] (but the paper makes this very transparent)

A2: We acknowledge that we adopted the theoretical framework from [1] in our main Theorem 4.4 and have provided all necessary credits to [1] in our paper. However, some points need to be clarified:

• Different Scopes of Analysis: Their analysis and ours cover different aspects of ICL. [1] focuses on "task learning" and "task recognition" when both the training and testing functions are linear regression (an in-distribution setting), whereas our analysis concentrates on how ICL selects pretraining priors when faced with arbitrary OOD functions.
• Non-Trivial Extension of Theory: While we adopted the closed-form ICL prediction from [1], extending it to derive how the interaction between the training and test distributions affects prior selection is non-trivial.
• Novel Insights: We extract novel insights from the existing analysis of [1] (Lemma 4.3 in our paper): ICL prefers pretraining priors with input distributions closer to the test context. However, [1] does not explicitly explore this aspect.

---

Q3: 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.

A3: Following your advice, to show that the failure in the synthetic word classification task is mainly due to its OOD nature instead of some other factors that make it difficult to learn, we add an experiment in Appendix B.2 in the updated PDF, in which we train a GPT-2 to perform the same task. We find that the task can be well addressed after training. In this task, the $\boldsymbol{x}\_i$ and $\boldsymbol{y}\_i$ are generated in the same way as Sec. 3.3. The only modification is that we use a smaller predefined vector embedding $E'\in \mathbb{R}^{10000\times 20}$ instead of the $E\_{llama}\in \mathbb{R}^{32000\times 4096}$ used in the experiment in Sec. 3.3.

Q4: What exactly is shown in the last row of both of these figures?

A4: We apologize for the lack of clarity. The final rows of original Fig. 8 and 9 (Fig. 9a and 9b in the updated PDF) present the test error of the pretrained ICL model using the closed-form prediction derived in Lemma 4.2., i.e. $\|\left\langle \boldsymbol{x}\_{i+1}, \sum\_{m=1}^M \tilde{\pi}\_m \tilde{\boldsymbol{w}}\_m\right\rangle-y\_{T+1}^*\|^2$.

---

Q5: Why are the losses so high here throughout?

A5: The high test loss values can be attributed to the following reasons:

• The test error is calculated using the squared loss without taking the square root.
• In Figure 8, the absolute values of the centers of the pretraining linear regression weights are large (e.g., $[\pm5,\pm5,\pm5]$), while the test ReLU NN weights are sampled from $\mathcal{N}(0,1)$, which results in relatively smaller magnitudes. This significant numerical discrepancy between the training and test function parameter distributions likely amplifies the gap in their outputs. A similar explanation applies to Figure 9, where the input distributions of the training and test sets differ significantly.

---

Q6: 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.

A6: Explanation: According to Lemma 4.2, the closed-form ICL prediction $M^*\left(\mathcal{S}\_i \oplus \boldsymbol{x}\_{i+1}\right)=\left\langle \boldsymbol{x}\_{i+1}, \sum\_{m=1}^M \tilde{\pi}\_m \tilde{\boldsymbol{w}}\_m\right\rangle$ is not the linear combination of the pretraining priors $\boldsymbol{w}\_m$. Instead, it is a combination of $\tilde{\boldsymbol{w}}\_m$, which represents the posterior of $\boldsymbol{w}\_m$ conditioned on the downstream context. As highlighted in [1], the shift from $\boldsymbol{w}\_m$ to $\tilde{\boldsymbol{w}}\_m$ reflects the "task learning" process of ICL, demonstrating that the model can refine its prediction when given the downstream context. This refinement explains why, in Fig. 9, the loss in the last row is lower than that in the first row. The loss in the first row is $\|\left\langle\boldsymbol{x}\_{T+1},\boldsymbol{w}\_m\right\rangle-y\_{T+1}^*\|^2$, while the last row measures the loss of using the refined $\tilde{\boldsymbol{w}}\_m$ to predict, i.e., $\|\left\langle \boldsymbol{x}\_{T+1}, \sum\_{m=1}^M \tilde{\pi}\_m \tilde{\boldsymbol{w}}\_m\right\rangle-y\_{T+1}^*\|^2$.

---

Q7: 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?

A7: Following your advice, we've rescaled the loss to log scale in Fig. 8 in the updated PDF (original Fig. 7). We can observe that the ICL performance is closest to the best pretraining function class.

---

Q8: Are the y-tick labels of Figure 7b+c correct

A8: Yes. The loss is large since the value of the labels of OOD functions cubic regression and linear+quadratic regression is high.

Dear Reviewer 7eJd,

We understand that you may be busy and might not have had the opportunity to review our rebuttal yet. With only three days remaining, we kindly remind you once again in the hope of discussing this further with you. We have invested significant time and effort into preparing a detailed response, as your feedback is incredibly important to us. Could you please take a moment to review it? Thank you so much for your time and consideration!

Thank you for your detailed response. I appreciate the additional experiment you ran and report in Appendix B.2 helping to resolve my concern on the task being generally too difficult. While I believe the presentation of the paper could be further improved by better guiding the reader through the many distinct sections and findings, I remain positive of this work and will maintain my score.

Minor point:

• Thank you for rescaling the loss to a log-scale in Figure 8. May I ask if there is a reason to not consistently do this throughout the paper?

Thank you for your careful reading and appreciation of our work! We'd like to address your concerns on the following points:

---

Response to Weakness 1: Our novel findings in Sec. 2: when evaluated on OOD tasks, ICL performs comparably to a corresponding model of the pretraining function class optimized by gradient descent (GD), given enough in-context examples. Additionally, we observe a double descent curve when the model is trained on (quadratic) linear regression and evaluated on OOD function classes. These findings suggest that ICL may only optimize within its pretraining function class rather than genuinely learn an OOD function.

Existing findings in Sec. 2: for ID tasks, given sufficient in-context examples, the test error can be minimized to near zero.

We have further clarified the novel contributions in Sec. 2 in the updated PDF. While merging Sec. 2 and Sec. 4 is a good suggestion, it would require significant restructuring of the paper, which may not align with the reviewing policy. Therefore, we maintain the current organization of the sections.

---

Response to Weakness 2: Sorry for the lack of clarity. Here we clarify the task setup of the retrieval task of Llama-2-7b (Fig. 6(b) in the original PDF, Fig. 13(b) in the new PDF). Concretely, we first randomly sample $C=10$ different vectors from $E_{llama}$ as $\boldsymbol{x}_i$ and compute $\boldsymbol{y}_i$ in the same way as the OOD classification task to get $S=[\boldsymbol{x}_1, \boldsymbol{y}_1,...,\boldsymbol{x}_C, \boldsymbol{y}_C]$. Then we repeat $S$ 10 times to construct the input sequence $S'=[S\oplus S\oplus...\oplus S]$, where $\oplus$ denotes concatenation operation. The goal is to predict the next token given a prefix of $S'$. To succeed in this task, the model has to retrieve the same token as the query token (the last $\boldsymbol{x}_i$ of $S'$) from the context and output $\boldsymbol{y}_i$. We add the clarification and the result of the same experiment using Llama-3-8b in Sec. 3.3. We move the original result of Llama-2-7b to Appendix B.3.

---

Response to Weakness 3 & Question 3: We do not consider the ability to retrieve OOD labels from the context as a genuine capability to learn OOD tasks because once the underlying task function is OOD, the model fails even if the target label appears in the context. This is supported by the experiment in Fig. 6 of the updated PDF (original Fig. 5). In this figure, the test function is OOD, so although the target label integer is likely to present in the context, the model is unable to retrieve it correctly.

Regarding your point that "models are capable of doing so if pre-trained on a diverse range of abstract label tasks," this is valid only when the underlying task is ID, as shown in Fig. 5 of the updated PDF (original Fig. 4). We have clarified it explicitly in Sec. 3.2.

---

Response to Weakness 4 and Question 4: The experiment in Fig. 6 (original Fig. 5) was designed to evaluate whether ICL can successfully address a task through retrieval when the ground-truth target label is highly likely to appear in the context but the underlying task function is OOD.

The usefulness of this experiment: It demonstrates that ICL can only solve a task by retrieving a similar label from the context if the underlying task is ID. When the task function is OOD, retrieval fails to provide a solution.

Reflections on real-world tasks: This finding highlights a 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 task is OOD and the underlying prediction rule is too complex for the LLM to learn.

---

Response to Question 1: Yes. In this experiment, we found that "2-layer ReLU NN" is consistently the most performant training function class across all four OOD test function classes, so we add the error of a 2-layer ReLU NN trained by GD. We find that the performance of the ReLU model trained by GD aligns well with the ICL performance of the GPT-2 trained on the same function class. This demonstrates that our findings in Sec. 2 still hold when the transformer is trained on a mixture of tasks.

---

Response to Question 2: According to the setup explained in Response to Weakness 2, the length of the repeated subsequence $S$ is 10, so only after 10 examples can the LLM correctly retrieve the target vector from the context.

We thank Reviewer KJfV for your careful reading and detailed comments! We'd like to address your concerns in the following points:

---

Q1: Presentation issues.

A1: Thank you for meticulously pointing out the presentation issues! We sincerely apologize for these oversights and would like to clarify the revisions we’ve made:

• In Eq. (1), yes, $f(\boldsymbol{x}\_i)$ should be $f(\boldsymbol{x}\_{i+1})$.
• In Fig. 2, "$y_i$ id" refers to the index of $y_i$, which corresponds to $I_{y_i}$. We've updated the title of Fig. 2 (Note: Fig. 2 is now Fig. 4 in the updated PDF.)
• Quotation marks: We’ve standardized all quotes by replacing single quotation marks (') with double quotation marks ("), ensuring consistency throughout the text.

---

Q2: I wonder if you have same number of training tasks for different range. ... Second, the setup is hard to understand, please improve the writing.

A2: Yes, all three models, regardless of their respective training ranges for $I_{y_i}$, were trained on 200,000$\times$64 sequences, where 200,000 represents the number of training steps and 64 corresponds to the batch size. The performance gain stems from the model's exposure to more retrieval rules (or equivalently, a larger variety of label vectors).

We sincerely apologize for any inconvenience caused during your reading. We have rewritten Sec. 3.1 to improve clarity and readability, with the revisions highlighted in blue.

---

Q3: Are previous works using similar design? If yes, please cite. If no, could you clarify why this setup?

A3: To the best of our knowledge, we haven't seen a similar design. The motivation behind repeating a similar synthetic word classification task in Fig. 6(a) (now Fig. 7) is to explore whether the observed phenomenon—that ICL struggles to learn tasks from an unseen distribution—applies to real-world pretrained LLMs.

---

Q4: 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?

A4: We acknowledge that it is difficult to determine whether a task was encountered during the pretraining stage of Llama-2. However, Llama-2's pretraining primarily focuses on natural language tasks rather than specialized mathematical or vector-based tasks. The Llama-2 paper [1] also does not mention training on tasks involving learning linear classifiers or vector-to-vector transformations. Therefore, it is likely that the task in Section 3.3—which requires learning a manually generated linear classification function—was unseen during the pretraining of Llama-2.

---

Q5: Limited evaluation setups.

A5: Following your advice, we add two natural-language experiments with Llama-3-8b to support our findings. Please refer to Sec. 2.3 and Sec. 4.3 of the updated PDF for details.

1. In Section 2.3, we demonstrate how the tendency of ICL to perform in-distribution (ID) predictions manifests in real-world LLMs. To this end, we designed a task involving predicting labels with letters reversed. The basic tasks include outputting antonyms, translating from English to French, etc. All the letters of the original labels are reversed. We found that in this task, a pretrained Llama-3-8b tends to output the reversed result of the query word rather than first predicting the correct label and then reversing it. Although both reversal tasks are uncommon, directly outputting the reversed version of a word is relatively more common than first reasoning and then outputting the reversed prediction. Therefore, this result reflects to some extent that LLMs, when performing ICL, are more inclined to make in-distribution predictions.
2. In Section 4.3, the LLM is required to perform an ambiguous sentence classification task through ICL. Each sentence can be classified based on one of three aspects: sentiment, type, or location. For example, the sentence "The groundbreaking discovery made by Japanese scientists has revolutionized renewable energy" can be categorized as ("positive", "science", "Asia"). In each ICL sequence, the label words are chosen from one of these aspects and are replaced with abstract strings. For instance, "positive," "neutral," and "negative" are mapped to "RqF," "IwZ," and "SdK," respectively. We observe that the LLM can identify the most appropriate classification criterion, verifying the algorithm-selection mechanism of ICL.

---

[1] Llama 2: Open Foundation and Fine-Tuned Chat Models

Dear Reviewer KJfV,

We understand that you may be busy and might not have had the opportunity to review our rebuttal yet. With only three days remaining, we kindly remind you once again in the hope of discussing this further with you. We have invested significant time and effort into preparing a detailed response, as your feedback is incredibly important to us. Could you please take a moment to review it? Thank you so much for your time and consideration!

Dear reviewer KJfV,

I hope this email finds you well. I wanted to kindly follow up regarding our previous response to your comments. If there are any further questions or clarifications needed, we would be more than happy to provide additional information or explanations to assist in the review process. Thank you for your time and consideration.

Dear Reviewer KJfV,

This is a gentle reminder that there are only three days left before the discussion period concludes. Your input is incredibly valuable to us.

In our previous responses, we have thoroughly addressed all the concerns you raised and have also included two carefully designed natural language experiments to further support our understanding of ICL's mechanism in OOD scenarios.

We sincerely hope to hear your thoughts and feedback on these updates. Thank you again for your time and effort in reviewing our work!