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

## Abstract

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

## 1 Introduction

Pretrained large language models (LLMs) can perform in-context learning (ICL) (Brown, ), where providing a few examples of input-output pairs and a query example improves the model’s ability to generate the desired output, compared to zero-shot predictions. Understanding ICL’s ability 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. (); Raventós et al. (); Zhang et al. (); Akyürek et al. () found that ICL can learn new linear regression weights after pretraining on a large set of weight vectors. Moreover, Pan (); Kossen et al. (); Vacareanu et al. () revealed that real-world LLMs like Llama-2 (Touvron et al., ) and GPT-4 (Achiam et al., ) are capable of solving artificially constructed tasks likely unseen in their pretraining data, such as a classification task with abstract labels.

However, another line of research (Yadlowsky et al., ; Ahuja & Lopez-Paz, ) has raised a contrasting view, showing that ICL struggles to generalize to OOD tasks where there are distributional shifts in either the input distribution $P(X)$ or the input-label mapping $P(Y|X)$. These findings raise several important questions:

> *Can ICL really learn new input-output mappings from the context? What underlying mechanism of ICL determines its performance on OOD tasks?*

In this work, we aim to consolidate previous findings by addressing these questions. First, we 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 functions encountered during pretraining, which could explain its failure on OOD tasks that significantly deviate from the training distribution. Furthermore, we reproduce the widely observed phenomenon that ICL can perform classification with abstract labels. We find that solving such tasks requires retrieving similar labels from the context, a capability that can be acquired through pretraining on analogous tasks. This implies that success in such tasks of ICL may not indicate 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. Building on the work of Lin & Lee (), we also provide a comprehensive theoretical framework for understanding the ICL mechanism.

Our contributions are summarized as follows:

- 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 2).
- 2.
We further investigate ICL’s ability to perform classification with unseen abstract labels. Although this appears to be evidence of ICL learning OOD tasks, 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 3.1) and only emerges when the testing function is in distribution (Section 3.2). Additionally, we find that a pretrained Llama2-7b fails to learn OOD functions through ICL in a synthetic word classification task (Section 3.3), further confirming ICL’s limitations in OOD scenarios.
- 3.
Finally, we explore the ICL’s behavior when trained on multiple tasks, and observe that the algorithm selection mechanism broadly occurs in OOD scenarios. We theoretically prove the low-test-error preference of ICL prediction, i.e., the ICL prediction prefers to implement the pretraining function with lower test error (Section 4.1). We also validate our theory with numerical experiments (4.2).

## 2 Transformers Implements Functions Classes Seen During Pretraining

Exploring the ICL performance on OOD tasks. To investigate the ICL performance on new tasks that are unseen during training, following Garg et al. (), we train a GPT-2 (Radford et al., ) from scratch on some simple functions and evaluate it on functions different from the training ones. Denote the Transformer model parameterized by $\theta$ as $M_{\theta}$. The pretraining objective is:

$$ $\min_{\theta}\frac{1}{T}\sum_{i=1}^{T}\mathbb{E}_{f\sim\mathcal{F}}[\|M_{ \theta}(\mathcal{S}_{i}\oplus\bm{x}_{i+1})-f(\bm{x}_{i})\|^{2}_{2}],$ (1) $$

where $\mathcal{S}_{i}=[\bm{x}_{1}\oplus y_{1}\oplus\bm{x}_{2}\oplus y_{2}\oplus...
\oplus\bm{x}_{i}\oplus y_{i}]\in\mathbb{R}^{d\times 2i}$ is the context of length $i$, $\oplus$ denotes concatenation. $\bm{x}_{i}\in\mathbb{R}^{d}$ are sampled from a standard normal distribution $\mathcal{N}(0,1)$ with dimension $d=20$. Let $y_{i}=f(\bm{x}_{i})$ represent the labels, with $\mathcal{F}$ denoting the hypothesis class to which $f$ belongs. We train three separate GPT-2 models on three different function classes $\mathcal{F}$: linear regression, quadratic regression (element-wise square followed by linear regression), and a 2-layer ReLU network (detailed descriptions are in Appendix A.1). We then evaluate their ICL performance on these three tasks. Note that even when the testing and training functions are i.i.d. sampled from the same task, the specific instances of the testing functions remain unseen during training. For comparison, we also train models within the corresponding $\mathcal{F}$ with gradient descent (GD) using the testing in-context examples (details in Appendix A.1).

Observations. We plot the test error on the three tasks in Figure 1 and observe that: 1) when evaluated on the same task $\mathcal{F}$ as pretraining, ICL can reach near-zero test error, as observed in previous works (Garg et al., ). 2) when evaluated on a new task, ICL performs similarly to the corresponding model of the pretraining function class optimized by GD, given enough in-context examples. This indicates that the ICL prediction implicitly implements function classes seen during pretraining. 3) The models trained on linear and quadratic regression exhibit a double descent error curve (Nakkiran, ), characterized by a high error when given exact $d$ examples and evaluated on a new task, which has been theoretically and empirically revealed under the noisy linear regression setting by Nakkiran () and Garg et al. (), respectively. This further demonstrates that ICL implements the linear regression pretraining task, as the double descent curve is a distinctive phenomenon unique to linear regression models.

Figure: (a)
Refer to caption: x1.png

Existing theoretical evidence. One recent work (Zhang et al., ) theoretically proved that a one-layer linear self-attention model (LSA, defined in Appendix B.1) pretrained on a linear regression task will still implement the linear predictor given downstream in-context examples of arbitrary new function classes, under some assumptions on the initialization of the Transformer weight matrices. We restate the Theorem 4.2 of Zhang et al. () as Lemma 2.1 below:

###### Lemma 2.1 .

(Theorem 4.2 of Zhang et al. (), informal) Let $\mathcal{D}$ be a distribution over $(\bm{x},y)\in\mathbb{R}^{d}\times\mathbb{R}$, whose marginal distribution on $x$ is $\mathcal{D}_{x}=$ $\mathcal{N}(0,\Lambda)$. Assume the test prompt is of the form $P=\left(\bm{x}_{1},y_{1},\ldots,\bm{x}_{T},y_{T},\bm{x}_{\text{query }}\right)$, where $\left(\bm{x}_{i},y_{i}\right),\left(\bm{x}_{\text{query }},y_{\text{query }}
\right)\stackrel{{\scriptstyle\text{ i.i.d. }}}{{\sim}}\mathcal{D}$. The prediction risk on the test query $y_{query}$ of an arbitrary task satisfies:

$$ $\mathbb{E}\left(\widehat{y}_{\text{query }}-y_{\text{query }}\right)^{2}= \underbrace{\min_{\bm{w}\in\mathbb{R}^{d}}\mathbb{E}\left(\left\langle\bm{w}, \bm{x}_{\text{query }}\right\rangle-y_{\text{query }}\right)^{2}}_{\text{Error of best linear predictor }}+const,$ $$

where $const$ is a constant depending on the downstream context, and the expectation is over $\left(\bm{x}_{i},y_{i}\right),\left(\bm{x}_{\text{query }},y_{\text{query }}
\right)\stackrel{{\scriptstyle\text{ i.i.d. }}}{{\sim}}\mathcal{D}$.

Lemma 2.1 serves as a shred of theoretical evidence that ICL can just implement the pretraining function class, while the role of the context examples is to optimize the model within the pretraining hypothesis space. We leave Assumption B.1 that Lemma 2.1 depends on in Appendix B.1.

###### Lemma 2.1 .

## 3 Learning Abstract Labels May Not Be A New-task-learning Ability

