In-context learning in presence of spurious correlations

## Abstract

Abstract Large language models exhibit a remarkable capacity for in-context learning, where they learn to solve tasks given a few examples.
Recent work has shown that transformers can be trained to perform simple regression tasks in-context.
This work explores the possibility of training an in-context learner for classification tasks involving spurious features.
We find that the conventional approach of training in-context learners is susceptible to spurious features.
Moreover, when the meta-training dataset includes instances of only one task, the conventional approach leads to task memorization and fails to produce a model that leverages context for predictions.
Based on these observations, we propose a novel technique to train such a learner for a given classification task.
Remarkably, this in-context learner matches and sometimes outperforms strong methods like ERM and GroupDRO.
However, unlike these algorithms, it does not generalize well to other tasks.
We show that it is possible to obtain an in-context learner that generalizes to unseen tasks by training on a diverse dataset of synthetic in-context learning instances.

## 1 Introduction

Large language models, such as GPT-3, have the ability of in-context learning (ICL), wherein they learn to solve a task given a few examples in the context (Brown et al., ).
The most significant aspect of in-context learning is that the learning happens during the forward pass on the context and query, without updating network parameters.
In order to study in-context learning in isolation, a number of studies considered training transformers (Vaswani et al., ) from scratch to solve simple learning tasks in-context.
In particular, Garg et al. () show empirically that transformers can be trained to perform in-context learning of simple regression functions, such as
dense or sparse linear functions, two-layer ReLU neural networks, and small decision trees.

Training on ICL instances can be seen as an instance of meta-learning (Schmidhuber, ; Naik and Mammone, ; Thrun and Pratt, ), where the goal is to learn a learning algorithm.
What exact algorithm is learned when training transformers on ICL instances is still an open problem.
Akyürek et al. () and Von Oswald et al. () show that transformers can implement a single gradient descent step of ordinary least squares and update the closed-form solution of ridge regression when a new example is added.
Additionally, they provide evidence that transformers trained on ICL instances of linear regression learn algorithms that closely match predictions of the known algorithms, such as gradient descent on ordinary least squares objective and ridge regression.
However, there is evidence that the learned algorithm may vary with model scale, depth, and pretraining task diversity (Akyürek et al., ; Raventós et al., ).
In particular,  Raventós et al. () demonstrate that in the setting of in-context learning of linear regression tasks with insufficient pretraining task diversity, the learned algorithm behaves like a Bayesian estimator with the pretraining task distribution as the prior, and hence fails to generalize well to unseen tasks.
Yadlowsky et al. () show that when trained on ICL instances where the regression function belongs to a union of distinct function classes, the learned algorithm fails to generalize beyond the pretraining function classes.
Ahuja and Lopez-Paz () show that in-context learning ability diminishes under strong distribution shifts.

In this work, we explore the limits of in-context learning further by testing it on challenging settings.
We deviate from the existing literature and consider visual classification tasks instead of regression tasks with simple function classes.
In particular, we consider classification tasks where some features are *spuriously correlated* with the label.
Such features are predictive of the label but are not causally related to it, due to which their correlation might not hold at test time.
A prominent example is the cow vs camel classification task, where the background often correlates with the label, as cows are typically photographed in pastures, while camels are typically photographed in deserts (Beery et al., ).
It is well-known that neural networks trained with gradient-based methods to minimize empirical risk can exploit spurious features, causing performance degradation under distribution shifts affecting these correlations (Torralba and Efros, ; Ribeiro et al., ; Gururangan et al., ; Zech et al., ; McCoy et al., ; Geirhos et al., , ; Xiao et al., ).

We start our analysis in the standard setting of having a single classification task with spurious features.
We consider the conventional approach of obtaining an in-context learner, wherein a transformer is trained on sequences of form $(x_{1},y_{1},\ldots,x_{k},y_{k},x_{k+1})$ to predict the label $y_{k+1}$ of the query example $x_{k+1}$.
We find that this conventional approach leads to models that do classification *ignoring the context*, essentially memorizing the task.
Furthermore, these models lack robustness to changes of the correlation between the label and spurious features.
In particular, we observe a significant performance drop when the query follows a distribution in which the label and spurious feature correlation is zero.
We propose an effective approach of addressing the task memorization issue.
Namely, we find that task memorization can be mitigated greatly by randomly permuting input embedding dimensions for each training sequence.
To address the issue of spurious features, we propose a novel way of forming ICL instances and a suitable transformer architecture, which work together to simulate distribution shift with respect to spurious features in the context.
Overall, our proposed techniques lead to strong in-context learners that outperform established methods such as 1-NN, empirical risk minimization (ERM), and GroupDRO (Sagawa* et al., ), suggesting that the in-context learner implements a more specialized algorithm.

