In-context Learning in Presence of Spurious Correlations
Authors: Hrayr Harutyunyan
Rafayel Darbinyan
Samvel Karapetyan
Hrant Khachatrian
Venue: Accepted by TMLR ()

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 in-weights learning 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.

Published in Transactions on Machine Learning Research (04/2026)
In-context Learning in Presence of Spurious Correlations
Hrayr Harutyunyan Google Mind
Rafayel Darbinyan ServiceTTiaan
darbingonnro@ggmail.com
Samvel Karapetyan YerevaNN, Yerevan State University
sameee@@yereeennn.ccmm
Hrant Khachatrian
bawnt@@ycrevemm.com
YerevaNN, Yerevan State University
Reviewed on OpenReview: https:// openreview. net/ forum? id= C9CSSTRR1A
Abstract
Large language models exhibit remarkable capacity for in-context learning, where they learn to solve tasks given few examples. Recent work has shown that transformers be trained to perform simple regression tasks in-context. This work explores the possibility of training 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. Morcover, when the meta-training dataset includes instances of only one task, the conventional approach leads to in-weights learning and fails to produce a model that leverages context for predictions. Based on these observations, we propose a novel technique to train such learner for 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 diverse dataset of synthetic in-context learning instances.
1 Introduction
Large language models, such GPT-3, have the ability of in-context learning (ICL), wherein they learn solve task given few examples in the context (Brown et al. 2020). The most significant aspect of in-context learning is that the learning occurs 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 /Vawaani et al., (017]) from scratch to solve simple learning tasks in-context. In particular, Garg et al. 2022) 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 ReL.U neural networks, and small docision trees.
Training on ICL instances be seen instance of meta-learning (Schmidhuber 1987 Naik & Mammone
1992 Thrun & Pratt 1998), where the goal is to learn learning algorithm. What exact algorithm is learned when training transformers on ICL instances is an open problem. Akyüüek et al. (2023) and Von Oswald el al. 2023) show that Iransformers can implement 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 lincar regression learn algorithms that closely match predictions of the known algorithms, such gradient descent on the ordinary least squares objective and ridge regression. However, there is evidence that the learned algorithm may vary with model scale,

---