### 3.1 Classification Tasks with Unseen Abstract Labels

Recent works (Pan et al., ; Kossen et al., ) have shown that LLMs can perform classification tasks in which the labels are ”abstract symbols” with no semantic meaning. For instance, in the SST-2 binary classification task, the labels ’positive’ and ’negative’ are substituted with abstract terms like ’foo’ and ’bar’, respectively. These tasks are likely not seen during pretraining. Pan et al. () refer to this ability of ICL to perform such classification as ”task learning” (TL). In this section, we explore whether the TL ability really reflects a new-task-learning capability of ICL or if it merely stems from the model having learned similar tasks during pretraining.

The retrieval ability can be gained by pretraining on a retrieval task with diverse input-label mappings. The classification of abstract labels can be approached by first retrieving an example with semantics similar to the query and then outputting the label of that example, as empirically demonstrated in previous research (Wang et al., ; Yu & Ananiadou, ). Therefore, the retrieval ability is a crucial prerequisite for performing abstract-label classification. To investigate whether ICL’s retrieval capability can emerge from training on similar tasks, we construct a retrieval task. Specifically, we generate a predefined word embedding $E\in\mathbb{R}^{N\times d}$ and randomly sample $\bm{x}_{i}\in\mathbb{R}^{d}$ from the first 5 rows of $E$. Let the index of $\bm{x}_{i}$ in $E$ be $I_{\bm{x}_{i}}$ (where $I_{x_{i}}\in[0,5)$). The labels $y_{i}$ are generated by mapping $I_{\bm{x}_{i}}$ to new indices $I_{y_{i}}\in[N]$ according to a mapping rule $I_{y_{i}}=R_{s}(I_{\bm{x}_{i}})=I_{\bm{x}_{i}}+s$ (with $s\in\mathbb{N}$), and then taking out $E_{I_{y_{i}}}$ as $y_{i}$. All in-context examples in a sequence share the same mapping rule $R_{s}$. To accomplish this task, the model must retrieve the same token as the query example from the context and output its subsequent token.

We train three models with three different ranges of $s$: $s\sim\mathcal{U}(50,150)$, $s\sim\mathcal{U}(50,250)$, and $s\sim\mathcal{U}(50,450)$ and evaluate on $s\sim\mathcal{U}(50,150)$, $s\sim\mathcal{U}(10,20)$, and $s\sim\mathcal{U}(500,600)$, where $\mathcal{U}$ denotes the uniform distribution. We plot the test error in Figure 2.

Figure: (a) Eval $s\sim\mathcal{U}(50,150)$
Refer to caption: x4.png

Observations. As observed in Figure 2, when the testing labels are in-distribution (left), all three models perform well. When testing labels are OOD, the more label tokens (random mappings) seen during training, the better the ICL performance is. This demonstrates that the retrieval ability can be obtained by training on a large number of similar retrieval tasks, even though the testing label tokens are unseen. This result may also provide new insights into how real-world LLMs acquire the in-context retrieval ability: by autoregressively pretraining on a huge corpus that may contain retrieval tasks, such ability emerges.

To further validate that the retrieval ability is evoked after trained on more random mappings, following Crosbie & Shutova (), we construct another retrieval task and visualize the prefix matching score of all attention heads of the three pretrained models in Figure 3. The prefix matching score is calculated by averaging the attention values from each token $t_{i}$ to the tokens after the same token as $t_{i}$ in earlier positions in the sequence, which correlates positively with the retrieval performance (Singh et al., ). In Figure 3, we observe that the model best at the retrieval task in Figure 2 exhibits more heads with high matching scores, further demonstrating it gains the retrieval ability by training on more retrieval sequences.

Figure: (a) PT $s\sim\mathcal{U}(50,150)$
Refer to caption: x7.png

The ability to perform linear regression and then retrieval can also be gained by pretraining on a similar task.
To further reproduce the emergence of the abstract label learning ability of real-world LLMs, we design a task that emulates the natural language classification with abstract labels like ’foo’ and ’bar’. The task function is defined as follows: $y_{i}=f(\bm{x}_{i})=E_{I_{\bm{x}_{i}}}$, where $I_{\bm{x}_{i}}=\text{floor}(0.4*(\bm{w}^{\top}\bm{x}_{i}))+s$, with $E$ being the predefined word embedding and $s\in\mathbb{N}+$ shared in the same sequence. Here, $\bm{x}_{i}$, $\bm{w}\sim\mathcal{N}(0,1)\in\mathbb{R}^{d}$. Each $\bm{x}_{i}$ is mapped to $y_{i}$ according to $\bm{w}$ and $s$. (^1^11In our experimental setup, given a sufficiently long context ($\approx 50$), the label of the query is highly likely to appear in the context, as the number of the possible classes is far less than the number of in-context examples.) In this task, estimating $\bm{w}$ and calculating $\bm{w}^{\top}\bm{x}_{i}$ simulates predicting the original label (’positive’ and ’negative’) based on the semantics in the natural language task, while retrieving the abstract labels from in-context examples that share the same $\text{floor}(0.4*(\bm{w}^{\top}\bm{x}_{i}))$ as the query from the context resembles identifying the abstract labels (’foo’ and ’bar’).

Again, we train three models on different ranges of mappings: $s\sim\mathcal{U}(100,200)$, $s\sim\mathcal{U}(100,1000)$, and $s\sim\mathcal{U}(100,2000)$, and evaluate on $s\sim\mathcal{U}(100,200)$, $s\sim\mathcal{U}(500,600)$, and $s\sim\mathcal{U}(3000,3100)$. The test error are plotted in Figure 4.

Figure: (a) Eval $s\sim\mathcal{U}(100,200)$
Refer to caption: x10.png

Observations. In Figure 4, the generalization ability to unseen labels also improves as the number of labels encountered during training increases. Notably, only the model trained with $s\sim\mathcal{U}(100,2000)$ performs well on the unseen labels. This suggests that as long as the LLM has been exposed to sufficiently many similar tasks with during training, it can effectively address arbitrary OOD labels retrievable from context through ICL. Therefore, the ability of ICL to perform abstract label classification may not serve as an evidence of learning new tasks.

### 3.2 Abstract Label Classification Can Only Be Achieved on In-Distribution Tasks

A retrieval task with OOD testing functions & observations. One might question whether, once a task can be solved through retrieval, ICL can generalize beyond the training function class. To investigate this, we replace the testing functions of linear regression with quadratic regression while preserving linear regression as the pretraining task. The results in Figure 5 show that the generalization doesn’t improves with training on more in-distribution functions. Combining observations from Figure 4, we conclude that only when the testing task is in distribution can ICL solve classification with unseen labels.

Figure: (a) Eval $s\sim\mathcal{U}(100,200)$
Refer to caption: x13.png

### 3.3 Real-world LLMs Mat Not Necessarily In-context Learn New Tasks

Evaluating Llama2 on an OOD synthetic word classification task. In this section, we assess whether real-world LLMs can tackle OOD tasks through ICL. We select the pretrained Llama-2-7b and evaluate it on a synthetic word classification task. To ensure the task is far from the pretraining distribution, we randomly sample $\bm{x}_{i}\in\mathbb{R}^{d}$ from the word embedding of Llama-2-7b (denoted as $E_{llama}$) and generate random linear mappings $\bm{w}\in\mathbb{R}^{C\times d}$ as task functions (where $C=10$). The label words are created by mapping $\bm{x}_{i}$ to one of the ten label vectors in $E_{llama}$ using $\bm{w}$. Experimental details are in Appendix A.3. To complete this task, the model must learn $\bm{w}$ in context.