Despite being trained on instance of a single task, the learned algorithm generalizes to other tasks *without* spurious features.
However, it fails to generalize to unseen tasks with spurious features.
For this reason, we next explore training an in-context learner that generalizes to unseen tasks with spurious features.
We create a dataset of in-context learning instances for various binary classifications tasks with varying spurious features.
We demonstrate the efficacy of the proposed techniques on this dataset too and find that it can be improved further by passing spurious feature annotations as input and injecting occasional queries requesting the label of a proceeding context example to promote learning induction heads.
The resulting model generalizes perfectly to unseen tasks, as long as the data generating process is similar.
However, generalization to unseen tasks with possibly different data generating process depends on the severity of the challenge posed by spurious features, indicating that the learned algorithm is more brittle to severe distribution shifts than conventional algorithms.
The source code for reproducing our experiments is available at [https://github.com/YerevaNN/incontext_spurious](https://github.com/YerevaNN/incontext_spurious).

We summarize our main contributions as follows.

- (i)
We show that the conventional approach of training an in-context learner is susceptible to presence of spurious features and also leads to task memorization in case of a single task.
- (ii)
We propose a suite of novel techniques of forming in-context training data to mitigate task memorization and increase robustness to spurious features, leading to in-context learners that outperform established learning algorithms.
- (iii)
We demonstrate that it is possible to obtain more general-purpose robust in-context learners by training on a diverse set of synthetic classification tasks involving spurious features.

Figure: (a) Naive approach
Refer to caption: extracted/5900723/figures/naive.png

## 2 In-context learning based on a single task

We start by considering the common setting of having a single classification task with spurious features.
For simplicity, we focus on label-balanced binary classification tasks in presence of a single binary spurious feature,
although what follows next applies to label-imbalanced multiclass settings as well.
Let $\mathcal{D}_{\mathrm{train}}$ be a set of training examples for the task, where each example is a triplet $(x,s,y)$ of input $x\in\mathbb{R}^{d}$, spurious feature value $s\in\left\{0,1\right\}$, and label $y\in\left\{0,1\right\}$.
Similarly, let $\mathcal{D}_{\mathrm{test}}$ be a set of test examples.
Importantly, we do not make any assumptions on the data generating process, except that $x$ has some information about $s$ and $s$ is predictive of $y$ on the training set, but their correlation does not hold on the test set.
For an example $(x,s,y)$, we define its *group* $g=2y+s$.
In a binary classification task with a single binary spurious feature, there are four groups.
Without loss of generality, we assume that for a majority of training examples we have that $y=s$.
Hence we refer to groups 0 and 3 as majority groups, while referring to groups 1 and 2 as minority groups.

Training a transformer to perform linear regression in-context requires millions of ICL training instances, even for small dimensional cases. For example, Garg et al. () use 32 million training instances for 20-dimensional inputs.
We next consider ways of generating so many ICL instances from a single task.

### 2.1 A naive approach of constructing ICL instances

The standard approach of constructing an ICL instance is to sample a subset of $n+1$ examples $\left\{(x_{i},s_{i},y_{i})\right\}_{i=1}^{n+1}$ from $\mathcal{D}_{\mathrm{train}}$ and form a sequence $S=(x_{1},\tilde{y}_{1},x_{2},\tilde{y}_{2},\ldots,x_{n},\tilde{y_{n}},x_{n+1})$, where $\tilde{y}_{i}\in\mathbb{R}^{d}$ is a fixed random representation of either $y_{i}$ or $g_{i}$
(this distinction will be elaborated later).
Then one trains a transformer $f_{\theta}:\cup_{k}\mathbb{R}^{k\times d}\rightarrow[0,1]$ to predict $y_{i}$ given $S_{i}\triangleq(x_{1},\tilde{y}_{1},\ldots,x_{i-1},\tilde{y}_{i-1},x_{i})$ (see Figure 1(a)), optimizing the following loss function:

$$ $\frac{1}{n+1}\sum_{i=1}^{n+1}\mathrm{CE}(y_{i},f_{\theta}(S_{i})),$ (1) $$

where $\mathrm{CE}(y,\widehat{y})=-y\log\widehat{y}-(1-y)\log(1-\widehat{y})$ is the binary cross-entropy loss.
We explore two options of setting $\tilde{y}_{i}$.
In the first option, we set $\tilde{y}_{i}$ to represent $y_{i}$ with a constant vector or its negative in $\mathbb{R}^{d}$. In this case we aim to obtain an in-context learner that is robust to spurious features without receiving spurious feature annotations as input.
ERM is one such learner that minimizes average loss on training examples and does not require spurious feature annotations.
In the second option, we set $\tilde{y}_{i}$ to represent $g_{i}$ as a sum of two constant vectors in $\mathbb{R}^{d}$, one representing the class and the other representing the spurious feature.
In this case we aim to obtain an in-context learner that does robust classification with respect to a specified spurious feature.
GroupDRO is one such learner that minimizes worst-group loss, therefore requiring spurious feature annotations at training time.

Unfortunately, the simple approach of (1) has several issues.
First, as the classification task is the same in all ICL instances, the model can ignore context examples and predict $y_{i}$ based solely on $x_{i}$, essentially memorizing the task.
Second, as all $n+1$ examples of a sequence $S$ are sampled from the training set and the spurious correlation holds for all of them, there is nothing preventing usage of spurious features in making predictions.
To confirm these two issues, we consider the Waterbirds dataset (Sagawa* et al., ), which is landbird vs waterbird image classification task where image background (sea or land) is correlated with the label in the training set (4,795 examples), but not in validation and test sets . A robust classifier should predict waterbird or landbird without relying on image background.
To separate out the representation learning challenge, we represent images with a pretrained and frozen DINOv2 ViT-B/14 distilled (Oquab et al., ).
This way each image is embedded in $\mathbb{R}^{768}$.
While using powerful pretrained representations increases overall performance under distribution shifts (Radford et al., ; Mehta et al., ), we note that it does not eliminate the problem of spurious correlations.
Representations obtained via large-scale self-supervised pretraining are likely rich enough to capture information about both the label and spurious feature.
Furthermore, many works have indicated that the main contribution to the out-of-domain generalization error comes from the classification head (rather than the representation learning module) and called for designing better methods of training the classification head (Galstyan et al., ; Menon et al., ; Kirichenko et al., ; Izmailov et al., ; Shi et al., ).

Figure: Figure 2: Majority-group and worst-group test accuracies on Waterbirds as a function of context size for the naive and proposed approaches with or without permuting input dimensions. Shaded regions show standard deviation across 5 training runs.
Refer to caption: x1.png

We train a causal decoder-only GPT-J transformer (Wang and Komatsuzaki, ) with 80M parameters on 2M in-context learning sequence with $n=512$ and $\tilde{y}_{i}$ representing labels, constructed from the training set of Waterbirds.
We use balanced sampling of classes and set the minority group proportion to 10% within each class.
We use the ADAM optimizer (Kingma and Ba, ) ($\beta_{1}=0.9$ and $\beta_{2}=0.999$) with 32 batch size and no weight decay.
The learning rate is selected from $\left\{3\cdot 10^{-5},6\cdot 10^{-5},10^{-4}\right\}$ based on average test performance over 5 runs.
Concretely, we evaluate on 8192 sequences where the context part is $n$ training examples, while the query is a sampled from the test set with equal group distribution. Exact metric definitions and missing details are provided in Appendix A.
Note that with 512 context length and 10% minority group ratio within each class, the expected value of the number of context examples from each of the 2 minority groups is about 25.
For reference, the smallest minority-group has only 56 examples in the Waterbirds training set.

Figure 2 plots majority-group and worst-group test accuracies as a function of context size $n$.
We see that naive approach results in models that ignore context – worst-group accuracy with 512 context examples is essentially the same as with 2 examples (see the *naive* curve).
This confirms the task memorization issue.
Figure 2 also shows that majority-group test accuracy of the naive approach is considerably higher compared to worst-group accuracy confirming the non-robustness issue.

### 2.2 The proposed approach of constructing ICL instances

Figure: Figure 3: Majority-group and worst-group test accuracies on Waterbirds-severe as a function of context size for the naive and proposed approaches with or without permuting input dimensions. Shaded regions show standard deviation across 5 training runs.
Refer to caption: x3.png

To address the task memorization issue, we propose to rotate image embeddings in each ICL instance independently, making it harder to memorize individual examples.
We found that generating random rotation matrices on fly is computationally expensive and slows down training.
We tried generating and storing 10K rotation matrices, but this resulted in less than 50M different training examples that were still possible to memorize to some extent.
A more effective and efficient alternative is to apply random permutations to image embedding dimensions (for brevity, this technique is denoted with *+P* in figures and tables).
We found this approach to be very effective in terms of inducing in-context learning (see *naive + P* in Figure 2).
We also see that the difference between majority-group and worst-group accuracies decreases, although an approximately 5 p.p. gap remains.

When training an ICL transformer, ideally, we would like to simulate the situation of making a test prediction based on a context of training examples.
Importantly, we would like to simulate the case where test distribution has balanced groups (i.e., the spurious correlation does not hold).
Given access to spurious feature annotations *for the training set*, we can simulate this scenario using only training examples.
In particular, we can form ICL instances of form $(x_{1},\tilde{y}_{1},\ldots,x_{n},\tilde{y}_{n},x_{n+1})$, where the context examples ($x_{1},\ldots,x_{n}$) are sampled in a way that the spurious feature is correlated with the label, while the query $x_{n+1}$ is sampled to have a uniform group distribution.
However, if we again optimize the loss of (1), for context lengths less than $n$, the network will be allowed to make predictions using the spurious feature, which is undesirable.
Please refer to Figure 13 of Appendix B for evidence of this.
Potential ways of addressing this issue is upweighting the final prediction loss in Equation 1 or upweighting predictions on minority examples.
In our preliminary experiments we found the former approach ineffective.
We did not experiment with the latter approach.

Instead, we propose a novel way of forming in-context learning instances and a modified transformer architecture that is suitable for such sequences.
In particular, we form sequences of form $S=(x_{1},\tilde{y}_{1},q_{1},x_{2},\tilde{y}_{2},q_{2},\ldots,x_{n},\tilde{y}_
{n},q_{n})$, where $(x_{i},\tilde{y}_{i})$ are context examples, while $q_{i}$ are queries, sampled with replacement from $\mathcal{D}_{\mathrm{train}}\setminus\{x_{1},\ldots,x_{n}\}$.
Importantly, $q_{i}$ are sampled with a uniform group distribution.
Redefining $S_{i}=(x_{1},\tilde{y}_{1},q_{1},\ldots,x_{i},\tilde{y}_{i},q_{i})$, we would like the prediction on $S_{i}$ to be the label of $q_{i}$.
When making a prediction on $q_{i}$, we want $q_{j}\ (j<i)$ to have no effect.
For this end we make two modifications.
First, we modify the causal attention matrix to disallow attending to query tokens, unless a query token is attending to itself.
Formally, if we enumerate tokens from $1$ to $3n$ and define $M_{i,j}$ to denote the attention mask for
token $i$ attending to token $j$, then we set

$$ $M_{i,j}=\begin{cases}0,&i<j,\\ 0,&i>j\text{\ \ and\ \ }j\equiv 0\mod 3,\\ 1,&\text{otherwise.}\end{cases}$ (2) $$

Second, we use modified token positions for computing positional encodings, in order to discount intermediate query tokens.
Namely, for the sequence $(x_{1},\tilde{y}_{1},q_{1},x_{2},\tilde{y}_{2},\ldots,x_{n},\tilde{y}_{n},q_{n})$, position indices are set to $(0,1,2,2,3,4,4,\ldots,2n-2,2n-1,2n)$.
Formally, enumerating tokens from $1$ to $3n$, the position index of the $i$-th token is set to $2\left\lfloor\frac{i-1}{3}\right\rfloor+(i-1)\mod 3$.
Please refer to Figure 1 for an illustration.
Hereafter, we refer to this approach as simply “proposed approach”.

Figure 2 compares the proposed and naive approaches with and without input dimension permutations.
Without random permutations, the proposed approach outperforms the naive approach marginally.
However, the same is not true with random permutations.
We found that image embeddings of DINOv2 have a bias towards representing objects more than backgrounds, alleviating the challenge posed by the spuriously correlated background in Waterbirds.
In fact, the linear probing accuracy of the spurious feature is just $\approx 82$%.
For this reason, we create a modified version of Waterbirds by adding a constant vector $\tilde{s}$ or $-\tilde{s}$ to image embeddings based on the spurious feature $s$.
We scale $\tilde{s}$ to have its norm equal to the average norm of image embeddings and verify that the linear probing accuracy of the spurious feature becomes 100%.
On this modified Waterbirds dataset, which we name Waterbirds-severe, we see a large separation between the naive and proposed approaches (see Figure 3).
We also see that without permutations, both naive and proposed approaches perform identically, indicating no robustness to the spurious correlation.
This is expected, because in the absence of in-context learning, we can think of the naive and proposed approaches as standard and reweighted empirical risk minimization with a complex classification head, respectively.
It has been observed that sample reweighting is not effective in overparameterized settings as all training examples will be perfectly fitted (Byrd and Lipton, ; Menon et al., ).

### 2.3 Comparison with conventional learning algorithms

Now that we have established the efficacy of the proposed technique, we compare it to a few established algorithms, such as 1-NN, ERM, and GroupDRO, that last of two have been historically hard to outperform (Gulrajani and Lopez-Paz, ; Koh et al., ).
Comparing to more existing methods designed for robustness to spurious correlations is outside of the goal of this work, namely studying limits of in-context learning.
In our comparisons, we follow the evaluation recipe used for the in-context learners.
Namely, we evaluate each baseline on 8192 sequences by training on the context part of the sequence and making a prediction on the single query.
More information about hyperparameters and model selection is presented in Appendix A.

Figures 5(a) and 5(b) compare the proposed and baseline approaches on Waterbirds and Waterbirds-severe respectively.
On Waterbirds, the proposed method outperforms ERM and GroupDRO on almost all context lengths, but is better than 1-NN only for short context lengths.
The good performance of 1-NN is due to the bias in DINOv2 representations.
On Waterbirds-severe, the proposed method outperforms the baselines at all context lengths.
From these results, we conclude that this in-context learner implements none of these algorithms.

It is worth noting that baseline worst-group accuracies at $n=512$ are actually *higher* than what we get when training on the entire dataset.
For example, on Waterbirds, 1-NN gets only 90.03 % worst-group accuracy, while ERM gets 84.23 $\pm$ 0.17 % and GroupDRO gets 92.43 $\pm$ 0.24 %.
This is due to balanced sampling of classes and setting the minority ratio to 10% withing each class, which is higher than the minority ratio of $\approx 5\%$ in the original Waterbirds dataset.
One can think of our resampling as a weaker form of down-sampling which has been found to be helpful in presence of spurious correlations (Nagarajan et al., ; Menon et al., ; Idrissi et al., ).

Figure: Figure 4: Worst-group test accuracies on CelebA for the proposed approach and conventional methods such as 1-NN, ERM, and GroupDRO. Shaded regions show standard deviation across 5 training runs. The “+G” marker indicates setting $\tilde{y}_{i}$ to represent $g_{i}$, i.e., passing context examples groups as input. Please refer to Section 3 for more information on this.
Refer to caption: x5.png

Additionally, we verify our findings on another popular dataset CelebA (Liu et al., ) designed for blond vs non-blond person classification, with sex being a spurious variable.
Unlike, Waterbirds, the spurious feature is asymmetric in CelebA, as blond and non-blond women are equally represented, while blond men are significantly infrequent compared to non-blond men.
In particular, we verify the two shortcomings of the conventional approach and demonstrate the efficacy of the proposed techniques compared to the baselines (please see Table 4 and Figure 4 of Appendix B).

### 2.4 Generality of the learned algorithm

Figure: (a) Waterbirds
Refer to caption: x6.png

Since we train in-context learners on ICL instances of a single task, a natural question arises whether the learned algorithm can generalize to unseen tasks.
Without permuting input dimensions, the model does not learn to do in-context learning.
Thus, we can not hope for any generality without permuting input dimensions.
We take the model obtained with the “Proposed + P” technique
and probe generality of its in-context learning by evaluating on various datasets.
We start by swapping the labels of two classes in Waterbirds at evaluation and observe $\approx 2$ p.p. overall accuracy drop and $\approx 5$ p.p. worst-group accuracy drop.
Despite the worsened performance, this indicates that the model treats class labels symbolically, which is remarkable given that the semantics of labels were constant during training.
However, when we evaluate on Waterbirds-severe, it gets 100% accuracy on the majority groups and *0% accuracy on minority groups*. Additionally, when we switch the task to predicting the background in the original Waterbirds dataset (now the class becomes a spurious feature), the overall test accuracy drops to 54.4%, while the worst-group accuracy drops to 9.3%.

It is worth noting that the learned algorithm is not completely useless for other tasks and works well in absence of spurious features, even on unseen tasks.
For example, evaluating on binary classification tasks derived from the CUB-200 (Welinder et al., ) dataset, from where the bird images of Waterbirds were taken, we get 99.7% accuracy at context size 100 (the accuracy is so high because most pairs of classes are easy to distinguish).
We also test on binary classification tasks derived from classes belonging to *Amphibia* and *Mammalia* supercategories of the iNaturalist (Van Horn et al., ) dataset.
At context length 512, the overall accuracy is 98.5%.

These OOD evaluation results indicate that the learned algorithm does something specific to the spurious feature of Waterbirds.
We hypothesize that it learns to ignore this particular spurious feature.
To test this, we evaluate on *group-balanced* Waterbirds sequences, with the task set to predicting background, and get 58.5% overall accuracy and 41.3% worst-group accuracy.
One potential way of improving generality and possibly also performance, is passing example groups as input, i.e., setting $\tilde{y}_{i}$ to represent $g_{i}$.
We did not observe performance improvements and increase of generality of the learned algorithm when passing groups as input (see the complete results in Tables 1 and 2 of Appendix B).
Thus, we conclude that when all ICL instances are derived from the task, the learned algorithm is inherently tied to the spurious feature of that task.

## 3 In-context learning based on a diverse set of tasks

Figure: (a) Before grafting
Refer to caption: extracted/5900723/figures/grafting_before.png

In Section 2, we showed that it is possible to obtain a good in-context learner for a given task, but it fails to generalize to tasks with different spurious features.
A better in-context learner should detect spurious features from context and make predictions without employing them.
In this section, we explore the possibility of obtaining such a learner by training on a diverse set of ICL tasks.
Since there exist few suitable datasets, we synthesize binary classification tasks with a single binary spurious feature, aiming to capture “structure” present in existing datasets.
In short, given a standard binary classification task, say cat vs dog classification, for a sampled minority of cats we overwrite some of their features with those of random dogs.
Similarly, we do an analogous operation for a sampled minority of dogs.
This way some cats share dog features and vice versa.
To create a diverse pool of in-context learning instances, we vary the two classes and the subset of grafted features.
Please refer to Figure 6 for an illustration of this grafting operation.

More concretely, we consider the iNaturalist dataset (Van Horn et al., ), which contains images from 5,089 natural fine-grained categories and filter out categories that have less than 500 images.
For testing purposes, from remaining 239 categories we set apart categories that belong to the supercategories *Amphibia* and *Mammalia*, along with 10% of random categories.
We denote the set of these 48 categories as $\mathcal{C}_{\mathrm{ood}}$, and the set of remaining 191 categories as $\mathcal{C}_{\mathrm{id}}$, which we use to create in-context learning instances for training.
For each category in $\mathcal{C}_{\mathrm{id}}$, we hold out half of the examples as in-distribution validation set.
To generate a single in-context learning instance, we sample two distinct classes from $\mathcal{C}_{\mathrm{id}}$ randomly and sample $n/2$ images from the training split of each class uniformly at random without replacement.
We then do the grafting operation, setting minority group ratio within each class to 10%.
We select the grafted features randomly, by first picking subset size $k$ uniformly at random from 0 to 199, and then sampling a random subset of embedding dimensions of size $k$.
With this we get $n$ examples that form the context part of the instance.
Abandoning the naive approach and focusing on the proposed one, for each class we sample $n/2$ queries from the remaining examples uniformly at random with replacement and do the grafting operation with 50% minority group ratio.

Figure: Figure 7: Majority-group and minority-group accuracies on the OOD test set of iNaturalist for the proposed approaches with or without permuting input dimensions and promoting induction heads. Shaded regions show standard deviation across 5 training runs.
Refer to caption: x8.png

Following the experiments in Section 2, we train the same transformer with the proposed approach on 4M ICL instances with $n=400$ context examples.
We use the same optimizer and sweep the learning rate in the same range, selecting the best value based on the average *minority-group accuracy* (defined exactly in Appendix B) on instances where both categories belong to $\mathcal{C}_{\mathrm{ood}}$ and thus were not observed during training.
The results presented in Figure 7 indicate a major difference compared to the results in the single-task regime – namely, the proposed approach learns to do in-context learning to some extent without permuting embedding dimensions.
As expected, we see much better performance with permuted embedding dimensions.
Notably, comparing majority-group and minority-group accuracies of the proposed approach with permutations, we see almost no sign of reliance on spurious features.

#### Promoting emergence of induction heads.

In-context learning ability has been linked to induction heads, which are specific type of circuits found within large language models that implement the operation of looking back over the sequence for finding previous instances of the current token and copying what comes after that (Olsson et al., ).
Inspired by this, we propose a data preparation technique that promotes learning of induction heads.
With probability $p$, we replace each intermediate query independently with a random example from the proceeding part of the context.
Note that this type of “hinting” is not possible in the naive approach and is enabled by the introduction of intermediate queries.
In all experiments with this technique enabled, we just set $p=0.25$.
We observed that training of typical runs escapes the initial loss plateau faster with this technique (in about 3k iterations compared instead of about 10k iterations).
Moreover, we see modest performance gains in iNaturalist experiments (see Figure 7, where *+I* stands for this technique).

Figure: Figure 8: Minority-group accuracy on the OOD test set of iNaturalist for the best proposed approach with or without passing group information as input. Shaded regions show standard deviation across 5 training runs.
Refer to caption: x10.png

#### Passing example groups as input.

In contrast to the findings in the single-task setting of Section 2, we observed that setting $\tilde{y}_{i}$ to represent group improves the proposed approach, even on top of permitting input dimensions and promoting induction heads.
One case of this is presented in Figure 9, while more cases can be found in the complete results presented in Appendix B.
For brevity, we mark passing groups as inputs with *+G* in figures and tables.

#### Comparison with conventional learning algorithms.

Similar to the experiments in Section 2, we compare the best variant of the proposed approach (G + P + I) to 1-NN, ERM, and GroupDRO.
Results presented in Figure 9 show that the learned algorithm is on-par with or outperforms the baselines starting at context length 32.
The results at context lengths below 20 are not as informative, because the way we implemented the grafting operation implies that no examples are grafted when there are less than 10 examples in a class.

Figure: Figure 10: Worst-group test accuracy on Waterbirds for the best variant of proposed approach trained on iNaturalist and for methods such as 1-NN, ERM, and GroupDRO. Shaded regions show standard deviation across 5 training runs.
Refer to caption: x12.png

#### Generality of the learned algorithm.

To test the generality of the learned algorithm, we report evaluation results on Waterbirds (Figure 11) and Waterbirds-severe (Figure 11).
We see that the learned algorithm outperforms baselines on Waterbirds and is as good as we got by training on Waterbirds itself.
However, the learned algorithm fails completely on Waterbirds-severe, while the baselines give meaningful results starting at context length 32.
We hypothesize that the challenge posed by the spurious features in Waterbirds-severe is significantly more severe compared to that in iNaturalist.
By varying the norm of the added background vector, we interpolate between Waterbirds and Waterbirds-severe, and we see good generalization until the norm of the added vector is $\approx$40% of the average embedding norm (see Figure 12).

Figure: Figure 12: Majority-group and worst-group test accuracies of a proposed model (G + P + I) trained on iNaturalist, but evaluated on a modified variants of Waterbirds where we add a vector representing the spurious feature (background). The x-axis is the relative norm of the added vector compared to the average Waterbirds image embedding norm. Relative norm of 0 corresponds to Waterbirds, while relative norm of 1 corresponds to Waterbirds-severe. Shaded regions show standard deviation across 5 training runs.
Refer to caption: x14.png

## 4 Related work

In this section, we discuss more related work in addition to the ones discussed earlier.

#### In-weights vs in-context learning.

We observe two modes of learner behavior in our experiments.
In the first mode, the learner acts like a standard supervised classifier, ignoring context examples.
This mode appears when training on ICL instance of a single task without permuting input embedding dimensions.
In the second mode, the learner does proper in-context learning.
Our experiments indicate that both permuting embedding dimensions and increasing the number of training tasks are reliable ways of steering the model towards the in-context learning mode.
The former is akin to the method of randomly projecting inputs proposed by Kirsch et al. () for obtaining general-purpose in-context classifiers.
Prior work has made a distinction between these two models of learning, naming them in-weights and in-context learning.
In particular, Chan et al. () demonstrate that certain distributional properties of data, such as long-tail of class frequencies and bursty distribution of context example classes, can promote in-context learning when meta-training on few-shot classification instances.
Singh et al. () show that in-context learning behavior is not persistent and decays away with overtraining, indicating a trade-off between in-weights and in-context learning mechanisms.
Moreover, they find that this in-context learning skill decay can be prevented by applying weight decay of embeddings and MLP layers, slowing down in-weight learning.
Anand et al. () make similar observations about these two modes of learning and propose active forgetting of token embeddings as an effective way of steering towards the in-context learning mode.

#### Many shot ICL.

One ancillary finding of this work is that transformers can be trained to do in-context learning of visual classification tasks when good image embeddings are provided.
This is remarkable because the input dimensionality we considered is much higher than what was considered in the pioneering works of Garg et al. () and Akyürek et al. () (784 vs 20).
Furthermore, we observe predictable performance gains from longer context sizes.
The number of “shots” we consider (up to 512 examples) is well beyond what is typically considered in ICL works (up to a few dozen of examples).
Our findings are complementary to those of Agarwal et al. (), Jiang et al. (), and Li et al. () who find that multimodal large language models, such as Gemini-1.5 Pro and GPT-4o, can benefit from large number of in-context demonstrations (up to 1000 demonstrations).

#### In-context learning for out-of-distribution generalization.

Closest to our work are the works that propose to make use of in-context learning for out-of-distribution generalization.
Han et al. () test multimodal large language models (MLLMs) on a variety of visual classification tasks.
They propose to leverage in-context learning abilities of MLLMs to improve performance on specialized domains and on tasks with significant corruptions.
However, they only consider the case where both context examples and query are from the target domain.
Zhang et al. () make similar observations, but additionally study robustness of in-context learning to distribution shifts, such as domain shifts, label shifts, and spurious correlations.
They find that in-context learning is highly susceptible to label shifts and presence of spurious correlations.
Finally, Gupta et al. () propose to address the problem of domain generalization (Muandet et al., ) by training an in-context learner that can take examples from a domain/environment and adapt to that domain in-context.

## 5 Discussion and conclusion

Spurious correlations can appear in various domains, often in complex and unexpected ways.
For example, a medical device might make a false positive or worse, a false negative prediction due to the specifics of the patient and training data.
Given the increasing reliance on in-context classification, particularly with the advancements in large language models, understanding and mitigating the negative effects of spurious correlations is important for ensuring the reliability and safety of these models.

We showed that it is possible to train an effective in-context learner tailored to a particular classification task with spurious features.
We did this by introduce two key techniques: (a) permuting input embedding dimensions and (b) forming ICL sequences with intermediate queries simulating distribution shift.
We provided evidence that the learned algorithm is highly competitive on the task it was trained on.
However, we found that while it generalizes to other tasks without spurious features, it does not work for tasks with other spurious features.
Understanding this failure mechanistically and exploring techniques for enabling better generalization are key future research directions.

We next explored training on synthetic ICL instances of diverse tasks and showed that it is possible to obtain an in-context learner that generalizes to unseen tasks, even with different data generating processes.
We established the usefulness of two more techniques: (c) passing example groups as input and (d) promoting learning of induction heads by occasionally querying past context examples.
We believe there is a room for improving in-context learning via improved strategies of choosing intermediate queries and possibly optimizing worst-group loss.
Understanding why the learned algorithm fails under extreme distribution shifts and why variants with permutations fail more (see Figure 12) is an interesting question to explore.
Another interesting direction to explore is to find out what exact algorithm is learned in the process of training on diverse tasks.
Based on the results presented in this work, we conclude that the learned algorithm is neither 1-NN, ERM, or GroupDRO.

Our work has several limitations.
First of all, training a transformer-based in-context learner with high-dimensional image embeddings is computationally costly (see Appendix A for information on compute resources), although it is faster than the baselines at inference.
For this reason, we did not explore more datasets and pretrained image embeddings.
We believe main conclusions of our work will be unchanged.
Second, we experimented with only one model size, width, and depth.
Larger models might behave differently (Wei et al., ).
Third, in our iNaturalist experiments, we considered only one “type” of spurious features.
It is likely that this choice has significant effect on the learned algorithm and its generality.
Future research should explore more ways of synthesizing spurious features and consider varying severity of the challenge posed by spurious features.
The latter can be done by considering multiple spurious features, introducing label imbalance, varying magnitude of spurious correlations, and varying the margin spurious features provide.
Finally, we acknowledge that the proposed approach requires spurious feature annotations, which is typically costly to obtain.
However, they are required only at training time, *not at inference*,
As we have shown, this limitation can be addressed by creating synthetic data, in which case spurious annotations are readily available.

## Acknowledgments

We would like to thank anonymous reviewers for valuable feedback. Samvel Karapetyan’s work is supported by Yandex Armenia Scholarship.

## Appendix A Further experimental details

#### Baselines.

For empirical risk minimization as a baseline, we tune 2 hyperparameters: learning rate (0.01 or 0.001) and number of epochs (100 or 200).
For GroupDRO we additionally tune its parameter that controls adaptiveness of group weights (0.01, 0.1, or 1) and we also try an optional strong L2 regularization (1.0 weight decay), as it has been observed to be useful for small datasets [Sagawa* et al., ].

#### Transformer-based methods.

In all transformer-based approaches, we train a causal decoder-only GPT-J transformer with 80M parameters that has 6 transformer layers with 8 multi-head attention, 768 model dimensionality, and 3072 hidden dimensionality.
When training on iNaturalist, we add a layer normalization [Ba et al., ] on transformer input, as we expect input norms to change when we evaluate on Waterbirds-based datasets.
The transformer input sequence in the proposed approach consists of 3 types of tokens: context image embeddings, query image embeddings, and label/group annotations.
While the network can rely on positions and content to distinguish image embeddings from annotations, we found it to be helpful to encode token types explicitly.
We do this by setting the first 3 dimensions of a token to be a one-hot vector representing token type (context image embedding, query image embedding, or annotation).
When permuting dimensions, we do the permutation before encoding token types to keep the location of token-type information consistent.
In our preliminary experiments and development, we used $n=128$ context length.
Apart from improved performance, we did not observe significant qualitative differences when we switched to larger context lengths for final experiments.

#### Evaluation and model selection.

For all transformer-based approaches and baselines, we do a grid search to find the best combination of hyperparameters.
In particular, we train each configuration with 5 different random seeds and selected one with the highest average test performance.
Importantly, for baseline methods model selection is done for each context length independently, while for transformer-based methods model selection is done once with respect to the test performance at maximum context length observed during training.
All evaluations are done on 8192 sequences, where the first $n$ examples are sampled from the corresponding train set while the query is sampled from the test set with a balanced group distribution.
Finally, even when training transformers on permuted image embeddings, we do not apply permutations during evaluation.
In all figures throughout this work, shaded regions show standard deviation across the 5 training runs.

Note that the most principled model selection approach would be selecting models based on a metric calculated on a dataset similar to the training set (e.g., a held-out part of training set), rather than the test set.
For example, in the case of experiments on Waterbirds or Waterbirds-severe, the principled approach would be to select based on performance on sequences where the context part is sampled from the training set, while the final query is sampled from a held-out validation set with balanced group distribution.
We tried this way of model selection and did not observe significant changes.
In the case of experiments on iNaturalist, the principled approach would be to select based on performance on sequences where the context part is sampled from the training set, while the final query is sampled from the hold-out part the training set.
We observed that this in-distribution metric is always around 99.5%-100%, and can be non-informative for model selection.
This is a typical scenario in OOD generalization (see for example  [Gulrajani and Lopez-Paz, ] or [Wenzel et al., ]).

#### Definitions of metrics.

Given a set of predictions on Waterbirds or Waterbirds-severe, worst-group accuracy is defined as the lowest accuracy of predictions among the 4 groups.
Note that worst-group accuracy is not applicable to iNaturalist, as different ICL sequences correspond to different classification tasks and hence form different groups.
For this reason, we introduce minority-group and majority-group accuracies.
Given a triplet $(C,q,\widehat{y})$, where $C$ is a context, $q$ is query, and $\widehat{y}$ is a prediction on $q$, we call $\widehat{y}$ a minority (majority) prediction, if $q$ is among the least (most) represented group(s) of the context $C$.
Given a list of triplets $(C,q,\widehat{y})$, we define minority (majority) group accuracy as the accuracy among minority (majority) predictions.

#### Compute resources.

We used NVIDIA A100 GPUs with 40GB memory to train transformer-based methods.
The network we considered is small enough to fit on one GPU with batch size 32 when $n=400$ (iNaturalist experiments) and batch size 24 when $n=512$ (Waterbirds and Waterbirds-severe experiments).
We did mixed 16-bit training to save compute and did not notice any quality degradation.
A single training takes around 12 hours for iNaturalist experiments and around 18 hours for Waterbirds experiments.
We used a mix of CPUs and weaker GPUs to train baselines, as they are not computationally as demanding.

## Appendix B Additional results

In addition to the figures presented in the main text, here we provide the exact experimental resources for multiple transformer-based and baseline approaches, some of which were not included in the main text due to space constraints.
Recall that +P means permuting input dimensions, +I means promoting learning of induction heads, and +G means passing example groups as input to in-context learning transformers.

Figure: Figure 13: Majority-group and worst-group test accuracies on Waterbirds-severe as a function of context size for the naive approach with a single modification of making the last example (query) group-balanced. Shaded regions show standard deviation across 5 training runs. As expected, at intermediate context lengths this method performs similar to the naive approach, but is much better at the training context length.
Refer to caption: x16.png

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

Table 1 presents worst-group accuracies on the test set of Waterbirds for 3 sets of approaches: (a) in-context learners trained on Waterbirds itself, (b) in-context learners trained on iNaturalist, and (c) baselines.
Similarly, Table 2 presents worst-group accuracies on the test set of Waterbirds-severe for 3 sets of approaches: (a) in-context learners trained on Waterbirds-severe itself, (b) in-context learners trained on iNaturalist, and (c) baselines.
As RoPE-based transformers are not good at length extrapolation [Press et al., ], we do not attempt evaluating models trained on iNaturalist with context size 400 on 512-long sequences of Waterbirds or Waterbirds-severe.
Finally, Table 3 presents minority-group accuracy on out-of-distribution classes of iNaturalist for two sets of approaches: (a) in-context learners trained on iNaturalist itself and (b) baselines.

#### Experiments on CelebA.

To further verify our main findings presented in Section 2, we conduct experiments on another popular visual classification tasks CelebA [Liu et al., ].
In CelebA, the task is to classify blond vs non-blond persons, with sex being a spuriously correlated variable.
Notably, the spurious correlation is asymmetric, in the sense that blond and non-blond women are almost equally represented, while blond men are much less represented compared to non-blond men.
We follow the design of Waterbirds experiments in our CelebA experiments, with the only difference that we set the group distribution of context examples to $(0.25,0.25,0.05,0.45)$, where group 0 are non-blond men, group 1 are non-blond women, group 2 are blond men, and group 4 are blond women.
Table 4 presents worst-group accuracies on the test set of CelebA for 2 sets of approaches: (a) in-context learners trained on CelebA itself and (b) baselines algorithms.
As in our Waterbirds experiments, we see that it is essential to permute input embeddings *and* to form ICL sequences in the proposed fashion.
Unlike Waterbirds, comparing “Proposed + P” with “Proposed + P + G” we see that providing spurious annotations in-context provides significant gains.
Figure 4 demonstrates that both of these approaches outperform 1-NN, ERM, and GroupDRO.

**Table 1: Complete results on Waterbirds. Reported numbers are average worst-group test accuracies, along with the their standard deviation. The top half of in-context learners were trained on Waterbirds itself, while the ones in the bottom half were training on iNaturalist.**
| Method / Context size | 4 | 8 | 16 | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Naive | 87.02<br>(0.79) | 84.52<br>(1.00) | 85.14<br>(0.42) | 84.82<br>(0.89) | 83.41<br>(0.75) | 84.45<br>(1.04) | 85.08<br>(1.15) | 84.82<br>(1.26) |
| Naive + P | 70.92<br>(1.18) | 75.32<br>(1.11) | 80.66<br>(0.68) | 83.24<br>(0.35) | 86.87<br>(0.62) | 89.87<br>(0.85) | 91.94<br>(0.75) | 92.60<br>(0.59) |
| Proposed | 87.91<br>(1.29) | 85.63<br>(2.20) | 86.51<br>(2.17) | 85.42<br>(1.73) | 85.12<br>(2.34) | 85.86<br>(2.22) | 86.72<br>(1.89) | 86.89<br>(1.82) |
| Proposed + I | 88.18<br>(1.07) | 85.89<br>(1.31) | 86.68<br>(1.02) | 86.01<br>(1.39) | 84.82<br>(1.02) | 85.92<br>(1.23) | 86.07<br>(1.27) | 86.46<br>(1.57) |
| Proposed + P | 68.44<br>(2.40) | 73.46<br>(2.53) | 80.00<br>(2.06) | 83.71<br>(2.15) | 87.30<br>(1.79) | 90.02<br>(1.16) | 92.11<br>(1.65) | 91.95<br>(1.20) |
| Proposed + P + I | 68.05<br>(1.51) | 72.47<br>(1.80) | 78.97<br>(1.12) | 82.58<br>(0.68) | 86.39<br>(0.69) | 90.00<br>(0.60) | 91.78<br>(0.56) | 92.17<br>(0.86) |
| Proposed + G | 88.74<br>(1.01) | 87.00<br>(1.60) | 87.62<br>(1.58) | 86.86<br>(1.31) | 86.18<br>(1.33) | 86.91<br>(0.98) | 87.26<br>(1.11) | 86.95<br>(1.21) |
| Proposed + G + I | 88.89<br>(0.53) | 87.49<br>(0.69) | 87.70<br>(0.74) | 86.90<br>(0.95) | 86.03<br>(0.71) | 86.64<br>(0.72) | 87.29<br>(0.77) | 87.35<br>(1.00) |
| Proposed + G + P | 68.47<br>(2.32) | 73.74<br>(2.00) | 79.21<br>(1.68) | 82.85<br>(1.33) | 86.55<br>(1.17) | 89.98<br>(0.72) | 92.00<br>(0.82) | 93.05<br>(0.40) |
| Proposed + G + P + I | 68.24<br>(1.88) | 73.78<br>(1.67) | 80.23<br>(0.94) | 83.02<br>(1.22) | 86.94<br>(1.31) | 89.89<br>(0.91) | 92.46<br>(1.00) | 92.69<br>(1.15) |
| 1-NN | 65.29<br>(1.23) | 72.53<br>(1.11) | 79.15<br>(1.16) | 82.81<br>(0.63) | 87.49<br>(1.18) | 90.00<br>(1.05) | 91.96<br>(0.51) | 93.40<br>(0.27) |
| ERM | 63.04<br>(1.22) | 70.76<br>(1.01) | 77.32<br>(1.16) | 83.04<br>(1.09) | 85.95<br>(1.38) | 87.20<br>(0.77) | 88.10<br>(0.98) | 88.48<br>(0.45) |
| GroupDRO | 64.61<br>(1.79) | 71.52<br>(0.73) | 77.81<br>(1.19) | 83.45<br>(1.57) | 87.34<br>(1.42) | 88.30<br>(0.91) | 89.79<br>(0.81) | 91.12<br>(0.62) |
| Naive | 69.77<br>(1.37) | 77.98<br>(1.51) | 79.23<br>(0.83) | 81.20<br>(1.35) | 82.57<br>(1.52) | 83.85<br>(1.56) | 84.21<br>(1.19) | - |
| Naive + P | 66.47<br>(1.17) | 73.12<br>(1.44) | 77.85<br>(1.74) | 81.76<br>(1.49) | 86.36<br>(0.86) | 88.02<br>(1.25) | 89.68<br>(0.77) | - |
| Proposed | 69.75<br>(5.51) | 77.51<br>(3.01) | 79.20<br>(2.11) | 81.39<br>(1.49) | 82.04<br>(1.29) | 83.51<br>(0.97) | 84.63<br>(0.80) | - |
| Proposed + I | 70.73<br>(1.42) | 77.10<br>(1.76) | 78.90<br>(1.49) | 80.86<br>(1.74) | 82.22<br>(1.72) | 84.22<br>(1.45) | 84.69<br>(1.47) | - |
| Proposed + P | 66.09<br>(1.49) | 73.71<br>(1.17) | 78.33<br>(0.69) | 82.75<br>(0.83) | 86.32<br>(0.52) | 88.85<br>(0.72) | 89.98<br>(1.35) | - |
| Proposed + P + I | 65.51<br>(2.16) | 70.91<br>(2.32) | 75.94<br>(3.04) | 81.51<br>(1.90) | 86.41<br>(1.50) | 89.39<br>(0.98) | 91.08<br>(0.75) | - |
| Proposed + G | 70.98<br>(2.52) | 78.41<br>(1.25) | 79.67<br>(1.26) | 81.59<br>(1.42) | 82.42<br>(1.28) | 83.91<br>(1.64) | 84.31<br>(1.31) | - |
| Proposed + G + I | 71.94<br>(2.70) | 78.56<br>(1.65) | 80.62<br>(1.66) | 82.31<br>(1.76) | 83.52<br>(1.57) | 84.52<br>(1.32) | 85.35<br>(1.20) | - |
| Proposed + G + P | 67.55<br>(0.78) | 73.79<br>(0.33) | 78.32<br>(0.93) | 82.56<br>(1.31) | 86.01<br>(1.09) | 89.40<br>(1.22) | 90.99<br>(1.15) | - |
| Proposed + G + P + I | 69.18<br>(2.76) | 74.13<br>(2.06) | 79.18<br>(1.81) | 83.17<br>(0.85) | 87.25<br>(0.37) | 90.67<br>(0.80) | 92.23<br>(0.69) | - |

**Table 2: Complete results on Waterbirds-severe. Reported numbers are average worst-group test accuracies, along with the their standard deviation. The top half of in-context learners were trained on Waterbirds-severe itself, while the ones in the bottom half were training on iNaturalist.**
| Method / Context size | 4 | 8 | 16 | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Naive | 83.04<br>(1.92) | 80.78<br>(1.58) | 80.78<br>(1.85) | 79.43<br>(2.77) | 80.50<br>(2.43) | 80.29<br>(2.30) | 81.67<br>(2.25) | 82.02<br>(2.72) |
| Naive + P | 10.89<br>(2.71) | 28.61<br>(4.98) | 46.23<br>(4.17) | 58.40<br>(2.46) | 67.13<br>(2.34) | 74.28<br>(2.25) | 77.18<br>(3.11) | 77.49<br>(4.08) |
| Proposed | 82.64<br>(1.56) | 81.01<br>(2.23) | 81.90<br>(1.80) | 81.36<br>(1.69) | 81.94<br>(1.91) | 81.70<br>(1.62) | 82.35<br>(1.72) | 82.09<br>(2.15) |
| Proposed + I | 83.23<br>(1.30) | 80.76<br>(1.93) | 81.65<br>(2.38) | 81.46<br>(2.11) | 81.63<br>(2.40) | 81.34<br>(2.01) | 81.46<br>(2.32) | 82.24<br>(3.49) |
| Proposed + P | 61.94<br>(8.91) | 68.23<br>(5.53) | 75.94<br>(3.13) | 81.93<br>(1.53) | 85.76<br>(2.03) | 88.36<br>(1.30) | 90.01<br>(1.98) | 90.20<br>(2.65) |
| Proposed + P + I | 64.01<br>(4.05) | 72.22<br>(4.43) | 78.45<br>(2.79) | 82.00<br>(2.20) | 85.86<br>(1.64) | 88.13<br>(1.39) | 90.09<br>(1.73) | 90.59<br>(1.54) |
| Proposed + G | 82.02<br>(3.37) | 81.15<br>(3.56) | 83.11<br>(1.84) | 81.22<br>(2.08) | 81.30<br>(1.72) | 81.90<br>(1.62) | 82.48<br>(1.62) | 82.44<br>(1.34) |
| Proposed + G + I | 82.61<br>(3.42) | 80.48<br>(2.69) | 81.20<br>(3.55) | 80.13<br>(3.18) | 81.09<br>(2.86) | 80.84<br>(2.47) | 81.61<br>(2.36) | 81.84<br>(2.51) |
| Proposed + G + P | 59.11<br>(2.89) | 64.44<br>(5.67) | 71.30<br>(3.74) | 79.46<br>(0.83) | 85.21<br>(1.54) | 88.60<br>(1.36) | 90.65<br>(1.01) | 91.38<br>(1.14) |
| Proposed + G + P + I | 64.26<br>(5.81) | 70.05<br>(4.01) | 77.76<br>(1.77) | 82.38<br>(1.66) | 86.56<br>(0.88) | 89.09<br>(1.02) | 90.75<br>(0.96) | 90.82<br>(0.73) |
| 1-NN | 5.44<br>(0.60) | 4.50<br>(0.43) | 3.49<br>(0.21) | 27.92<br>(0.54) | 45.04<br>(0.88) | 52.58<br>(1.39) | 61.74<br>(0.48) | 71.20<br>(0.58) |
| ERM | 6.81<br>(0.44) | 4.35<br>(0.26) | 1.87<br>(0.24) | 35.30<br>(1.55) | 29.52<br>(1.49) | 45.84<br>(1.00) | 65.35<br>(0.53) | 75.69<br>(0.88) |
| GroupDRO | 7.42<br>(0.57) | 5.26<br>(0.35) | 2.75<br>(0.29) | 17.62<br>(0.65) | 45.47<br>(1.18) | 65.13<br>(1.06) | 78.57<br>(0.77) | 86.89<br>(0.57) |
| Naive | 48.18<br>(3.52) | 49.39<br>(3.28) | 48.71<br>(6.49) | 52.58<br>(4.56) | 54.10<br>(6.04) | 56.41<br>(5.03) | 56.86<br>(4.75) | - |
| Naive + P | 0.88<br>(0.45) | 0.06<br>(0.05) | 0.00<br>(0.00) | 0.13<br>(0.29) | 0.19<br>(0.43) | 0.13<br>(0.29) | 0.02<br>(0.04) | - |
| Proposed | 49.04<br>(2.76) | 53.39<br>(4.74) | 54.82<br>(8.82) | 59.44<br>(10.75) | 61.04<br>(12.23) | 62.26<br>(12.31) | 63.77<br>(12.38) | - |
| Proposed + I | 48.45<br>(6.15) | 52.44<br>(11.15) | 54.74<br>(10.69) | 58.67<br>(11.38) | 60.37<br>(9.19) | 62.42<br>(8.69) | 63.27<br>(9.15) | - |
| Proposed + P | 1.88<br>(0.56) | 0.27<br>(0.21) | 0.06<br>(0.10) | 0.08<br>(0.13) | 0.30<br>(0.60) | 0.15<br>(0.28) | 0.03<br>(0.06) | - |
| Proposed + P + I | 2.27<br>(0.74) | 0.66<br>(0.49) | 0.15<br>(0.14) | 1.18<br>(1.09) | 2.50<br>(2.49) | 1.14<br>(0.88) | 0.50<br>(0.20) | - |
| Proposed + G | 50.00<br>(5.03) | 52.31<br>(5.05) | 53.69<br>(4.54) | 57.87<br>(3.33) | 59.11<br>(3.16) | 60.33<br>(3.36) | 62.30<br>(3.01) | - |
| Proposed + G + I | 51.78<br>(5.76) | 53.87<br>(6.15) | 55.07<br>(6.20) | 60.15<br>(7.40) | 60.73<br>(8.02) | 62.40<br>(8.01) | 61.86<br>(7.77) | - |
| Proposed + G + P | 1.52<br>(0.69) | 0.16<br>(0.13) | 0.00<br>(0.00) | 0.10<br>(0.20) | 0.04<br>(0.05) | 0.03<br>(0.05) | 0.36<br>(0.73) | - |
| Proposed + G + P + I | 1.59<br>(0.17) | 0.23<br>(0.16) | 0.08<br>(0.10) | 0.50<br>(0.69) | 1.91<br>(3.05) | 2.19<br>(3.67) | 2.34<br>(4.00) | - |

**Table 3: Complete results on iNaturalist. Reported numbers are average minority-group accuracies on the OOD test set of iNaturalist, along with the their standard deviation.**
| Method / Context size | 4 | 8 | 16 | 32 | 64 | 128 | 256 | 400 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Proposed | 91.80<br>(0.39) | 93.20<br>(0.29) | 93.71<br>(0.35) | 94.58<br>(0.22) | 95.01<br>(0.42) | 95.27<br>(0.42) | 95.30<br>(0.40) | 94.89<br>(0.27) |
| Proposed + I | 92.88<br>(0.31) | 93.82<br>(0.37) | 94.61<br>(0.56) | 95.36<br>(0.45) | 95.76<br>(0.40) | 95.90<br>(0.44) | 95.94<br>(0.18) | 95.06<br>(0.54) |
| Proposed + P | 92.04<br>(0.22) | 92.90<br>(0.30) | 94.80<br>(0.32) | 96.64<br>(0.30) | 97.65<br>(0.20) | 98.39<br>(0.27) | 98.49<br>(0.14) | 98.55<br>(0.23) |
| Proposed + P + I | 92.15<br>(0.28) | 92.97<br>(0.30) | 94.67<br>(0.28) | 96.86<br>(0.21) | 97.80<br>(0.29) | 98.46<br>(0.20) | 98.54<br>(0.11) | 98.61<br>(0.25) |
| Proposed + G | 92.48<br>(0.45) | 93.27<br>(0.72) | 93.88<br>(0.43) | 94.91<br>(0.63) | 94.99<br>(0.38) | 95.29<br>(0.45) | 95.13<br>(0.33) | 94.64<br>(0.43) |
| Proposed + G + I | 92.59<br>(0.33) | 93.80<br>(0.23) | 94.18<br>(0.38) | 95.50<br>(0.33) | 95.82<br>(0.41) | 95.83<br>(0.34) | 95.82<br>(0.55) | 95.28<br>(0.60) |
| Proposed + G + P | 91.90<br>(0.17) | 92.84<br>(0.19) | 94.69<br>(0.15) | 97.28<br>(0.31) | 98.29<br>(0.13) | 98.70<br>(0.19) | 98.85<br>(0.19) | 99.00<br>(0.11) |
| Proposed + G + P + I | 92.28<br>(0.10) | 93.25<br>(0.09) | 94.93<br>(0.22) | 97.73<br>(0.07) | 98.44<br>(0.20) | 98.99<br>(0.09) | 99.04<br>(0.14) | 99.06<br>(0.07) |
| 1-NN | 92.08<br>(0.64) | 94.56<br>(0.39) | 95.84<br>(0.16) | 97.17<br>(0.23) | 97.84<br>(0.12) | 98.49<br>(0.20) | 98.55<br>(0.23) | 98.80<br>(0.21) |
| ERM | 89.67<br>(0.43) | 92.98<br>(0.30) | 94.65<br>(0.17) | 96.17<br>(0.24) | 96.88<br>(0.23) | 97.70<br>(0.21) | 98.15<br>(0.17) | 98.43<br>(0.11) |
| GroupDRO | 91.20<br>(0.55) | 93.79<br>(0.39) | 95.33<br>(0.18) | 97.39<br>(0.20) | 97.85<br>(0.20) | 98.46<br>(0.13) | 98.91<br>(0.20) | 99.01<br>(0.18) |

**Table 4: Complete results on CelebA. Reported numbers are average worst-group test accuracies, along with the their standard deviation. All in-context learning were train on CelebA itself.**
| Method / Context size | 4 | 8 | 16 | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Naive | 24.88<br>(2.03) | 24.56<br>(2.26) | 25.80<br>(1.98) | 25.14<br>(2.11) | 23.85<br>(1.91) | 25.62<br>(1.63) | 25.84<br>(2.10) | 26.20<br>(1.42) |
| Naive + P | 20.72<br>(2.21) | 17.27<br>(2.38) | 12.43<br>(2.04) | 14.85<br>(2.33) | 13.56<br>(1.60) | 16.17<br>(2.83) | 20.16<br>(4.14) | 26.03<br>(5.13) |
| Proposed | 25.83<br>(1.77) | 25.42<br>(1.66) | 26.85<br>(1.52) | 25.80<br>(1.60) | 25.18<br>(1.06) | 26.65<br>(2.11) | 26.89<br>(1.41) | 27.53<br>(1.31) |
| Proposed + P | 26.90<br>(4.56) | 26.54<br>(1.46) | 27.00<br>(2.72) | 37.30<br>(1.55) | 47.29<br>(2.67) | 54.66<br>(2.67) | 60.48<br>(2.55) | 68.45<br>(1.99) |
| Proposed + G | 23.87<br>(1.50) | 24.67<br>(1.51) | 25.60<br>(1.30) | 24.95<br>(1.12) | 24.42<br>(1.17) | 25.77<br>(1.24) | 25.70<br>(0.81) | 26.55<br>(1.39) |
| Proposed + G + P | 26.71<br>(3.57) | 32.06<br>(4.62) | 39.21<br>(6.06) | 46.13<br>(5.38) | 53.41<br>(3.73) | 59.37<br>(3.07) | 64.39<br>(1.84) | 69.58<br>(1.72) |
| 1-NN | 35.87<br>(1.48) | 37.63<br>(0.86) | 36.08<br>(1.13) | 37.86<br>(0.45) | 38.28<br>(0.80) | 36.40<br>(0.32) | 36.90<br>(0.99) | 37.81<br>(0.36) |
| ERM | 30.70<br>(1.05) | 28.93<br>(0.64) | 26.37<br>(0.66) | 30.75<br>(0.40) | 34.29<br>(0.93) | 38.64<br>(1.29) | 45.18<br>(1.39) | 49.92<br>(1.28) |
| GroupDRO | 35.32<br>(0.88) | 34.64<br>(0.99) | 30.24<br>(1.01) | 37.49<br>(0.60) | 47.11<br>(0.59) | 54.56<br>(0.66) | 56.11<br>(0.60) | 61.47<br>(0.93) |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: See the weakness part.
soundness: 3
strengths: 1. The paper addresses an important and interesting problem in in-context learning, as spurious correlations are a common issue in real-world scenarios.
The proposed method is generally sound and is validated by experimental results.
summary: This paper investigates the problem of in-context learning in the presence of spurious correlations. The authors first examine the single-task setting and demonstrate that conventional approaches are vulnerable to spurious correlations. To address this issue, they propose several techniques to improve generalization in the presence of spurious correlations for in-context learning. However, they find that these methods do not generalize well to multiple tasks. Consequently, they develop additional techniques for training on diverse datasets and demonstrate the effectiveness of their method through experiments.
weaknesses: 1. A major issue, although acknowledged in Section 5, is the requirement of spurious feature annotations. This is a strong assumption in practice and may significantly limit the applicability of the proposed method.
2. In Section 2.2, an intuitive method would be to sample all in-context samples $x_i$ such that they all have a uniform group distribution. This could mitigate the issue of spurious correlations. Could the authors explain why they did not adopt this method? What are the experimental results of this approach?
3. The paper would benefit from more intuitive figures to explain the settings and proposed methods. Given the complexity of these settings, visual aids would help readers better understand them.
4. The authors construct a simulated dataset, Waterbird-Severe, to validate the effectiveness of their method. However, this dataset may not be realistic and contains a spurious correlation that is unrealistically strong. As a result, it would be more convincing if the authors conducted experiments on more realistic datasets.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Some simulations and theoretical analysis or proof can make the paper stronger. 
2. It might worth to explore this idea with more benchmarks and experiments in visionLLM or LLM setup? ICL are frequently used in those models with prompts and I am curious if this method can improve the capability of ICL there, which can greatly enhance the impacts.
soundness: 2
strengths: 1. This paper proposes an interesting idea and the method is novel. 
2. Conducted analysis and experiments to demonstrate the idea and effectiveness.
3. Overall, the paper is well-written.
summary: In this paper, authors explored the limit of current ICL framework by present more challenging setup in image classification with presence of spurious correlation in the features. They showed that existing approaches can lead to poor performance under distribution shift due to memorization issue. They proposed a new way of forming new ICL instances that can outperform baselines and mitigate this issue.
weaknesses: 1. Impact of the paper is limited given current setup and experiments. 
2. Experiments are very limited and need more benchmarks to show this approach is generalizable to other cases.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Could the authors provide a visual analysis to help readers intuitively understand the specific spurious features on which the existing methods and the proposed methods are effective or ineffective?
soundness: 2
strengths: 1. The paper introduces a new perspective on ICL by focusing on the impact of spurious correlations, which is an important yet underexplored area in the context of large language models.
2. The authors provide a thorough analysis of the problem, backed by empirical evidence from experiments on various datasets, including Waterbirds, CelebA, and iNaturalist.
3. The proposed method to mitigate task memorization and increase robustness to spurious features are innovative and show promising results in improving ICL, especially when the number of context examples is limited.
4. The paper not only focuses on in-context learning for a specific task but also addresses the generalization of the learned algorithm to unseen tasks, which is a critical aspect of real-world applicability.
5. The paper is well-organized, with a clear presentation of the problem, methodology, experiments, and results, making it easy to follow the authors' line of reasoning.
summary: The paper explores the challenges and limitations of training large language models for in-context learning (ICL) on classification tasks that involve spurious features. The authors argue that conventional ICL approaches are susceptible to spurious correlations and can lead to task memorization rather than leveraging context for predictions. To address these issues, the paper proposes novel techniques for training ICL models that are more robust to spurious features and less prone to task memorization. The proposed methods involve permuting input embedding dimensions and constructing ICL instances with intermediate queries to simulate distribution shifts. The paper demonstrates that these techniques can lead to in-context learners that match or outperform established algorithms like ERM and GroupDRO on certain tasks. Furthermore, the authors show that by training on a diverse dataset of synthetic ICL instances, it is possible to obtain a more general-purpose in-context learner that can generalize to unseen tasks involving spurious features.
weaknesses: My primary concern with this paper lies in the limited support for the effectiveness of the proposed methods: "(a) passing example groups as input and (d) promoting learning of induction heads by occasionally querying past context examples."  The paper lacks a strong theoretical foundation for these methods, and the experimental validation is somewhat restricted.

Specifically, I have the following observations:

1. The experiments consider only a single model structure, depth, and width. The impact of these architectural hyperparameters on the effectiveness of the proposed methods remains unexplored.  Investigating different architectures could reveal whether the observed benefits are consistent across a wider range of models.

2. The experimental results are confined to three datasets: Waterbirds, CelebA, and iNaturalist. While these datasets are commonly employed in the Out-of-Distribution (OOD) generalization literature, they are relatively simple and may not fully capture the complexity of spurious correlations encountered in real-world applications. Evaluating the proposed methods on more diverse and challenging datasets would strengthen the claim of their effectiveness.

3. The paper does not address the possibility of data leakage, which could inadvertently inflate the reported performance.  A thorough analysis of potential data leakage sources and mitigation strategies is necessary to ensure the validity of the results.

To enhance the paper, I recommend the authors consider the following:

1. Provide a theoretical foundation: Develop a theoretical framework that explains why the proposed methods are expected to mitigate spurious correlations. This could involve analyzing the underlying mechanisms through which these methods operate and providing formal guarantees or bounds on their performance.

2. Increase experimental diversity: Expand the experimental evaluation to include a wider range of model architectures, hyperparameters, and datasets. This would provide a more comprehensive assessment of the effectiveness and generalizability of the proposed methods.

3. Address potential data leakage: Thoroughly investigate and discuss potential sources of data leakage in the experimental setup. Implement appropriate mitigation strategies and report any remaining limitations.

By addressing these issues, the authors can provide a more compelling argument for the effectiveness of their proposed methods in addressing spurious correlations.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. When producing Waterbirds-severe dataset, how is $\tilde s$ computed? Is it added before or after the permutation if using $+P$ technique.
2. For the $+I$ technique, is there any evidence that the performance gain is attributed to the forming of induction heads? This is not very obvious to me since the pair $(x_i,\tilde y_i)$ does not contribute to the loss function explicitly and the label of $q_{i-1}$ does not appear in the sequence. 
3. In the last paragraph, the authors claims the proposed approach does not require spurious feature annotations at inference time. Am I missing anything? I feel using $\tilde y$ to represent group ($+G$) technique requires the annotations at inference time since the inference sequence is also using the same form of $\tilde y_i$.
4. The questions in the weakness part.
soundness: 3
strengths: This paper has a clear structure and interesting findings. I appreciate that authors also conduct experiments in the setting of different spurious correlations and test the trained model on unseen tasks.
summary: This paper studies in-context learning image classfication tasks in presence of spurious correlations. The authors find that:
* permuting the image embeddings prevents transformers from memorizing the task and promotes in-context learning;
* inserting intermediate queries along with permuting embeddings enhences robustness to spurious correlations;
* using proceeding examples as queries improves the performance on both majority and minority groups; 
* passing group values to the context improves the performance on minority groups.

Except for Waterbirds dataset, these findings are also verified on (i) the Waterbirds-severe dataset where the spurious features are enhenced; (ii) a varying spurious correlations dataset adapted from iNaturalist. However, the transformer trained on dataset in (ii) fails to be robust against the spurious correlations in the dataset in (i).
weaknesses: 1. Motivation for the proposed approach is unclear. I feel the effcacy of the proposed approach should be attributed to the fact that $q_i$ are sampled from the uniform group distribution and the autoregressive loss is only taken on these queries (Please correct me if I am wrong). This seems a bit werid to me -- under the spurious correlation setting, is it appropriate to assume that we have access to the uniform group distribution at the training stage? If one has access to the uniform group distribution data, why not simply use a training sequence consisting of these data as in-context examples?

2. Lacks of in-depth analysis and interpretation. For example, in section 2.4 the authors obtain interesting results by stwiching the evaluation data to (i) Waterbirds-severe; (ii) background prediction on Waterbirds; (iii) background prediction on Waterbirds with group-balanced context. There is one vague hypothesis saying the model learns to ignore the specific spurious feature in the pretrain data, which lacks further analysis and evidence. The trained model gets around 50% overall test accuracy and 9% worst-group accuracy on predicting the backgroud, which justifies your this hypothesis. However, the results on Waterbirds-severe and group-balanced Waterbirds seem to suggest the model still *in-context learns* to use the background feature. It feels to me the model learns some "biased initialization" (prefer to ignore specific feature) after pretraining on the spurious correlated data with the proposed approach, and then adjusts accroding to the in-context examples on top of this "biased initialization". Discussions along similar directions are lacked. Similar issues for the generality part in section 3. I feel more in-depth discussions would be beneficial.

3. Clarity issue: the paragraph "More concretely, ....." in section 3 is hard to parse. Would be helpful if there's an illustration for this data generating procedure. Also some other minor problems, please see questions.

Meta-review: This paper explores how to make in-context learning (ICL) robust against spurious correlations in classification tasks. The authors find that standard ICL approaches are vulnerable to spurious features and tend to memorize tasks rather than truly learn from context. To address this, the authors propose several techniques including permuting input embeddings and constructing ICL sequences with intermediate queries to simulate distribution shifts, which outperform baseline methods like ERM and GroupDRO on specific tasks. They show it's possible to train ICL models that generalize to unseen tasks by training on diverse synthetic data (this isn't novel and has been shown previously, eg. see [1]). The main strengths are the paper's novel techniques for improving ICL robustness and demonstration of both single-task and multi-task generalization capabilities. The reviewers pointed out weaknesses in limited theoretical foundation for why the proposed methods work, using relatively simple experimental settings with a single model architecture and limited datasets, and strong assumptions about having access to spurious feature annotations during training. 

[1] https://arxiv.org/abs/2207.04179

Decision: Reject