Published in Transactions on Machine Learning Research (04/2026)
y
Zeacafommar fo
Pranadormas fo
1
.12
3/2
81
322
M2
Q2
(a) Naive approach
(b) Proposed approach
Figure 1: In-context learning transformer architectures of the naive and proposed approaches. The proposed approach allows arbitrary query tokens after each learning example. Token positions and the attention mask are modified so that these intermediate queries have no effect on other tokens.
depth, and pretraining task diversity (Akyürek et al. 2023 Raventós et al. 2024 Goddard et. al. 2025 In particular, [Ravenós et al. (004)) 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. (2023) show that when trained on TCL, 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 Lopez-Paz (2023) show that in-context learning ability diminishes under strong distribution shifts.
In this work, explore the limits of in-context learning further by studying it in
challenging setting. In particular, motivated by results indicating that in-context learning is susceptible to shorteuts and spurious
correlations (Tang et al. 2023 Zhou el al. 2024: Song et al. 2024), we pose the following question:
Can we oblain an effective classificntion algorithm that is robust to spurious features by training an in-contert learner on a suitable meta-training dataset, rather than designing the learning algorithm manually?
To address this question in isolation, we consider visual classification tasks where some features are spuriously correlated with the label. Such features are predictive of the label but are not causally related it, due which their correlation might not hold at test time. A prominent example is the vs. camel classification task, where the background often correlates with the label, since cows are typically photographed in pastures, while camels are typically photographed in deserts (Beery et al. 2018). It is well known that neural networks trained with gradient-basse based methods to minimize empirical risk can exploit spurious features, causing
performance degradation under distribution shifts affecting these correlations [Toralba & Efros
2011
Ribeiro
et al.
2016 Gururangan el al. 2018 Zech el al. 2018
McCoy al. 2019 Geirhos el. al. 2019 2020
Xiao
et al.
2021
We start analysis in the standard setting of having single classification task with spurious features. We consider the conventional approach of obtaining an in-context learner, wherein a transformer is trained sequences of form (Z1. Tk. Mk, T++1) 1) to predict the label Mk++1 of the query example I4+1- We find that this conventional approach leads to in-weights learning (Chan et al. 2022), wherein models perform classification ignoring the context, essentially memorizing the task. Furthermore, these models lack robustness to changes in 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 effective approach to reduce in- in-weights learning and improve in-context learning. Namely, we find that in- weights learning can be greatly mitigated by randomly permuting input embedding dimensions for each training sequence. To address the issuc of spurious features, we propose novel way of forming ICL instances and a suitable transformer architecture, which work logether 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" al. 2020), suggesting that the in-context learner implements more specializod algorithm.
2

---

Published in Transactions on Machine Learning Research (04/2026)
95
90
Worst-group accuracy
Worst Worst-group group accurac
85
Naive
Naive P
Proposed
Proposed
Proposed
Proposed
60
2
4
8
16
32
64
128
256
512
2
4
8
16
32
64
128
256 512
examples
context examples
(a) Waterbirds
(b) Vaterbirde-severe
Figure Worst-group test accuracy on Waterbirds and Waterbirds--sevre as a function of context size for the naive and proposed methods with without permuting input dimensions. Shaded regions show standard deviation across training runs.
Despite being trained on instances 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 datasot of in-context learning instances for various binary classification tasks with varying spurious features. We demonstrate the efficacy of the proposed lechniques on this dataset too and find that it be improved further by passing spurions feature annotations input and injecting occasional queries requesting the label of preceding context example to promote learning induction heads. The resulting model gencralizes perfectly to unseen tasks, as long as the data generating process is similar. However, generalization to unseen tasks with possibly different dala 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 httpps///ttthhh.ccc///////aa//rs
We summarize main contributions follows.
(i) We show that the conventional approach of training an in-context learner is susceptible presence of spurious fcatures and also leads to in-weights learning in case of single task.
(ii) We propose a suite of novel techniques of forming in-context training data to reduce in-weights learning 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-purpos robust in-context learners by training on a diverse set of synthetic classification tasks involving spurious features.
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 single binary spurious feature, although what follows next applies to label-imbalanced multiclass settings well. Let D be set of training examples for the task, where each example is triplet (æ, s,y) of input Rd spurious feature value s E {0, 1}, and label E {0, 1}. Similarly, let Drest be set of test examples. Importantly, we do not make any assumptions on the data generating process, except that x has some information about a and 8 is predictive of on the training set, but their correlation does not hold on the test set. For an example (I, s, y), we define its group 2y + s. In binary classification task with single binary spurious feature, there are four groups. Without loss of generality, we assume that for majority of training examples we have y = s. Hence, refer to groups and as majority groups, while referring to groups 1 and 2 minority groups
3

---

Published in Transactions on Machine Learning Research (04/2026)
Training transformer to perform linear regression in-context requires millions of ICL training instances, even for small dimensional cases. For example, Garg et al. (2022) use 32 million training instances for 20-dimensional inputs. Next, we consider ways of generating so many ICL instances from a single task.
2.1 A naive approach of constructing ICL instances
The standard approach to constructing an ICL instance is to sample a subset of n examples {(a; 8,, ki from Derain and form a sequence S = (x1, 22, .,,..., Tw, Fn, In+1), where ú: E R is a fixed random representation of either y: or gi (this distinction will be elaborated on later). Then one trains a transformer fo U URR + [o. 1] to prodict y given S - (x1,0) 12.0h..... 2;-1, Di-1, (sec Figure lap, optimizing the following loss function:
1 Σ
CE fo().))
(1)
where CF(y, 3) = - -yogū
log(1 D) the binary cross-entropy loss. We explore two options of setting Dt- In the first option, we set y/ represent yi with a constant vector or its negation in gd. In this case we aim to obtain an in-context learner that is robust to spurions features without receiving spurious feature annotations as input. ERM is one such learner that minimizes average loss training examples and does not require spurious feature annotations. In the second option, we set JJ to represent Ji as a sum of two constant vectors in I&, 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 (Ip has several issues. First, as the mapping from inputs to labels is the same for all ICL training instances, the model can do in-weights learning, wherein the model encodes this stable mapping in its weights and classifies queries without using context examples. In other words, the learned algorithm does no in-context learning and predict ye based solely on Ij, essentially memorizing the task. Second, as all n 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., 2020), which is landbird versus waterbird image classification task where image background (sca land) correlated with the label in the training set (4,795 examples), but not in the 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 pretrained and frozen DINOv2 ViT-B 14 distilled (Oquab) et al. 2023). This way each image is embodded in 7688 While using powerful pretrained representations increases overall performance under distribution shifts (Radford et al. 2021 Mehta et al., 2022), we note that it does not eliminate the problem of spurious correlations. Representations obtained via large-scale self- supervisod 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. 2022 Menon et al. 2021 Kirichenko et al,
2023 Izmailov et al. 2022 Shi et al.) 2023
We train a causal decoder-only GPT-J transformer Wang & Komatsuzaki 2021 with 80M parameters on 2M in-context learning sequences with n 512 and ğ representing labels, constructed from the training 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 & Ba 2014 (B1 = 0.9 and Bn 0.999) with 32 batch size and no weight decay. The learning rate is selected from {3 10 6 10-5, ,1100-1 _.004 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 Section 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.
4

---

Published in Transactions on Machine Learning Research (04/2026)
Figure 2a plots worst-group test accuracy as function of context size n. We see that the 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 in in-weights learning issue. Figure 17a 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
To prevent in in-weights learning and induce in-context learning instead, potentially enabling generalization to other tasks, one approach is to increase the number of training masks/mappings. We propose rotating input embeddings in cach ICL instance independently as simple approach of deriving many tasks from a single source lask. When training with rotated input embeddings, every ICL sequences receives a different permutation and thus represents a different mapping from inputs to labels. This naturally encourages in-context learning, where the model uses context examples to infer the mapping on the fly, instcad of encoding a single mapping. We found that generating random rotation matrices on-the-fly is compulationally expensive and slows down training. We tried generating and storing 10K rotation matrices, but this resulted in fewer 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 by +P in figures and tables; please see Figure 12 for an illustration of this technique). We found this approach to be very effective in inducing in-context learning (see naive + P in Figure 2a). We also see that the difference between majority-grous and worst-group accuracies decreases, although an approximately 5 p.p. gap remains.
When training an ICL transformer, idcally, we would like to simulate the situation of making test prodiction based on context of training examples. Importantly, we would like to simulate the case where the test distribution has balanced groups (i.e., the spurions 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 (I1, Ii,..., In, In+1), where the context examples (æ1...., 2x) are sampled in such a way that the spurious feature is correlated with the label, while the query Fn+1 is sampled to have a uniform group distribution. However, if we again optimize the loss of [ for context lengths less than n, the network will be allowed to make predictions using the spurious fcature, which is undesirable. Please refer to Figure 19 of Section C for evidence of this. To address this, one can compute loss only on the query token, ignoring all intermediate predictions. Unfortunately, this approach leads to reduced sample efficiency and slower learning.
Instead, we introduce a notion of intermediate queries, which allows us to simulate making predictions on test examples sampled from a balanced group distribution at all context lengths at once. Given training examples {xx,, )=) we sample test examples {9:}= from Dirain \ (xI,..., with uniform group distribution. We name these test examples {q:W=n intermediate queríes. Then, form sequence S = (#1. 21, f2. -2..... :n. In). Our goal is to do a single forward pass on this sequence S and obtain the predictions of form 0i _ fo((21. U1, 22, Dz...., Ii, Ds, 9:)) for all i E [n]. To this end we make two modifications, essentially making inpresenlations of context lokens 31, Vi, Tn, In agnostic to the presence of the intermediate queries .1...., 9n- First, we modify the causal attention matrix to disable attention to query tokens, unless query token is attending to itself. See Figure [1 for an illustration for n = 3. Formally, if we enumerate tokens from to 3n and define Ma as the atlention mask for token i attending to token j, then we set
0, i i<> j.
Ma,- =
0, i>j and j = 0 mod 3,
(2)
1, otherwise.
Second, we use modified token positions for computing positional encodings, in order to discount interme- diate query tokens. Namcly, for the sequence (æ1, U1 Q1, Z2 Ø2. D2...., In, In qn), position indices are set to (0, 2, 4,4,... 2n 2,2n 2n 1, 2m). Formally, enumerating tokens from to 3n, the position index of the i-th token is set. to 2 [ + 1) mod 3. Please refer to Figure 1 for an illustration.
5

---

Published in Transactions on Machine Learning Research (04/2026)
Finally, once the predictions on the intermediate queries are obtained, optimize the average loss the query examples:
CE 0t).
(3)
n2
where Dg. is the label of query 4s- Hereafter, we refer to this approach as simply "proposed approach".
Figure 2a compares the proposed and naive approaches with and without input dimension permutations. Without random permutations, the proposed approach marginally outperforms the naive approach. However, the same is not true with random permutations. We found that image embeddings bias toward representing objects more than backgrounds, alleviating the challenge posed by the spuriously correlated background in Waterbirds. In fact, with 512 training examples, the linear probing accuracy of the spurious feature is only 65%, while that of the label is = 95%. For comparison, with ResNet-50 (He et al. 2016) embeddings, the linear probing accuracies of the spurious feature and label are R 81% and x 85% respectively.
DINOv2 have a
For this reason, we create a modified version of Waterbirds by adding a constant vector å or -8 to image embeddings based on the spurious feature s. We scale to have its norm equal to the average norm of image embeddings and verify that the linear probing accuracy of the spurions feature becomes 100%. On this modified Waterbirds dataset, which we call Waterbirds-sseere,, we see large separation between the naive and proposed approaches (see Figure 2b). We also see that without permutations, both the naive and proposed approaches perform identically, indicating 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. Sample reweighting has been observed to be ineffective in overparameterized settings, as all training examples will be perfectly fitted (Byrd & Lipton 2019 Menon et al. 2021
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, of which the latter two have been historically hard to outperform (Gulrajani & Lopez-Pat Paz 2021 Koh et al.] 2021). A comparison to more methods designed for robustuess to spurious correlations is outside of the goal of this work, namely, studying limits of in-context learning. In our comparisons, each baseline on 8192 sequences by training the context part of the sequence and making prediction on the single query. More information about hyperparameters and model selection is presented in Section A
follow the evaluation recipe used for the in-context learners. We evaluate
Figures 3a and 35 compare the proposed and bascline approaches on Waterbirds and Waterbirrss-savera respectively. On Waterbirds, the proposed method outperforms ERM and GroupDRO on almost all context. lengths, but better than 1-NN only for short context lengths. The good performance of 1-NN due to the bias in DINOv2 representations. On Paterbirds-sevent the proposed method outperforms the baselines at all context lengths. From these results, we conclude that in-context learners obtained with the proposed approach implement none of these algorithms.
It should be noted that worst-group accuracies of baselines 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 0.17 % and GroupDRO gets 92.43 0.24 %. This is due to balanced class sampling and sctting the minority ratio to 10% within each class, which higher than the minority ratio of R 5% in the original Waterbirds dataset. One can think of our resampling as weaker form of down-sampling which has been found to be helpful in presence of spurious correlations (Nagarajan et al.
2021 Menon et al. 2021 Tdrissi et al. 2022
Additionally, we verify our findings on one more image classification task CelebA (Liu et al. 2015 and on a natural language inference task MultiNLI (Williams et al. 2018a). CelebA is blond vs non- non-blon 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 MultiNLI, given pair of sentences, a premise and hypothesis, the task
6

---

Published in Transactions on Machine Learning Research (04/2026)
95
Praposed
90
1-NN
ERM
80
85
GroupDRO
80
60
75
Worst
65
1-NN
60
55
GroupORO
16
32
64
128
256
512
16
32
64
128
256
512
Number examples
examples
(a) Waterbirds
(b) Waterbirds-severe
Figure Worst-group test accuracies on Waterbirds and Waterbirds-severe for the proposed approach and conventional methods such 1-NN, ERM, and GroupDRO. Majority-group accuracies are reported in Figure 18 of Section C
is determine whether the hypothesis is entailed by, neutral with, or contradicts the premise. Prior work has observed that there is a spurious correlation between contradictions and the presence of the negation words nobody, no, never, and nothing (Gururangan al. 2018). For our experiments, we frame a binary classification task entails neutral contradicts, with binary spurious feature, which is 1 if and only if the hypothesis has one of the four negation words. For both datasets, we verify the two shortcomings of the conventional approach and demonstrate the efficacy of the proposed techniques compared to the baselines. The results for CelebA are presented in Table 4 and Figure 20a while the results for MultiNLI are presented in Table and Figure 20b
2.4 Generality of the learned algorithm
Since we train in-context learners on ICL instances 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, cannot expect any generality without permuting input dimensions. We take the model obtained with the "Proposed P" approach and probe its in-context learning generality by evaluating and observe ss 2 p.p. overall accuracy drop and & 5 p.p. worst-group accuracy drop. Despite the worsened performance, this indicates that the model treats class labels symbolically, which is remarkable as labels had consistent semantics during training. However, when we evaluate on Waterbirds-severe it gets 100% accuracy on the majority groups and 0% accuracy on minorily groups. Additionally, when we switch the task to predicting the background in the original Waterbirds dataset (now the class becomes spurious feature), the overall test accuracy drops to 54.4%, while the worst-group accuracy drops to 9.3%.
various tasks. We start by swapping the labels of two classes in Waterbirds at evaluation
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.) 2010) dataset, from where the bird images of Waterbirds were taken, we get 99.7% accuracy at context size 100 (the accuracy so high because most pairs of classes are casy distinguish). We also test on binary classification tasks derived from classes belonging to Amphihia and Mammallia supercategories of the iNaturalist (Van Horn et. al.) 2018) 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. Additionally, we do a forward pass on 1024 ICL sequences and collect final query representations various layers of the transformer. We then do a linear probing (512 examples for probe training and 512 for validation) to measure predictability of the background
7

---

Published in Transactions on Machine Learning Research (04/2026)
Prodict background from each layer
02-
Diveril arterrery
43-
32-
33-
Layer number
Figure 4: Linear probing accuracy of the background variable at varions layers of in-context learner trans- formers trained on Waterbirds.
variable. We find that the "Proposed + P" approach reduces background information effectively as we sweep from input to the final layer, while the "Naive" fails to reduce background probing accuracy (see Figure 4
One potential way of improving generality and possibly also performance is passing example groups as input, i.c,, sctting Ii to represent 9:- We did not observe performance improvements or an increase in generality of the learned algorithu when passing groups as input (see Tables 1 and of Section (C) Thus, we conclude that when all ICL instances are derived from a single 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
In Scction 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 sampled minority of cats we overwrite some of their fcatures with those of random dogs. Similarly, we do an analogous operalion for a sampled minority of dogs. In this way some cats share dog features and vice versa. To create a diverse pool of in in-context learning instances, we vary the two classes and the subset of grafted features. Please refer to Figure 16 for an illustration of this grafting operation.
More concretely, we consider the iNaturalist datasel (Van Horn el al. 2018), which contains images from 5,089 natural fine-grained categories and filter out categories that have fewer than 500 images. For testing purposes, from remaining 239 categorics we set apart the ones belonging to Amphibia and Mammalia supercategories, along with 10% of random categories. We denote the set of these 48 categories Cood, and the set of remaining 191 categories Cia, which we use to create ICT. instances for training. For each category in Cid, we hold out half of the examples as in-distribution validation set. To generate an ICL instance, we randomly sample two distinct classes from Cia and sample n/2 images from the training split of each class uniformly at random without. replacement. Please refer to Figure 15 for an illustration of our preprocessing of iNaturalist. We then do the grafting operation, sctting the minority group ratio within cach class to 10%. We select the graßted features randomly, by first picking subset size k uniformly at random from 0 to 199, and then sampling 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 grasting operation with 50% minority group ratio.
Following the experiments in Section 2 we train the same transformer with the proposod 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 Section C
8

---

Published in Transactions on Machine Learning Research (04/2026)
99
Proposed
Proposed
Propesed 1
Proposed P
Minority garrityggruuu group accuracy
Propesed
bilnorty-group group accuracy
98
97
Propesed + 1
96
96
95
94
93
93
2
8
16
32
64 128 256 400
16 32
64
128 256 400
Number examples
Number of context examples
Figure 5: Minority-group accuracy on the OOD test set of iNaturalist for the proposed ap- proach with or without permuting input dimen- sions and promoting induction heads.
Figure 6: Minority-group accuracy on the OOD test set of iNaturalist for the best proposed ap- proach with or without passing group information as input.
Proposed
G
90
1-NN
Minority group accuracy
woost:group accuracy
FRM
GroupDRO
96
75
70
92
Proposed
1-NN
GroupDRD
2
8
16
32
64
128
256 400
2
4
16
32
64
128
256
Number context examples
Number of context examples
Figure 7: Minority-grone group accuracy on the OOD test set of iNaturalist for the best variant of pro- posed approach and conventional methods such as 1-NN, ERM, and GroupDRO.
Figure 8: Worst-group test accuracy on Waterbirds for the best variant of proposed ap- proach trained on iMaturalist and for methods such as 1-NN, ERM, and GroupDRO.
on instances where both categories belong to Good and thus were not observed during training. The results presented in Figure indicaat major difference compared the results in the single-task regime, namely, the proposed approach learns do in-context learning 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 (Figure 23 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 a specific type of circuit 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. 2022). Inspired by this, propose data preparation technique that promotes learning of induction heads. With probability p, replace each intermediate query independently with a random example from the precoeding part of the context (scc Figure 14 for an illustration of this technique). Note that this lype 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 about 10k iterations). Moreover, we see modest performance gains in iNaturalist experiments (see Figure
5 where +7 stands for this technique).
9

---

Published in Transactions on Machine Learning Research (04/2026)
80
Propesed
P+1
1-NN
ERM
GroupDRO
Mors-ereupp
20
10
8
16
32
64
128
256
Number examples
Figure 9: Worst-group test accuracy Waterbirds-sever: for the best variant of proposed approach trained on iNaturalist and for methods such as 1-NN, ERM, and GroupDRO.
Proposed
100.0
majonty-group accuracy
99.5
99.0
60
98.5
98.0
40
97.5
Worst
97.0
20
96.5
Propesed +++1
0.2 Relative norm of the added background vector
0.05
1.0
Relative norm the added background vector
Figure 10: Majority-group and worst-group test. accuracies of proposed model (G + P + I) trained on iNaturalist, but evaluated modified variants of Waterbirds where add vector representing the spurions 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 corresponds to Waterbirds-severe. Shaded regions show standard deviation across training runs.
Passing example groups as input. In contrast to the findings in the single-task setting of Section 2 we observed that setting ÿi to represent group improves the proposed approach, even in addition of permuting input dimensions and promoting induction heads. One case of this presented in Figure f while more cases C For brevity, we mark passing groups inputs with +G in figures and tables. Please see Figure 13 for an illustration.
can be found in the complete results presented in Section
Comparison with conventional learning algorithms. Similar to the experiments in Section we compare the best variant of the proposed approach (G + P to I-NN, ERM, and GroupDRO. The results presented in Figure 7 show that the learned algorithm is on par with or outperforms the baselines starting at context. length 32. The results context lengths below 20 are not as informative, since our implementation of the grafting operation implics that no examples are grafted when there are less than 10 examples
class.
Generality of the learned algorithm. To test the generality of the learned algorithm, report evaluation results Waterbirds (Figure and Waterbirds-sene sereirrdds-eeeree (Figure g We see that the learned algorithm outperforms basclines on Waterbirds and performs comparably to the model trained on Waterbirds itself. However, the learned algorithm fails completely on Waterbirde-severs while the baselines give meaningful results starting at context length 32. We hypothesize that the challenge posed by the spurious 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 good generalization until the norm of the added vector is R400 of the average embedding norm (see
features Waterbirds-sevent
10

---

Published in Transactions on Machine Learning Research (04/2026)
Figure 10 It is possible that the learned algorithm specializes to spurious features planted with the grafting operation, but fails to handle a strong, dense, additive spurious feature.
4 Related Work
In-weights vs in-context learning. We observe two modes of learner behavior in our experiments. Tn 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. 2022) for obtaining general-purpose in-context classifiers. Prior work has made a distinction between these two modes of learming, naming them in-weights and in-context learning. In particular, Chan et al. (0022 demonstrate that certain distributional properties of data, such as long-tail of class frequencies and bursty distribution of context example classes, can promote in-contexi learning when meta-training on few-shot classification instances. [Singh al. (2024) 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- in-weights learning. [Anand et al. (2024) make similar observations about these two modes of learning and propose active forgetting of token embeddings as an effective way of steering lowards the in-conlext 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. (022)) and Akyürek et al. (2023) (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 lypically considered in ICL works (up to a few dozen examples). Our findings are complementary to those of [garral et al. (2024), Jiang et al. (0024), ; and Li et al. (2024) 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 and 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. (2023) test multimodal large language models (MLLMs) on a variety of visual classification tasks. They propose to leverage in-context lcarning 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. (2024) make similar observations, but additionally study robustness of in-context learning 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. This is in agreement with the findings of dAhuja & Lopez- Pites-Pax Paz brittle and don't generalize under severe distribution shifts, Finally, Gupta et al. 2024) propose to address the problem of domain generalization (Muandet et al., 2013 by training an in-context learner that can take examples from a domain/ reviromenv and adapt to that domain in-context.
2023 and our work that in-context learners are
Task Diversity. The important of task diversity for out-of-distribution generalization is well-studied outside of the in-context learning literature. In transfer learning setting, Tripuraneni et al. (2020) derive theoretical results highlighting the importance of task diversity in addition to samples per task in obtaining transferable representations. Hendryeks et al. (2020) find that pretrained then fine-tuned models generalize to oul-of-distribution examples better than only fine-tuned models, connecting the success to pretraining task diversity. Ramanujan et al. (2023) provide more direct evidence of this connection for image classification. Task diversity has also been one of the primary drivers of success of instruction tuning (Wei et al.] 2022). Similarly, the excellent ditto--distribttion generalization performance of the CLIP (Radford et al.) 2021) has been attributed to its pretraining dataset diversity (Fnng et al. 2022 Similar observations have been made
11

---

Published in Transactions on Machine Learning Research (04/2026)
for training in-context learners (Raventós et al., 2024 Goddard et al., 2025 Raparthy et al. 2024 Our
findings of Section 3 complement this literature and suggest that for training general in-context learner robust to spurious features, one likely needs dataset of diverse tasks with diverse spurious features.
5 Discussion and conclusion
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 introducing 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 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 room for improving in-context learning via improved strategies of choosing intermediate queries and possibly optimizing worst-group loss. Understanding why the learned a algorithm fails under extreme distribution shifts and why variants with permutations fail more (see Figure 10 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 neither 1-NN, ERM, GroupDRO.
Our work has several limitations. First, training a transformer-based in-context learner with high-dimensional image embeddings is computationally costly (see Section A for information on compute resources), although it is faster than the baselines during inference. For this reason, we did not explore more datasets and pretrained image embeddings. We believe the main conclusions of our work will remain unchanged and provide an experiment on CelebA with larger network Section C Second, we experimented with only one model size, width, and depth. Larger models might behave differently (Wei et al. 2023). Third, in our iNaturalist experiments, we considered only one "type" of spurious features. This choice is likely to have significant. effect on the learned algorithm and its generality. Future research should explore more ways of synthesizing spurious features and consider varying the severity of the challenge posed by spurious features. The latter can be done by considering multiple spurious features, introducing label imbalance, varying the magnitude of spurious correlations, and varying the margin spurious features provide.
Finally, we acknowledge that the proposed approach of training robust in-context learners requires spurious feature annotations, which are typically costly to obtain. As we have shown, this limitation can be addressed by creating synthetic data, in which case spurious annotations are readily available. At inference time, the learned algorithm does not require spurious annotations if it is trained with J, set to represent Y4 (i.e., ERM-like algorithm), but requires them when it is trained with passing example groups as input (i.c., ÿ set to represent 95, GroupDRO-like algorithmn). It is important to note that as we consider classification problems where the learner receives training data only from a single environment, spurious annotations are generally necessary to disambiguate core and spurious features.
Broader Impact Statement
Spurious correlations can appear in various domains, often in complex and unexpected ways. Given the increasing reliance on few-shot prompting in large language models, understanding and mitigating the negative effects of spurious correlations is important for cnsuring the reliability and safcty of these models. We hope that our work can lead to new methods and training strategies for making LLMs more robust to spurious correlations present in the context.
12

---

Published in Transactions on Machine Learning Research (04/2026)
Acknowledgments
We would like to thank anonymous reviewers for valuable feedback. Samvel Karapetyan's work is supported by Yandex Armenia Scholarship. The research was partially supported by the Higher Education and Science Committee of MESCS RA (Research project No. 25DD-1B082).
References
Rishabh Agarwal, Avi Singh, Lei M Zhang, Bernd Bohnet, Stephanie Chan, Ankesh Anand, Zaheer Ab- bas, Azade Nova, John Co-Reyes, Eric Chu, et al. Many-shot in-context learning. arXiv preprint arXiv:2404.11014 2024.
Kartik Ahuja and David Lopez-Paz. A closer look at in-context learning under distribution shifts. arXiv preprint ar xrx::2335.16700, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023.
Suraj Anand, Michael A Lepori, Jack Merullo, and Ellie Pavlick. Dual process learning: Controlling use of in-context vs. in-weights strategies with weight forgetting. arXiv preprint arXiv:240600005. 2024.
Jimmy Lei Ba, Jamic Ryan Kiros, and Gcoffrey E Hinton. Layer normalization. arXiv preprint N60:1607.06455 2016.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition lerra incognila. In Proceedings of the European conference on computer vision (ECCV), pp. 456 473, 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Jonathon Byrd and Zachary Lipton. What is the effect of importance weighting in deep learning? In International conference on machine learning, pp. 872 881. PMLRR 2019.
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. Advances in Newral Information Processing Systems, 1118878-18891. 2022.
Alex Fang, Gabriel Ilharco, Mitchell Wortsman, Yuhao Wan, Vaishaal Shankar, Achal Dave, and Ludwig Schmidt. Data determines distribuutonal robustness contrastive language image pre-training (clip). Tn International conference on machine learning, pp. 6216 6234. PMLR, 2022.
Tigran Galstyan, Hrayr Harutyunyan, Hrant Khachatrian, Greg Ver Steeg, and Aram Galstyan. Failure modes of domain generalization algorithms. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pottern Recognition, pp. 19077-19086, 2022.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-3093 2022.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. Imagenet-trainee CNNs are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations, 2019.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaclis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shorteut learning in deep neural networks. Nature Machine Intelligence, 2(11): 665-673, 2020.
13

---

Published in Transactions on Machine Learning Research (04/2026)
Chase Goddard, Lindsay M. Smith, Vudtiwat Ngampruetikom, and David J. Schwab. When can in-context learning generalize out of task distribution? In Forty-second International Conference on Machine Learning, 2025.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2021.
Sharut. Gupta, Stefanie Jegelka, David Lopez-Paz, and Kartik Ahuja. Context is environment. In The Twelith International Conference on Learning Representations, 2024.
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A. Smith. Annotation artifacts in natural language inference data. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Humon Language Technologies, Volume 2 (Short Papers), pp. 107-112, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10.18653/11 Nt.18853/v1/N188 /
Zhongyi Han, Guanglin Zhou, Rundong He, Jindong Wang, Xing Xie, Tailin Wu, Yilong Yin, Salman Khan, Lina Yao, Tongliang Liu, ct al. How well does gpt-4v (ision) adapt to distribution shifts? a preliminary investigation. Xiv preprinl orXiv:2312.07424, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on compuler vision and pollern recognilion, pp. 770 778, 2016.
Dan Hendryeks, Xiaoyuan Liu, Eric Wallace, Adam Dziedzic, Rishabh Krishnan, and Dawn Song. Pretrained transformers improve out-of-distribution robustness. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.). Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2744 2751, Online, July 2020. Association for Computational Linguisties. doi: 10.18653/v1/2020, acl-main.244.
Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Papez-Paz. Simple data balancing achieves competitive worst-group accuracy. In Conference on Causel Learning and Reasoning, pp. 336-30 351. PMLR, 2022.
Pavel Iæmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On feature learning in the presence of spurious correlations. Advances in Newral Information Processing Systems, 35:38516 38532, 2022.
Yixing Jiang, Jeremy Irvin, Ji Hun Wang, Muhammad Ahmed Chaudhry, Jonathan H Chen, and Andrew Y Ng. Many-shot in-context learning in multimodal foundation models. arXiv preprint arXiv:2405.00 09798, 2024.
Diederik P Kingmn and Jimmy Ba. Adam: A method for stochastic optimization. ar preprint arXiv:1412.6980, 2014.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer ro-training is sufficient for robustness to spurious correlations. In The Eleventh International Conference Learning Representations, 2023.
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. arXiv preprint arXiv:2212.04456 2022.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Trena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, pp. 5637 5664. PMLR, 2021.
Tianle Li, Ge Zhang, Quy Due Do, Xiang Yue, and Wenhu Chen. Long-context llms struggle with long in-context learning. arXiv preprint Xrx:224404. 02060, 2024.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
14

---

Published in Transactions on Machine Learning Research (04/2026)
Tom McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 3428 3448, Florence, Italy July 2019. Association for Computational Linguistics. doi: 10.18653/ 10.18653/v1/P199 1t/1//11111111/--
Raghav Mehta, Vitor Albiero, Li Chen, Ivan Evtimov, Tamar Glaser, Zhiheng Li, and Tal Hassner. You only need a good embeddings extractor to fix spurious correlations. arXiv preprint Tfi:2212.0025:: 2022.
Aditya Krishna Menon, Ankit Singh Rawat, and Sanjiv Kumar. Overparametterisatoo and worst-case generalisation: friend or foe? In International Conference on Learning Representations, 2021.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International conference machine learning, pp. 10-18. PMLR, 2013.
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. In International Conference on Learning Representations, 2021.
Devang K Naik and Richard J Mammone. Meta-neural networks that learn by learning. In /Proceedings 1992] IJCNN International Joint Conference on Neural Nebworks, volume 1, pp. 437-442. IEEE, 1992.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield- Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Cirewits Thread, 2022. htasss//ttasfoomme-eee Icarning-and awwaannnnniiinnnnnnonooo hhattccttioooocssseeeeeeeaeassee index.html.
pub/2022/ httb//2222/imconteer /in-tettext--
Maxime Oquab, Timothée Darcet, Theo Moutakanni, Huy V. Vo, Mare Szafranice, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Russell Howes, Po-Yao Huang, Hu Xu, Vasu Sharma, Shang-Wer Wen Li, Wojcicch Galuba, Mike Rabbat, Mido Assran, Nicolas Ballas, Gabriel Synnaeve, Ishan Misra, Herve Jegou, Julien Mairal, Paurick Labatut, Armand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision, 2023.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Tlya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8748 8763. PMLR, 18 24 Jul 2021.
Vivek Ramanujan, Thao Nguyen, Sewoong Oh, Ali Farhadi, and Ludwig Schmidt. On the connection between pre-training data diversity and fine-tuning robustness. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp- 66426 66437.63777
Curran Associates, Inc., 2023.
Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikael Henaff, and Roberta. Raileanu. Learning to solve new sequential decision-making tasks with in-context learning, 2024.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems, 36, 2024.
Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin. ' why should trust you?" explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1135-1144 1144, 2016.
15

---

Published in Transactions on Machine Learning Research (04/2026)
Shiori Sagawa*, Pang Wei Koh*, Tatsunoi B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Represenlations, 2020.
Jürgen Schmidhuber. Exolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-... hook. PhD thesis, Technische Universität München, 1987.
Yuge Shi, Imant Daunhawer, Julia E Vogt, Philip Torr, and Amartya Sanyal. How robust is unsupervised representation learning to distribution shift? In The Fleventh International Conference on Learning Representations, 2023.
Aaditya Singh, Stephanic Chan, Ted Moskovitz, Erin Grant, Andrew Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers. Advances Newral Information Processing Systems, 36, 2024.
Rui Song, Yingji Li, Lida Shi, Fausto Giunchiglia, and Hao Xu. Shorteut learning in-context learning: A survey. Xiv preprint orXiv:2411 11.02018 2024.
Ruixiang Tang, Dehan Kong, Longtao Huang, and Hui Xue. Large language models can be lazy learners: Analyze shorteuts in in-context learning. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguisties: ACL 2023, pp. 4645 4657, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.1 (0.18653/v1/2088 URL https://aclarhooogyy. httgg2223..indings... findings-acl. 284/
Sebastian Thrun and Lorien Pratt. Learning to learn: Introduction and overview. In Learning to learn, pp. 3 17. Springer, 1998.
Antonio Torralba and Alexei Efros. Unbiased look at dataset bias. In CVPR 2011, pp. IEEE, 2011.
Nilesh Tripurancni, Michael Jordan, and Chi Jin. On the theory of transfer learning: The importance of task diversily. Advances in newral information processing systems, 33:7852 7862, 2020.
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 8769 8778, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp.
PMLR, 2023.
Ben Wang and Aran Komatsuzaki. GPT-J-6B: A 6 Billion Paramcter Autoregressive Language Model. jax, May 2021.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Lc. Finctuned language models are zero-shot learners. In International Conference on Learning Representations, 2022.
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arxis preprint arXiv:2303.088. 038466, 2023.
P. Welinder, S. Branson, T. Mita, C. Wah, F. Schroff, Belongie, and P. Perona. CCLD Birds 200. Technical Report CNS-TR-2010-001, California Institute of Technology, 2010.
16

---

Published in Transactions on Machine Learning Research (04/2026)
Florian Wenzel, Andrea Dittadi, Peter Gehler, Carl-Johann Simon-Gabrien Max Horn, Dominik Zietlow, David Kernert, Chris Russell, Thomas Brox, Bernt Schiele, et al. Assaying out-of-distribution generalization in transfer learning. Advances in Neural Information Processing Systems, 35:7181 7198, 2022.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapler of the Associalion for Compulational Linguislics: Human Language Technologies, Volume (Long Papers), pp. 1112 1122. Association for Computational Linguisties, 2018a.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long Papers), http://allleb..r////attttaoooooa
pp. 1112 1122. Association for Computational Linguistics, 2018b. URL N18-11011
Kai Yuanqing Xiao, Logan Engstrom, Andrew llyas, and Aleksander Madry, Noise or signal: The role of image backgrounds in object recognition. In International Conference on Learning Representations, 2021.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities in transformer models. arXiv preprint arXf::811.00871, 2023.
John R. Zech, Marcus A Badgeley, Manway Liu, Anthony B Costa, Joseph J Titano, and Eric Karl Oermann. Variable generalization performance of a deep learning model to detect pneumonia in chest radiographs: a 2018
cross-sectional study. PLoS medicine,
Xingxuan Zhang, Jiansheng Li, Wenjing Chu, Junjia Hai, Renzhe Xu, Yuqing Yang, Shikai Guan, Jiazheng Xu, and Peng Cui. On the out-of-distribution generalization of multimodal large language models. arXiv 2024.
preprint ar 2rXiv24402.003
Yuqing Zhou, Ruixiang Tang, Ziyu Yao, and Ziwei Zhu. Navigating the shortcut maze: A comprehensive analysis of shortcut learning text classification by language models. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Associalion for Compulational Linguislics: EMNLP 2024, pp. 2586 2614, Miami, Florida, USA, November 2024. Association for Computational Linguisties. doi: 10. 10.18653/v1/2024 http:///222..iiinseeeeeeeee findings-emmlp.145 URL httppp:/aaclannoooooooooo/////// otg/2004.....2/00 ougg/2024.fidnnnn-mmmlll. findings-emnlp... 1466
A Further experimental details
Baselines. For empirical risk minimization as baseline, we tune lyperparaaeters: learning rate (0.01 or additionally tune its parameter that controls
0.001) and number of epochs (100 or 200). For GroupDRO adaptiveness of group weights (0.01, 0.1, 1) and we also try optional strong L2 regularization (1.0 weight decay), as it has been obscrved to be useful for small datascts (Sagawa* et al. 2020).
Transformer-based methods. In all transformer-basedd 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 al.] 2016) on transformer input, as we expect input norms to change when we evaluate on Waterbirds-baser 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 positions and content to distinguish image embeddings from annotations, 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-hol vector representing loken lype (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.
17

---

Published in Transactions on Machine Learning Research (04/2026)
Evaluation and model selection. For all transformer-based approaches and baselines, do grid search to find the best combination of hyperparameters. In particular, we train each configuration with 5 different random seeds and select the one with the highest average test performance. Importantly for baseline methods model selection is done for each context longth independently, while for transformer- based mothods model selection done once with respect to the lest 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 permuted image embeddings, during evaluation. In all figures throughout this work, shaded regions show standard deviation across the 5 training runs.
do not apply permutations
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 exporiments on Waterbirds or Waterbirds-severee the principled approach would be select, based on performance on sequences where the context part is sampled from the training set, while the final query is sampled from 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 always around 99. 99.5%-100% and can be -on-informative for model selection. This is a typical scenario in OOD generalization (see for example (Gulrajani & Lopez-Pat 2021 or (Wenzel et al. 2022
Definitions of metrics. Given a set of predictions on Waterbirds or Vaterbbirds-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 triplet (C, q, ū), where C' is context, q is query, and is prediction on q, we call ÿ a minority (majority) prediction, if q is among the least (most) represented group(s) of the context C. Given a list of triplets (C, D), 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 one GPU with batch size 32 when 400 (iNaturalist experiments) and batch size 24 when n = 512 (Waterbirds and Vaterbirda-seven 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.
experiments). We did
B Additional illustrations
In this section, we present illustrations of the main techniques and data preparation steps involved in this work.
Figure 11 plots the attention mask M of (2 in the proposed approach for 3.
12 presents an illustration of the proposed technique of permuting image embedding dinensions (denoted by +P throughout the paper). Figure 13 presents an illustration of the proposed approach with
Figure
passing example groups as input (denoted by +G throughout the paper). Figure I4 presents an illustration of the proposed approach with promoting emergence of induction heads (denoted by +I throughout the paper). Figure 15 presents an illustration of our preprocessing of the iNaturalist datasct. Figure 16 presents an illustration of the grafting operation for creating spurious features.
18

---

Published in Transactions on Machine Learning Research (04/2026)
211
81
41
372
1/2
42
223
93
93
-11
J2
912
02
53
03
Figure 11: The attention mask in the proposed approach with n 3 context examples. Green cells indicate allowed attention pairs, i.c., Mia 1.
ICL instance #1: embedding dimension permutation =2, 4, 6, 5)
3
E2
f/2
000
E 1
C2
ICL instance #2: embedding dimension permutation = (1, 6, 5, 3, 4, 2)
32
2
J/2
000
C1
.2
Figure 12: An illustration the proposed technique permuting image embedding dimensions (denoted by +P throughout the paper). Note that in each ICL, instance we sample a new permutation, but the same permutation is used to permute dimensions of all image embeddings within one ICL instance. Whenever we use the proposed approach of forming ICL sequences (sce Figure 1b), the dimensions of intermediate queries are also permuled.
(2yl-
1)u
(2x]
1)
@1
(2g2
1) + (2x2
1je
Ú1
1/2
92
Figure 13: An illustration of the proposed approach with passing example groups as input (denoted by +G throughout the paper).
F1
y1
!1
U2
T/2
x
.33
/3
.22
U+
44
aoe
Figure 14: An illustration the proposed approach with promoting emergence of induction heads (denoted by +I throughout the paper). Intermediate queries that randomly selected to be one of the previous context examples are shown in green.
19

---

Published in Transactions on Machine Learning Research (04/2026)
Plantae
Insects
Reptilla
Category
Training
ID validation
validation
Fungi
Mollusca
Animalia
Mammalia
Actinopterygii
Chromista
Protozoa
Arachnida
Amphibia
Held-out supercategorles
Figure 15: An illustration of our preprocessing of the iNaturalist dataset.
5 image embeddings from class
image embeddings from class B
Grafting features
y=0 y=0 5-0 2
s-0
(a) Before grafting
(b) After grafting
Figure 16: An illustration of the grafting operation for creating spurious features. The figure (a) depicts two classes examples, each having 5 examples given by 12-dimensional embeddings. In this example, the grafting operation selects embedding dimensions to become spurious features. For this end, these features of examples and of class A are swapped with those of examples 2 and 4 of class respectively. Figure (b) depicts the embeddings after the grafting operation.
20

---

Published in Transactions on Machine Learning Research (04/2026)
C Additional results
In addition to the figures presented in the main text, here we provide the exact experimental results for multiple transformer-basel and baseline approaches, some of which were not included in the main text due to space constraints. Recall that +P means permuting imput dimensions, +I means promoting learning of induction heads, and +G means passing example groups as input to in-context learning transformers.
Table presents worst-group accuracies on the test set of Waterbirds for 3 sets of approaches: (a) in-context. learners trained Waterbirds itself, (b) in-context learners trained iNaturalist, and (c) baselines. Similarly, Table 2 presents worst-group accuracies on the test set of Waterbirds-severr for 3 scts of approaches: (a) in-context learners trained on Waterbirds-sevvrr itself, (b) in-context learners trained on iNaturalist, and (c) baselines. As RoPF-based transformers are not. good at length extrapolation (Press et al. 2021). we do not attempt evaluating models trained on iNaturalist with context size 400 on 512-long soquences of Waterbirds or Waterbirds-seven
Finally, [Talle 3) presents minority-group accuracy on out-of-distribution classes of iNaturalist for two sets of approaches: (a) in-context learners trained iNaturalist itself and (b) baselines.
Figure presents majority- group accuracies on Waterbirds and Waterbirds for the naive and proposed methods. Figure 18 presents majority- group accuracies on Waterbirds and Waterbirds-seven for the proposed approach and conventional methods.
100
95
90
90
Majority grourit-ggouupppo
Majority
75
Naive
Naive Naive
Naive P
60
Proposed
Proposed
Proposed
2
16
32
64
128
256
512
16
32
64
128
256 512
examples
examples
(a) Waterbirds
(b) Waterbirds-severe
Figure 17: Majority-group accuracies on Waterbirds and Waterbirds-severe as function of context size for the naive and proposed approaches with or without permuting input dimensions. Shaded regions show standard deviation across training runs.
Figure 19 demonstrates that important to simulate distribution shift at all intermediate query locations, not just for the main query. This makes sure that the learned network make robust predictions for all intermediate context sizes.
2 we conduct Experiments on CelebA. To further verify our main findings presented in Section experiments on CelebA (Liu et al. 2015). Here the task is to classify blond vs non-blond persons from face images, with sex being a spuriously correlated variable. Notably, the spurious correlation is asymmetric, the sense that blond and non-blond women are almost equally represented, while blond mon are much less represented compared
non-blond men. We follow the design of Waterbirds experiments in our CelebA experiments, with the only difference that set the group distribution of context examples to (0.25,0 0.25, 0.05, 0.45), where group 0 are non-blond men, group are non-blond women, group 2 are blond men, and group 4 are blond women. "Table 4 presents worst-group accuracics on the test set of CelebA for 2 sets of approaches: (a) in-context learners trained on CelebA itself and (b) baseline algorithms. As in our Waterbirds experiments, we see that it is essential to permute input embeddings and to form ICL sequences the proposed fashion. Unlike Waterbirds, comparing "Proposod P with "Proposod + P + G we see that providing spurious annotations in-context provides significant gains, Figure 20a demonstrates that both of these approaches outperform 1-NN, ERM, and GroupDRO.
21

---

Published in Transactions on Machine Learning Research (04/2026)
100
100
95
Majority group
group group
80
Majority
70
75
Proposed
1-NN
60
1-NN
GroupORO
GroupORO
2
8
16
32
64
128
256
512
2
8
16
32
64
128
256
512
examples
examples
(a) Waterbirds
(b) Waterbirds-severe
Figure 18: Majority-group test accuracics on Waterbirds and saterbirds-severe for the proposed approach and conventional methods such as I-NN, ERM, and GroupDRO, Shaded regions show standard deviation across 5 training runs.
Naive
80
Naive balanced query
Majarit group accuracy
Worst accurscy
60
94
40
92
90
20
Naive balanced query
88
32
64
128
32
64
12B
Number of context examples
Number of context examples
Figure 19: Majority-group and worst-group test accuracies on Vaterbirds-severt as a function context: size for the naive approach with a single modification of making the last example (query) group-balanced. Shaded regions show standard deviation across training runs. As expected, at intermediate context lengths this method performs similarly to the naive approach, but is much better at the training context length.
Experiments MultiNLI. Beyond visual classification tasks, also conduct experiments natural language inference (NLI) task. In particular, consider the MultiNLI dataset (Williams al. 2018b where given pair of sentences, a premise and hypothesis, the task is to determine whether the hypothesis is entailed by, neutral with, or contradicts the premise. Prior work has observed that there is spurious correlation between contradictions and the presence of the negation words nobody, no, never, and nothing (Gururangan) et al. 2018 For our experiments, we frame a binary classification task enlails or neutrul us. contradicts, with binary spurious feature, which is if and only the hypothesis has one of the four negation words. We follow the design of Waterbirds experiments and set the group the group distribution of context examples the test set of MultiNLI for 2 sets
to (0.45, 0.05, 0.05, 0.45). Table 6 presents worst-group accuracies of approaches: (a) in-context. learners trained on CelebA itself and (b) baseline algorithms. Similar to Waterbirds and CelebA results, we see that the proposed approach of intermediate queries with balance group distribution helps to decrease the reliance on the spurious feature. However, in contrast to Waterbirds and CelebA experiments, permuting input embeddings has smaller effect on preventing in-weight learning. Figure 20b demonstrates the 'Proposed + P' method outperforms 1-NN, ERM, and GroupDRO.
Experiments with a larger network. To verify that our findings generalize to larger models, we repeat
CelobA experiments but with transformer architecture of 12 layers with 12 multi-head attention (instead of layers with multi-head attention). Due memory increase, we decrease the batch size from 24 8. Besides these two changes, keep all other experimental details the same. The complete results presented
22

---

Published in Transactions on Machine Learning Research (04/2026)
Proposed
Proposed + P 1-NN
Worst group accuracy
Worst-group accuracy
ERM
GroupDRO
70
Proposed
30
1-NN
GroupDRO
2
4
8
16
32
64
128
256
512
2
4
8
16
32
64
128
256
512
examples
examples
(a) CelebA
(b) MultiNLI
Figure 20: Worst-group test accuracies on CelebA and MultiNLI for the proposed approach and conventional methods such as 1-NN, ERM, and GroupDRO. Shaded regions show standard deviation across 5 training runs.
in Table 5 are qualitatively the same compared to the smaller network case (Talle 4), with the difference that the results of transformer-based entrics are lower. Furthermore, the standard deviation of the +P approaches is significantly higher, indicating difficulties in optimization. We hypothesize that this is due to reusing learning rate and training length that were tuned for the smaller network with times larger batch size.
90
95
Majority group accuracy
90
Worst-g Worst-gruup group accurac
85
80
80
70
70
65
65
Evaluating with randomly rotated embeddings
Evaluating with randomiy rotated embeddings
8
16
32
64
128
256
512
128 Number examples
8
16
32
64
256
512
context examples
Figure 21: Majority-group and worst-group test accuracies on Waterbirds as function of context size for a "Proposed p" run evaluated on ICL sequences with randomly rotated input embeddings. Largely unchanged evaluation results fail to confirm that there is any data leakage when input embeddings permuted during training.
On data leakage in single task regime. In the single task setting of Section 2 there is a potential for data leakage, not in the sense that individual examples might be leaked (we always evaluate on unseen examples), but in the sense that the learner effectively observes more data from the single task than its context. length at evaluation. Indeed, when we do not permute input embeddings, we observe task memorization (i.e., data leakage) and the model does very well at evaluation with even close to empty context. To verify that there is no data leakage when we enable permuting input embeddings (+P), we take one of the "Proposed Pd runs trained Waterbirds and evaluate it on ICL sequences where input embeddings of each sequence are rotated with random rotation matriz. As the set of permutation matrices is measure-zero subset of general rotation matrices, we expect that in case of data leakage we would observe degraded performance, as the model would be expecting randomly permuted embeddings of some memorized embedding space. In results presented in Figure 21 we see that under this new evaluation the results are the same (up to statistical noise). failing confirm that there any data leakage when input emboddings are permuted during training. Finally, note that data Icakage not a concern in the multiple task setting Section because we evaluate unseen tasks such as Waterbirds and Waterbirds-severe
either unseen categories of iNaturalist
23

---

Published in Transactions on Machine Learning Research (04/2026)
95
90
90
85
85
80
75
Majority
70
Words
65
65
Group
60
sroup-Salanced
arouppbaalanee contexts permute
Group proup-aalanced contexts permute
32 Number examples
16
64
128
32 Number examples
16
64
128
Figure 22: Majority-group and worst-group test accuracios on Waterbirds as function of context size for the naive approach trained and evaluated on group-balanced contexts. The training is done on ICL sequences with 128 context examples. Shaded regions show slandard deviation across training runs.
99
Propesed Proposed +1
98
Proposed P
97
Proposed ++1
96
96
95
94
Majority
Minority
94
Propesed +1
93
92
Proposed P 1
4
8
16
32
64
128
256400
2
4
8
16 32
64
128
256 400
Number context examples
Number context
Figure 23: Majority-group and minority- 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 training runs.
Experiments with group-balanced contexts. As noted in Section the proposed approach of training in-context learner requires spurious annotations. Given access to spurious annotations, one can simply train in-context learner on sequences with balanced groups. While in-context learners obtained this way will not be useful for new tasks for which we do not have spurions annotations (and thus cannot form group-balanced contexts), it is still useful to compare how well this approach does in the single task setting of Section 2 For this end, we train in-context. learners on balanced-group sequences consisting of 128 Waterbirds examples. This way each group is represented with 32 context examples. Note that in our main Waterbirds experiments with 512 context examples but group-imbalanced contexts, the minority groups are represented with even fewer, 25 examples. As the group-balanced sampling context breaks the correlation between the label and spurious feature, only consider the naive approach of forming ICL sequences (Figure 2a.. The results presented in Figure 22 show that, expected, group-balanced sampling improves worst-group accuracy. The naive approach, which again leads to in-weights learning, reaches 86.08 + 1.87 worst-group accuracy with 128 group-balanced context. examples, compared to 84.82 1.26 worst-group accuracy on 512 group-imlallnnce context examples (see Table I). This positive effect of downsampling has been also observed in standard (not in-context) training settings (Nagarajan et al.) 2021 Menon al. 2021 Idrissi et al., 2022). Furthermore, we again see that the proposed technique of permuting embedding dimensions induces strong in-context learning and reaches 90.44 + 1.10 worst-group accuracy with 128 group-balanced context examples. This is just bit below the 91.95 + 1.20 worst-group accuracy get training "Proposed + P on 512 group-imbalanced
context examples (see Table 1).
24

---

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

---

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

---

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

---

Published in Transactions on Machine Learning Research (04/2026)
Table 4: Complete results on CelebA. Reported numbers are average worst-group test accuracies, along with their standard deviation. All in-context learners were Irained on CelebA itself.
Method / Context size
8
16
32
64
128
256
512
24.88
24.56
25.80
25.14
23.85
25.62
25.84
26.20
Naive
(2.03)
(2.26)
(1.98)
(2.11)
(1.91)
(1.63)
(2.10)
(1.42)
20.72
17.27
12.43
14.85
13.56
16.17
20.16
26.03
Naive P
(2.21)
(2.38)
(2.04)
(2.33)
(1.60)
(2.83)
(4.14)
(5.13)
25.83
25.42
26.85
25.80
25.18
26.65
26.89
27.53
Proposed
(1.77)
(1.66)
(1.52)
(1.60)
(1.06)
(2.11)
(1.41)
(1.31)
26.90
26.54
27.00
37.30
47.29
54.66
60.48
68.45
Proposed + P
(4.56)
(1.46)
(2.72)
(1.55)
(2.67)
(2.67)
(2.55)
(1.99)
23.87
24.67
25.60
24.95
24.42
25.77
25.70
26.55
Proposed + G
(1.50)
(1.51)
(1.30)
(1.12)
(1.17)
(1.24)
(0.81)
(1.39)
26.71
32.06
39.21
46.13
53.41
59.37
64.39
69.58
Proposed + G P
(3.57)
(4.62)
(6.06)
(5.38)
(3.73)
(3.07)
(1.84)
(1.72)
35.87
37.63
36.08
37.86
38.28
36.40
36.90
37.81
1-NN
(1.48)
(0.86)
(1.13)
(0.45)
(0.80)
(0.32)
(0.99)
(0.36)
30.70
28.93
26.37
30.75
34.29
38.64
45.18
49.92
ERM
(1.05)
(0.64)
(0.66)
(0.40)
(0.93)
(1.29)
(1.39)
(1.28)
35.32
34.64
30.24
37.49
47.11
54.56
56.11
61.47
GroupDRO
(0.88)
(0.99)
(1.01)
(0.60)
(0.59)
(0.66)
(0.60)
(0.93)
Table 5: Complete results on CelebA, but with larger network of 120m paramcters, consisting of 12 layers (instead of 6 layers) with 12 multi-head attention (instead of 8 heads). Reported numbers are average worst-group test accuracies, along with their standard deviation. All in-context learners were trained on CelebA itself.
Method / Context size
8
16
32
64
128
256
512
24.43
21.79
23.68
23.02
23.99
23.18
22.62
20.22
Naive
(0.57)
(0.86)
(0.58)
(0.49)
(0.84)
(1.04)
(0.83)
(1.00)
21.34
14.64
12.56
13.35
13.76
15.73
17.69
21.67
Naive + P
(1.40)
(0.66)
(0.74)
(1.54)
(2.58)
(2.41)
(3.27)
(4.55)
22.90
20.86
23.11
21.93
23.35
21.82
21.75
19.54
Proposed
(2:30)
(2.40)
(2.66)
(2.50)
(2.80)
(3.11)
(2.32)
(2.06)
35.13
31.54
30.89
35.19
41.63
47.81
51.60
55.53
Proposed + P
(5.19)
(1.73)
(4.40)
(5.28)
(6.74)
(9.40)
(11.00)
(11.70)
23.08
20.89
22.74
21.73
22.70
21.24
21.50
18.90
Proposed + G
(2.47)
(2.66)
(2.54)
(3:23)
(3.25)
(3.45)
(2.61)
(2.01)
31.44
32.09
36.54
43.67
49.62
54.77
58.75
61.74
Proposed + G
(4.27)
(5.24)
(1.90)
(2.72)
(2.93)
(4.24)
(5.94)
(6.10)
35.87
37.63
36.08
37.86
38.28
36.40
36.90
37.81
I-NN
(1.48)
(0.86)
(1.13)
(0.45)
(0.80)
(0.32)
(0.99)
(0.36)
30.70
28.93
26.37
30.75
34.29
38.64
45.18
49.92
ERM
(1.05)
(0.64)
(0.66)
(0.40)
(0.93)
(1.29)
(1.39)
(1.28)
35.32
34.64
30.24
37.49
47.11
54.56
56.11
61.47
GroupDRO
(0.88)
(0.99)
(1.01)
(0.60)
(0.59)
(0.66)
(0.60)
(0.93)
28

---

Published in Transactions on Machine Learning Research (04/2026)
Table 6: Complete results on MultiNLI. Reported numbers are average worst-group test accuracies, along with their standard deviation. All in-context learners were trained on MultiNLI itself.
Method Context size
2
4
8
16
32
64
128
256
512
84.15
84.58
84.74
84.13
83.92
85.20
84.38
84.30
84.05
Naive
(0.30)
(0.37)
(0.37)
(0.45)
(0.25)
(0.24)
(0.32)
(0.33)
(0.39)
85.30
87.29
87.43
87.56
87.50
88.05
87.35
87.99
87.54
Naive + P
(1.58)
(0.95)
(1.05)
(0.77)
(0.57)
(0.38)
(0.27)
(0.20)
(0.34)
86.62
86.02
87.45
87.54
86.99
86.51
86.11
87.35
86,89
Proposed
(0.65)
(0.33)
(0.58)
(0.42)
(0.26)
(0.27)
(0.29)
(0.19)
(0.30)
85.84
87.22
87.43
87.75
87.70
87.86
87.31
87.56
87.73
Proposed + P
(1.12)
(0.55)
(0.48)
(0.24)
(0.43)
(0.62)
(0.53)
(0.46)
(0.30)
87.15
86.23
87.23
87.37
87.02
86.51
86.47
87.69
87.50
Proposed + G
(0.82)
(0.65)
(0.42)
(0.35)
(0.27)
(0.58)
(0.49)
(0.51)
(0.55)
85.41
87.32
87.57
87.81
87.74
87.97
87.56
87.78
87.68
Proposed + G G+P
(1.38)
(0.76)
(0.34)
(0.56)
(0.50)
(0.30)
(0.27)
(0.31)
(0.12)
64.49
82.41
85.84
86.39
86.01
86.04
86.44
85.90
86.53
1-NN
(0.90)
(0.81)
(0.84)
(0.76)
(0.40)
(0.93)
(0.81)
(0.43)
(0.29)
76.40
81.71
85.10
86.82
86.76
87.30
87.04
87.10
87.55
ERM
(0.65)
(0.58)
(1.13)
(0.54)
(0.54)
(0.78)
(0.55)
(0.48)
(0.41)
76.54
81.56
85.45
86.75
86.55
87.16
87.14
87.10
87.23
GroupDRO
(0.42)
(0.75)
(1.10)
(0.47)
(0.30)
(0.55)
(0.37)
(0.14)
(0.49)
29


Reviews:
[{"review_id": "QwgLY3bgFb", "reviewer": "Reviewer_vawz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies in-context learning (ICL) in the presence of spurious input features (e.g., in image classification where background features correlate with the foreground label). It is worth noting that the notion of ICL used here differs from the standard usage in large language models: in this work, the model is explicitly trained on sequences containing in-context demonstrations, rather than being used as a frozen model.\n\nThe authors first conduct experiments on the Waterbirds classification task, using image features extracted from a pretrained vision model. They observe that standard ICL training leads to in-weights learning: the model ignores the in-context examples and instead simply learns the task in its parameters. To mitigate this, they introduce random permutations of input embedding dimensions (with a consistent permutation applied across all examples within an ICL instance). This forces the model to rely on contextual information and successfully induces in-context learning.\n\nThe second proposed modification inserts an additional \"query\" after each demonstration example. These intermediate queries are drawn from a distribution where spurious correlations do not hold. The attention mask is changed so that these queries do not influence the context seen by later tokens. The model is trained to correctly classify each intermediate query, thereby simulating a test-time distribution with no spurious correlations.\n\nOn the original Waterbirds dataset, this second modification does not improve performance. The authors attribute this to the strength of the spurious signal being too weak in the pretrained features. They therefore construct a modified \"Waterbirds-severe\" version by artificially amplifying the spurious background signal, and in this setting the proposed method shows performance gains.\n\n**Weaknesses:**\n\n1. Limited experimental scope. For the single task setup, the approach is evaluated on essentially a single dataset, and even within that dataset the main proposed modification only succeeds after artificially amplifying the spurious features. This makes the experimental setting feel somewhat synthetic and raises concerns about generality.\n\n2. Dependence on spurious feature annotations. The method requires explicit labels for spurious features. Such annotations are rarely available in real-world applications, which limits the practicality of the approach.\n\n3. Lack of intuition for why perturbing features improves performance. The core justification for random feature permutation is underdeveloped. It is unclear why simply perturbing input dimensions, which effectively makes the task harder, should lead to better performance than in-weights learning (with enough in-context examples), especially when only a single underlying task exists.", "claims_and_evidence_yn": "Yes", "claims_explanation": "While the contribution is modest, the presentation is clear and the experimental results adequately support the claims made in the paper.", "audience_yn": "Yes", "audience_explanation": "Yes, the topic of ICL in the presence of spurious features is of clear interest to the ML community.", "requested_changes": "1. Include additional datasets for the single task setup to compare the naive baseline with the proposed approach. Those datasets should preferably be ones without the synthetic modifications introduced in the paper, to strengthen the empirical validation and demonstrate broader applicability.\n\n2. Provide an explanation, even at an intuitive level, for why random feature perturbations lead to improved performance compared to using clean features (especially when there is only a single underlying task). This would help clarify why in-weights learning is suboptimal in this setting and why inducing in-context learning is beneficial.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dvGBxQ1lWY", "reviewer": "Reviewer_vLLM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the in-context learning of classification tasks in the setting where spurious correlations lead to poor accuracy on underrepresented subgroups. The standard approach of generating ICL instances is shown to induce in-weights learning, and a method based on intermediate query tokens is proposed as an alternative. Under this data generation procedure, and with regularization techniques such as permutation, in-context learners are shown to be robust to spurious correlations. Unlike robustness algorithms such as GroupDRO, it is observed that these in-context learners adapt to task structure and so do not generalize well out-of-distribution; training on a diverse set of ICL instances is proposed as a method to mitigate this.\n\nSummary of strengths:\n1. This paper’s empirical investigation is comprehensive and rigorous.\n2. The paper has potential to bring two distinct research communities closer together.\n3. The insight that ICL aggressively leverages task structure, leading to worse OOD generalization than conventional robustness algorithms, is particularly interesting and could motivate future research.\n\nSummary of weaknesses:\n1. There are a few methodological decisions that would benefit from further justification (see Requested Changes).\n2. Inclusion and discussion of various references could be improved (see Requested Changes).", "claims_and_evidence_yn": "Yes", "claims_explanation": "Overall, the claims in the paper are well-substantiated. Experiments are performed on several different benchmark datasets at varying levels of granularity and difficulty; the experimental designs are detailed and clear; the analysis is comprehensive and addresses natural questions which may arise while reading the paper; the appendix handles extensions such as larger networks, group-balancing, and checking data leakage.", "audience_yn": "Yes", "audience_explanation": "A strength of this paper is its potential to bring two distinct research communities closer together, namely in ICL/meta-learning and group robustness. Each community would likely be interested in this paper: the ICL/meta-learning community for its challenging generalization of standard ICL, observations on how ICL adapts to task structure, and application to relevant theories such as induction heads; the group robustness community for its new problem setting using existing datasets and comparison between ICL and conventional robustness algorithms. I expect future lines of research to address open problems raised by this paper, e.g., whether in-context learners can adapt to problems with unknown group structure, missing spurious annotations, or multiple correlations.", "requested_changes": "Critical for acceptance: None\n\nWould strengthen the work:\n1. It’s unclear to me why visual classification tasks were chosen as the primary benchmarks for ICL under spurious correlations given the proclivity of autoregressive Transformers to sequential language-based tasks. I would appreciate further justification on this point. Moreover, while not strictly necessary for acceptance, I believe it would improve the work to evaluate ICL on language benchmarks like CivilComments or MultiNLI.\n2. While the permutation method for regularizing ICL memorization is reasonable, why were standard data augmentation techniques not used given the computer vision setting? For example, instead of computing and storing 10K rotation matrices as mentioned in the text, one could rotate (or flip, jitter, etc) the image itself and compute its DINO-ViT embedding. Also, it would be interesting to comment on whether the necessity of regularization for robustness in ICL has any relation to the necessity of regularization in GroupDRO.\n3. Given the bias of DINO image embeddings towards objects instead of backgrounds, the paper proceeds with increasing the difficulty of Waterbirds by introducing Waterbirds-Severe. However, an alternative would have been to utilize a different image encoder, such as ResNet, which is known to be biased towards image backgrounds. Was this method considered or investigated?\n4. The strategy to create Waterbirds-Severe, i.e., adding a constant vector to image embeddings, is interesting and clearly works. However, the more common method to increase difficulty in the spurious correlations literature is to increase the strength of the correlation, or in the most extreme case to have a complete correlation, where the model does not see, e.g., any waterbirds on land at train time. An investigation of this point, or justification for its absence, would improve the paper.\n5. Inclusion and discussion of various references could be improved:\n    * The impact of task diversity on OOD generalization is well-studied outside of the ICL literature. It may be worth discussion of, e.g., [1, 2, 3] and other related papers.\n    * Missing citation and discussion of papers on task diversity and OOD generalization in ICL: [4, 5]\n    * I would appreciate a more detailed contextualization with [6] in the “ICL for OOD generalization” subsection.\n\n[1] Hendrycks et al. Pretrained Transformers Improve Out-of-Distribution Robustness. ACL 2020.\n\n[2] Tripuraneni et al. On the Theory of Transfer Learning: The Importance of Task Diversity. NeurIPS 2020.\n\n[3] Ramanjuan et al. On the Connection between Pre-training Data Diversity and Fine-tuning Robustness. NeurIPS 2023.\n\n[4] Goddard et al. When can in-context learning generalize out of task distribution? ICML 2025.\n\n[5] Raparthy et al. Generalization to New Sequential Decision Making Tasks with In-Context Learning. ICML 2024.\n\n[6] Ahuja and Lopez-Paz. A Closer Look at In-Context Learning under Distribution Shifts. ArXiv 2023.", "broader_impact_concerns": "I have no broader impact concerns.", "additional_comments": "1. I believe the Akyurek et al citation has the wrong year.\n2. Please explain the +P and +I notations in the figure captions and/or reference to where they are defined in the main text.\n3. Is the group structure where groups 0 and 2 are majorities while groups 1 and 3 are minorities, referenced in the beginning of Section 2, only for Waterbirds or is this defined for general datasets? It also only holds for binary classification, and doesn’t make sense in situations where the spurious correlation applies to only one class (e.g., CelebA hair-color).\n4. It would be nice to give the “Proposed” method a name for future papers to reference.\n5. The first two paragraphs of Section 2.4 would benefit from a figure; I found it difficult to keep track of the benchmarks, baselines, and accuracies.", "extra_scores": {}, "extra_text": {}}, {"review_id": "7VHCvUQQNr", "reviewer": "Reviewer_2z2i", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "**Summary**\n---------------\nThis paper investigates the challenge of training Transformer-based in-context learners that are robust to spurious correlations. The authors demonstrate that conventional training approaches fail, leading to \"in-weights learning\" where the model memorizes the task instead of leveraging the context, and remains susceptible to shortcut features. To overcome this, they propose a novel training methodology that involves permuting input embedding dimensions to prevent memorization and forming ICL sequences with intermediate, group-balanced queries to simulate a distribution shift. Their experiments show that this approach can yield a specialized in-context learner for a single task that outperforms strong baselines like ERM and GroupDRO. Furthermore, by training on a diverse dataset of synthetic tasks, they successfully create a more general-purpose learner that can identify and ignore spurious features in unseen tasks. However, the work also critically highlights that this learned algorithm is brittle, failing completely when faced with spurious correlations that are significantly more severe than those seen during training, thereby outlining a key limitation in its ability to handle strong distribution shifts.\n\n**Strengths:**\n---------------\n1. The experimental results on Waterbirds and Waterbirds-severe are strong.\n2. The authors don't oversell their results. They dedicate significant effort to understanding and reporting on the \"brittleness\" of their learned algorithm, which is a crucial and sobering finding for the field.\n\n**Weakness**\n---------------\n1. The proposed method, particularly the use of group-balanced queries, requires access to spurious feature labels during the training phase. This is could be a limitation, as such annotations are often costly or unavailable.\n2. According to Figure 2(a), the performance gains appear to come almost entirely from permuting the input embeddings (+P), with the proposed intermediate query mechanism offering little-to-no additional benefit. The authors' argument is that the image embeddings of DINOv2 have a bias toward representing objects more than backgrounds. While this specific experiment is convincing, the paper's argument for the general utility of intermediate queries would be further strengthened by demonstrating this separation on existing real-world benchmarks with naturally strong spurious correlations, rather than relying on a single synthetically hardened one.", "claims_and_evidence_yn": "Yes", "claims_explanation": "The claims made in the submission are, for the most part, supported by accurate, convincing, and clear evidence. The authors do a good job of providing specific experiments to back up their arguments.", "audience_yn": "Yes", "audience_explanation": "Yes. Since the advent of large language models, understanding the mechanism and limits of in-context learning has become a central research question. The findings of this paper would be of high interest to a substantial portion of TMLR's audience.", "requested_changes": "For weakness 2, I suggest the users to rerun the Waterbirds experiment using weaker, pre-trained image embeddings. If your hypothesis is correct, these less-robust embeddings should be more susceptible to the spurious background feature. In this setting, we would expect the performance gap between Naive + P and Proposed + P to widen considerably, thus providing strong, direct evidence for the utility of intermediate queries when the representation itself is not robust.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)