For comparison, we also evaluate the ICL performance of Llama-2-7b on a retrieval version of this task, where we sample $C$ classes of tokens from $E_{llama}$ to generate random sequences. The objective is to retrieve the same tokens as the query token from the context and output its subsequent token. The results of these two tasks are presented in Figure 6.

Figure: (a)
Refer to caption: x16.png

Observations. From Figure 6, we observe that the ICL performance on the synthetic classification task is close to random guessing (10% accuracy), while performance on the retrieval task is significantly better. Considering that the input and label distributions of the two tasks are very similar, we have reason to believe that LLMs struggle to learn new input-output mappings from context; instead, ICL appears to be more adept at retrieval tasks.

## 4 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. (); Yadlowsky et al. () 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, we will show that this algorithm-selection phenomenon of ICL persists even when evaluated on OOD tasks, regardless of the distribution of the testing functions, and provide a theoretical characterization of the algorithm-selection mechanism.

The Model pretrained on a single task vs. the model pretrained on multiple tasks. In Figure 7, we compare the performance of GPT-2 models trained on a single task—linear regression (LR), quadratic regression (QR), 2-layer ReLU network (ReLU NN) regression—against the model trained on all three tasks when encountering four kinds of OOD tasks. The results are in Figure 7.

Figure: (a)
Refer to caption: x18.png

Observations. From Figure 7, we observe that the ICL performance of the model trained on mixed tasks is comparable to the performance of the model trained on a single task with the lowest test error on the evaluation task. ICL seems to automatically select the best pretraining functions according to the downstream context.

### 4.1 Theoretical Analysis

A mixed Gaussian pretraining dataset of multiple tasks. In this section, we theoretically analyze the algorithm selection mechanism of ICL on OOD tasks, based on the theoretical framework of Lin & Lee (). Consider a noisy linear regression pretraining dataset with the inputs and task weights following the mixed Gaussian distribution:

###### Assumption 4.1 .

(Mixed Gaussian pretraining data) Input distribution: $P(\bm{x}|\bm{\mu})=\mathcal{N}(\bm{x}|\bm{\mu},\sigma_{x}^{2}\bm{I})$, label distribution: $P(y|\bm{x},\bm{w})=\mathcal{N}\left(y|\langle\bm{x},\bm{w}\rangle,\sigma_{y}^{
2}\right)$. The input means and task weights are sampled from a mixed Gaussian distribution: $P(\bm{\mu},\bm{w})=\sum_{m=1}^{M}\pi_{m}\mathcal{N}(\bm{\mu};\bm{\mu}_{m},
\sigma_{\mu}^{2}\bm{I})\cdot\mathcal{N}\left(\bm{w};\bm{w}_{m},\sigma_{w}^{2}
\bm{I}\right)$, where $\sum_{m=1}^{M}\pi_{m}=1$, $0<\pi_{m}<1$ and $\left\|\bm{\mu}_{m}\right\|=\left\|\bm{w}_{m}\right\|=1,\forall m$. Define $\delta_{\mu}=\frac{\sigma^{2}_{\mu}}{\sigma^{2}_{x}}$ and $\delta_{w}=\frac{\sigma^{2}_{w}}{\sigma^{2}_{y}}$. Each training sequence $\mathcal{S}_{T}=[\bm{x}_{1}\oplus y_{1}\oplus...\oplus\bm{x}_{T}\oplus y_{T}]$ is constructed by first sampling the input mean and the task weight according to $P(\bm{\mu},\bm{w})$ and then sampling $\bm{x}_{i}$ and $y_{i}$ according to $P(\bm{x}|\bm{\mu})$ and $P(y|\bm{x},\bm{w})$, respectively. Denote this pretraining distribution as $P_{tr}$.

The lemma below states that the closed-form prediction of the model trained on the pretraining data under Assumption 4.1, given the testing context, remains a Gaussian mixture of the reweighted pretraining task weights.

###### Lemma 4.2 .

(Corollary 2. of Lin & Lee (), closed-form ICL prediction of the pretrained model) Denote the model $M^{*}$ that minimizes the risk on the pretraining data of Assumption 4.1, i.e., $M^{*}\in\arg\min\frac{1}{T}\sum_{i=0}^{T-1}\underset{S_{i}\sim P_{tr}}{\mathbb
{E}}\left[\left\|M\left(\mathcal{S}_{i}\oplus\bm{x}_{i+1}\right)-y_{i+1}\right
\|^{2}\right]$, then the prediction on any sequence $\mathcal{S}_{i}\oplus\bm{x}_{i+1}$ by $M^{*}$ is as follows: $M^{*}\left(\mathcal{S}_{i}\oplus\bm{x}_{i+1}\right)=\left\langle\bm{x}_{i+1},
\sum_{m=1}^{M}\tilde{\pi}_{m}\tilde{\bm{w}}_{m}\right\rangle$.
where $\tilde{\pi}_{m}$, and $\tilde{\bm{w}}_{m}$ depending on both the pretraining task and the downstream context example are given in Lemma 1 of Lin & Lee ().

Based on the closed-form ICL prediction, we now analyze how the downstream context affects $\tilde{\pi}$, which determines how ICL selects the pretraining functions. First, we introduce Lemma 4.3 that characterizes the ratio of the reweighted weight of two pretraining tasks:

###### Lemma 4.3 .

(Appendix H.1 of Lin & Lee ()) Consider any two different pretraining component $\alpha$ and $\beta$, given a testing context $\mathcal{S}_{T}\oplus\bm{x}_{T+1}$ and the well-pretrained model $M^{*}$, the ratio between the weights of the two task priors $\tilde{\pi}_{\alpha}/\tilde{\pi}_{\beta}$ in $M^{*}$’s ICL prediction can be decomposed into two terms: $\tilde{\pi}_{\alpha}/\tilde{\pi}_{\beta}=\frac{\pi_{\alpha}}{\pi_{\beta}}\exp
\left(\Psi_{\bm{\mu}}(\alpha,\beta)+\Psi_{\bm{w}}(\alpha,\beta)\right)$, where

$$ $\displaystyle\Psi_{\bm{\mu}}(\alpha,\beta)=\left(\sum_{i=1}^{T+1}\left\|\bm{ \mu}_{\beta}-\bm{x}_{i}\right\|^{2}-\sum_{i=1}^{T+1}\left\|\bm{\mu}_{\alpha}- \bm{x}_{i}\right\|^{2}\right)/\left(2\sigma_{x}^{2}\left(1+(T+1)\delta_{\mu} \right)\right).$ (2) $$

Further, assuming the testing in-context examples $\bm{x}_{i}\sim\mathcal{N}(\bm{\mu}^{*},\tau_{x}^{2}\bm{I})$, if $\|\bm{\mu}_{\beta}-\bm{\mu}^{*}\|^{2}-\|\bm{\mu}_{\alpha}-\bm{\mu}^{*}\|^{2}\geq
0$ holds, then as the context length $T\rightarrow\infty$, the first term $\Psi_{\bm{\mu}}(\alpha,\beta)\rightarrow(\|\bm{\mu}_{\beta}-\bm{\mu}^{*}\|^{2}
-\|\bm{\mu}_{\alpha}-\bm{\mu}^{*}\|^{2})/2\sigma^{2}_{\mu}\geq 0$.

However, Lin & Lee () didn’t analyze how the second term $\Psi_{\bm{w}}(\alpha,\beta)$ 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 $\Psi_{\bm{w}}(\alpha,\beta)$ converges to a non-negative value when the pretraining function class $\alpha$ performs better on the downstream context than $\beta$.

###### Theorem 4.4 .

(ICL prediction favors the pretraining function with low error on the context) Given the context $\bm{S}_{T}$, if the empirical risk of implementing a function of the pretraining task $\alpha$ is less than that of $\beta$, i.e., $\frac{1}{T}\sum_{i=1}^{T}|\bm{w}_{\beta}\bm{x}_{i}-y_{i}|^{2}-|\bm{w}_{\alpha}
\bm{x}_{i}-y_{i}|^{2}\geq 0$,
then, under some mild Assumptions C.2 on the distribution of $\bm{S}_{T}$, we have $\Psi_{\bm{w}}(\alpha,\beta)\geq 0$.

Combining Lemma 4.3, if the downstream inputs $\bm{x}_{i}$, $\bm{x}_{i}\sim\mathcal{N}(\bm{\mu}^{*},\tau_{x}^{2}\bm{I})$ and $\|\bm{\mu}_{\beta}-\bm{\mu}^{*}\|^{2}-\|\bm{\mu}_{\alpha}-\bm{\mu}^{*}\|^{2}\geq
0$ hold, then as $T\rightarrow\infty$, we have $\tilde{\pi}_{\alpha}/\tilde{\pi}_{\beta}\geq\pi_{\alpha}/\pi_{\beta}$.

Summary of the algorithm-selection mechanism.
4.3 and Theorem 4.4 together elucidate the algorithm-selection mechanism of ICL. According to Lemma 4.2, the ICL prediction of the model pretrained on the mixed Gaussian data will be a reweighted combination of the pretraining task vectors $\bm{w}_{i}$. Whether the ratio between the weights of two pretraining tasks, $\tilde{\pi}_{\alpha}/\tilde{\pi}_{\beta}$, given a downstream context, exceeds the original ratio $\pi_{\alpha}/\pi_{\beta}$ depends on two factors: 1) whether the pretraining input distribution of $\alpha$ is closer to the downstream input distribution than that of $\beta$; 2) whether the task function of $\alpha$ induces lower test error in downstream context than that of $\beta$. When both conditions are met, we have $\tilde{\pi}_{\alpha}/\tilde{\pi}_{\beta}\geq\pi_{\alpha}/\pi_{\beta}$, indicating that ICL prefers $\alpha$ over $\beta$ in its predictions.

Discussion of the setup of our theory.
Notably, our theoretical result doesn’t assume a Transformer architecture, while previous theoretical works of understanding ICL often adopt Transformers with oversimplified assumptions on their parameters or structures (Ahn et al., ; Zhang et al., ; Huang et al., ; Collins et al., ). Additionally, our analysis shows that models pretrained on the ICL tasks can implement algorithm selection during ICL inference following Lin & Lee (). In contrast, prior work on algorithm selection (Bai et al., ) only shows that a specific set of parameters in a simplified ReLU Transformer can enable algorithm selection. However, the parameter construction is complex and somewhat tricky, and there is no theoretical or experimental guarantee that Transformers exhibiting algorithm selection will necessarily implement these parameters.

###### Assumption 4.1 .

###### Lemma 4.2 .

###### Lemma 4.3 .

###### Theorem 4.4 .

### 4.2 Empirical Validation of the Algorithm-Selection Mechanism of ICL

Now we validate our theoretical findings regarding ICL’s algorithm-selection mechanism in OOD tasks by conducting numerical experiments following Lin & Lee (). In Figure 8 and 9, the training data is a Gaussian mixture with four components (see Assumption 4.1), while the test function is a two-layer ReLU network (Appendix A.1). Both the training and the test data are in ICL format. We compute the test error of using each pretraining task function to predict the downstream function (the first row of Figure 8 and 9), the weights for each pretraining function during ICL inference (the second row), and the test error of the pretrained ICL model (the third row). We evaluate five different noise levels ($\delta_{x}=\delta_{w}\in\{1/81,9/1,1,9,81\}$) and consider two settings described below.

Low-test-error preference of ICL. To validate Theorem 4.4, we ensure that the distributional distances between the inputs of each training task and the test data remain consistent. Specifically, all $\bm{x}_{i}$ in both training data and test data are sampled from $\mathcal{N}([0,0,0]^{\top},\sigma^{2}_{x}\bm{I})$. The task weights for different pretraining tasks vary, as detailed in the top half of Table 1. In this setup, only the test error of the pretraining functions influences algorithm selection. From the top two rows of Figure 8, we can observe a clear negative correlation between the ICL performance and the test error of the task weight. This result supports Theorem 4.4, confirming that ICL prefers the pretraining functions with a low test error in the downstream context. Also, it’s consistent with the observations in Figure 7.

**Table 1: Experiment setting of Figure 8 and Figure 9. ’PT’ and ’DS’ are short for ’pretraining’ and ’downstream’, respectively.**
| Experiment | DS inputs | PT task id | PT input distribution | PT task functions | PT-DS input distance |
| --- | --- | --- | --- | --- | --- |
| Figure 8 | $\mathcal{N}([0,0,0]^{\top},\sigma^{2}_{x}\bm{I})$ | 1 | $\mathcal{N}([0,0,0]^{\top},\sigma^{2}_{x}\bm{I})$ | $\mathcal{N}([5,5,5]^{\top},\sigma^{2}_{w}\bm{I})$ | 0 |
| 2 | $\mathcal{N}([0,0,0]^{\top},\sigma^{2}_{x}\bm{I})$ | $\mathcal{N}([-5,5,5]^{\top},\sigma^{2}_{w}\bm{I})$ | 0 |  |  |
| 3 | $\mathcal{N}([0,0,0]^{\top},\sigma^{2}_{x}\bm{I})$ | $\mathcal{N}([-5,5,-5]^{\top},\sigma^{2}_{w}\bm{I})$ | 0 |  |  |
| 4 | $\mathcal{N}([0,0,0]^{\top},\sigma^{2}_{x}\bm{I})$ | $\mathcal{N}([-5,-5,-5]^{\top},\sigma^{2}_{w}\bm{I})$ | 0 |  |  |
| Figure 9 | $\mathcal{N}([-4,-4,-4]^{\top},\sigma^{2}_{x}\bm{I})$ | 1 | $\mathcal{N}([5,5,5]^{\top},\sigma^{2}_{w}\bm{I})$ | $\mathcal{N}([1,1,1]^{\top},\sigma^{2}_{w}\bm{I})$ | 15.59 |
| 2 | $\mathcal{N}([-5,5,5]^{\top},\sigma^{2}_{w}\bm{I})$ | $\mathcal{N}([1,1,1]^{\top},\sigma^{2}_{w}\bm{I})$ | 12.77 |  |  |
| 3 | $\mathcal{N}([-5,5,-5]^{\top},\sigma^{2}_{w}\bm{I})$ | $\mathcal{N}([1,1,1]^{\top},\sigma^{2}_{w}\bm{I})$ | 9.11 |  |  |
| 4 | $\mathcal{N}([-5,-5,-5]^{\top},\sigma^{2}_{w}\bm{I})$ | $\mathcal{N}([1,1,1]^{\top},\sigma^{2}_{w}\bm{I})$ | 1.73 |  |  |

Figure: Figure 8: Empirical validation of the low-test-error preference of ICL. The test error of the four pretraining functions in this experiment varies (see the first row), while the distributional distances between the inputs of each training task and the test data remain consistent. In the first two rows, a clear negative correlation between the test error of pretraining task $i$ and the task weight $\tilde{\pi}_{i}$ can be observed.
Refer to caption: x22.png

Similar-input-distribution preference of ICL. We also empirically validate Lemma 4.3 in Figure 9. In this case, the distributional distances between the input of different pretraining tasks and that of the test context vary, while the test errors of different pretraining functions are almost the same (detailed setup is in the bottom half of Table 1). As shown in the bottom two rows in Figure 9, the task weight $\tilde{\pi}_{i}$ is positively correlated with the similarity between the training and testing input distribution. This is consistent with Lemma 4.3 which demonstrates that ICL prefers to select the pretraining function whose input distribution is close to the downstream one.

Figure: Figure 9: Empirical validation of the similar-input-distribution preference of ICL. The test error of the four pretraining functions remains almost consistent (see the first row), while the distributional distances between the inputs of the four training tasks and the test data vary. The distances of different tasks are ordered from largest to smallest as $1>2>3>4$. From the first two rows, we observe that the task weight $\tilde{\pi}_{i}$ is negatively correlated with the distance between the training and testing input distribution.
Refer to caption: x23.png

## 5 Comparison with Related Works

Abilities to learn new tasks of ICL.
Besides studies indicating that ICL can learn new weights of linear regression (Garg et al., ; Raventós et al., ; Zhang et al., ; Akyürek et al., ), other research has found that LLMs can tackle tasks that are unlikely to have been encountered during pretraining. For example, Pan () showed that LLMs perform better than random guessing on classification tasks with meaningless labels. Kossen et al. () demonstrated that ICL can identify authorship based on writing style in private communication messages not included in the pretraining corpus. Additionally, Vacareanu et al. () 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 encountered tasks from a similar distribution in their pretraining corpus. For instance, the LLMs could have been pretrained on a corpus containing authorship identification tasks or on statistical data encompassing different functions. Our work does not claim that ICL cannot generalize to new task instances; rather, it highlights the limitation in generalizing to an unseen input-label distribution.
Additionally, Yadlowsky et al. () finds that ICL struggles to generalize to testing function classes that are unseen during training (e.g., convex combinations or extreme versions of the pretraining functions). They didn’t delve into how ICL behaves on OOD data, while we reveal that it implements the pretraining functions.

The algorithm-selection mechanism of ICL. Recent works by Bai et al. (); Wang et al. () have uncovered the algorithm selection phenomenon, demonstrating that Transformers pretrained on both linear regression and classification tasks perform well when presented with the context of either task during ICL inference. Theoretically, they showed that a Transformer with specific parameters can achieve algorithm selection. Yadlowsky et al. () empirically found that ICL selects the optimal pretraining function class after observing in-context examples from a 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, we empirically and theoretically demonstrate 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.

## 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 a near-optimal solution within the pretraining function classes. However, ICL performs well in retrieval tasks where the shift in the input-label mapping is only caused by replacing the in-context label tokens with new ones while the underlying function distribution retains. We also examine ICL’s performance on OOD tasks after pretraining on multiple tasks. Our theoretical and empirical analysis reveals ICL’s preference for low-test-error functions, i.e., ICL tends to implement pretraining function classes with low test error in the test context. This finding, alongside previous work (Lin & Lee, ), highlights two key factors that determine how ICL will implement the prediction function based on the testing context and pretraining tasks: the distance between the training and testing input distributions, and the ability of a pretraining function to solve the test task.

## 7 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 3 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 for smaller models like Llama2-7b. 2) The models are 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.

## 8 Reproducibility

In the main text and Appendix A, we’ve stated all setups for reproducing our experimental results. For the theoretical part, we’ve included the assumptions (Assumption C.2) and proofs in Appendix C.

## Appendix A Experimental Details

### A.1 Experimental Details in Section 2 and Section 4.2

Definitions of the function classes. The function classes in Figure 1 and Figure 7 are:

- •
Linear regression: $y_{i}=\bm{w}^{\top}\bm{x}_{i}$, where $\bm{w}$, $\bm{x}_{i}\in\mathbb{R}^{d}$ and $\bm{w}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$.
- •
Quadratic regression: $y_{i}=\bm{w}^{\top}(\bm{x}_{i})^{2}$, where $\bm{w}$, $\bm{x}_{i}\in\mathbb{R}^{d}$ and $\bm{w}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$, $(\bm{x}_{i})^{2}$ denotes the element-wise square of $\bm{x}_{i}$.
- •
2-layer ReLU network regression: $y_{i}=\bm{w}_{1}^{\top}\text{ReLU}(\bm{w}_{2}\bm{x}_{i})$, where $\bm{w}_{1}\in\mathbb{R}^{d^{\prime}}$, $\bm{w}_{2}\in\mathbb{R}^{d^{\prime}\times d}$, and $\bm{x}_{i}\in\mathbb{R}^{d}$. $\bm{w}_{1}$, $\bm{w}_{2}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$.
- •
Square root linear regression: $y_{i}=\bm{w}^{\top}\sqrt{\bm{x}_{i}}$, where $\bm{w}$, $\bm{x}_{i}\in\mathbb{R}^{d}$ and $\bm{w}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$, $(\bm{x}_{i})^{2}$ denotes the element-wise square root of $\bm{x}_{i}$.
- •
Cubic linear regression: $y_{i}=\bm{w}^{\top}(\bm{x}_{i})^{3}$, where $\bm{w}$, $\bm{x}_{i}\in\mathbb{R}^{d}$ and $\bm{w}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$, $(\bm{x}_{i})^{2}$ denotes the element-wise cube of $\bm{x}_{i}$.
- •
Linear+quadratic regression: $y_{i}=\bm{w}_{1}^{\top}(\bm{x}_{i})^{2}+\bm{w}_{2}^{\top}\bm{x}_{i}$, where $\bm{w}_{1}$, $\bm{w}_{2}$, $\bm{x}_{i}\in\mathbb{R}^{d}$ and $\bm{w}_{1}$, $\bm{w}_{2}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$.
- •
2-layer Sigmoid network: $y_{i}=\bm{w}_{1}^{\top}\text{Sigmoid}(w_{2}\bm{x}_{i})$, where $\bm{w}_{1}\in\mathbb{R}^{d^{\prime}}$, $\bm{w}_{2}\in\mathbb{R}^{d^{\prime}\times d}$, and $\bm{x}_{i}\in\mathbb{R}^{d}$. $\bm{w}_{1}$, $\bm{w}_{2}$, $\bm{x}_{i}\sim\mathcal{N}(0,1)$.

Baseline models in Figure 1. The models of each pretraining hypothesis class are implemented by training a neural network that yields functions of that hypothesis class. For example, a linear regression weight $w$ can be implemented by a single linear layer. The models are optimized using SGD with learning rate 1e-3 for 1000 steps.

### A.2 Experimental Details in Section 3.1

We now provide additional details regarding the experiments of Figure 3
. Following Crosbie & Shutova (), we generated a dataset consisting of 100 sequences of random tokens, each containing repeated sub-sequences. The task is to 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.

### A.3 Experimental Details of Section 3.3

We uniformly sample 1000 word vectors $\bm{x_{i}}\in\mathbb{R}^{d}$ from the word embedding $E\in\mathbb{R}^{N\times d}$ of the pretrained Llama-2-7b, where $N=32000$ and $d=4096$. Then we generate a task weight $\bm{w}\in\mathbb{R}^{d^{\prime}\times C}$ that only takes the first $d^{\prime}$ dimensions of $\bm{x}_{i}$ (denoted as $\bm{x}_{i}[:d^{\prime}]$) to compute a probability distribution over $C$ classes: $p_{i}=\bm{x}_{i}[:d^{\prime}]^{\top}\bm{w}\in\mathbb{R}^{C}$. Thus, each $\bm{x}_{i}$ is mapped to one of the $C=10$ classes. Next, we sample the label vectors $y_{i}\in\mathbb{R}^{d}$ from $E$ according to the index of maximum probability: $y_{i}=E_{\arg\max_{j}p_{i}[j]+s}$, where $s\in\mathcal{N}+$ is a offset. We set $d^{\prime}=30\ll d=4096$ to reduce the complexity of the task. Hence, $\bm{x}_{i}$ are classified into $C$ labels words $E[s:s+C]$. The predicted token of $\bm{x}_{i}$ is computed as: $\arg\max_{j}\hat{p}_{i}[j],~{}j\in\{s,s+1,...,s+C-1\}$, where $\hat{p}_{i}$ is the output of the last linear layer of Llama-2-7b given $\bm{x}_{i}$.

## Appendix B Additional Theoretical Details

### B.1 Additional Details of Lemma 2.1

In this section, we provide the necessary details of the theoretical setting of Zhang et al. ().

The linear self-attention (LSA) model considered in the Theorem 4.2 of Zhang et al. () (Lemma 2.1) is defined as follows:

$$ $f_{\mathrm{LSA}}(E;\theta)=E+W^{PV}E\cdot\frac{E^{\top}W^{KQ}E}{N},$ (3) $$

where $E$ is the input embedding defined as follows:

$$ $E=E(P)=\left(\begin{array}[]{ccccc}\bm{x}_{1}&\bm{x}_{2}&\cdots&\bm{x}_{N}&\bm {x}_{\text{query }}\\ y_{1}&y_{2}&\cdots&y_{N}&0\end{array}\right)\in\mathbb{R}^{(d+1)\times(N+1)}.$ (4) $$

$W^{PV}$ is obtained by merging the projection and value matrices into a single matrix, and $W^{KQ}$ is attained by merging the query and key matrices into a single matrix. $N$ is the context length.

Now we present the assumption on the attention weights of the linear-attention model in Lemma 2.1.

###### Assumption B.1 .

(Assumption 3.3 in Zhang et al. (), initialization). Let $\sigma>0$ be a parameter, and let $\Theta\in\mathbb{R}^{d\times d}$ be any matrix satisfying $\left\|\Theta\Theta^{\top}\right\|_{F}=1$ and $\Theta\Lambda\neq 0_{d\times d}$. We assume

$$ $W^{PV}(0)=\sigma\left(\begin{array}[]{cc}0_{d\times d}&0_{d}\\ 0_{d}^{\top}&1\end{array}\right),\quad W^{KQ}(0)=\sigma\left(\begin{array}[]{ cc}\Theta\Theta^{\top}&0_{d}\\ 0_{d}^{\top}&0\end{array}\right)$ $$

The training objective is to minimize the population risk of the linear regression task:

$$ $L(\theta)=\lim_{B\rightarrow\infty}\widehat{L}(\theta)=\frac{1}{2}\mathbb{E}_{ w_{\tau},\bm{x}_{\tau,1},\cdots,\bm{x}_{\tau,N},\bm{x}_{\tau,\text{ query }}} \left[\left(\widehat{y}_{\tau,\text{ query }}-\left\langle w_{\tau},\bm{x}_{ \tau,\text{ query }}\right\rangle\right)^{2}\right],$ (5) $$

where $w_{\tau}\sim\mathcal{N}(0,I_{d})$, $\bm{x}_{\tau,i}$, $\bm{x}_{\tau,query}\sim\mathcal{N}(0,\Lambda)$, $\widehat{y}_{\tau,\text{ query }}$ is the prediction of the LSA model.

###### Assumption B.1 .

## Appendix C Proof of Theorem 4.4

We restate Theorem 4.4 as the Theorem C.1 below, and present the assumption it depends on.

###### Theorem C.1 .

(ICL prediction favors the pretraining function with low error on the context) Given the context $S_{k}$, if the empirical risk of implementing a function of the pretraining task $\alpha$ is less than that of $\beta$, i.e., $\frac{1}{T}\sum_{t=1}^{T}|\bm{w}_{\beta}\bm{x}_{i}-y_{i}|^{2}-|\bm{w}_{\alpha}
\bm{x}_{i}-y_{i}|^{2}\geq 0$,
then, under some mild Assumptions C.2, we have $\Psi_{\bm{w}}(\alpha,\beta)\geq 0$.

Combining Lemma 4.3, if the downstream inputs $\bm{x}_{i}$, $\bm{x}_{i}\sim\mathcal{N}(\bm{\mu}^{*},\tau_{x}^{2}\bm{I})$ and $\|\bm{\mu}_{\beta}-\bm{\mu}^{*}\|^{2}-\|\bm{\mu}_{\alpha}-\bm{\mu}^{*}\|^{2}\geq
0$ hold, then as $T\rightarrow\infty$, we have $\tilde{\pi}_{\alpha}/\tilde{\pi}_{\beta}\geq\pi_{\alpha}/\pi_{\beta}$.

###### Assumption C.2 .

(Assumption on the distribution of the downstream context examples.) Assume that: the minimum eigenvalue of the covariance matrix of any in-context example $\bm{x}_{i}$ satisfies $\lambda_{\text{min}}(\bm{x}_{i}\bm{x}_{i}^{\top})\geq 1$;
$(\bm{I}+T\delta_{w}\bm{I})(\bm{I}+\delta_{w}\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_{i}
^{\top})^{-1}=\bm{I}$;
$\frac{1}{T}\sum_{i=1}^{T}2(\bm{w}_{\alpha}-\bm{w}_{\beta})^{\top}\bm{x}_{i}y_{
i}\frac{1}{T}\sum_{j=1}^{T}\left(\bm{x}_{j}^{\top}\bm{x}_{i}\frac{y_{j}}{y_{i}
}-\bm{x}_{i}^{\top}\bm{x}_{i}\right)\geq 0$

###### Proof.

According to Lemma 1 of Lin & Lee (),

$$ $r(\alpha,\beta)=\frac{\tilde{\pi}_{\alpha}}{\tilde{\pi}_{\beta}}=\frac{\pi_{ \alpha}C_{0}c_{\alpha}^{\mu}c_{\alpha}^{w}}{\pi_{\beta}C_{0}c_{\beta}^{\mu}c_{ \beta}^{w}}=\frac{\pi_{\alpha}}{\pi_{\beta}}\exp\left(\Psi_{\mu}(\alpha,\beta) +\Psi_{\bm{w}}(\alpha,\beta)\right).$ (6) $$

In the Appendix H.1 of Lin & Lee (), they have proved that when the context length $T\rightarrow\infty$, under the first condition in Assumption C.2, $\lim_{T\rightarrow\infty}\Psi_{\mu}(\alpha,\beta)=\geq 0$.

Now we prove that when the empirical risk on the in-context example of pretraining task function $\alpha$ is no more than that of $\beta$, the second term $\Psi_{\bm{w}}(\alpha,\beta)\geq 0$.

$$ $\displaystyle\Psi_{\bm{w}}(\alpha,\beta)$ (7) $\displaystyle=\log\left(\frac{\exp\left(-\frac{\left\|\bm{w}_{\alpha}\right\|^ {2}-\left\|\bm{w}_{\alpha}+T\delta_{w}\overline{\bm{w}}\right\|_{\left(\bm{I}+ T\delta_{w}\bm{\Sigma}_{\bm{w}}\right)^{-1}}^{2}}{2\sigma_{w}^{2}}\right)}{ \exp\left(-\frac{\left\|\bm{w}_{\beta}\right\|^{2}-\left\|\bm{w}_{\beta}+T \delta_{w}\bm{w}\right\|_{\left(\bm{I}+T\delta_{w}\bar{\Sigma}_{\bm{w}}\right) ^{-1}}^{2}}{2\sigma_{w}^{2}}\right)}\right)$ $\displaystyle=\frac{\left\|\bm{w}_{\beta}\right\|^{2}-\left\|\bm{w}_{\beta}+T \delta_{w}\overline{\bm{w}}\right\|_{\left(\bm{I}+T\delta_{w}\overline{\bm{ \Sigma}}_{\bm{w}}\right)^{-1}}^{2}}{2\sigma_{w}^{2}}-\frac{\left\|\bm{w}_{ \alpha}\right\|^{2}-\left\|\bm{w}_{\alpha}+T\delta_{w}\overline{\bm{w}}\right \|_{\left(\bm{I}+T\delta_{w}\overline{\bm{\Sigma}}_{\bm{w}}\right)^{-1}}^{2}}{ 2\sigma_{w}^{2}}$ $\displaystyle=\frac{\|\bm{w}_{\beta}\|^{2}-\left\|\bm{w}_{\beta}+\delta_{w} \sum_{i=1}^{T}\bm{x}_{i}y_{i}\right\|_{\left(\bm{I}+T\delta_{w}\overline{\bm{ \Sigma}}_{\bm{w}}\right)^{-1}}^{2}}{2\sigma_{w}^{2}}-\frac{\|\bm{w}_{\alpha}\| ^{2}-\left\|\bm{w}_{\alpha}+\delta_{w}\sum_{i=1}^{T}\bm{x}_{i}y_{i}\right\|_{ \left(\bm{I}+T\delta_{w}\overline{\bm{\Sigma}}_{\bm{w}}\right)^{-1}}^{2}}{2 \sigma_{w}^{2}}$ $\displaystyle=\frac{\|\bm{w}_{\beta}\|^{2}-\left\|(\bm{w}_{\beta}-\frac{\sum_{ i=1}^{T}\bm{x}_{i}y_{i}}{T})+(\bm{I}+T\bm{I}\delta_{w})\frac{\sum_{i=1}^{T}\bm {x}_{i}y_{i}}{T}\right\|_{\left(\bm{I}+T\delta_{w}\overline{\bm{\Sigma}}_{\bm{ w}}\right)^{-1}}^{2}}{2\sigma_{w}^{2}}$ $\displaystyle-\frac{\|\bm{w}_{\alpha}\|^{2}-\left\|(\bm{w}_{\alpha}-\frac{\sum _{i=1}^{T}\bm{x}_{i}y_{i}}{T})+(\bm{I}+T\bm{I}\delta_{w})\frac{\sum_{i=1}^{T} \bm{x}_{i}y_{i}}{T}\right\|_{\left(\bm{I}+T\delta_{w}\overline{\bm{\Sigma}}_{ \bm{w}}\right)^{-1}}^{2}}{2\sigma_{w}^{2}}$ $\displaystyle\overset{(a)}{=}\|\bm{w}_{\beta}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y_ {i}}{T}\|^{2}_{\bm{I}-\left(\bm{I}+T\delta_{w}\overline{\bm{\Sigma}}_{\bm{w}} \right)^{-1}}-\|\bm{w}_{\alpha}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y_{i}}{T}\|^{2}_ {\bm{I}-\left(\bm{I}+T\delta_{w}\overline{\bm{\Sigma}}_{\bm{w}}\right)^{-1}}$ $\displaystyle\overset{(b)}{=}\|\bm{w}_{\beta}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y_ {i}}{T}\|^{2}_{\frac{\delta_{w}\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_{i}^{\top}}{1+ \delta_{w}\sum_{i=1}^{T}\bm{x}_{i}^{\top}\bm{x}_{i}}}-\|\bm{w}_{\alpha}-\frac{ \sum_{i=1}^{T}\bm{x}_{i}y_{i}}{T}\|^{2}_{\frac{\delta_{w}\sum_{i=1}^{T}\bm{x}_ {i}\bm{x}_{i}^{\top}}{1+\delta_{w}\sum_{i=1}^{T}\bm{x}_{i}^{\top}\bm{x}_{i}}}$ $$

where equation $(a)$ is due to the third condition in Assumption C.2, equation $(b)$ is by applying the Sherman–Morrison formula. Since $\frac{\delta_{w}}{1+\delta_{w}\sum_{i=1}^{T}}\geq 0$, to prove that $\Psi_{\bm{w}}(\alpha,\beta)\geq 0$, we only need to show that

$$ $\|\bm{w}_{\beta}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y_{i}}{T}\|^{2}_{\sum_{i=1}^{T} \bm{x}_{i}\bm{x}_{i}^{\top}}-\|\bm{w}_{\alpha}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y _{i}}{T}\|^{2}_{\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_{i}^{\top}}\geq 0.$ (8) $$

Further, we can derive that the term $\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}-\bm{x}_{i}y_{i}\|^{2}_{\bm{x}_{i}\bm
{x}_{i}^{T}}-\|\bm{w}_{\alpha}-\bm{x}_{i}y_{i}\|^{2}_{\bm{x}_{i}\bm{x}_{i}^{T}}$ below is non-negative by using the condition 2 in Assumption C.2:

$$ $\displaystyle\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}-\bm{x}_{i}y_{i}\|^{2}_{ \bm{x}_{i}\bm{x}_{i}^{T}}-\|\bm{w}_{\alpha}-\bm{x}_{i}y_{i}\|^{2}_{\bm{x}_{i} \bm{x}_{i}^{T}}$ (9) $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}(\bm{w}_{\beta}-\bm{x}_{i}y_{i})^{\top} \bm{x}_{i}\bm{x}_{i}^{T}(\bm{w}_{\beta}-\bm{x}_{i}y_{i})-(\bm{w}_{\alpha}-\bm{ x}_{i}y_{i})^{\top}\bm{x}_{i}\bm{x}_{i}^{T}(\bm{w}_{\alpha}-\bm{x}_{i}y_{i})$ $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}(\bm{w}_{\beta}+\bm{w}_{\alpha}-2\bm{x} _{i}y_{i})^{\top}\bm{x}_{i}\bm{x}_{i}^{T}(\bm{w}_{\beta}-\bm{w}_{\alpha})$ $\displaystyle\underset{(c)}{\geq}\frac{1}{T}\sum_{i=1}^{T}(\bm{w}_{\beta}+\bm{ w}_{\alpha}-2\bm{x}_{i}y_{i})^{\top}(\bm{w}_{\beta}-\bm{w}_{\alpha})$ $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}^{\top}\bm{x}_{i}-y_{i} \|^{2}-\|\bm{w}_{\alpha}^{\top}\bm{x}_{i}-y_{i}\|^{2}\underset{(d)}{\geq 0}$ $$

where the inequality $(c)$ holds since according to the condition 2 in Assumption C.2, $\bm{x}_{i}\bm{x}_{i}^{T}-\bm{I}$ is positive semi-definite, and the inequality $(d)$ holds since the population downstream risk of $\alpha$ is lower than $\beta$. Therefore, to prove inequality (8), we just need to prove that the l.h.s. of inequality (8) multiplying $\frac{1}{T}$ is not less than $\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}-\bm{x}_{i}y_{i}\|^{2}_{\bm{x}_{i}\bm
{x}_{i}^{T}}$ in Equation (9):

$$ $\frac{1}{T}\left(\|\bm{w}_{\beta}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y_{i}}{T}\|^{2 }_{\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_{i}^{\top}}-\|\bm{w}_{\alpha}-\frac{\sum_{i= 1}^{T}\bm{x}_{i}y_{i}}{T}\|^{2}_{\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_{i}^{\top}} \right)\geq\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}-\bm{x}_{i}y_{i}\|^{2}_{ \bm{x}_{i}\bm{x}_{i}^{T}}-\|\bm{w}_{\alpha}-\bm{x}_{i}y_{i}\|^{2}_{\bm{x}_{i} \bm{x}_{i}^{T}}.$ (10) $$

First, let’s simplify the l.h.s of inequality (10):

$$ $\displaystyle\frac{1}{T}\left(\|\bm{w}_{\beta}-\frac{\sum_{i=1}^{T}\bm{x}_{i}y _{i}}{T}\|^{2}_{\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_{i}^{\top}}-\|\bm{w}_{\alpha}- \frac{\sum_{i=1}^{T}\bm{x}_{i}y_{i}}{T}\|^{2}_{\sum_{i=1}^{T}\bm{x}_{i}\bm{x}_ {i}^{\top}}\right)$ (11) $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}(\bm{w}_{\beta}-\frac{\sum_{j=1}^{T}\bm {x}_{j}y_{j}}{T})^{\top}\bm{x}_{i}\bm{x}_{i}^{\top}(\bm{w}_{\beta}-\frac{\sum_ {j=1}^{T}\bm{x}_{j}y_{j}}{T})-(\bm{w}_{\alpha}-\frac{\sum_{j=1}^{T}\bm{x}_{j}y _{j}}{T})^{\top}\bm{x}_{i}\bm{x}_{i}^{\top}(\bm{w}_{\alpha}-\frac{\sum_{j=1}^{ T}\bm{x}_{j}y_{j}}{T})$ $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}^{\top}\bm{x}_{i}-\frac {1}{T}\sum_{j=1}^{T}\bm{x}_{j}^{\top}\bm{x}_{i}y_{j}\|^{2}-\|\bm{w}_{\alpha}^{ \top}\bm{x}_{i}-\frac{1}{T}\sum_{j=1}^{T}\bm{x}_{j}^{\top}\bm{x}_{i}y_{j}\|^{2}$ $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}(\bm{w}_{\beta}^{\top}\bm{x}_{i})^{2}-( \bm{w}_{\alpha}^{\top}\bm{x}_{i})^{2}+2(\bm{w}_{\alpha}-\bm{w}_{\beta})^{\top} \bm{x}_{i}\frac{1}{T}\sum_{j=1}^{T}\bm{x}_{j}^{\top}\bm{x}_{i}y_{j}.$ $$

Then we simplify the r.h.s. of inequality (10):

$$ $\displaystyle\frac{1}{T}\sum_{i=1}^{T}\|\bm{w}_{\beta}-\bm{x}_{i}y_{i}\|^{2}_{ \bm{x}_{i}\bm{x}_{i}^{T}}-\|\bm{w}_{\alpha}-\bm{x}_{i}y_{i}\|^{2}_{\bm{x}_{i} \bm{x}_{i}^{T}}$ (12) $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}(\bm{w}_{\beta}^{\top}\bm{x}_{i})^{2}-( \bm{w}_{\alpha}^{\top}\bm{x}_{i})^{2}+2(\bm{w}_{\alpha}-\bm{w}_{\beta})^{\top} \bm{x}_{i}\bm{x}_{i}^{\top}\bm{x}_{i}y_{i}$ $$

Subtracting Equation (LABEL:proof:eq-2) from Equation (LABEL:proof:eq-1), we get

$$ $\displaystyle\frac{1}{T}\sum_{i=1}^{T}2(\bm{w}_{\alpha}-\bm{w}_{\beta})^{\top} \bm{x}_{i}\frac{1}{T}\sum_{j=1}^{T}\bm{x}_{j}^{\top}\bm{x}_{i}y_{j}-2(\bm{w}_{ \alpha}-\bm{w}_{\beta})^{\top}\bm{x}_{i}\bm{x}_{i}^{\top}\bm{x}_{i}y_{i}$ (13) $\displaystyle=\frac{1}{T}\sum_{i=1}^{T}2(\bm{w}_{\alpha}-\bm{w}_{\beta})^{\top }\bm{x}_{i}y_{i}\frac{1}{T}\sum_{j=1}^{T}\left(\bm{x}_{j}^{\top}\bm{x}_{i} \frac{y_{j}}{y_{i}}-\bm{x}_{i}^{\top}\bm{x}_{i}\right).$ $$

applying the condition 4 in Assumption C.2, we get the final conclusion.

∎

###### Theorem C.1 .

###### Assumption C.2 .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: See the comments in the weakness section.


Update:

Sorry for the late response. Most of my concerns were properly addressed. I have raised my score to 6.
soundness: 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. 
2. 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.
3. Besides GPT-2, they use Llama2-7b, a pre-trained language model for evaluation.
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.
weaknesses: 1. Presentation issues. 

(a) In Eq.1, why the prediction is supervised by f(x_i)? Shouldn’t it be f(x_{I+1}) 

(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. 

(c)Please correct the way how you are using the quote. See examples in the caption of Fig.3.

2. 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. 


3. 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?


4. 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?


5. 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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
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?
2. What causes the sudden increase in performance after context length size 10 in Figure 6B
3. 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.
4. Can you please explain the motivation behind the experiment in Figure 5 better.
soundness: 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.

**Quality**: The paper does thorough, well thought out experiments and provides new theoretical insight.

**Clarity**: The paper is clearly written, well structured and easy to understand.

**Significance**:  The question asked by the paper are significant as it further clarifies the mystery of ICL especially regarding limitations to OOD tasks.
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:
1. 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.
2. 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.
3. 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.
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.
- The explanation of the retrieval version of the Llama-2-7B task was hard to understand, it would be helpful to clarify that section
- 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).
- 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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
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?
- 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.
- 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.
- 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?
soundness: 2
strengths: - Characterizing the out-of-distribution generalization ability of in-context learning is an important problem with potential impact
- The experiments are well-designed to investigate the respective hypotheses
- The paper does a good job at referencing and contextualizing existing related work
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.
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.
- The theoretical results heavily borrow from [1] (but the paper makes this very transparent)

[1] Lin, Ziqian, and Kangwook Lee. "Dual operating modes of in-context learning." arXiv preprint arXiv:2402.18819 (2024).

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: 1. Which findings in particular do the authors believe are genuinely surprising to most deep learning researchers?
2. What findings from the paper can’t be explained through the theoretical Bayes-optimal prediction framework for ICL?
soundness: 3
strengths: 1. The experimental setups are quite easily to follow and well-explained.
2. 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.
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.
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.
2. 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.
3. 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.

Meta-review: This paper investigates the behavior of in-context learning (ICL) on out-of-distribution tasks through empirical analysis and theory development. The reviewers generally found merit in the paper's experimental methodology and results, particularly highlighting the well-designed synthetic tasks and comprehensive evaluation approach. The paper received ratings ranging from 5 to 8, with reviewers acknowledging both its empirical contributions in characterizing ICL's behavior on OOD tasks and its theoretical insights into how ICL implements functions from its pretraining distribution, though half of the reviewers found the presentation not unified and one felt this paper was subsumed by existing Bayesian ICL theories. Nevertheless, with an appropriate revision for clarity and comparison with the Bayesian framing of ICL which is prevalant, I assess that this work would be a good fit for this ICLR.

Decision: Accept (Poster)