Context is Environment
Authors: Sharut Gupta
Stefanie Jegelka
David Lopez-Paz
Kartik Ahuja
Venue: ICLR 2024 poster (2024)
Keywords: Domain Generalization; In-Context Learning
Primary area: general machine learning (i.e., none of the above)
TL;DR: Researchers in domain generalization should consider environment as context, and harness the adaptive power of in-context learning. Researchers in LLMs should consider context as environment, to better structure data towards generalization.

Abstract: Two lines of work are taking the central stage in AI research. On the one hand, the community is making increasing efforts to build models that discard spurious correlations and generalize better in novel test environments. Unfortunately, the hard lesson so far is that no proposal convincingly outperforms a simple empirical risk minimization baseline. On the other hand, large language models (LLMs) have erupted as algorithms able to learn in-context, generalizing on-the-fly to eclectic contextual circumstances that users enforce by means of prompting. In this paper, we argue that context is environment, and posit that in-context learning holds the key to better domain generalization. Via extensive theory and experiments, we show that paying attention to context$\unicode{x2013}\unicode{x2013}$unlabeled examples as they arrive$\unicode{x2013}\unicode{x2013}$allows our proposed In-Context Risk Minimization (ICRM) algorithm to zoom-in on the test environment risk minimizer, leading to significant out-of-distribution performance improvements. Furthermore, training with context helps the model learn a better featurizer. From all of this, two messages are worth taking home. Researchers in domain generalization should consider environment as context, and harness the adaptive power of in-context learning. Researchers in LLMs should consider context as environment, to better structure data towards generalization. Code is available at https://github.com/facebookresearch/ICRM.

Published as a conference paper at ICLR 2024
CONTEXT IS ENVIRONMENT
Sharut Gupta
Stefanie Jegelka MIT CSAIL stef Crefjeemit.eem edu
David Lopez-Paz, Kartik Ahuja Meta Al {dlp,
Meta AI, MIT CSAIL sharut @mit.mmii........m edu
ABSTRACT
Two lines of work are taking center stage in Al research. On the one hand, the community is making increasing efforts to build models that discard spurious correlations and generalize better in novel test environments. Unfortunately, a hard lesson so far that no proposal convincingly outperforms a simple empirical risk minimization baseline. On the other hand, large language models (LLMs) have erupted as algorithms able to learn in-context, generalizing on-the-fly to the eclectic contextual circumstances that users enforce by prompting, We argue that context environment, and posit that in-context learning holds the key to show better domain generalization. Via extensive theory and experiments, that paying attention to context- unlabeled examples as they arriv--aloow our proposed In-Context Risk Minimization (ICRM) algorithm zoom-in the test environment risk minimizer, leading to significant out-of-distribution performance improvements. Furthermore, training with context helps the model learn a better featurizer. From all of this, two messages are worth taking home: researchers in domain generalization should consider environment as context, and harness the adaptive power of in-context learning. Researchers in LLMs should consider context as environment, better structure data towards generalization. Code is available at https: ///gtttttt//b// wtt.thhww..oooceooeeeoooe.eeeeee wwt.///cccoooooeerrrccccoaa
INTRODUCTION
One key problem in AI research is to build systems that generalize across a wide range of test envi- ronments. In principle, these algorithms should discard spurious correlations present only in certain training environments, and capture invariant patterns appearing across conditions. For example, we would like to build self-driving systems that, while trained on certain weather conditions, levels of traffic, and driving rules, can perform satisfactorily in new circumstances. Unfortunately, this has proven challenging instance, these models often fail to drive in unseen weather conditions (Lech- ner et al., 2022), creating immediate hazards. Despite its importance, how to perform well beyond the distribution of the training data remains a burning question. In fact, major international conferences offer well-attended workshops dedicated to the issue (Wald et al., 2023), and news articles remind us of the profound societal impact of the failures of MI. systems (Angwin al., 2016).
Research efforts have so far led to domain generalization algorithms that fall into two broad categories. On the one hand, invariance proposals (Ganin et al., 2016; Peters al., 2016; Arjovsky et al., 2019), illustrated in Figure la, discard all environment-specife information, thus removing excessive signal about the problem. On the other hand, marginal transfer proposals (Blanchard et al., 2011; Li et al., 2016a; Zhang et al., 2020; Bao and Karaletsos, 2023), illustrated in Figure 1b, summarize observed inputs each environment level. So far, the bitter lesson is that no algorithm geared towards -ut-tf-distribution generalization convincingly outperforms simple empirical risk minimization (ERM) baseline, which pools the data from environments, when evaluated across standard real-world benchmarks (Gulrajani and Lopez-Paze 2020; Gagnon-Aude Audet et al., 2023; Yao et al., 2022). Has the generalization project hit a dead end?
coarse embedding. thus diluting important signal at the example
In parallel, large language models (OpenAI, 2023; Touvron et al., 2023, LLMs) are taking the world by storm. LLMS are next-token predictors built with transformers (Vaswani et al., 2017) and trained on enormous amounts of natural language. One impressive capability of LLM systems is their ability to learn in-context, that is, to generalize on-the-fly to the eclectic circumstances that users enforce by prompting (Brown et al., 2020). For example, trained I.LM would complete the sequence

---

Published as a conference paper at ICLR 2024
classifier
yr
yr
classifier
1
)(2))
transformer
2
21
æi
22g
(a) Invariance DG.
(b) Marginal transfer DG.
(c) In-context DG (ours).
Figure I: Three frameworks for domain generalization (DG), predicting the target yf from the input z in test environment c. Depicted blue, the last example x;-1 contains relevant features for the current prediction. (a) Invariance DG discards all of the previously observed information from the test environment, removing too much predictive signal. (b) Marginal transfer DG summarizes all of the previously observed test inputs as coarse embedding. diluting predictive signal found the example level. (b) Our in-context DG directly observes all of the previous test inputs, allowing the search of meedle--n-theehaastee thy-thastack" signals, such the relevant one, i.e., -²-1-
"France-Paris Italy-Rome Spain-" with the sequence "Madrid," effectively learning, from the input itself, that the user demanding a capital prediction task. When interacting with LLMss one feels closer towards solving the puzzle of out-of disttiibiiiiitoonnnn (OOD) generalization. Could LLMs hold key piece to the OOD puzzle?
This paper suggests positive answer, establishing strong parallel between the concept of en- viromment in domain generalization, and the concept of context in next-token prediction. In fact, different environments describe varying contextual circumstances such as time, location, experimental intervention, and other background conditions. On the one hand, describing environments as context opens the door using powerful next-token predictors teetthe-shelf, with their adaptability to learn in-context, to address domain generalization problems. This allows us to move from coarse domain indices fine and compositional contextual descriptions, helpful amortize learning across similar environments. On the other hand, using context environment can help LLM researchers to use various domain generalization methods such as distributionally robust optimization (Sagawa et al., 2019; Xie et al., 2023, DRO) across varying contexts.
Based on these insights, we propose a natural algorithm, In-Context Risk Minimization (ICRM), illustrated in Figure lc. Given examples (af, .f) from environment propose address out-of- distribution prediction as in-distribution context-basee based prediction, training a machine:
yf ha
environment context
While the requested prediction yf concerns only the input xf the machine can now pay attention to the test experience far, extracting relevant environment information from instance and distributional features. Our theoretical results show that such in in-context learners can utilize context to zoom-in on the empirical risk minimizer of the test environment, achieving competitive performance. Further, we show that the extended input-context feature space in ICRM can reveal invariances that FRM-based algorithms ignore. standout feature ICRM its capability to improve feature learning through context-based training enabling ICRM to outperform counterparts even the absence of context. Our extensive experiments demonstrate the efficacy of ICRM, and extensive ablations dissect and deepen understanding of it.
-idt-----ittiuttionn
We organize the rest of the exposition as follows. Section reviews the fundamentals of domain generalization, centered around the concept of environment. Section explains the basics of next- token prediction, with an emphasis on learning from context. Section knits these threads together propose ICRM, framework to learn from contexts from multiple environments, supported by
2

---

Published as a conference paper at ICLR 2024
host of theory. Section 6 showcases the efficacy of our ideas in variety of domain generalization benchmarks, and Section 7 closes the exposition with some topics for future discussion.
THE PROBLEM OF DOMAIN GENERALIZATION
The goal of domain generalization (DG) is to learn a predictor that performs well across a set of domains or environments & (Muandet et al., 2013). Environment indices E & list different that may occur due to time, location, experimental
versions of the data collection process-variatiin interventions, changes in background conditions, and other contextual circumstances leading to distribution shifts (Arjovsky et al., 2019).
During training we have access to collection of triplets D {1 yr, J Each triplet contains a vector of features z;, a target label yi. and the index of the corresponding training environment e, E Eu E. Each example (z's y4) is sampled independently from joint distribution P Y). Using the dataset D. we learn a predictor h that maps features to labels, while minimizing the worst risk across the set of all environments E:
= arg min max (h), eEE'
where (h) pex())- Y) is the risk of the predictor h in environment e, as measured by the expectation of the loss function & with respect to the environment distribution pe.
As one example, we could train self-driving model classify images x; into label ys indicating the presence of a pedestrian. Each training example (I;, y/) is hereby collected from one of the training cities e; &i.. with its own weather conditions. The goal of Equation (2) is obtain predictor that correctly classifies x in new cities E & observed during test time. This has proved to be challenging (Lechner et al., 2022), predictors often exhibit penurious performance in unseen weather conditions.
Domain generalization challenging because we do not have access test environments during training time, rendering Equation (2) challenging to estimate. Therefore, to address the DG problem in practice, researchers have proposed various algorithms that make different assumptions about the invariances shared between & and Ete. In broad strokes, domain generalization algorithms fall in one of the two following categories. based invariance (Muandet et al., 2013; Ganin et al., 2016; Peters et al., 2016; Arjovsky et al., 2019), illustrated in Figure 1a, regularize predictors h(zf) to not contain any information about the environment e. This however results in removing lot of signal about the prediction task. In the second category, domain generalization algorithms based on marginal transfer (Blanchard et al., 2011; Li et al., 201 6a; Zhang al., 2020; Bao and Karaletsos, 2023) extract environment-specee information. These methods implement predictors ha of), where of - Σ ((z)) coarsely summarizes the environment e in terms of previously observed instances. Different choices for include kernel functions (Blanchard et al., 2011, MTL), convolutional neural networks (Zhang et al., 2020, ARM), and patch embeddings (Bao and Karaletsos, 2023, Context- ViT). Alas, all of these alternatives in the second category dilute relevant features found in individual examples. For example, the size of the representation é would have grow linearly with the size of the training data to describe aspects corresponding
the first category, domain generalization algorithms
small group of examples, such as extreme value statistics.
As a result, and despite all efforts, no proposal so far convincingly outperforms simple empirical risk minimization baseline (Vapnik, 1998, ERM) across standard benchmarks (Gulrajani and Lopez-Paz, 2020: Liggoon-Auder et al., 2023; Yao et al., 2022), Effectively, ERM simply pools all training data together and seeks the global empirical risk minimizer:
ht arg min
Σ e))= (h). eEE,
Does the efficacy of ERM suggest that environmental information is useless? We argue that this is not the case. The key to our answer resides in a recently discovered emergent ability of next-token predictors, namely, in-context learning.
3

---

Published as a conference paper at ICLR 2024
paradigm
training data
testing data
estimates
ERM T, V
xee xe
P(YIX) P(Y (X)) P(Z++1 Ze.... Z1) P(YIX, ~ P (Y X)
IRM
3. y;
LLM
and context 2jet and context of = (15))
ICRM
2, U.C.
x
Table I: Different learning paradigms discussed in this work, together with their training data and testing data formats, as well as the estimated predictors. In our ICRM, we amortize the current input x² and its context c containing previously experienced unlabeled examples from the same environment e', and "zoom-in" (~+) to the appropriate local risk minimizer.
NEXT-TOKEN PREDICTORS AND IN-CONTEXT LEARNING
In next-token prediction, we aim learn the conditional distribution
P(Z = 2++1 Z+ = 2t, ... Z1 21);
describing the probability of observing the token 1t41 after having observed the sequence of tokens (.1,..., 21). The quintessential next-token prediction task is language modeling (Bengio al., 2000), where the sequence of tokens represents a snippet of natural language text. Most LLMs estimate Equation (4) via transformer 1 R 1(z;;; 1e- (Vaswani et al., 2017).
LLMs exhibit a certain ability, termed in-context learning (ICL., relevant to our interests. ICL is the ability describe and learn about learning problem from the sequence of tokens (typically labeled (z, y) pairs) itself, called the context or prompt. As an example, trained LLM would complete the sequence "France- Paris Italy-Rome Spain-" with the sequence "Madrid," demonstrating its ability to infer from a few input-output pairs that the user is demanding a capital prediction task. To illustrate this ability for contexts without labels, consider the two following sequences:
"You are talking to a teenager. Write a poem on gravitational fields."
context ey
"You are talking to Physics graduate.
Write a poem on gravitational fields."
context 02
As widely observed, LLMs answer differently to these two sequences, producing two poems, say Y1 and y2. each adapted to the assumed audience. While nothing unexpected is happening here at the sequence level-the model simply produces high-likelihood continuation to each of the two prompts-we we observe degree of compositional generalization, because the LLM provide different but correct answers to the same question 2 = Ø2 when presented under two contexts ca and C2 By addressing the general task of in-distribution language modeling. L.Ms can attain significant put-of-distribunt reminiscent of meta-learning (Schmidhuber, 1987; Finn al., 2017), yet seamlessly accommodates to contexts without labels, and does not require updating the parameters of the model. ICL is also similar in spirit test-time adaptation (Wang al., 2020, TTA); however, TTA often requires updating model parameters with an externally hand-crafted objective.
abilities in a multitude of specific tasks-Suse as writing poems. ICL is
Notably, ICL emerges without supervision. The training corpus does not contain any explicit division between questions and their context beyond the natural order of the words within each snippet of language the training data. However, since we train the machine to produce an enormous amount of completions, some of which start with partially overlapping contexts, the predictor has the opportunity to amortize learning to a significant degree i.e. use the trained model to generalize across unseen distributions rather than explicitly optimizing separate model for each distribution. While the machine may have never observed the context č = "You are now speaking to a teenager," its semantic similarity C1 above -pllspplssse other similar contexts where the word speaking appears-- may endow OOD generalization. This is the desired ability to generalize over environments described in the previous section, which remained out of reach when using coarse domain indices.
4

---

Published as a conference paper at ICLR 2024
4 ADAPTIVE DOMAIN GENERALIZATION VIA IN-CONTEXT LEARNING
Our exposition has so far laid out two threads. First, Section 2 motivated the need for domain generalization algorithms capable of extracting relevant environment-speccinn features, at both the example and distributional levels. To this end, we have argued to move away from coarse environment indices, and towards rich and amortizable descriptions shared in new circumstances. Second, Section suggests understanding context as an opportunity to describe environments in precisely this manner. We now knit these threads together with a protocol to address domain generalization with in-context learners.
In- in-Context Risk Minimization (ICRM, Figure 1c):
Collect dataset of triplets D {(z, Vt. e;)) as described in Section 2. Initialize context-based predictor ý h(r; c), tasked with predicting the label y associated to the input 2, as supported by the context
During training, select E Er at random. Draw examples from this environment at random, construct one input sequence (zf....,æf)) and its associated target sequence vf). Update the context-based predictor to minimize the auto-regressive loss Σ (() where the context is ey (2 (fi.-..,,$$ 25-1)- for all j = ... and cf = 0.
During test time, a sequence of inputs from the test environment e' E Egg. We predict of = h(æz", ,c,) for æÿy where the context ef' = (2f (=, x5-1) for all = 2,... l', and ef' = 0.
......., x4) arrives for prediction, one by one, all
A few critical remarks about the above proposal are in order. The idea of using contextual information to aid the prediction through attention mechanisms in transformers has been used in carlier works on neural processes (Kim et al., 2019; Nguyen and Grover, 2022), prior data fitted networks (Müller et al., 2021) and more recent works that study the mechanisms underlying in-context learning (Garg al., 2022; Akyürek et al., 2022). These works leverage labeled contextual data. Our proposal embraces the challenge of domain generalization and only uses unlabeled data from the environments both train and test time. The most natural way to construct contexts is use past samples that appear in the natural order in which data was collected (e.g.. video). Since existing DG datasets do not provide such a refined ordering, we build contexts using environment indices that are more readily available. The proposal also requires the data at test time to be sampled from the same or slowly changing environments. While the proposal is strongly inspired from LLMs that both pay attention current query and the contextual information, there are differences namely we predict the label of the input and not the next x in the sequence.
Next, we develop theoretical guarantees on the behavior of ICRM. The results below concern the joint distribution of ((X,, Xx), (YI..... .y.....Yy Y.), E)) where each X,, Y, is an independent draw from environment E with distribution (xx (X,Y). Y). For query X,, the context preceding it is C,- = (Xi,... Xj-1) and the environment underlying this context is E. To orient ourselves around these results, we recall three predictors featured in the exposition so far. First, the global empirical risk minimizer over the pooled training data, denoted by hi in Equation (3), estimates P(Y XX) Second, the environment risk minimizer estimates P(Y X, E). Third, our in-context risk minimizer, denoted by
h arg min h
Σ E E M...G.GDDMMM,
estimates the conditional expectation E(Y X, The sequel focuses on the binary cross-entropy loss E. Our first result shows that, in the absence context, ICRM zooms-out to behave conserva- tively.
Proposition 1 (ot) In the absence of context, ICRM behaves as the global empirical risk minimizer across the support of the training environments, i.e., h(- () = ht(().
The above result is built on the insight that ICRM is Bayes optimal at all context lengths and ERM is Bayes optimal for context 0. Having established the connection between ICRM and ERM in
5

---

Published as a conference paper at ICLR 2024
the absence of any context, now study the benefits of ICRM in the presence of sufficiently long contexts. The following result shows that, when provided with context from a training environment E Er. our ICRM zooms-in and behaves like the appropriate environment risk minimizer, as shown in Table 1. We assume that P(Y = 1 X = x. F = e) is parametrized by a function h* (1c 0%), where 0% describes features of the environment relevant to the query x. for all there exists an ideal amortization functiio b that takes as input the query X and context Ce preceding it-both sampled from environment Æ'-and approximates 0f Formally, the sequence of random variables 6(X, C+) indexed by converges almost surely to the random variable 0g.
&. We also assume
Theorem 1 (Full iid zoom-in). Let h- (x, 0%) describe P(Y 1 X I, E e) for all e E E. Further, we assume the existence of an amortization function 6(X, C.) 3 0g. Then, ICRM zooms-in on the environment risk minimizer and achieves cross-entropy loss over the training distribution
lim H(Y X,C,) H(Y X) 2-000
Further, if TV; X) > O, ICRM has better performance than the global risk minimizer.
Theorem states that ICRM converges to empirical risk minimizer of the environment under infinitely long contexts. Next, we show that ICRM can partially zoom-in on the appropriate environment risk minimizer even with contexts of length of one.
Theorem 2 (Partial iid zoom-in). Suppose the joint distribution ((X1....Xp (Y (Y....., Yi), F) is Markov with respect to a Bayesian network. The query X and the environment E are statistically dependent and form the Markov blanket of Then ICRM partially zooms-in on the environment risk minimizer, improving over the performance of the global empirical risk minimizer in terms of the cross-entropy loss. Further, the improvement is strictly monotonic context length t.
Next, move to the out-of distiiibbitin setting where the test environments can be different from the training environments. To provide theory for a domain generalization result, we must place some assumptions on the data generation process. In particular, and for all E E, let = y, e N(ng, Σÿ), and x + g(z) where the latent variables z are sampled conditional on the label y and environment e from Gaussian distribution with mean and covariance depending on (y, e), and are then mixed by a map to generate the observations x. We summarize the environment as parameter vector % [P% p²z EX) ye (0.1) where p the probability of label y in environment e. Our next result shows that ICL algorithms that learn h(;; c) exhibit robust behavior under distribution shifts. In contrast, such guarantees are not known for algorithms that generate predictors of the form h(x).
Define 6, to be permutation of % that swaps its two components. We construct the Voronoi cells corresponding to the points in the union of sets {Te)eeE,, and {& )e The set of points in the Voronoi cells corresponding environments. Next, we show that there exists an ICL algorithm, which takes the data from multiple environments as input and outputs a predictor that takes current query and context as input, whose output predictors perform well in novel test environments even those that are sufficiently far away from the training environments, so long as they are in the Voronoi cells of training environments.
the set of points {Te)eeé.. define the Voronoi cells of the training
Theorem 3 (Full OOD zoom-in). Consider data triplets (I, v),e) generated from z Nuu Eγ) and x + 9(2). Ve E E, where is the identity map (see Appendix A for extension to general diffeo- morphism g). There exists an ICL algorithm that in the limit of infinitely long contexts produces Bayes optimal predictions for all the test environments in the Voronoi cells of the training environments.
5 ICRM UNDER THE LENS OF INVARIANCE
Common advice in domain generalization recommends following the invariance principle to learn robust predictors (Peters et al., 2016; Arjovsky et al., 2019). One simple version of the invariance principle is to "select those inputs leading to stable predictors across training environments." At first sight, one could argue that the proposed ICRM does not adhere to such principle, as it is adapting to environment-specie specific information provided in the form of context. As we shall now illustrate, ICRM's implementation of ERM on the extended input-context feature space reveals invariant predictors that vanilla implementation of ERM the standard feature space fails to find. To see this, consider linear least-squares regression problem mapping two dimensional inputs x = (x',x²) into a target under environments E m as y x² + 3 22 + e, where n E F ee. the pair (a. (3)

---

Published as a conference paper at ICLR 2024
invariant regression coefficients, and is independent noise term. We make one simplifying assumption for pedagogic purposes. During training, we provide ICRM training directly with the relevant extended feature space (x',2² MM (2.)))) (2), instead of requiring the algorithm to learn such representation from general-form sequential context.
In this setup, ICRM learns predict using a xx + 0-
pe B p In contrast, ERM
trains a linear model on (x* (²") and predicts using â x² + B x². The main point is: if B o 0 and cov(X', x)) * O, then ā * a, and the error of ERM in a new environment grows with the variance of x². On the other hand, ICRM estimates the true invariant coefficient a. and the resulting error independent of variance of x², even in the absence of context during test time. As result, ICRM exhibits better outtof-distribution performance than ERM without any contextual information at test time. For derivation and generalization of these claims, see Appendix A.
We believe that ICRM, and more generally ICL, provide a novel view on invariance. On the one hand, prior DG algorithms advocated to remove features as guide to reveal invariance. On the other hand, in-context learners suggest that extending features with context affords invariance otherwise unnoticed. This needs further clarification: while the process of zooming-in to environment risk minimizer does not provide us with an invariant predictor over the original feature space, the process of zooming-in is often an invariant mechanism over the extended feature space. These points reminiscent of the concept of "fragility" in the philosophy of causation (Menzies and Beebee, 2020). Does smoking cause cancer? Not invariably across all contexts or environments. Yet, smoking does cause cancer iivariably-acros all contexts or environments-wher extending the feature space as to include additional causes such as diet, genetic predispositions, and the number of smoked cigarettes. The ever wer-grooingg collection of causes approaches what Mill (1856) called the total cause, a large context sharpening invariance the expense of constraining the diameter of the environment. In the extreme, when constraining the environment to contain only one smoker, the outcome of lung cancer disease invariably follows.
EXPERIMENTS
To evaluate the efficacy of ICRM, our experiments address the following questions:
1. How does ICRM fare against competitive DG algorithms, for different context sizes?
2. What is the impact of model architecture on ICRM's gains?
3. Can ICRM search for query relevant signals in the context?
In our experiments, we compare ICRM against marginal transfer methods such as Adaptive Risk Minimization (Zhang et al., 2020, ARM), and test-time adaptation proposals such as TENT (Wang et al., 2020). As a strong baseline, we also include ERM in our experimental protocol. To ensure a fair comparison across different algorithms for each dataset, we use a standardized neural network backbone (ConvNet or ResNet-50 depending on the dataset) as described in Appendix C.4. For ICRM, the same backbone is used to featurize the input, which is then processed by the decoder-only GPT- 2 (Radford et al., 2019). During both training and inference for ICRM, data in a sequence is sampled from the same environment. For fair comparisons, hyper-paramenn parameter tuning, and testing (Gulrajani and Lopez-Paz, 2020), details in Appendix C.4. We assess these methods across six image classification benchmarks, featuring diversity shift FEMNIST (Cohen et al., 2017) contains MNIST digits and handwritten letters from individual writers environments. Rotated MNIST concerns varied rotational angles as environments. Tiny ImageNet-C and CIFARIO- (Hendrycks and Dietterich, 2019) introduce diverse image corruptions to create multiple environments. WILDS Camelyon 17 (Koh et al., 2021) studies tumor detection with data from multiple hospitals as distinct environments and Imagenet-R R (Hendrycks et al., 2021) contains various renditions (e.g., paintings, embroidery, etc.) of ImageNet object classes as domains. More details are provided in Appendix C.3. We consider addressing correlation shifts (Ye et al., 2022) as a future work of our paper, as involves scenarios where the test domain contains naturally occurring subpopulations time-based shifts without clear domain separation
adhere DomainBed's protocols for training,
6.1 ADAPTATION TO DISTRIBUTION SHIFT
To study the adaptation of other approaches to distribution shifts, we report the average performance on four datasets across three independent runs of the entire sweep for test context lengths of 0, 25, 50,
7

---

Published as a conference paper at ICLR 2024
Table 2: Average/worst OOD accuracy for different context lengths, for Adaptive Risk Mini- mization (ARM), Empirical Risk Minimization (ERM), Test Entropy Minimization (TENT) and our ICRM on FEMNIST, Rotated MNIST, WILDS Camelyon17 and Tiny-mmaeNNetCC.
Data method
Average test accuracy
Worst case test accuracy
FEMNIST
25
75
100
25
50
100
ARM
49.5
83.9
84.4
84.7
84.6
23.6
59.5
60.7
57.0
58.8
TENT
78.1
77.9
81.2
82.5
83.3
55.2
57.2
63.3
65.9
67.2
ERM
79.3
79.3
79.3
79.3
79.3
59.0
59.0
59.0
59.0
59.0
ICRM
78.7
87.2
87.4
87.5
87.8
59.8
69.3
70.6
70.6
70.6
Rotated MNIST
0
25
50
75
100
25
50
75
100
ARM
36.5
94.2
95.1
95.3
95.5
85.3
87.2
87.9
87.9
TENT
94.1
88.0
93.8
94.3
80.2
88.5
88.5
80.2
81.3
ERM
94.2
94.2
94.2
94.2
94.2
80.8
80.8
80.8
80.8
80.8
ICRM
93.6
96.1
96.2
96.2
96.2
82.5
88.5
88.5
88.8
88.8
WILDS Camelyon17
25
50
75
100
25
50
75
100
ARM
61.2
59.5
59.7
59.7
59.7
TENT
67.9
81.8
87.2
89.4
89.4
ERM
68.6
68.6
68.6
68.6
Samee average accuracy
ICRM
92.0
90.7
90.8
90.8
90.8
Tiny Iet-CNet-C
25
75
100
25
50
100
ARM
30.8
31.0
31.0
31,0
31.0
8.2
8.3
8.2
8.3
8.2
TENT
31.7
1.6
1.7
2.0
2.1
9.4
1.2
1.4
1.6
1.6
ERM
31.8
31.8
31.8
31.8
31.8
9.5
9.5
9.5
9.5
9.5
ICRM
38.3
39.2
39.2
39.2
19.2
19.5
19.4
75, and 100 samples. The results for other datasets and more DG algorithms are reported in Table and Table 8. As Table 2 shows, ICRM outperforms all methods across context lengths, except at null context length on MNIST datasets, where ERM exceeds by 1%. Further, our gains persist over both the worst group and average accuracy across testing environments. Figure 5 zooms into the model's performance between no-context and 25 context samples, highlighting the consistent superiority of in the
ICRM even with small contexts. Additionally, ICRM demonstrates gains in performance absence of test context. Specifically for both WILDS Camelyon17 and Tiny ImageNet-C. ICRM outperforms baselines despite not leveraging any context from the test environment. We hypothesize that ICRM training still benefits from contexts as to find contextual features that ERM ignores.
6.2 UNDERSTANDING THE IMPACT OF ARCHITECTURE
To dissect the performance gains potentially arising from ICRM's transformer architecture, we explore two additional competitors. First, we train an ERM baseline, ERM* using an identical architecture to ICRM, but without context. Second, we train an ARM baseline, ARM*, where the input and context coarse summary are sent to GPT-2 such that the model now attends to the summary through attention layers.
Table 3: Worst group OOD test accuracies for +RM and ERM+ contrast to their base algorithms, ARM and ERM across FEMNIST, Rotated MNIST, WILDS, Camelyon17, and Nen-
Dataset
ARM
ARM
ERM
ERM +
# Context Samples
0
100
0
100
100
100
FEMNIST
23.6
58.8
51.7
62.0
59.0
59.0
53.3
53.3
Rotated MNIST
28.2
87.9
71.4
81.1
80.8
80.8
$1.9
81.9
WILDS Camelyon17
61.2
59.7
55.8
55.0
68.6
68.6
50.1
50.1
Tiny ImageNet-C
8.2
8.2
1.9
1.9
9.5
9.5
8.3
8.3
Table 3 presents the performance of both ERM* and ARM relative to ERM and ARM, across the four datasets. ARM demonstrates superior zero-shot performance over ARM both FEMNIST and Rotated MNIST. However, ARM maintains a performance advantage over *RM across varying counts of in-context samples WILDS Camelyon17 and Tiny ImageNet-C,, with pronounced
8

---

Published as a conference paper at ICLR 2024
difference on the latter. Similarly, ERM either matches outperforms ERM* on all four datasets. Therefore, even with similar architectures, prior protocols fall short of the proposed ICRM.
a
2
3
Query
0.12
0.06
0.32
0.10
0.15
0.10
0.05
0.11
0.00
0.00
Query
0.00
Volley ball
Bench Wagon
Academic
Hog
Walking stick
Plate
Volleyball
Volleyball
Gown
Query Image
Attention scores for random test sequences, for ICRM on FEMNIST (top) and Tiny ImageNet-C (bottom), with the class label shown below. Samples following the query have an attention score of 0.0 because of the causal attention mechanism.
Figure
6.3 INVESTIGATING ATTENTION IN ICRM
As discussed in Section 2, ICRM can learn amortization function by paying attention to the input query and context. To understand this better, we construct random data sequence from the test environment and analyze attention scores between each example in this context and a novel input query. Figure 2 illustrates attention scores from single head for query image (marked in blue) for FEMNIST and Tiny ImageNet-C. Note that samples following the query have attention score of 0.0 because of the causal attention mechanism. The top row reveals that the model selectively attends images featuring at least curved arcs (marked in green) while paying little attention to partial circle (highlighted in red). Similarly, in the bottom row, the model effectively discerns individuals across samples within the sequence and also indicates semantic understanding similarity.
DISCUSSION
We introduced In-Context Risk Minimization (ICRM), framework to address domain generalization as context-based prediction. ICRM learns in-context about environmental features by paying attention to unlabeled instances as they arrive. In such away, ICRM dynamically zooms-in in on the test environment risk minimizer, enabling competitive out-of of-distribttion generalization.
ICRM provides new perspective invariance. While prior work on DG focused information removal as a guide to generalization, ICRM suggests that extending the feature space with the relevant environment information affords further invariance. By addressing the general problem of context-based prediction in-distribution, we amortize the performance over a multitude of specific custt--diitribution tasks. More generally, by framing DG next-token prediction, our approach can be adapted to fully exploit data in natural order (such as in video or text, ordered by time and position), Léon Bottou once said, Nature does
more closely mimicking the human learning experience- not shuffle data. That said, we view extending ICRM to scenarios where the environment contains naturally occurring subpopulations or time-based shifts without clear separation as exciting future direction. As a word of caution, we must conduct research to guarantee that in-context learners do not "zoom-in" on toxic spurious correlations with high predictive power in certain environments. We close with quote from Andersen et al. (2022), for whom zooming-in
refers cognitive agent's 's ability intelligently ignore irrelevant information and in those aspects the world are relevant to their goals. The relevance realization framework suggests that the hrain achieves this feat by attempting to balance the competing goals of remaining efficient in the current environment while also being resilient in the face of environmental perturbations.
Paralleling the examples from Andersen et al. (2022), we would like to further understand how next- token prediction and in-context learning serves powerful mechanism to amortize and dynamically navigate trade-offs such as such as sffeciency-resilinnnyy exploration-exxploitatio specialization- generalization, and focusing-diversifying..

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
SG and SJ acknowledge funding from the Office of Naval Research grant N00014-20-1-2033 (MURI ML-SCOPE) and NSF award CCF-2112665 (TILOS AI Institute). We thankful to Martin Arjovsky, Léon Bottou, Elvis Dohmatob, Badr Youbi Idrissi, Maxime Oquab, and Ahmed Touati for their valuable feedback and help.
REFERENCES
Kartik Ahuja and David Lopez-Paz. A closer look at in-context learning under distribution shifts. arXiv, 2023.
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In ICML., 2020.
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean- Shan-Christophe Gagnon-Audet, Yoshua Bengio, loannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of- distribution generalization. NeurIPS, 2021.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learnin?? investigations with linear models. arXiv preprint arXiv:221 1,.666,, 2022.
Brett P Andersen, Mark Miller, and John Vervaeke. Predictive processing and relevance realization: exploring convergent solutions to the frame problem. Phenomenology and the Cognitive Sciences, 2022.
Julia Angwin, Jeff Larson, Surya Mattu, and Lauren Kirchner. ProPublica, 2016.
Machine bias.
URL https://www. hwtp://wwww......../o.aaaaaa.ee. machine bias risk--isseeeeeeeeee assessments -in criminal sentencing.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv, 2019.
Robert B Ash and Catherine A Doléans-Dade. Probability and measure theory. 2000.
Yujia Bao and Theofanis Karaletsos. Contextual Vision Transformers for Robust Representation Learning. arXiv, 2023.
Yoshua Bengio, Réjean Ducharme, and Paseal Vincent. A neural probabilistic language model. NeurlPS, 2000.
Gilles Blanchard, Aniket Anand Deshmukh, Ürun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. JMLR, 2011.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. NeurIPS, 2020.
Shiyu Chang, Yang Zhang, Mo Yu, and Tommi S Jaakkola. Invariant rationalization. In JCMI, 2020.
Yining Chen, Elan Rosenfeld, Mark Sellke, Tengyu Ma, and Andrej Risteski. Iterative feature matching: Toward provable domain generalization with logarithmic environments. NewrIPS, 2022.
Gregory Cohen, Saced Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. In IJCNN, 2017.
Cian Eastwood, Alexander Robey. Shashank Singh, Julius Von Kügelgen, Hamed Hassani, George J Pappas, and Bernhard Schölkopf. Probable domain generalization via quantile risk minimization. NeurlPS, 2022.
Cian Eastwood, Shashank Singh, Andrei Liviu Nicolicioiu, Marin Vlastelica, Julius von Kügelgen, and Bernhard Schölkopf. Spuriosity didn't kill the classifier: Using invariant predictions to hamesss spurious features. arXiv preprint arXiv.2307.09933, 2023.
10

---

Published as a conference paper at ICLR 2024
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learinng for fast adaptation of deep networks. In ICML, 2017.
Jean-Christophe Gagnon-Audes Kartik Ahuja, Mohammad Javad Darvishi Bayazi, Pooneh Mousavi, Guillaume Dumas, and Irina Rish. WOODS: Benchmarks for out-of-distribution generalization in time series. 7MLR, 2023.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. JMLR, 2016.
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-30598, 2022.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes. In International conference on machine learning, pages 1704-1713. PMLR, 2018a.
Marta Gamelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Tch. Neural processes. arXiv preprint arXiv:1807.01622. 2018b.
Ishaan Gulrajani and David Lopez-Paz. search of lost domain generalization. arXiv, 2020.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv, 2019.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, al. The many faces of robustness: A critical analysis of -ut-of-distribution generalization. In Proceedings of the IEEE/CVF international conference on computer vision, pages 3340-8349 2021.
Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On feature learning in the presence of spurious correlations. NeurlPS, 2022.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Enforcing predictive invariance across structured biomedical domains, 2020.
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. In A/STATS, 2020.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint 4rXiv:1901.05761, 2019.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv, 2014.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training sufficient for robustness to spurious correlations. arXiv, 2022.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In ICMI., 2021.
Masanori Koyama and Shoichiro Yamaguchi. Out-of-distribution generalization with maximal invariant predictor. arXiv, 2020.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapola- tion. arXiv, 2020.
Sébastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Jacostt-JJliinn Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica. In Conference on Causal Learning and Reasoning, 2022.
11

---

Published as a conference paper at ICLR 2024
Mathias Lechner, Ramin Hasani, Alexander Amini, Tsun-Hsuar Wang, Thomas A Henzinger, and Daniela Rus. Are all vision models created equal? a study the open-loop to closed-loop loop causality gap. arXiv, 2022.
Ya Li, Mingming Gong, Xinmei Tian, Tongliang Liu, and Dacheng Tao. Domain generalization via conditional invariant representations. In Proceedings of the AAAI conference artificial intelligence, 2018.
Yanghao Li, Naiyan Wang, Jianping Shi, Jiaying Liu, and Xiaodi Hou. Revisiting batch normalization for practical domain adaptation. arXiv, 2016a.
Yanghao Li, Naiyan Wang, Jianping Shi, Jiaying Liu, and Xiaodi Hou. Revisiting batch normalization for practical domain adaptation. arXiv preprint wrXiv:1603.04770 2016b.
Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. arXiv, 2020.
Maggie Makar, Ben Packer, Dan Moldovan, Davis Blalock, Yoni Halpern, and Alexander D'Amour Causally motivated shortcut removal using auxiliary labels. AISTATS, 2022.
Peter Menzies and Helen Beebee. Counterfactual Theories of Causation. In The Stanford Encyclope- dia of Philosophy. 2020.
John Stuart Mill. A System of Logic, Ratiocinative and Inductive, volume 1. 1856.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In ICML, 2013.
Jens Müller, Robert Schmier, Lynton Ardizzone, Carsten Rother, and Ullrich Köthe. Learning robust models using the principle of independent causal mechanisms. arXiv, 2020.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. arXiv preprint arXiv:2 2021.
Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-awan meta learning via sequence modeling. arXiv preprint arXiv:2207.04179 2022.
OpenAl. GPT-4 Technical Report. arXiv, 2023.
Giambattist Parascandolo, Alexander Neitz, Antonio Orvieto, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary. In ICLR, 2021.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society Series B: Statistical Methodology, 2016.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Alexandre Rame, Corentin Dancette, and Matthicu Cord. Fishr: Invariant gradient variances for out-of-distributtion generalization. ICML., 2022.
Alexander Robey, George Pappas, and Hamed Hassani. Model- based domain generalization. arXiv, 2021.
Mateo Rojas-Carulla, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant models for causal transfer learning. JMLR. 2018.
Walter Rudin. Principles of mathematical analysis. 1953.
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv, 2019.
12

---

Published as a conference paper at ICLR 2024
Olawale Salaudeen and Oluwasanmi Koyejo. Target conditioned representation independence (tcri); from domain-invenn invariant to domain comain-general representations. arXiv preprint arXiv:2212.11342 2022.
Jürgen Schmidhuber. Evolutionary principles in self-referential learning, learning how to learn: the meta-meta-... hook. PhD thesis, Technische Universität München, 1987.
Steffen Schneider, Evgenia Rusak, Luisa Eck, Oliver Bringmann, Wieland Brendel, and Matthias Bethge. Improving robustness against common corruptions by covariate shift adaptation. Advances in neural information processing systems, 33: 1539-115511 2020.
Zheyan Shen, Jiashuo Liu, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards generalization: A survey, arXiv, 2021.
Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. arXiv preprint arXiv:2104.09937, 2021.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In ECCV Workshops, 2016.
Damien Teney, Ehsan Abbasnejad, and Anton van den Hengel. Unshuffling data for improved generalization. arXiv, 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
Vladimir Vapnik. Statistical learning theory. Wiley, 1998.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurlPS, 2017.
Victor Veitch, Alexander D' Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual invariance to spurious correlations in text classification. NeurIPS, 2021.
Johannes Von Oswald, Eyvind Niklasson, Fttore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, PMLR, 2023.
Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. NeurIPS, 2021.
Yoav Wald, Gal Yona, Uri Shalit, and Yair Carmon. Malign overfitting: Interpolation can provably preclude invariance. arXiv preprint arXiv:2211.15724, 2022.
Yoav Wald, Claudia Shi, Aahlad Puli, Amir Feder, Limor Gultchin, Mark Goldstein, Maggie Makar, Victor Veitch, and Uri Shalit. Workshop on spurious correlations, invariance and stability. ICML, 2023. URL https: wtw.e.c......itttieeemmeee /0tp////tta/////3/////k///////// wwookkhhho/219999
Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. arXiv, 2020.
Haoxiang Wang, Haozhe Si, Bo Li, and Han Zhao. Provable domain generalization via invariant- feature subspace recovery. In ICML,, 2022.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V Le. Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining, arXiv, 2023.
Sidney J Yakowitz and John D Spragins. On the identifiability of finite mixtures. The Annals of Mathematical Statisties, 1968.
Shen Yan. Huan Song. Nanxiang Li. Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.00677, 2020.
13

---

Published as a conference paper at ICLR 2024
Huaxiu Yao, Caroline Choi, Bochuan Cao, Yoonho Lee, Pang Wei W Koh, and Chelsea Finn. Wild-time: benchmark of in-the-wild distribution shift over time. NeurIPS, 2022.
Nanyang Ye, Kaican Li, Haoyue Bai, Runpeng Yu, Lanqing Hong, Fengwei Zhou, Zhenguo Li, and Jun Zhu. Ood-bench: Quantifying and understanding two dimensions of out-of-distributioo generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 7947-7958, 2022.
Marvin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: Learning adapt domain shift, NeurIPS, 2020.
Yihua Zhang, Pranay Sharma, Parikshit Ram, Mingyi Hong, Kush Varshney, and Sijia Liu, What is missing IRM training and evaluation? challenges and solutions. arXiv, 2023.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
Theorems and Proofs
15
A.1 Proof of Proposition
15
A.2 Proof of Theorem
16
A.3 Proof of Theorem 2
21
A.4 Proof of Theorem
23
A.5 Extension of Theorem 3
24
................................
A.6 Comparing ICRM and ERM under the lens of invariance
27
A.7 Illustration of failure of existing MTL methods
Related work
30
Supplementary experimental details and assets disclosure
30
C.1
Assets
................................
30
C.2 Hardware and setup
30
C,3 Datasets
30
C.3.1 Federated Extended MNIST (FEMNIST)
................................
30
C33.2
Rotated MNIST
31
C.3.3
WILDS Camelyon17
................................
31
C.3.4
CIFARIO-C and Tiny ImageNet-C
31
C.3.5
ImageNet
C.4
Experimental protocols
32
D Additional experiments
33
D.I Adaptation curves of various algorithms
33
D.2 Domain generalization accuracies per algorithm and dataset.
................................
34
D.2.1 Adaptation to distribution shift
34
D.2.2 Robustness of ICRM in the absence of environment labels
37
D.2.3 Understanding the impact of architecture
37
D.3 Comparison of ICRM with In-Context Learning
D.4
Investigating the features learned by ICRM
A THEOREMS AND PROOFS
A.1 PROOF OF PROPOSITION 1
Lemma 1. ICRM is Bayes optimal at all context lengths. Suppose E is the binary cross-entropy loss and the labels Y are binary. The optimal in-context learner h (equation 5) satisfies the following condition, i.e., for each k E [t]
BE Fk, Cs ck),
(6)
15

---

Published as a conference paper at ICLR 2024
for almost all Tk) in the support of training distribution except set of measure zero, and where the expectation is over conditional on [ck, In other words, the in-context learner is Bayes optimal at each context length.
Proof. In this result, we consider the problem of binary classification, Suppose h(TR ox) is the predicted probability of class = conditional on 2y and Ok. Define h(xx: o&) = [M(*ks Cx), 1
h(zg; cx) describing the probability of both the classes.
From equation 5, recall that the objective of ICRM is minimize
Σ
(7)
Consider one of the terms in the sum above [e(((h(xx; t Yt)] in this term, obtain
Substituting as the cross-entropy
((h(xx: Ck), Ck). H(YI C2) + KL (
Ck), then the second term in the above is zero and E [[E(h((a Ck): Y&)] equals H [Rx Ca). Since KL divergence is always non-negative, [Xx Ck) corresponds to the lowest value that can be achieved by IE [E((h(Xa; Yx). If h Cx) = P(Yx Xx. C&) for all k E [t], then each of the terms in the sum in equation 7 are minimized. As result, h(Xx; Ck) |X,, C²) for all [f] is a solution to equation 5.
If Cx)
Consider another minimizer h of equation 5 and define the corresponding distribution h' For each E [t]. the second term E [KL((PYx)XX) Cx) (Xx; Ck)] has to be zero for to be a minimizer.
O. then we claim that h' (IrR; (e) = P(Y = IA. Ck cx) for almost all (I&, cu) in the support of training distribution except over set of measure zero. If the probability measure associated with Xx, C₂ is absolutely continuous w.r.t Lebesgue measure, then this claim follows from Theorem 1.6.6 (Ash and Doléans-Dade, 2000). If the probability measure associated with Xg, C& absolutely continuous w.r.t counting measure, then this claim trivially follows.
If E (LL( C (Xx;
We proved the above result for classification and cross-entropy loss for measures over X, C' that are either absolutely continuous w.r.t Lebesgue measure or the counting measure. It is easy extend the above result for regressions and least square loss; see Lemma 1 in Ahuja and Lopez-Paz (2023).
Proposition 1 (Zu) In the absence of context, ICRM behaves as the global empirical risk minimizer across the support of the training environments, i.e., h(- 0) = h+).
Proof. From Lemma 1, it follows that (((2; = B| = Jı. Ck ct). The solution to empirical risk minimization (x())) E = xT), where the expectation is computed over the training distribution of Y conditional on x. When the context is empty, then we have h(x; 0) = E|Y|X1 x] hi (xx) for almost all x in the support of training distribution except over set of measure zero.
A.2 PROOF OF THEOREM 1
Before stating the proof of Theorem 1, we provide an example of an ideal amortization map b(-).
Example of ideal amortization map.
Consider the example from equation ??. where y = a xx
3- 22 + s. P(Y = x x, e) pp ax 3p2)) where Pe is the probability density of noise. Observe that P(Y = y|X = x, E = is parametrized in terms of μa and the sequence of random variables b(X, C.) - E X3 converge almost surely to n where X3 is the second
component of X2 and C. (X1,".. X
16

---

Published as a conference paper at ICLR 2024
For ease of exposition, start with the case when all the concerned random variables X, Y, Ct E, b(X, Ct), where X is the current query and is its label and C₁ is the context preceed- ing it sampled from environment E, and b(+) is the ideal amortization map, are sacrete-valued with finite support. Subsequently, we study more general settings.
Theorem 1 (Full iid zoom-in). Let h (x, 0%) describe P(Y =1 X x, = for all e E E. Further, we assume the existence of an amortization function b(X, Ct) 3 6- Then, ICRM zooms-in on the environment risk minimizer and achieves cross-entropy loss over the training distribution
lim H(Y XX H(Y X)
Further, if [YY X) O, ICRM has better performance than the global risk minimizer.
Proof. As stated above, in this proof, work with discrete-valued X, Y, C., E,b(X, C.) that also have finite support, where X is the current query and Y is its label and C2 is the context preceeding it sampled from environment E, and b(-) is the ideal amortization map. Subsequently, we study more general settings.
Since each (X,, Y2) is sampled independently given training environment E, we can conclude C, E) = 0. Therefore,
I(Y;C|JXEE) =0 = HYYX,E.)
Observe that for all t Z+
H H( H(Y & H(YIX, S H(YIX,
(8)
where 24 is the set of all positive integers. The first inequality in the above follows from the fact that conditioning reduces entropy. For the second inequality, we use the following property. Consider U, V two random variables and define W a(V). Observe that [UU; W[V) = = H(UVV)
(UVV, W) & (W).
Since the inequality above equation 8 holds for all t, we obtain
HY lim C.) < lim H(Y]XX, 10000
(9)
In the above, we use the following property. If an I Dn, Vn E Z, and lim,-+oo On and lim bn exist, then lim, -o On < lim,-roe bn- In what follows, lim HY||X, C.) and lim,- HY b(X, C.)) exist. First observe that HY|[X, C.+1) S H(Y|X, C.) for all / as result the sequence is decreasing bounded below by 0 and thus from monotone convergence theorem (Rudin, 1953) lim H(Y X, G,) exists. Next, we will show that lim; H (Y) b(X, Ct)) H(Y E). We will then combine it equation obtain what we intend to prove, i.c., lim, oo HY C.) = HYY|X,
will show that both the limits
For each X = and E in the support of training distribution, we argue that b(X, C2) 9 05. Suppose this was not true. This implies that the probability that P(limy+0o b(X, C.) # 0%
E e) 3 0. Since X 2, E occurs with finite probability (as X and E are discrete-valued and x, e in the support) say a, then o33 fraction of sequences of b XX C.) do not converge to 8% which contradicts the assumption that b(XC)) " g%.
Consider a (z, 0) from the support of (X, (x, where X is the current query and E is the envi- ronment from which X and context preceeding is sampled. Let us consider the distribution VOY)Xb((X,CC
P(Y = -/x = = (,, = 0) -Y=
(10)
P(X MXX=,,MXXC.)=
We simplify lim P(Y X, b(x, C.)) below.
lim P(Y X ,,,
lim P(Y y 6( =)
(11)
100
lim = - [
17

---

Published as a conference paper at ICLR 2024
We show that the limits of the numerator and denominator exist (and non-zero for the denominator) and we simplify these separately below.
lim P(Y = y, X = x. b(x, C+) = =0) = lim t0000
Σ P(Y = X = e,, 0) t
Σ P(Y y x, E= e) lim
P(X x, x, e, ((,C,) =
(12)
1000
Σ P(Y=
x, x, E lim P(b(X, C.) -| x. E. = y1 y| = E P
1-000
In the simplification above, we firstly used the fact that we can interchange sum and limits, this is true because only takes finitely many values. In the simplification above, we also use the fact Y L C, Since b(X, C+) converges 8% almost surely, the distribution lim, P(b(X, C.) = 0 = z, E = e) takes a value one if f 0% and zero otherwise. As a result, the above expression becomes
lim P(Y b(X
E P(Y P x E= x, e). eEE..0
(13)
where E, is the set of all the environments observed conditional on X with 0 = 0. Observe
that all the environments in E have the same P(Y = X|X = x, E = e) given by h (1( 0). We can write
lim P(Y = xxX,,,, = 0) = (x,
E
P(X= =1x=z E=)
(14)
cEE.,
We simplify 4om P(X x, ((,, Ct) = 6) in similar manner to obtain
P(X = x. 6(x C+) = 0) Σ 100
lim
E(e=
(15)
cEE.,0
Observe that the denominator is positive and not zero because x, 0 is in support of X, S. We use equation 14 and equation to obtain
so P(Y = 11X,,,,,,,
lim P(Y = 1|x = xz b(x, C.) = 0) 1+000
lim4-100 P(X b))
(16)
h*(2,0) EeEE,
Exee, P(X x= E,EEce) (x,
Therefore,
lim P(Y = 1|X x, b( = - = Ex t000
(17)
where is any environment in E,,,, i.c., it is in the support of data sampled with X and that also
satisfies 0% = 0.
lim M(,C)) = Σ lim P(X = x. MX,( = lim H(Y = xx x, 1-000
a.,
Σ E
Σ
P(X == a)
lim
(+3
6EE.66
(18)
18

---

Published as a conference paper at ICLR 2024
In the above simplification, we again swap limits and sum because the summation is over a finite
set of values. From equation 17, it follows that lim HY [X = I, b(X, C2) = 0) = HY
r, E e), where is any environment in Ex,o- We use this in the above to get
lim HY b(x,
P(
E
x E E H =I,
2,0
BEE..'
= = P(X = x, F. = 8) =( E
(19)
P(X = x.
x x. o)
= H)Y
2,6
We combine the above with equation 9 to obtain lim, H( = H(Y E). Finally, observe that if I(Y; E XY||X, E), ICRM impoves over ERM that attains a cross-entropy loss of H(][X).
HY| E) < H(|J)). Since lim,-soo H(Y
This completes the proof.
We now extend the argument to setting beyond discrete random variables. In particular, we consider settings where X, E, b(X, Ct) can be either discrete or continuous random variables. In the notation to follow, we use dP to denote the Madon-Niooyym derivatives. For discrete random variable, the Radon-Nikodym derivatives correspond to the standard probability mass function and for continuous random variables it would correspond to standard probability density functions. We operate under some regularity assumptions. We assume that the support of E has finite volume and the support of (X, b(X, C.)) has a finite volume for all t. Further, we assume that the Kadon-Nikody. the joint dP x, E e, b(X, C.) 8) is bounded above. While much of the proof that follows is same as the previous proof, we repeat the arguments for completeness.
derivative of
Theorem 4. Let h* (I, 82) describe dP = X x. E é) for all E. Further, we assume the existence of an amortization function b(X, C.) $ gg. Then, ICRM zooms-in on the environment risk minimizer and achieves a cross-entropy loss the training distribution
lim 1-0
H(Y (,) = H(Y X)
Further, if I E X) > O, ICRM has better performance than the global risk minimizer.
Proof. Since each (X,, is sampled independently given a training environment E, we can conclude I(Y; C. X, E) 0. Therefore,
I(Y; C,|, E) = 0 =
= HY(I,E,
Observe that for all t E 2+
H(Y|X, E) = H(YIX,E,C) & H(YIX, & UXX]X,6(XXC.)
(20)
where Z+ is the set of all positive integers. The first inequality in the above follows from the fact that conditioning reduces entropy. For the second incquality, we use the following property. Consider U, V two random variables and define W a(V). Observe that (U; WIV) = H(UVV)
HUVV W) S (W).
Since the inequality above equation 20 holds for all we obtain
H(Y|X,E)
lim H(Y|X,C,) < lim H(Y]X, C.). t-200
(21)
In the above, we use the following property. If an < bn, Vn E 2, and lim On and lom by exist, then lim On < lim,- bn- In what follows, we will show that both the limits (,,,((XC)) exist. First observe that HY| C1+1) S
lim X Ci) and lim H(Y|X, C,) for all as result the sequence is decreasing bounded below by 0 and thus from Monotone convergence theorem lim (Y C.) exists. Next, we will show that
19

---

Published as a conference paper at ICLR 2024
lim H( Ci) H(Y]X We will then combine it with equation 21 to obtain what we intend to prove, i.., lim; H = H(,
and E in the support except over set of probability measure zero, argue that 6(X, C.) $ 0%. Suppose this was not true. Define I' to be the set of values of x., e for a.e.
For each X
which 6(X, C.) A 0g. Let P((X, E) E [) > and the probability that P(litu,-oe 6(X, C.) + X E) E r O. If this is true then Polime+oo b(X, C.) (%) 0 contradicts the fact that b(X,C) C.) g Therefore, P((X, E) E T) =
Consider (ar 0) from the support of (X, 0%) except from r, where X is the current query and E is the environment from which X and context preceeding it is sampled. Let us consider the distribution dP(Y)X, b(X, C)).
dP(Y y/x =x,b(X,xC)) = 0) dP(Y = =,
(22)
We simplify lim,- dP(Y = y|X' = b(x, = 0) below.
lim dP(Y y| x, 6(x, C.) = =0 1i = y, xx, =
(23)
1000
, dP(X= = xx,((X,))) =
We simplify the numerator and the denominator of the above separately.
lim dP(Y- y, = 2,, - = = lim t0000
-
dP(Y y, X = x, xx e, eX = 1000
/ dP(Y y|xx z, E E e)
lim dP(X 2, == e, b(X, ())=)))) =0)
1+:0
[ dP(Y u|X = E. = e) x, E =) lim dP(b(X, C.) = 0| x. E = e)
(24)
In the above, we use dominated convergence theorem (Ash and Doléans-Dade 2000) to swap limit and the integrals (to dominated convergence theorem, we use the fact that the dP(X = z, e, b(X,C) Gt) 0) is bounded and support of E has a finite volume). In the simplification above, also the fact Y L C F. Since b(X, C₁) converges to 0% almost surely, the distribution lim dP(b(X, C.) = 0 = x, E = e) evaluates to probability one when 0 = 0= and is zero otherwise. As result, the above expressions become
{ Jeee,. dP(Y = =,, = =x,E= E= e).
lim dP(Y = =y, X x, (( 0) 10000
(25)
where E, is the set of all the environments observed conditional on X with 0 = 8. Observe that all the environments in E,,9 have the same dP(Y = 1 = x, E = e) given by h" (x, 0). Similarly,
lim xx, 100
JeeE..
dP(X = x, c)
(26)
As result, we can write
lim dP(Y = = -X, = 0) = (x;))) JeEE.. (400
= xx
20

---

Published as a conference paper at ICLR 2024
We use this to obtain
lim dP(Y 1|X xx,((, 0) 0) ci dP(Y 1-000
lime- dP(X= = (x)
h²(2,)) SPEE2 dP(X = =2,
(27)
(s,
SeEEr, dP(X = I, E =c
Therefore,
lim dP(Y - y|X = x:0==,,bXXX- 100
= y = x, =).
(28)
where is any environment that is in the support of data sampled with X and that also satisfies 0 O.
lim H(Y]X,,(XXCI)) -f lim dP = x. b(, = =) lim 2x = 11000
10,0 60000
80000
- =Joo ( dp(x=2,E= x. z,
H(Y x 0)
1000
(29)
In the above, we use dominated convergence theorem to swap the limits and integrals (Recall that
dP(X x, e, b(, C.) 0) is bounded say by say s and the volume of the support of E is bounded say by sp. As result, dP(X = x, b( C.) HY sip log(2).). From equation 28, it follows that lim;- H(Y x, b(X, Ct) = 0) H(Y
x. b, C.) 0) S
E e), where is any environment in E, We use this in the above to get
lim H(Y|X, b(x, Ct)) = -f.0 (Lf JEEEZO 4000
=
dP(X = x, d xx
=
dP(X x, E a) =(
=
2,8 JEEE2o
(30)
dP(X= xx x E))
We combine the above with equation 21 obtain lim, H( H(Y F). Finally, observe that if T(Y; E > = HY| E) < H((I)). Since lim H(( C.) [Y),,, E), ICRM impoves over ERM that attains a cross-entropy loss of (Y)))).
This completes the proof.
A.3 PROOF OF THEOREM 2
Theorem 2 (Partial iid zoom-in). Suppose the joint distribution (.XI...-X)), (Y....., YJ), E) is Markov with respect to a Bayesian network. The query X and the environment E are statistically dependent and form the Markov blanket Y. Then 1CRM partially zooms-in on the environment risk minimizer, improving over the performance of the global empirical risk minimizer terms of the cross-entropy loss. Further, the improvement is strictly monotonic in context length t.
Proof. Let us consider the setting where the context is of length one. We denote the current query as with corresponding label Y and environment F. The example in the context X which has corresponding label ÿ and it shares the same environment E. Recall that as part of the context, the learner only sees X and not Y. Both Y and E are real-valued scalars and is a dimensional vector.
21

---

Published as a conference paper at ICLR 2024
Following the assumption in the theorem, the distribution of (X, XXYYX E) is Markov with respect to a Bayesian network. We first establish that E cannot be a child of any variable in the directed acyclic graph (DAG). The assumption (X,Y) - (X, Y) implics X 1 XX] E and Y T E Suppose E is a child variable of Y. Due to the symmetry, (X, Y, E) and (X, Y, E) follow the same distribution. As result, E' is also a child variable of Y, which implies Y L Y (since E is a collider on the path from Y to Y). This contradicts Y T YIE. Suppose E' is a child variable of some component of X say X Due to symmetry, E' is also a child variable of X which implies X y X E. This contradicts X T X E. Therefore, E cannot be a child of any of the variables in the DAG.
X,Y,E)
Since both X and E form the Markov blanket of Y. there are two possible cases. Either E is directly connected to Y or E is connected to Y through some element of X.
In the first case, E can only have an arrow into Y and not the other way around as F is not child of any other node. Let us consider the setting when E is one of the parents of and denote it as E + Y. Since X (X) ) is on the Markov Blanket of ()), we claim that each component of X is either parent of Y or a child of Y. Suppose this was not the case. This implies that there exists component of X say X', which is on the Markov Blanket as parent of E. But that would make E a child of Y. However, E' cannot be child variable shown above. As result, each component of X is either a parent or a child of Y. We now consider two subcases.
Let us consider the setting when there exists a child X of Y. Observe that X is a child of Ý and it has a path to E and as a result it has a path to Y. This path from elements of X' to Y passes through E. This path has colliders and does not contain any element of X on it (We show this case in Figure 3a). As result, Y L X Thus [(;; X|)) > 0 (use chain rule of mutual information).
Let us consider the other setting when each X is parent of Y (shown in Figure 3b). this case, E has to have path to some element of X, say X as otherwise E 1 X, which contradicts the assumption that E L X. Consider the path X to E to Y. Observe that this path is not blocked. As result, [Y; XX) > 0.
Let us consider the other possibility when Y is connected to E through X. Here the only way this possible is if some element of X say X' is a child of Y and E is a parent of that element (as shown in Figure 3c). Therefore, we know that X is connected to Y through E and X
Observe that this path from X to Y is not blocked as X is collider. Therefore, T(Y; XX)) > 0. We showed the result so far assuming that the context length was one. Suppose that the context has k 1 examples denoted as Cs = [X -I. The chain rule of mutual information tells us [Y; Cx) = I(Y; Xx- X)= +
IX, The proof above already demonstrates that the first term I(Y; x ()) is strietly positive. Since mutual information is non-negative, we can conclude that I(Y; (X) > O.
Next, we want to argue that entropy strictly reduces as context length increases. In other words,
C())X,Cb) < ((YIX,Cx-1) + I(Y; Xx(XX >0.
We want to show Y L Xx, Ck-1). In the proof above, we had three cases shown in Figure 3. In each of these cases, we argued that the path from X₁ to Y is not blocked. Even if we condition on contexts Ck-1 this continues to be the case. In the first two cases, the path from N₂ to Y is direct and does not contain any element from the conditioning set. In the third case, the direct path involves a collider X from the conditioning set and thus is also not blocked. As a result, Y L XxI(X, Ck-1). This completes the proof.
Remark on the Theorem 2 It is possible to extend Theorem 2 to the case when only a subset of X and E' form the Markov blanket. Observe that the analysis of Case a) and Case c) in Figure 3a, Figure 3c does not change. The analysis of Case b) is more nuanced now. In Case b), we used the fact that E is connected to X that is on the Markov blanket. This need not be the case if only a subset of X is on the Markov blanket. Suppose XME denote the set of X that are on the Markov Blanket. If E is connected any member of XMB, the same analysis as Case b) continues to hold. Consider the
22

---

Published as a conference paper at ICLR 2024
E
Y
Y
Y
E
X
A
Y
Y
X2
X
(a) Case
(b) Case
(c) Case
Figure 3: Illustrating the different key cases for Theorem 2.
case when E' is connected to some other member of that not in XMB Denote this member as X'. Observe that the same element X' from X will have direct path into Y through E that is not blocked. As result, even in this case conditioning X helps.
A.4 PROOF OF THEOREM 3
Theorem 3 (Full ood zoom-in) Consider data triplets (x. y. e) generated from 2 ~ N( U Σ2) and t- g(2)..or for all environments E E. where g the identity map. There exists an ICL algorithm that in the limit of infinitely long contexts produces Bayes optimal predictions for all the test environments that fall the Voronoi cells of the training environments.
Proof. The learning algorithm works follows. For each e, pair in the training data, define the set E's as De.y Maximize the likelihood of De.y assuming that the underlying distribution is Gaussian, This can be stated
Pr SW
arg min pe
Σ RED
[| n (1))))111 log(det(E?))
The solution to the above are standard sample mean based estimators of means and covariance. Also, use sample mean based estimator to estimate the probability of each class in environment e and denote as p.. Define te [(pe, po. 122), (fl, Aie E.)]. The model test time works as follows.
We are given samples De at test time from some environment E &tc- Estimate the param- eters of Gaussian mixture model with mixture components to maximize the likelihood of observing D We denote the estimated parameters as f.r [pe, Me' Ee', Pel Pe,, Eee. Define a permutation f₂r as Ber [pe' Pe', See Pe', , Me', E.e).
Find the closest environment to the estimated parameters the training set.
min eEEr
miin Rell. | Sel13
(31)
Suppose ē is the closest training environment that solves the above. If 0er is closer to 9/ than Se, then Per, Me', Ee correspond to the label and pe', Enr correspond to the label 1. For the query x, the probability assigned to label is
| 11?
(3-1)-1
P
cxx)
|| 2 [1s + e
10 -8- (8-))1
e
If Bar is closest to this environment, then Pet, Me's Ene correspond to the label and
Per, Re, E. is the label 0. For the query z, the probability assigned to label 0 is 1 - c(x).
For the training environments, in the limit of infinitely long contexts the estimated parameters take exact values, i.e., te Te, for all &
23

---

Published as a conference paper at ICLR 2024
Ye
Ye
Veronoi cell of training environment
de
Figure 4: Illustration of Voronoi cells of training environment.
For the test environment, the true of parameters that generate the data are Ye's where 7e
(w' ue Ee),, (pb ub, [. Define the permutation of % de ((p', mb, F..), (pe., 1 E2
There can be two types of test environments. One in which the mean and covariance for both classes identical. The method above assigns probability of 2 to both the classes, which is the Bayes optimal prediction. Let us consider the latter environments, where the class conditional parameters for not the same. In the limit of infinitely long contexts test time, there two possible values 0- can take, either far = Ye' or 0e = dg, This follows from identifiability of Gaussian mixtures, Yakowitz and Spragins (1968).
Consider the first case, 0gr Ye In this case, the equation 31 becomes
min eEEtr
(min(l1ze - Tell- |1 - Ye
Suppose some environment è solves the above optimization. Following the assumption in we know that Yer falls in the Voronoi region of some V/2 and thus Y is closer to Ya than 8ę (see Figure 4). As result, pe n2,,, E. is associated with class O, which is actually correct and thus the final predictor would match the Bayes optimal predictor for the test environment. In the second case, 0, = 8p, Therefore, 3r' "e' and pb. P. E would be correctly associated with class one thus leading to Bayes optimal predictions. This completes the argument we set out to prove.
We now brietly explain how the method fails if test parameter is outside the Voronoi cell of training parameters. Suppose 0gr Ye' but Ye' is in Voronoi region of some 8_. In this case, Ber would be closest Y and ph,, 16 Id, would be incorrectly associated with class 1. This shows that beyond the Voronoi region the proposed algorithm fails.
A.5 EXTENSION OF THEOREM 3
In the previous theorem, we assumed that g is identity. We now describe how the result can be extended to general non-linear mixing maps g. For this result, we leverage the theoretical results from identifiable variational autoencoders (i-VAE) (Khemakhem et al., 2020).
short review of identifiable variational autoencoders (Khemakhem et al., 2020) We are provided with observations x's that are generated from latent variable using injective map g.
24

---

Published as a conference paper at ICLR 2024
where 2: g(2). The theory of i- VAV provides with method and the conditions under which the underlying true latent variables 2 can be identified up to permutation and scaling. In i-VAEs, is assumed that along with each sample as u. For our results, auxiliary information is available to us in the form of the environment index and the lahel of the data point. In the theory of i- i-VA,, the distribution of the latent variables assumed to follow conditionally factorial exponential distribution stated as follows.
provided with auxiliary information, which they term
(rr(=/u)= = II Qi(z)) M.(u) exp
Σ (2,
(32)
A.,(u)) are the parameters of the distribution that vary with Qi is base measure and M, is normalizing constant. We concatenate Tj's and X,S across latent dimensions to make construct di dimensional vectors denoted as X(u) and Tz)). Thus the data generation process is summarized
where T; IT....... .T.ss) are the sufficient statistics, (2) (u,,
2 ~ PT,a--lu,, x g{)),
(33)
where g, À the parameters. We now revisit the data generation process that we consider and explain how it falls under the umbrella of the data generation processes considered in i-VAE. For all e E E,
zlv,e e ~ N( EE).
(34)
x g(z);
where the latent variables are sampled conditional on the label and environment e from Normal distribution whose mean and covariance depend both y, Define X' as the image of i.e., X' gRd)). We further assume that the covariance matrix a diagonal structure as stated below. Assumption 1. Each Σn is diagonal matrix.
Since Σ. is diagonal matrix, we denote the ith diagonal element (i)) Similarly, the hth component of μ² is denoted as (i). Observe that the distribution of conditional on y, belongs the family conditionally factorial exponential distributions studied in i-VAE (Khemakhem
L. M, e) = ()))))) A., (y, e) = (x)(s)
et al., 2020). we substitute (z,), = V2n 2 e) ( (z) z and (2,( =2, then obtain the distribution of described by equation 34.
(a]
Definition 1. We define an equivalence relation between sets of parameters of the model as follows.
(g.T. (g. T, + EA,c T(g (x)) = ( + c. E X.
(35)
If A is invertible, then we denote the relation by If A is a block permutation matrix, then we denote by rpp
We now state some key results from (Khemakhem et al., 2020).
Theorem 5. Assume that the data is sampled from the data generation in equation 33 according to with parameters (g, T, (TT).. A). Assume the following holds
The mixing function g is injective
The sufficient statistics Ts.j are differentiable almost everywhere, and (Trs)ısisk are linearly independent on any subset of X of measure greater than zero.
There exists dk 1 distinct points -.. -. , Udlk such that the matrix
L ==(Aa]) A(uo), A(uax) X(wa))
of size dk dk is invertible.
25

---

Published as a conference paper at ICLR 2024
then the parameters (g. T) are ~A identifiable.
Theorem 6. Assume the hypotheses of the Theorem 5 holds, and k 2 2. Further assume:
The sufficient statistics Tig are twice differentiable.
The mixing function g is C'-diffeomorphism.
then the parameters (g. T, X) are ~ identifiable.
We can leverage the above two theorems (Theorem 5, Theorem 6 and Theorem 4 from Lachapelle al. (2022)) and arrive at the following corollary for the Gaussian data generation process from equation 34.
Theorem 7. If the data generation process follows follows equation 34, where g is a C2 didffeemmorphism. Suppose there exist 2d + points u (yo, co), ,u2d (yzd, G2a) in the support of (y. 6) observed in training distribution such that
(N(u1) X(uo), ... A(w2a) {(wo))
invertible. If P..r.A--l,ee Pg. for all y, in the support of (y, e) in the training distribution, then ž = Allz + where ž = g (x) and 2 g (x).
Proof. We equate the probability of observations under two models g. T, X and g. T, À for each y, e. Consider z ~ PT. (-ly, e) and the corresponding x = g(z). These x's follow Pg. x (-1,, e) since Po.T,x(-lv, e) Pa.r. e). Define ă = g (x) and these ă follow Pr. . (-1,, e). We can
write a(z), where = g
Observe (=(2)y,e) e) = ps (=(((2)
and
log (2lyk; ek) log
ex)) + log det (Da(=));
1o (z|yo, co) log (p=(a(z) yo eo)) + log det (et()) (36) log (zlys, ex) log (p=(=luo, ea)) = log (p= ex)) log (a=(a(=) Bo. Go))
Substituting the exponential form we obtain that T(z) [A(U,, ex) A(yo, eo))] = T2)) [Ā(UE, cx) - A(yo)
If we use sufficient variability conditions, we obtain T(2) = AT( + c. We now use the fact that sufficient statistics T==) = (z, 2) are minimal to conclude that
T(z) AT +
where A is invertible. In the above, the line of reasoning used in in the proof of Theorem 4 in (Lachapelle et al., 2022).
After this point, leverage Theorem 6 to conclude that
T=(2,) AT3 +
We can expand the above to write
D
Note that the above relationship holds for all E Z. If 2 depends on 2² then 2 would be degree four polynomial in 2; and it would be equated to a degree 2 polynomial 2, stated in the RHS. This cannot be true for all 2; in the support. As result, 2 is scalar multiple of 2, Since for every i there is such a j. it follows that ž = AlI= + r.
26

---

Published as a conference paper at ICLR 2024
Theorem 8. (Zoom-in [ood]) Consider the data generation process in equation 34. We make few additional assumptions on the data generation stated below.
Each Σy is diagonal matrix
There exist 2d + points u (yo, eo), TM u22d (yzdd e2d) in the support of (y. e)
observed in training distribution such that
(X(u1) - Muu)). ... M(uza) x(vo))
is invertible.
g C'-diffeomorphiss
Under the above assumptions, can guarantee that there exists in-context learning algorithm that the limit of infinitely long contexts generates Bayes optimal predictions for all the test environments that fall in Voronoi cells of training parameters weighted by certain vector.
Proof. The training proceeds as follows. Train an autoencoder on training data under the constraint that the output of the encoder follow a Gaussian distribution with independent components conditional on each This is stated as the following minimization.
g. PM Ey arg
min Ell f(x) 2)11 F.F. E2)
Σ KL (p(())))) (p E.) U.E
(37)
where = = fxx), p=(- y, is the distribution of ž. The first term is standard reconstruction loss and the second term is the KI. divergence between distribution of z and a Normal distribution with independent components. Also, estimate the class probabilities for each environment and denote them as p. Similar the proof of Theorem 3 define te [60% pag (e.1)
The model at test time works as follows. We first use the trained encoder f and generate ž for test time inputs. After this the model operates in exactly the same way on e's as the proof of Theorem 3. Basically the output of encoder takes place of raw x' in the procedure described in proof of Theorem 3.
The assumptions in this theorem along with following i) E follows Gaussian distribution with independent components, ii) g(5) follows distribution of x conditional on y, e for each y, e, implies use the previous result in Theorem 7 to conclude that ž Allz + Observe that also follows a Gaussian distribution with independent components conditional on each y, e. In the limit of infinitely long contexts, 9e is equal to scaled means of original training environments and covariances also scaled componentwise according to the transform AII. We can now apply the previous Theorem 3 on z's follows. the parameters of the test environment are in the Voronoi cell of the train distribution of 2'8, then the procedure described above continues to generate Bayes optimal predictions in those environments.
A.6 COMPARING ICRM AND ERM UNDER THE LENS OF INVARIANCE
The label y is related x² and mean of x environment as follows.
y ax + 3y +
(38)
ERM learns a linear model on two dimensional feature vector x = (x 2 The closed form solution for linear regression is A p. where A E[XX which is assumed to be invertible, and
E[XY]. The covariance matrix of X is defined as Σ = [o] 012 012 o
Proposition 2.
Let E[x = for all E. If is invertible, 3 0. 012 + 0. 12 + 0 for some E Exr, then the coefficient estimated by ERM for x. is not the same as the invariant coefficient
27

---

Published as a conference paper at ICLR 2024
Proof. We compute p first.
(B aE[x
X
(39)
of
012 + _
where d = )
Next, we compute A.
of 012
A = =
(40)
012 02 +
The solution ERM is
a''
0 +
12
o
3'
[o + 8)07
o2
12 of
01 + 2
(41)
Simplifying the above, we obtain the coefficient for 1 to be
122 (1
a' a
(42)
o (1 0
Owing the assumptions, B * 0, 12 + 0 and 3 for some obtain that the second term in the above is not zero. As a result, the estimate computed by ERM for a is biased.
Proposition Let E[X,[E = e] for all e E. If is invertible, 8 + 0, 012 + 0, 1² + 0. The error of ERM in test environment increases in of
Proof. The error of ERM is given as
B[(aX' + +1 - -xx3))+'33]++
(43)
= -a-aa ( - where o² is the variance of the noise variable If we take the derivative of the above emoor w.r.t of. we obtain (a a²) which is positive. This completes the proof.
-333)))-33333BB8 233'E((12))) -(())))------------0-- 2 og
ICRM learns linear model on (x+ ,0, ",,,,,,,, (.. 12)) We study two settings to analyze the error of ICRM at test time. If at test time, the model has seen sufficiently long contexts, then it knows the means corresponding to x¹ and x² and the model achieves the test error of o². On the other hand, the context is empty, then also note that the expected eтror of the model is 8 301 || (assuming the model uses default value of zero for the mean in the absence of any context), where μ²² is the mean of x in environment e'. Since the error of ICRM in the absence of any context is independent of variance of at, the error of ERM can be much worse than that of ICRM in this setting as well.
Extending the above example beyond linear settings.
Let us consider general setting.
y u(x 12 te, x² r(p 1),
(44)
where u(-) and v(-) are maps (potentially non-linear), E and are independent zero mean noise variables. Following the same line of thought as the above example. ICRM learns non-linear model on (I1, 22, pf, pb) and learns [| 1 x² ur v.] From equation 44, it follows that
Y 1. (x2 (X) 122))
|E(Y|ar aot..o x² 1e 12l
Y xt n3 u(x .
28

---

Published as a conference paper at ICLR 2024
From the above it follows that ICRM learns u(x', 142). In comparison, consider standard ERM learns a non-lineen model on (x x²). Consider the DAG corresponding to setting equation 44. We assume that the joint distribution described in equation 44 is Markov w.rt to the following DAG X + Y + le > X². As a result, Y X²|x². This follows from the fact there a path to X2 through FE and is not blocked by X From Y / X X it follows that ERM learns predictor that relies on both x² and x² Therefore, ICRM learns the right invariant model and does not rely on x² and ERM relies on spurious feature x²
A.7 ILLUSTRATION OF FAILURE OF EXISTING MTL METHODS
In this section, we provide a simple example to show the failure mode of marginal transfer learning (MTL) methods that based averaging Te] En.ec Φ(-) to summarize information about the environment. These methods can be summarized take the following form:
Σ (()))
(45)
We are only going to consider maps that are differentiable.
Example. Suppose we want to learn the following function
w(r,c)
I(x < x.),, Σ #CEC
(46)
c
where Iy is the ith input in the context and x is the current query and I(-) is indicator function that takes value of one if the argument inside is true and zero otherwise. We claim that if
Ensec ((x;), = w(x, c) for all E R, E Rle, then the output dimension of O grows in
context length lc|. Suppose this was not the case. If Φ's output dimension smaller than |c]. then & cannot be differentiable bijection. As result, there exists contexts and c' of same length for which En.cc (xx;)) = En.er ((x)). We argue that there exists an x such that w(r, c) # w(x d').
This would lead to contradiction as En,eo ((x)), w( e) for all
Without loss
of generality, suppose that the smallest value of context e is smaller than that in context e'. If a is larger than smallest value of but lesser than smallest value of d', then w(x, d') 1 on the other hand w(z c) < lcl
We can translate the insight from the above example into more general settings. Consider any map w2z c), that satisfies the following property. For no two distinct contexts and c' w(z, w(x.e)
for all x E I.. Maps of the form f प Exyec (x;)),z can only approximate such w's provided dimension of & grows in length
We explain how the above example can be described by iitention-basee architectures with much x
fewer parameters. First take the current query æ and transform it through a linear map = 1
and
transform the past context values through a linear map well obtain transform for x; to 2
-10 0 and thus T =QQ
We set the Query Q and Key K matrices such that QTK = 0
(+++ Instead of softmax, if we pass the output through a sigmoid, we obtain ( T Q Kā:) Ite nea 27. If sufficiently large, then this approximates [[x < x.)). Therefore, one layer linear
attention with sigmoid and sufficiently high achieves the target, i.e., Ez,ec aT (T TQQ K 2 Enec 1(x < xx).
29

---

Published as a conference paper at ICLR 2024
B RELATED WORK
brief tour of domain generalization. Muandet et al. (2013) developed kernel methods to learn transformations such that the distance between the feature distributions across domains is minimized and the information between the features and the target labels is preserved. The pioneering work of Ganin et al. (2016) proposes method inspired from generative adversarial networks to learn feature representations that are similar across domains. Sun and Saenko (2016) developed a method based on natural strategy match the means and covariances of feature representations across domains. Li et al. (2018) went a step further to enforce invariance on the distribution of representations conditional on the labels. In parallel line of work, led by Peters et al. (2016); Rojas- Carulla al. (2018); Arjovsky et al. (2019), the proposals sought to learn representations such that the distribution of labels conditional on the representation are invariant across domains. These works were followed by several interesting proposals to enforce invariance (Teney et al., 2020; Krueger et al., 2020; Ahuja al., 2020; Jin et al., 2020; Chang al., 2020; Mahajan et al., 2020; Koyama and Yamaguchi, 2020;
Müller et al., 2020; Parascandolo et al., 2021; Robey et al., 2021; Wald et al., 2021; Chen et al., 2022;
Wang et al., 2022: Zhang et al., 2023; Eastwood et al., 2022; Rame et al., 2022; Veitch et al., 2021;
Makar et al., 2022; Wald et al., 2022: Salaudeen and Koyejo, 2022; Eastwood et al., 2023) wwhich is an incomplete representative list. See Shen et al. (2021) for a more comprehensive survey of these works. Most of the above works have focused on learning features that enable better generalization. Recently there been an intriguing line of work from Kirichenko et al. (2022); Izmailov et al. (2022) that shifts the focus from feature learning last layer retraining. These works show that under certain conditions (e.g., avaiability of some data that does not carry spurious correlations) one can carry out last layer retraining and achieve significant out-of-distribution performance improvements.
In the main body of the paper, we already discussed the other prominent line of work in domain generalization on marginal transfer learning, where the focus and learn environment specific relationships. This line of work was started by the notable work of Blanchard al. (2011) and has been followed up by several important proposals such as Zhang et al. (2020); Bao and Karaletsos (2023).
leverage the distributional features
Context-supported prediction frameworks. Existing works have exploited contextual information to develop a variety of prediction frameworks. The works on neural processes and conditional neural processes (Garnelo et al., 201 8b;a) combined the uncertainty estimation capabilities of Gaussian processes with function aprpoximation capabilities of neural networks and showed promising results on meta-lcarning. These works were later improved through transformer based archictectures in atten- tive neural processes (Kim et al., 2019; Nguyen and Grover, 2022). Context-basse architectures have also been used to study in-context learning in (Garg et al., 2022; Akyürek et al., 2022; Von Oswald et al., 2023). These works use labeled data in the context enable adaptation. In contrast, our work adheres to the constraints of domain generalization and only leverages unlabeled data.
C SUPPLEMENTARY EXPERIMENTAL DETAILS AND ASSETS DISCLOSURE
C.I ASSETS
We do not introduce new data in the course of this work. Instead, we use publicly available widely used image datasets for the purposes of benchmarking and comparison.
C.2 HARDWARE AND SETUP
Each experiment was performed on 8 NVIDIA Tesla V100 GPUs with 32GB accelerator RAM for single training run. The CPUs used were Intel Xeon E5-2698 v4 processors with 20 cores and 384GB RAM. All experiments use the PyTorch deep-learning framework.
C.3 DATASETS
C.3.1
FEDERATED EXTENDED MNIST (FEMNIST)
Building on the Extended MNIST (EMNIST) dataset, which includes images of handwritten uppercase and lowercase alphabets along with digits, FEMNIST (Zhang et al., 2020) enriches
30

---

Published as a conference paper at ICLR 2024
this data by attributing each data point to its originating writer. This extension associates each 28 (8822ssize image in the dataset to one of the 62 classes. our setup, each writer serves as a distinct environment. We evaluate the performance of each method based on both worst-case and average accuracy across a set of 35 test users, who are distinct from the 262 training users and 50 validation users. Unlabelled data from environment in this dataset could provide about the writing style of the user and disambiguate data points.
C.3.2 ROTATED MNIST
We employ customized version of the MNIST dataset in Zhang al. (2020). The dataset contains images rotated in increments 10 degrees, ranging from to 130 degrees. Each degree of rotation constitutes a separate environment, effectively acting distinct value. The training set for the two most extreme rotations, 120 and 130 degrees, contains only 108 data points each. For rotations between 90 and 110 degrees, each environment includes 324 data points. The total training set comprises 32,292 points. For evaluation, test images are generated from the MNIST test set, and are duplicated for each environment. Performance metries include both worst-case and average accuracy across these testing domains. Analogous to FEMNIST, unlabeled samples from an environment within this dataset can assist in distinguishing images that may seem similar due their rotated orientations.
C.3.3 WILDS CAMELYON17
We the Camelyon17 17 dataset, part of the WILDS benchmark (Koh al., 2021), which features image patches derived from whole-slide lymph node sections of patients with potential metastatic breast cancer. Each patch is labeled indicate the presence absence of tumor. In our experimental design, each participating hospital is treated distinct environment. The dataset is partitioned in alignment with the official WILDS configuration: three hospitals contribute to the training set, fourth is designated for validation, and the remaining hospital's is data is used for testing.
C.3.4 CIFARIO-C AND TINY IMAGENET-C
Adapting the methodology from Hendrycks and Dietterich (2019), introduce 56 distinet distortions to the training set, treating each as a separate environment. For evaluation, we use a non-overlapping set of 22 test distortions, largely differing in nature from those used in training, For Tiny ImageNet-C, each 64x 64-sized distorted image is associated with one of the 200 classes in the dataset. The same set corruptions are employed augment the CIFARIO dataset, resulting in 32 x 32-sized images for the CIFAR10-C dataset. This setup permits an investigation into whether exposure to distortions during training equips the model to better manage novel distortions during testing. We assess performance through both worst-case and average accuracies across these test distortions.
C.3.5 IMAGENET R
ImageNet-R comprises a diverse range of artistic and creative content, encompassing art, cartoons, deviant art, graffiti, embroidery. graphics, origami, paintings, patterns, plastic objects, plush objects, sculptures, sketches, tattoos, toys, and video game interpretations of ImageNet classes. This dataset consists of renditions for 200 ImageNet classes, totaling 30,000 images. For our experiments, we utilize images from categories of cartoons, paintings, stickers, graphics, sculptures, sketches, tattoos, toys, and video games for training. Validation is conducted using images from embroidery, miscellaneous, and graffiti categories, while the test environments incorporate images from art, deviant art, and origami categories. This segregation introduces extreme real-world distribution shifts, where the data features during testing differ significantly from those observed during training.
31

---

Published as a conference paper at ICLR 2024
C.4 EXPERIMENTAL PROTOCOLS
To ensure a fair comparison across different algorithms for each dataset, we use a standardized neural network backbone. The details for these architectures are provided in Table 4 and Table 5. We use the ConvNet architecture as outlined in Zhang et al. (2020).
For ICRM, the same backbone is used featurize the input, which is then processed by the decoder- only Transformer (Vaswani et al., 2017) architecture from the GPT-2 Transformer f'amily (Radford al., 2019). Our model is standardized have 12 layers, attention heads, and 128-dimensional embedding space across all datasets. Linear layers are employed to map both the input sequence to the transformer's latent embedding and the model's predicted output vector to the output label. For training ICRM on larger datasets like ImageNet R. CIFAR 1O-C, WILDS Camelyon 17 and Tiny ImageNet-C, we start with a ResNer50 model pre-trained on ImageNet (as shown in Table 4) and freeze all batch normalization layers before fine-tuning.
We adopt the same Context Network as used in ARM, specifically retaining their choice of output channels one for smaller datasets like FEMNIST and Rotated MNIST, and three for the others.
For TENT, all reported metrics are based on its episodic version, where the model reset to its trained state after processing each batch. This ensures a fair comparison with other methods. Additionally, during testing. the model's ' parameters are updated for 10 steps using stochastic gradient descent by minimization test entropy across all datasets.
Table 5: ConvNet architecture for (Zhang et al., 2020). We use 2x2 kernels and "same" padding.
Table 4: Network architectures for each dataset.
Architecture
Dataset
# Layer
ICRM
Others
Conv2D (in=d, out=128) BatchNorm2d (dim=129)
FEMNIST Rotated MNIST
ConvNet Transformer
GPT2
2
ConvNet
ReLU
4
Max Pooling (2)
CIFARIO-C
5
Conv2D (in=128, out=128)
Camelyon17
ResNet-50 GPT2
ResNet-50
6
BatchNorm2d (dim=128)
Tiny ImageNet-C
Transformer
ReLU
ImageNet R
8 Max Pooling (2)
9 Global wveaae-poooling
We list all hyperparameters, their default settings, and search boundaries for random sweeps in Table 6. The maximum context length, or support, is fixed at 100 for all algorithms. All models are optimized using the Adam optimizer (Kingma and Ba, 2014). To ensure a fair comparison, we perform random search of 5 trials across the hyperparameter range (refer to Table 6) for each algorithm. The model with the highest validation set accuracy is selected for each run. We then report the average of this number across three independent runs the entire sweep, and its corresponding standard error.
Table 6: Hyperparameters, their default values and distributions for random search.
Condition
Parameter
Default value
Random distribution
learning rate
0.0001
Uniformm( -5.33.5) 10 Uniform( -6-2)
ResNet
weight decay
learning rate
0.0001
10 Unitorm -4.5.
not ResNet
weight decay
0
10 Uniiormm( -6,-2)
32

---

Published as a conference paper at ICLR 2024
D ADDITIONAL EXPERIMENTS
D.1 ADAPTATION CURVES OF VARIOUS ALGORITHMS
0.9
Wost Coop Tea. Noc.co
Averge 1axt Aecurayy
0.7
0.4
0.8
ABM
TENT
Ifsconteet coomplee
Invornical coompla
West Glouy Tol Aocunee
Avenge Test Accurayy
0.00
0.8
AXX
EBM
0.4
exomples
exomppa
Worst Group That Accurscy
0.9
Awerage Accuracy
0.9
ICRM
TENT
10
20
40
50
10
20 In-context examples
30
40
In-context
Wont Croup Ten Nocese
0.20
Atrenge Test Avcurec
0.15
0.3
0.10
0.2
0.05
0.1
E83
ICRM
TENT
0.00
0.0
10
20
30
40
50
10
20
30
40
IInscvnical commple
AIncconical coumpla
Figure 5: Accuracy adaptation curves for worst accuracy (left) and average accuracy (right) across test environment as function of increasing count of context samples. Showing results in order for (EM)IST((op). RotatedMNIIST WILDS Camelyon 17 and Tiny ImageNet-C C(bottom). The average and worst-case accuracy plots for WILDS Camelyon17 17 are identical since the dataset contains only single test environment.
33

---

Published as a conference paper at ICLR 2024
D.2 DOMAIN GENERALIZATION ACCURACIES PER ALGORITHM AND DATASET
D.2.1 ADAPTATION TO DISTRIBUTION SHIFT
In our experiments, we compare ICRM against marginal transfer methods such Adaptive Risk Minimization (Zhang et al., 2020, ARM), test-time adaptation proposals such as TENT (Wang et al., 2020) and Empirical Risk Minimization (Vapnik, 1998, ERM). We also include comparisons with six additional baselines, including approaches that follow the invariance-based paradigm like Fish (Shi al., 2021) and IB-IRM (Ahuja et al., 2021). alongside those that use contextual information differ- ently, such as BN Adapt (Schneider et al., 2020; et al., 2016b) and Bayesian BN Adapt (Schneider al., 2020). BN Adapt replaces the global batch normalization statisties learned during training with test batch statistics at inference. On the other hand, Bayesian BN Adapt assumes the global statistics of the training data prior and linearly interpolates between these statistics and the test batch statistics during inference. Additionally, we evaluate methods that employ classic regularization tech- niques such Mixup (Yan al., 2020) and IB-ERM (Ahuja et al., 2021). Table 2 shows the average performance attained by these methods across four benchmark datasets. Further, Table 7 and Table 8 demonstrate the average and worst group -fitt--distiibutiin performance, respectively, accompanied by the corresponding standard errors. These statistics are computed across three independent runs of the entire sweep, wherein the model selected for evaluation is the one with hyper--prannne paprr-prraeeters yielding the highest validation accuracy.
34

---

Published as a conference paper at ICLR 2024
Table Average out-of-distribution accuracies along with their corresponding standard errors for various counts of context samples. The methods compared include Adaptive Risk Minimization (ARM). Empirical Risk Minimization (ERM), Test Entropy Minimization (TENT), Batch Norm Adaptation (BN Adapt), Bayesian Batch Norm Adaptation (Bayesian BN Adapt), Fish, IB-ERM, IB-IRM, Mixup and our method ICRM on FEMNIST, Rotated MNIST, WILDS Camelyon17, Imay--mageNet-C, ImageNet R and CIFARIO-C.
Dataset / algorithu
Average test accuracy (by # in-context examples)
FEMNIST
100
ARM
49.5 1.0
83.9
+ 0.5
84.4 + 0.5
84.7 * 0.6
84.6 - 0.3
TENT
78.1 1.2
77.9
1.2
81.2 0.9
82.54
8.
BN Adapt
78.3 0.6
76,9
1 1.4
80.3 00.3:.99
81.5 01.50.77
824
Bayesian BN Adapt
78.3 1.8
79.6
81.
82.9 8222
Fish
77.2
7.
0.
0.2
0.6
IB-ERM
79.0 1.5
79.0
79.0
79.0 1.
IB-IRM
79.0 0.4
79.0
79.0
79
79.0 +
Mixup
78.6
0.
0.
78
7..6
ERM
79.3 0.4
7.4
+.4
7933
79.3
+ 0.4
ICRM
78.7
87.2
8..4
0.2
87.8
0.2
Rotated MNIST
100
ARM
36.5 5.2
94.2
95.1
95.3 04
95.5 0.3
TENT
9..1=0.3
8.4
9.39
938
9.2
BN Adupt
94.6 0.8
87.0
91.5
93.7 1.2
94.3 1.0
Bayesian BN Adapt
94.6=
91.2
93.4 12
94.3 1.0
94.7
Fish
94.8 0.4
9.
94.8
94.8 04
94.8
0.4
IB-ERM
92.2 0.5
0.5
9.22
92.2
92.2 0.5
91.0-11
91.0
Mixup
93.6 0.0
9.06
93.6
93.6 0.0
93.6
FRM
94.2 0.3
94.2 + 0.3
0.3
94.2 + 0.3
94.2 0.3
ICRM
93.6
96.1 1
96.2 0.1
96.2
96.2
WILDS Camelyon17
0
25
50
75
100
59.7
TENT
= 7.6
81.8 1.1
87.2 + 1.1
89.4 1.1
89. 1.0
BN Adapt
67.5
82.0 0.3
87.4
0.3
89.7
0.3
89.9
0.3
Bayesian BN Adapt
67.5
82.0
87.3
89.6
02
89.7
0.3
Fish
53.9
53.9
53.9 + 2.8
53.9 + 28
53.9
2.8
IB-ERM
51.8
51.8
51.8
IB IRM
$3.9 1.3
$3.9 1.3
53,9
1.3
53.9
1.3
53.9
1.3
Mixup
62.8 5.7
62.8
62.8 2
62.8
= 5.7
62.8
5.7
ERM
68.6
68.6
68.6
ICRM
92.0 0,6
90.7 1 0.8
90.8 + 0.8
90.8
0.8
90.8
0.8
Tiny ImageNet-C
0
25
50
75
100
ARM
30.8
1..0
31.0 0.2
31.0 0.2
31.0 0.2
TENT
31.7
1.1
1.7
2.1
2.1
Adapt
1.1
Bayesian BN Adapt
31.7 31.7-0.X 0.8
2.2
21
2.3
2,4
Fish
33.7
33.7
33.7 3..7
33.7 33.700.8
33.7 0.8
IB-ERM
3.4
0.4
0.4
3.5
35.5 1.
IB-IRM
3.3
3.34
3.
354
35
Mixup
=.3
0.5
03
ERM
31.8 0.6
3188
31.8
318
1.8
318 0.6
ICRM
38.3 0.1
39,2 0.3
39,2 + 0.3
39.2 = 0.3
39,2 - 0.3
ImageNet
100
ARM
56.3 0.8
58.1 40.3
58.8
59.8 = 0.8
$9.0
TENT
58,9
10.1022
10.7
12.1
0.1
BN Adapt
58.9 0.5
9.1
1..
122
0.11
Bayesian BN Adapt
58.9 0.5
11990.11
12.3
139
14.
Fish
58.6 1.0
1.0
58.6
58.6
58.6
IB-ERM
58.5 0.5
5.55
58.5
5.5
58.5
0.5
IB-IRM
57.8 0.0
57.8
57.8 0.0
57.8
57.8
0.0
Mixup
58.8 0.8
5888
58.8
588
58.8 0.8
ERM
58.9 0.5
58.9
58.9 0.5
589
58.9 4.9 0.5
+
59.4
CIFARIO-C
100
ARM
65.9 1.3
66.0 1.3
66.0
66.0 1.
66.0
TENT
66.1 1.6
63.9 2.1
68.4
69.9 20
70.5
BN Adapt
66.1 1.6
63.9
68.4
099220
70.1
Bayesian BN Adapt
66.1
70.0
Fish
72.3 1.0
72.3
72.3
72.3 1.0
12.3100
IB-ERM
65.2 29
65.2 2.9
65.2
65.2 4 29
65.2
4-24-29 2.9
IB-IRM
64.3 2.6
64.3
64.3
64.3
2.6
64.3
Mixup
72.8 0.4
72.8 + 0.4
72.8 + 0.4
72.8 + 0.4
7288
0.4
66.1 1.6
66.1
+ 1.6
66.1 +
66.1
66.1
1.6
ICRM
70.6 0.2
71.0
0.2
71.0
0.2
71.0
8 0.2
71.0
0.3
35

---

Published as a conference paper at ICLR 2024
Table Worst environment out-of-distrriution accuracies along with their corresponding stan- dard errors for various counts of context samples. The methods compared include Adaptive Risk Minimization (ARM), Empirical Risk Minimization (ERM), Test Entropy Minimization (TENT), Batch Norm Adaptation (BN Adapt), Bayesian Batch Norm Adaptation (Bayesian BN Adapt), Fish, IB-ERM, IB-IRM, Mixup and our method ICRM on FEMNIST, Rotated MNIST, WILDS Camelyon17, Ima-eeeaeNeetCC ImageNet R and CIFARIO-C.
Dataset/ / algorithu
Worst case test accuracy (by # in-context examples)
FEMNIST
0
100
ARM
23.6 1.7
59.5
3.5
60.7 + 3.8
$7.0 + 7.3
58.8 + 4,0
TENT
55.2
57.2 -
63.3 1.3300
65.9
0.6
67.2
BN Adapt
52.7 6.2
56.2
6.19
64.7 6
6
Bayesian BN Adapt
54.3 2.6
60.4 1.2
64.7 0.9
655
66.3 1.2
Fish
52.8 1.2
52.8
52.8
52.8
IB-ERM
58.6 3.4
3.4
5..6
58.6
5.
IB-IRM
57.3 2.6
57.3
57.3
$7.3
Mixup
$7.0
1.
$7.0
57.0
57.0
ERM
39.0 0.2
59.0
59.0 +
59.0 + 0.2
59.0 0.2
ICRM
59.8
69.3 0.0
70.6 2.3
70.6
70.6 = 0.7
Rotated MNIST
100
ARM
28.2 2.1
85.3 1.6
87.2 1.0
87.9 1.0
87.9 0.9
TENT
80.
88.5
0.9
80.2 =
8
BN Adupt
80.5 2.4
70.9
76.9
79.8
2.7
Bayesian BN Adapt
80,5
75.4
79.2
80.7
2.8
Fish
83.2 1.9
83.2
83.2
83.2 + 1.9
83.2 112+1.9
IB-ERM
72.0 0.9
72.0
72.0
72.0 0.9
72.0
Mixup
81.2+ 0.7
81.2 0.7
81.2+ 0.7
81.2 0.7
81.2 0.7
FRM
80.8 1.1
80.8 1.1
80.8 + 1.1
80.8
1.1
80.8 1.1
ICRM
82.5 0.5
88.5 1 0.5
88.5 1: 0.5
88.8
0.5
88.8
WILDS Camelyon17
0
25
50
75
100
59.7
TENT
=
81.8 + 1.1
87.2 + 1.1
89.4 1.1
89.4
1.0
BN Adapt
67.5
82.0 0.3
87.4
0.3
89.7
0.3
89.9
0.3
Bayesian BN Adapt
67.5
82.0
87.3
89.6
0.2
89.7
Fish
53.9
53.9 2.8
53.9 + 2.8
53.9 + 288
$3.9
2.8
IB-ERM
51.8
51.8
51.8
IB IRM
53.9 1.3
$3.9 1.3
53,9 1 1.3
53.9
1.3
$3.9
1.3
Mixup
62.8 5.7
62.8
62.8 5.7
62.8
= 5.7
62.8
5.7
ERM
68.6
68.6
68.6 1: 7.8
68.6
7.8
ICRM
92.0 0.6
90.7 1 0.8
90.8 + 0.8
90.8
0.8
90.8
0.8
Tiny ImageNet-C
0
25
50
75
100
ARM
8.2 0.3
8.3 0.3
8.2
8.3 8.3: 0.3
0.3
TENT
1.2+
1.0
1.6
1.6 0.0
1.6
Adapt
9.4
1.3
Bayesian BN Adapt
9.4
1.6
1.6+
1.8
1.8 +00+000
Fish
11.1=011
1.11
1.11
1.11
11.1
IB-ERM
0.6
1.6
1.
1.8
0.6
IB-IRM
1..9=0.4
15.9
1..9
15.9
1..9
Mixup
113=005
1,
0..
0.5
0.5
ERM
9.5
9,5
9,5
9.5
9.5
9.5
ICRM
18.8 0.2
19,2
19,5
19.5
0.1
19.4 0.2
ImageNet
0
100
ARM
47.4 L.I
4..3
47.2 1.9
49.81.
474 1.0
TENT
48.0
8610.1
8.40.11
8910.1
9.1
BN Adapt
48.0 1.0
8.1
8.5
8.900.1
9.1
Bayesian BN Adapt
48.0 1.0
10.5002
103
107
109
Fish
46.0
2.10
46.0
46.0
46.0
IB-ERM
47.
4.3
4..2
472
4..2
IB-IRM
472
4..2
47.2
47.2
47.2
Mixup
479
2.1
47.9
47.
47.9 2.1
ERM
48.0 1.0
48.0
48.0 1.0
48.0 - 1.0
48.0 1.0
CIFARIO-C
100
ARM
39.3 1.7
39.3 1.7
39.4
39.3
= 1.7
394 1.7
TENT
39.8
2.1
48.9
49.7
2.0
52.6
BN Adapt
39.8 2.5
43.8
45.1
47.8
200 2.0
48.6
Bayesian BN Adapt
39.8
49.6
Fish
49.9 1.5
49.9
49.9
4999
+
IB-ERM
4499
44.443.4
44.9
34 449
449
34
IB-IRM
43.3
43.3
43.3
43.3
43.3
Mixup
53.9 2.4
53.9 2.4
53.9 2.4
53.9
24
53.9
244
39.8
39.8 2.5
39.8
2.5
39.8
25
39.8
ICRM
54.6 0.4
56.0 1 0.5
55.8
0.5
55.8
0.5
55.9
0.5
36

---

Published as a conference paper at ICLR 2024
D.2.2 ROBUSTNESS OF ICRM IN THE ABSENCE ENVIRONMENT LABELS
As outlined in Section 4, the training regimen of ICRM assumes dataset D (Ti Vi, ex )==1 collected under multiple training environments E &r. However, in scenarios lacking such domain separation during training, does ICRM continue show an edge over ERM baselines? To study this question, we modify the sampling strategy: rather than constructing context vectors containing examples from one environment, we construct context vectors containing iid samples from of the environments pooled together. To continue to test for out-of out-o--distributin generalization, however, we evaluate the performance on examples from a novel test environment. We term this modified approach ICRM-Mix.
Table and Table 10 contrasts the performance of ICRM with ICRM-Mix. ICRM consistently outperforms ICRM-Mix across varying counts of in-context samples on both FEMNIST and Rotated MNIST. Surprisingly. ICRM-Mix and ICRM perform similarly on WILDS Camelyon17 and Tiny ImageNet-C. Consider setting where the model benefits the most attending to examples from the same class or related classes. If classes are distributed uniformly across domains, then ICRM and ICRM-mix are bound to perform similarly. Consider another setting where the model benefits the most by attending to environment-specinn examples such as characters drawn by the same user. In such a case, ICRM and ICRM-mix have very different performances.
Table 9: Average oftto-iriittitttinn test accuracies along with their corresponding standard errors for ICRM and ICRM-Mix across FEMNIST, Rotated MNIST, WILDS Camelyon17 17 and Tiny- Imn-emmaeNet- C. ICRM-Mix trains on sequences with samples drawn i..d. from the unified dataset comprising various environments.
Dataset / algorithm
Average test accuracy (by # in-context examples)
FEMNIST
25
50
75
100
ICRM
78.7 + 0.5
87.2 + 0.4
87.4 + 0.5
87.5
87.8 0.2
ICRM-Mix
77.6 + 0.8
81.1 + 0.2
81.1 + 0.2
80.9 + 0.3
80.9 + 0.1
Rotated MNIST
0
25
50
75
100
ICRM
93.6 + 0.2
96.1 + 0.1
96.2 + 0.1
96.2 + 0.1
96.2 + 0.1
ICRM-Mix
88.9 + 1.4
92.6 + 0.3
92.7 + 0.2
92.6 + 0.3
92.7 0.2
WILDS Camelyon17
25
50
75
100
ICRM
92.0 + 0.6
90.7 + 0.8
90.8 + 0.8
90.8 + 0.8
90.8 0.8
ICRM-Mix
92.9 = 0.3
90.7 + 0.6
90.8 + 0.5
90.7 = 0.5
90.7 + 0.5
Tiny ImageNet-C
0
25
50
75
100
ICRM
38.3 + 0.1
39.2 + 0.3
39.2 : 0.3
39.2 + 0.3
39.2 0.3
ICRM-Mix
38.4 + 0.2
39.3 + 0.2
39.3 + 0,2
39.3 + 0.2
39.3 + 0.2
Imagenet R
25
50
75
100
ICRM
57.4 - 0.4
59.7 + 0.4
59.6 + 0.6
59.4 + 0.4
60.5 + 0.3
ICRM-Mix
54,9 + 1.0
54.9 + 1.0
57.8 + 1.0
57.8 + 1.0
57.8 + 1.0
CIFARIO-C
0
25
50
75
100
ICRM
70.6 - 0.2
71.0 + 0.2
0.2
71.0 0.2
71.0 1: 0.3
ICRM-Mix
69.2 + 0.2
69.4 + 0.3
69.4
+ 0.3
69.4 0.3
69.4 + 0.3
D.2.3 UNDERSTANDING THE IMPACT OF ARCHITECTURE
Table 3 presents the average performance of both ERM+ and ARM relative ERM and ARM, across the four datasets. Further, Table 11 and Table 12 demonstrate the average and worst group out-of-distribution performance of these approaches, respectively, along with the corresponding standard errors. These statistics are computed across three independent runs of the entire sweep. wherein the model selected for evaluation the one with hyper-parameters yielding the highest validation accuracy.
37

---

Published as a conference paper at ICLR 2024
Table 10: Worst environment outto--distribution test accuracies along with their corresponding standard errors for ICRM and ICRM-Mix across FEMNIST, Rotated MNIST, WILDS 1amelyon17 and -mageNeege ICRM-Mix trains on sequences with samples drawn i.i.d from the unified dataset comprising various environments.
Dataset algorithm
Worst case test accuracy (by # in-context examples)
FEMNIST
25
50
75
100
ICRM
59.8 + 0.7
69.3 + 0.0
70.6
70.6
70.6
ICRM-Mix
57.5 + 1.4
62.7 + 1.1
65.0 + 0.3
64.1 + 1.5
62.9 + 2.3
Rotated MNIST
0
25
50
75
100
ICRM
82.5 + 0.5
88.5 + 0.5
88.5 + 0.5
88.8 + 0.5
88.8 + 0.4
ICRM-Mix
68.8 + 3.8
77.1 + 0.7
76.8 + 0.9
76.4 + 0.9
76.6 0.9
WILDS Camelyon17
25
50
75
100
ICRM
92.0 + 0.6
90.7 + 0.8
90.8 + 0.8
90.8 + 0.8
90.8 0.8
ICRM-Mix
92.9 = 0.3
90.7 + 0.6
90.8 = 0.5
90.7 = 0.5
90.7 = 0.5
Tiny ImageNet-C
0
25
50
75
100
ICRM
18.8 + 0.2
19.2 + 0.1
19.5 + 0.2
19.5 + 0.1
19.4 + 0.2
ICRM-Mix
18.7 - 0.2
19.2 + 0.2
19.4 + 0.1
19.5 + 0.1
19.4 0.1
Imagenet R
25
50
75
100
ICRM
45.4 0.7
48.0 + 0.2
47.2 + 0.8
46.9 + 0.4
50.6 + 1.3
ICRM-Mix
44.4 + 1.7
46.9 + 0.4
48.1 + 1.6
46.6 + 0.6
48.7 1.0
CIFARI0-C
0
25
50
75
100
ICRM
54.6 - 0.4
56.0 + 0.5
55.8 0.5
55.8
55.9 0.5
ICRM-Mix
53.3 = 0.0
54.2 1.1
54.2 1.1
54.3 + 1.1
54.2 + 1.1
38

---

Published as a conference paper at ICLR 2024
Table 11: Average cut-of-distribution test accuracies along with their corresponding standard errors for ARM and ERM* in contrast to their base algorithms, ARM and ERM across FEMNIST, Rotated MNIST, WILDS Camelyon17 and Tiny--.ageNer-C.
Dataset / algorithm
Average test accuracy (by # in-context examples)
FEMNIST
0
25
50
75
100
ARM
49.5 + 1.0
83.9 + 0.5
84.4 + 0.5
84.7 + 0.6
84.6 0.3
ARM
71.4
+ 1.2
83.4 + 0.2
84.0 + 0.2
83.8 + 0.2
83.5 + 0.1
ERM
79.3 + 0.4
79.3 + 0.4
79.3 + 0.4
79.3 + 0.4
79.3 + 0.4
ERM
77.4 1.3
77.4 + 1.3
77.4 1.3
77.4 1.3
77.4 + 1.3
Rotated MNIST
0
25
50
75
100
ARM
36.5 + 5.2
94.2 + 0.7
95.1 0.4
95.3 0.4
95.5 0.3
ARM
86.9 = 2.0
92.6 + 0.7
92.7 + 0.6
92.8 + 0.6
92.8 + 0.6
ERM
94.2 = 0.3
94.2 - 0.3
94.2 0.3
94.2
94,2 + 0.3
ERM+
94.3 = 0.4
94.3 + 0.4
94.3 0.4
94.3 0.4
94.3 0.4
WILDS Camelyon17
0
25
50
75
100
ARM
61.2 + 5.2
59.5 + 4.2
59.7
4.2
59.7
59.7
ARM
55.8 = 0.8
55.1 + 1.7
55.0 + 1.7
55.0 - 1.8
55.0 + 1.8
ERM
68.6 + 7.8
68.6
+ 7.8
68.6 + 7.8
68.6 + 7.8
68.6 + 7.8
ERM
50.1 = 0.1
50.1 + 0.1
50.1
0.1
50.1 + 0.1
50.1 0.1
Tiny ImageNet-C
25
50
75
100
ARM
30.8 + 0.2
0.2 31.0 0.2
31.0 0.2
31.0 0.2
31.0 0.2
ARM
5.5 + 0.2
5.7 0.2
5.7 = 0.2
5.7 0.2
5.7 = 0.2
ERM
31.8 + 0.6
31.8 + 0.6
31.8 + 0.6
31.8 1.6
31.8 + 0.6
ERM
29.7 1 0.3
29.7 + 0.3
29.7 0.3
29.7 0.3
29.7 0.3
Imagenet R
25
50
75
100
ARM
56.3 + 0.8
58.1 +0.3
58.8 + 0.8
59.8 1:::00 0.8 0.8
59.0 0.3
ARM
1.8 + 0.2
1.5 0.2
1.5 = 0.2
1.5
1.5
ERM
58.9 + 0.5
58.9
+ 0.5
58.9 + 0.5
58.9 0.5
58.9 + 0.5
+
57.0 0.4
57.0 + 0.4
57.0
57.0
57.0
CIFAR1O-C
25
50
75
100
ARM
65.9 + 1.3
66.0 + 1.3
66.0 :: 1.3
66.0 1.3
66.0 1.3
ARM
42.7 = 0.1
44.2 0.1
44.3 - 0.1
44.3 0.1
44.3 + 0.1
ERM
66.1 + 1.6
66.1 1.6
66.1 1.6
1.61
66.1 + 1.6
ERM+
66.5 + 1.2
66.5 + 1.2
66.5 + 1.2
66.5 1.2
66.5
39

---

Published as a conference paper at ICLR 2024
Table 12:
Worst environment out-of-distribution test accuracies along with their corresponding standard errors for ARM and ERM* in contrast their base algorithms, ARM and ERM across FEMNIST, Rotated MNIST, WILDS Camelyon17 and Tiny-ImagNNet-C.
Dataset / algorithm
Worst case test accuracy (by # in-context examples)
FEMNIST
0
25
50
75
100
ARM
23.6 + 1.7
59.5 + 3.5
60.7 3.8
57.0 7.3
58.8 4.0
ARM
51.7 + 2.2
63.0 + 2.1
64.0 + 0.8
60.7 + 1.6
62.0 + 0.8
ERM
59.0 0.2
59.0 + 0.2
59.0 0.2
59.0 0.2
59.0 + 0.2
ERM
53.3 = 2.7
53.3 + 2.7
53.3 = 2.7
53.3 + 2.7
53.3 + 2.7
Rotated MNIST
0
25
50
75
100
ARM
28.2 + 2.1
85.3 + 1.6
87.2 + 1.0
87.9 1.0
87.9 0.9
ARM
71.4 1 2.6
80.9 + 1.8
81.0 + 1.8
81.2
81.1 + 1.8
ERM
80.8 + 1.1
80.8 1.1
80.8 + 1.1
80.8 + 1.1
80.8 + 1.1
ERM+
81.9 = 0.7
81.9 + 0.7
81.9 0.7
81.9 + 0.7
81.9 0.7
WILDS Camelyon17
0
25
50
75
100
ARM
61.2 + 5.2
59.5 + 4.2
59.7 4.2
59.7 4.3
59.7
ARM
55.8 = 0.8
55.1 + 1.7
55.0 + 1.7
55.0 - 1.8
55.0 + 1.8
ERM
68.6 + 7.8
68.6
+ 7.8
68.6 + 7.8
68.6 + 7.8
68.6 + 7.8
ERM
50.1 = 0.1
50.1 + 0.1
50.1 0.1
50.1 0.1
50.1 0.1
Tiny ImageNet-C
O
25
50
75
100
ARM
8.2 + 0.3
8.3 + 0.3
8.2 = 0.3
8.3 + 0.3
8.2 = 0.3
ARM
1.9 + 0.1
1.9 0.1
1.9 0.1
1.9
1.9 + 0.1
ERM
9.5 + 0.4
9.5 0.4
9.5 - 0.4
9.5 +
9.5 = 0.4
ERM
8.3 + 0.3
8.3 + 0.3
8.3 - 0.3
8.3 0.3
8.3 0.3
Imagenet R
25
50
75
100
ARM
47.4 + 1.1
45.3 + 0.4
1.9 47.2 1.9
49.8 + 1.2
47.4 1.0
ARM
1.4 0.2
0.9
+ 0.3
1.0
0.4
0.8 0.3
1.3
ERM
48.0 + 1.0
48.0 + 1.0
48.0 + 1.0
48.0 1.0
48.0 + 1.0
+
45.3 1.0
45.3 + 1.0
45.3 1.0
45.3 1.0
45.3
CIFAR1O-C
25
50
75
100
ARM
39.3 + 1.7
39.3 + 1.7
39.4 + 1.7
39.3 : 1.7
39.4 1.7
ARM
24.5 - 0.7
24.7 + 0.6
24.8 + 0.5
24.8 0.5
24.8 + 0.6
ERM
39.8 + 2.5
39.8 2.5
39.8 2.5
2.5
39.8 + 2.5
ERM+
40.4 + 1.8
40.4 + 1.8
40.4 1.8
40.4
40.4
40

---

Published as a conference paper at ICLR 2024
D.3 COMPARISON OF ICRM WITH IN-CONTEXT LEARNING
The first and most popular conception of In-context learning (Brown et al., 2020) involves providing the model with contextual information, typically a few sample (x. y) pairs that represent specific "task". In contrast, our approach ICRM introduces an alternative perspective on in-context learning, where unlabeled inputs act as the contextual backdrop for a task, also known an environment. Note that, in order to benefit from in-context learning in domain generalization, the context itself must be a sequence of unlabeled inputs since test-time the learner only has access the unlabeled r's from the test environment and not the labels y's.
Our method, ICRM, can seamlessly adapt to supervised settings, functioning with input sequences containing both (z, y) pairs, instead of just z. We refer to this approach as Supervised ICRM or ICL. We evaluate both Supervised ICRM and ICRM on FEMNIST and Rotated MNIST datasets. As anticipated, Supervised ICRM demonstrates superior performance compared to ICRM. However, it not suitable for domain generalization settings where data labels are unavailable at inference.
Table 13: Average/worst OOD test accuracy for different context lengths, for ICRM and Supervised
ICRM on FEMNIST and Rotated MNIST. Supervised ICRM refers to ICRM trained on labeled input sequences containing (z, y) pairs as context.
Data / method
Average test accuracy
Worst case test accuracy
FEMNIST
0
25
50
75
100
0
25
50
75
100
ICRM
78.7
87.2
87.4
87.5
87.8
59.8
69.3
70.6
70.6
70.6
Supervised ICRM
79.0
87.8
87.7
88.2
87.9
61.2
72.2
73.5
74.5
74.9
Rotated MNIST
0
25
50
75
100
0
25
50
75
100
ICRM
93.6
96.1
96.2
96.2
96.2
82.5
88.5
88.5
88.8
88.8
Supervised ICRM
93.3
96.3
96.3
96.3
96.3
82.0
89.0
89.0
89.1
89.3
D.4 INVESTIGATING THE FEATURES LEARNED BY ICRM
Figure 2 presents attention maps for two randomly sampled sequences for FEMNIST and Tinu ImageNet-C datasets. Figure 6 shows similar visualization for other random sequences. Par- ticularly in the second row, attention is predominantly allocated to lines of length similar to that of the query (also in green). thereby disregarding shorter lines (shown in red). Similarly, the third row shows that the model, when presented with query image of a "train",, attends not only on other trains but also on a "bus"-indicating semantic understanding of similarity.
The key takeaway from this visualization is that ICRM effectively learns to attend to a select few samples in an input sequence. Interestingly, these samples either belong the same class or exhibit similar features, despite potentially belonging to different classes.
To gain deeper insights into the features learned by ICRM, we examine its capability to transition from broad domain indices to nuanced, compositional contextual descriptions of environments. This analysis crucial for understanding how ICRM facilitates amortization across similar environments. In particular, we extract embeddings from the penultimate layer of our trained model, for data from every environment within the training set. Subsequently, we train linear classifier to predict the corresponding environment from each embedding. We repeat the experiment for models trained using both ICRM and ICRM-Mix.
As illustrated in Figure 7, the linear model using embeddings from ICRM attains accuracy of up to 75% on FEMNIST and 98% on Rotated MNIST. This efficacy suggests that ICRM embeds representations that linearly separable with respect to their environmental origins. Additionally, ICRM not only exhibits superior accuracy but also a faster rate of convergence in comparison to M-XM This advantage is likely due ICRM's i.i.d data sampling from each environment. In
41

---

Published as a conference paper at ICLR 2024
2
Σ
!
2
2
3
Query
0.12
0.06
0.32
0.10
0.15
0.10
0.05
0.11
0.00
0,00
2
J
I
Query
0.08
0.15
0.10
0.04
0.21
0.03
0.04
0.15
0.19
0.00
Query Bullet Train School bus
0.15
0.23 Bullet train
0.00
Mantis
Rocking chair
Pizza
Bullet train
Query Volley ball
0.16
0.24
0.12 Hog
0.12
0.11 Plate
0.15 Volleyball
0.11 Volleyball
0.00
0.00
0.00
Rench Wagon
Academic Gown
Walkking stick
Query Image
Figure 6: Attention scores for random test sequences, for ICRM on FEMNIST (top two rows) and Tiny ImageNet-C (bottom two rows).
contrast, linear model trained embeddings from ICRM-Mix model also achieves significant accuracy, reaching 70% on FEMNIST and 91 % on Rotated MNIST. This further explains ICRM- Mix's robust performance in out-of-distribtiio generalization even in the absence of explicit domain separation during training, as analyzed in Appendix D.2.2.
0.8
0.8
Acancy
Accuncy
0.4
0.4
0.2
ICRM ICRM-Mi
ICRM-Mix
0.0
3000
Tminnin epochs
Training spocks
a
(b)
Figure 7: Evolution of the classification accuracy of the linear model trained on the retrieved embeddings as a function training epochs for (a) FEMNIST (b) Rotated MNIST
42


Reviews:
[{"review_id": "soHrhhvmme", "reviewer": "Reviewer_7Bph", "summary": "The authors introduce In-Context Risk Minimization (ICRM): an algorithm that leverages techniques from next-token predictors (i.e., \"in-context learning\") to learn a transformer-based prediction model that generalizes at test time to new unseen domains.  The authors present several theoretical results that provide intuition on the similarities and differences between their proposed approach and existing approaches in domain generalization. They also demonstrate the value of ICRM by benchmarking their algorithm against other domain generalization approaches on several datasets. \n\nMore abstractly, the authors argue that the success of their approach provides evidence that \"context is environment\": the intuitive idea considering examples from a similar environment as \"context\" can encourage the learned model to \"zoom in\" to control for environment-specific features, while also learning environment-agnostic invariant features.", "questions": "In addition to the above 3 weaknesses which are most important to my score, I also had several comments, questions, and suggestions about places I got confused when reading the paper that can be made more clear.  \n\n* nit: Why did you choose the title \"context is environment\"? I wonder if \"environment is context\" is more appropriate, as you're aiming to solve domain generalization using in-context learning, and treating environment as context.  \"Context is environment\" to me implies that you are labeling \"context\" as environment to feed as input to domain generalization algorithms. \n* nit, Intro: Is the expectation here that we have environment labels available during training time? – when you state that \"simple empirical risk minimization\" is a strong baseline, do you mean just training using ERM on the entire dataset in an environment-agnostic way? (As a reader, I would need to dig into the references you've cited here to clarify). EDIT:  I see in Section 2 that you mean a global ERM that \"pools the data together\".  Can you clarify this earlier in the intro?\n* nit, Intro: Can you make more clear, perhaps using an example, what it means to \"reveal\" an \"invariance\"?  By an \"invariance\", do you mean some semantic feature that is predictive of the true label $y_i^e$ across all environments $e$?  By \"reveal\", do you mean that you discover that your in-context learner models actually \"use\" the invariant features to make predictions, or something else?\n* Section 2: \"the size of the representation would have to grow linearly with the size of the training data to describe aspects corresponding to a small group of examples, such as extreme value statistics\". I had trouble interpreting this sentence.  What do you mean by \"size\" here – do you mean the dimensionality of each $\\phi_i^e$?  Why are more dimensions necessary for the representations to encode relevant information?  Why do we care about \"describing aspects corresponding to a small group of examples\" – is it because we have fewer samples from some environments relative to others?  Is the problem that we must use the same function $\\phi$ across all the environments?\n* Section 5: \"define $\\delta_e$ to be a permutation of $\\gamma_e$ that swaps its two components\".  What do you mean by \"two components\" – do you mean swaps the parameters for $y = 0$ and $y = 1$?\n* Section 5: I may be missing something, but I'm not sure if I understand the linear least-squares example given in Equation 7.  Is it fair in this case to compare ICRM to \"global ERM\", which in many cases will fail to learn to add environment-specific averages $\\mu_e^i$? Does this phenom (that the true model can be learned) not hold for simple per-environment ERM too?  I am unsure of how illustrative this example (that ICRM tends to learn the true $\\alpha$ vs. some $\\alpha'$) is for the overarching claim that \"ICRM learns invariant features\".\n* Section 5, last paragraph: I don't really understand the value of including the last sentence here (\"when constraining the environment to only one smoker, the outcome of lung cancer disease invariably follows\").  Are you trying to suggest that when we have fewer environments (a \"smaller diameter\"), that invariance is easier to achieve? Can you provide more practical context around this point?\n* Section 6: I am confused by why in Figure 2, you show images that appear (from my eye) to be identical to the query, and calculate the attention score (e.g., between the two rotated 2s or Js).  Can you clarify exactly how the images on the right of the line are \"augmented\"?  Are you certain that these augmentations were significant enough?  Does it make sense for the model to see two almost identical images in the set of test examples?\n* Section 6: nit: In Figure 2, can you write (in small font) what the true labels are for each image, rather than the integer indices? (or at least for the query images?)\n* Section 7: \"we enable learning machines to fully exploit data in natural order\". I question if you should emphasize this point in your conclusion, as you didn't actually present any of the data \"in natural order\" (i.e., chronologically from its creation date) in any of your experiments.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This paper is extremely timely and would be of great interest to researchers on domain generalization and group robustness. I strongly recommend that this paper is accepted.\n\n* The authors clearly motivate, present from first principles, and connect modern research in domain generalization and in-context learning.  Their notation and exposition in Sections 2 and 3 is thoughtful and clear.  Abstract concepts are made clear using examples, such as the self-driving car in Section 2 and sentence examples in Section 3.\n* The authors' proposed algorithm (ICRM) is intuitive to implement, and described clearly.\n* The authors' theoretical results in Section 3 (with the exception of the domain generalization result Theorem 3) are relatively intuitive, and also demonstrate the conditions under which ICRM is equivalent (Prop. 1) vs. beneficial (Theorem 1) over global ERM.\n* The authors provide detailed information about their experimental design and instantiation of each algorithm being benchmarked so that their experiments are clearly reproducible.\n* The authors' proposed method appears to result in significant performance gains over comparable domain generalization algorithms, which suggests its ability to have positive impact in practice.\n* I appreciate the authors' efforts to thoughtfully connect more abstract interpretations of how \"environment is context\" to existing literature from causality that studies similar phenomena (at the end of Section 5).  I found the smoking example to be an especially clear analogy that helped me better understand the motivation behind the approach.", "weaknesses": "I am happy to consider raising my score if the authors address the below concerns.\n1. **Choice of datasets for evaluation + including experiments on \"harder\" datasets with spurious correlations that vary across environments**. \n  * My basic sense of the datasets that the authors chose to benchmark the value of their domain generalization method on is that there is arguably not *that* high variance in what semantic features are present across environments (e.g., simple rotations or corruptions are somewhat artificial shifts that are detached from reality). To really buy the author's claim that this approach is relevant for domain generalization, I would love to see experimental results on prominent spurious correlation benchmarks, such as those used in [1] (e.g., CivilComments, WaterBirds, CelebA).  I understand that several of these benchmarks only have 2 total test environments (so are perhaps unsuitable for domain generalization specifically), but can you attempt to include additional results on a spurious correlation benchmark (or justify why it's not appropriate)? I'm curious if on more difficult benchmarks, the model learned by ICRM may undesirably \"zoom-in on toxic spurious correlations\".\n2. **Weakening claims about \"what ICRM has learned\"**.  \n  * Unless the authors attempt to peek under the hood (using post-hoc explanation methods or by designing controlled experiments) to probe at what features the ICRM models have \"learned\", I suggest that the authors make clear that some of their claims that their models learn \"invariant\" or \"contextual features\", are hypotheses/intuition (rather than assertions of the truth). ex: in 6.1, you state that the gains are \"because ICRM training still benefits from contexts as to find contextual features that ERM ignores\".  I think this is a bit too strong (what are the \"contextual\" invariant features for your experiments using real data, for example?).  Perhaps soften to \"we hypothesize that this is because…\"\n3. **More discussion of how to use ICRM in practice**, i.e., the nuances of what test set data to present as belonging to the same \"environment\" (/context) vs. different environments.  IMO this approach is perfect for domain generalization settings where the \"environment\" is literally the institution in which the algorithm is being deployed (ex: I begin using it in a new hospital), but has way more challenges when we consider \"environments\" as naturally occurring subpopulations or time-based shifts.  Can you please include additional discussion within your paper about how one would decide what examples to group together as being in the same \"environment\" in practice? I also thought it was a really cool result that you don't need to observe labels in order to include examples in the same \"context\" that should be emphasized further in such a practical discussion.\n\n\n[1] https://proceedings.mlr.press/v177/idrissi22a/idrissi22a.pdf", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xs6P222RJm", "reviewer": "Reviewer_1GJP", "summary": "This paper introduces a new method for domain generalization that utilizes data previously seen in the test set as contextual information to improve generalization accuracy.", "questions": "* The method proposed in this paper was inspired by some ideas in in-context learning, but its relationship with in-context learning in LLM is not significant. The description in the third paragraph of the introduction can easily lead to misunderstandings.\n\n* A concern regarding Theorem 1: Why does the inequality hold? Even though we assume that data can be independently sampled from x, it doesn't necessarily imply that the distribution of the sequence formed by the samples is independent of the variables. Therefore, The assumption in Theorem 1 seems not correct. Are there any relevant references that can confirm the existence of this independence?\n\n* Does function h have different parameters for different environments? Theorem 1 indicates that h has different parameters depending on the environment. If this is the case, which set of parameters should be used during testing?\n\n* The current motivation of this article lies in the idea that applying context information in domain generalization can get better generalization results. Furthermore, this paper argues that previous methods did not consider this aspect. However, there is an issue in assuming that testing data comes from the same domain (or we know the domain partition of test data) when discussing the generalization problem. If it samples from different domains, this approach may potentially mislead the model to produce incorrect results, as the function h learned during training somehow combines information from c and x to make decisions.\n\n* Although the experiments have yielded many promising results, they are only compared to ERM and two baselines from three years ago. The baseline's performance is even lower than ERM, although the author initially acknowledges that most methods may not necessarily outperform ERM. However, the choice of baselines appears insufficiently comprehensive.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The writing in the article flows smoothly, and the experiments seem to yield very promising results.", "weaknesses": "The combination of domain generalization and in-context learning is an ambitious idea. However, the theoretical and experimental discussions in this paper are not sufficient. \n\nI find that the motivation for integrating these two concepts does not fully convince me. It appears that this paper only utilizes information from observed data samples x, which I believe is not entirely consistent with the current concept of in-context learning in LLM because a sequence formed solely from observed training data x may not be sufficient as a demonstration to the model. Authors could consider providing sequences of x, y pairs as context, rather than just a sequence of x. \n\nWhile this paper seems to get promising results, its practical applicability may be limited because it does not account for the possibility of the testing set data samples being a mixture of unknown domains. To address this issue, it might be worth exploring the identification of domain-specific information or introducing additional information as prompts (context). \n\nIn fact, the method proposed in this paper is quite similar to the example presented in Figure 1.b, with the key difference being that this paper's method does not involve averaging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3WJVoBcvGn", "reviewer": "Reviewer_NswY", "summary": "This paper proposes a new method for domain generalization called In-Context Risk Minimization (ICRM), a method which generalizes to a new domain given a sequence of unlabeled examples from that domain. The authors motivate ICRM by referring to the LLM literature, in particular their emergent ability to perform well OOD. They prove that ICRM can \"zoom-in\" on good/optimal risk minimizers given a sample/infinite data. Empirically, they demonstrate that ICRM outperforms baselines on several DG datasets and perform supplemental experiments to ablate and interpret these results.", "questions": "- Some people would argue that LLMs don't actually perform well OOD - rather, they are trained on a large enough dataset that instances where they have to perform truly OOD are much rarer. This might be good to address in the motivation/intro\n- In general, I find the two concrete examples in the paper more confusing that clarifying (gravitational fields poem, ICRM example in Sec 5): I think clarifying these would go a long way to my understanding of the method and paper", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- I think this is an ambitious paper and connects two interesting areas within OOD-type research - there are definitely interesting insights gained from the new setup and the idea of receiving a sequence of new-domain examples\n- empirical results are strong, demonstrate a good use of unlabelled data\n- experiments are pretty thoroughly done, I appreciate the supplementary studies in Table 3, Fig 2, and Fig 5\n- the new method is described clearly and the theoretical results seem useful: in particular Theorem 3 and its extension seem important in terms of plugging this work into the iVAE work. I also like the specification of the limitations in terms of Voronoi cells of the training environments as a concept", "weaknesses": "- I find myself somewhat confused by the analogy between LLMs and ICRM - the paper makes it seem as those these should map 1:1 but I can't quite make it clear to myself, perhaps the authors can clarify. It's not clear what the sequence of Xs that arrive correspond to in LLMs, since they are listed as being selected at random at training time: this means they can't be language tokens, and if they are unrelated sequences I don't see how they correspond with the notion of context laid out in the \"gravitational fields poem\" example (these are also correlated - context sequences in the LLM space are usually not thought of as random wrt what follows them). I think ICRM is interesting nonetheless but I have trouble getting the intuition straight because of this\n- I think that the relationship to various types of distribution shift should be highlighted a little bit more here: for instance, it seems like ICRM will be vulnerable to shifts in P(Y | X) across environments, since it only receives unlabelled examples from the new domain at test time. I think this is true of many DG methods, and I think this is what the \"Voronoi cells\" specification in Thm 3 is constraining, which I think is nice. However, I think this weakness is somewhat obscured by the discussion of context and the analogy to LLMs (e.g. gravitation fields poem example): in LLMs the \"context\" as it is generally considered is quite rich and can help with generalization to compound ie Y | X shifts, but ICRM as I understand it does not have that capability.\n- in Sec 5, I don't find the example particularly clarifying: I may be missing something but it seems unfair to give ICRM access to the \\mu variables, and not give them to ERM. It's possible this is a reasonable comparison to make but I think it needs to be explained a bit more in that case.\n- I don't understand how ICRM allows models to take advantage of \"natural order\", since the examples are drawn randomly at training time, correct?\n- I find it a little puzzling that the performance of ICRM decreases with added context in Camelyon17 - is there an explanation for this?\n- It's strange to me that ICRM would outperform other methods even at 0 context: isn't the advantage of the method the ability to use context?\n\n\nSmaller questions:\n- Intro: Minor point but important I think - the Angwin et al paper (Propublica COMPAS investigation) is not an example of an investigation of the impact of out of distribution model failures, but rather the impact of unbalanced model error rates. There are many studies of societal impacts of OOD failures and I think it would be better to cite one of those\n- Thm 3: \"there exists an ICL algorithm\" - what is an ICL algorithm? I don't think this is defined as a mathematical object\n- end of Sec 5: I think you mean that lung cancer is \"deterministic\" in the single smoker case, rather than it \"invariably follows\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qWw8kl6I3d", "reviewer": "Reviewer_cBZm", "summary": "The authors propose a domain generalization (DG) method called In-Context Risk Minimization (ICRM).\nICRM uses a transformer architecture that takes multiple datapoints as input and predicts their labels.\nIn DG, inputs come from different environments for which the distribution of inputs and labels change.\nDG assumes that the environment index is known for each datapoint at training and test time, i.e. we can group datapoints by their environment.\nIn ICRM, each set of 'in-context' inputs belongs to the same environment.\nAt test time, this allows ICRM to learn _in-context_ to adapt to a particular set of test inputs.\nThe authors argue this is advantageous to previous work in DG, which either learns models that aggregate test points into a single embedding or learn invariant predictors that ignore dependencies between test observations.\n\nThe authors provide theoretical evidence that ICRM can achieve better loss in iid and ood settings than naive context-free empirical risk minimization.\nThey provide experiments on FEMNIST, Rotated MNIST, WILDS Camelyon17, and Tiny ImageNet-C where they show ICRM outperforms the following methods: ARM, TENT, and ERM.", "questions": "O) \"Define δe to be a permutation of γe that swaps its two components\" → What are the two components (it seems to me there are three?) and what does it mean to swap them?\n\nP) How do the computational costs of ICRM and the baselines compare?\n\nQ) \"minimizing the worst risk across\" → Later you write that you sample environments at random during training. Does that not mean that you are actually minimizing the average risk across environments?\n\nR) \"to minimize the auto-regressive loss\" --> Why is this loss auto-regressive? It does not depend on the model's past predictions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "I am convinced that ICRM is advantageous to methods that do not rely on context or aggregate all context points into a single embedding.\n\nThe linear regression toy example they provide nicely illustrates the advantages of ICRM over invariant methods.\n\nThe paper is largely well-written and the description of the ICRM method is clear.", "weaknesses": "Unfortunately,  I cannot recommend acceptance of the submission in its current state.\n\nA) My biggest concern is the strong resemblance the proposed approach bears to neural processes (NPs) (Garnelo et al., 2019), and in particular attentive neural processes NPs (Kim et al., 2019); ANPs), as well as other related work, e.g. from Müller et al. (2022) and Kossen et al. (2021).  This line of related work is not discussed at all in the current draft. Chiefly, these works are related because, excluding the original NP paper, they all propose transformer architectures that process multiple datapoints at the same time to improve predictions. Neural Processes are framed around 'function learning', where 'functions' are analogous to 'environments' for ICRM. For a set of context points $\\{(x_i, y_i)\\}_{i\\in\\mathcal{C}}$ drawn from a 'function' they make predictions for the labels $y_t$ of test points $x_t$, $x\\in \\mathcal{T}$ of that function. In the attentive neural process, attention is used to predict test labels by attending to the context points – strikingly similar to ICRM! \n\nI would strongly encourage the authors to highlight ANPs (and related work) as important prior work, and dial down claims of novelty when appropriate.\nI do believe that the application of ANP-like architectures to domain generalization is interesting and sufficiently novel, however, the current omission of relevant prior work is not acceptable.\n\nNote that there are some differences between prior work and ICRM.\nFor ANPs, the $x_i$ are often single features, e.g. the pixels of an image or (multi-dimensional) observations drawn from a Gaussian Process.\nHowever,  ICRM itself does not process high-dimensional inputs itself, and instead relies on ConvNet/ResNet-50 embeddings: these could also be used as pre-processing with ANPs.\n(Note that Neural Processes have been scaled to image classification before by Requeima et al. (2019))\n\nB) Perhaps the biggest difference between prior work and ICRM is, that ICRM does not use labels of examples as input.  This could be a restriction of ICRM: it can only model p(X|C) of the context. If there are any shifts in p(Y|X, C) between environments, ICRM cannot model these. For ANPs, this is not a problem, as inputs are (x, y) pairs. I would argue, ANPs are a generalisation of the proposed ICRM architecture.\nDoes domain generalization always preclude access to (even a few) of the test environment labels? Relatedly, does DG always assume that p(Y|X, E) does not shift? Even if both of these are true,  ANPs can still be used without test-context labels as these can be iteratively predicted. Conversely, if either is not true, then it seems like ANP-like architectures are clearly advantageous.\n\n\n* Kim, Hyunjik, et al. \"Attentive neural processes.\"  ICLR (2019). https://openreview.net/forum?id=SkE6PjC9KX.\n* Garnelo, Marta, et al. \"Neural processes.\" ICML 2018 workshop on Theoretical Foundations and Applications of Deep Generative Models. https://arxiv.org/abs/1807.01622\n* Müller, Samuel, et al. \"Transformers can do bayesian inference.\" ICLR (2022). https://openreview.net/forum?id=KSugKcbNf9.\n* Kossen, Jannik, et al. \"Self-attention between datapoints: Going beyond individual input-output pairs in deep learning.\" NeurIPS (2021). https://openreview.net/forum?id=wRXzOa2z5T.\n* Requeima, James, et al. \"Fast and flexible multi-task classification using conditional neural adaptive processes.\" NeurIPS (2019).\n\n\nC) My second biggest concern is with the results of the experimental evaluation. Figure 5 and Table 2 show that ICRM does not consistently improve performance when increasing the number of in-context demonstrations. In particular, for WILDS Camelyon17 and Tiny ImageNet-C the results are concerning. For Tiny ImageNet-C performance only improves very slightly when including additional in-context observations: crucially, the zero-shot ICRM performance outperforms all competing methods at _any_ number of in-context observations. Most concerning is WILDS Camelyon17: here, performance _decreases_ when including in-context examples – performance is maximal when the context is empty! \n\nIt seems to me as though the main selling point of ICRM, improving performance by including test-time examples in the context, only comes true for FEMNIST. The benefits of including datapoints in the context with ICRM are not as large as the authors proclaim. I did not see any discussion of this in the paper, and would ask the authors to provide one.\n\nD) Thanks for providing the ERM+ baseline, which uses \"an identical architecture to ICRM, but without context\". Comparing the numbers of Tables 3 and 2, I was surprised to see that ICRM at context size 0 significantly outperforms ERM+ across tasks. Do you have any insights as to why this would be the case? What differences are left between ICRM and ERM+?\n\nE) \"[...] even in the absence of test context. Specifically for both WILDS Camelyon17 and Tiny ImageNet-C, ICRM outperforms baselines despite not leveraging any context from the test environment. This is because ICRM training still benefits from contexts as to find contextual features that ERM ignores\" --> This is unclear to me – how can ICRM benefit from context if there is no context? Why does ICRM outperform the other baselines in the zero-shot setting?\n\nF) I think your definition of 'in-context learning' is misleading throughout the draft and a significant departure from the meaning of in-context learning (ICL) as introduced by Brown et al. (2020). This is a recurrent problem throughout the submission. For one, ICRM is about learning from _unlabelled_ examples. ICL is about learning novel tasks from a few-shot set of _labelled_ examples. For example, the 'poem' example the authors give on page 4 is inappropriate: neither does it provide multiple examples, nor does it provide labels. A more appropriate example for ICL for poem writing would be a set of input-output pairs, e.g. a list of target audiences and matching poems for them.\n\nG) \"The key to our answer resides in a recently discovered emergent ability of next-token predictors, namely, in-context learning.\"  I agree with you that an interesting aspect of ICL in LLMs is that it seems to 'emerge', at least the model is not knowingly trained to learn ICL. However, for ICRM, you _train_ the model to rely on the context, i.e. very similar to the setup of Neural Processes that predates ICL. Thus, I find it highly misleading to suggest ICRM relies on any emergent abilities.  Further, ICRM does _not_ preform auto-regressive next-token prediction: your inputs are images and your outputs are labels, opposed to input tokens at time t and output tokens at time t+1. One could even argue that what ICRM does, does not meet common expectations for in-context learning, as there are no labelled in-context examples.\n\nH) You frequently write that ICRM can  \"fully exploit data in natural order\". Are you using causal attention or positional embeddings? If not,  your Transformer architecture should be equivariant to the order of the inputs. Further, the prediction on the test input should then be _invariant_ to the order of the context examples. I therefore find your statements about 'nature does not shuffle data' misleading. In fact, I believe that, often, the order of the data should not affect predictions (cf. exchangeability). On the other hand, if you _are_ using causal attention or positional embeddings (which are part of the GPT-2 default architecture), I believe these might negatively affect your performance, and would like to see an ablation with them removed.\n\nI) \"The maximum context length, or support, is fixed at 100 for all algorithms\" → Do you sample uniformly between [0, 100] examples at training time, or do you train different ICRM models, one for each context length? Relatedly, you write that you train your model to minimize the 'auto-regressive loss'. Does this mean that, for each training input, you minimize the predictive loss for each of the training inputs, i.e. you perform 'teacher forcing', just like in LLM training?\n\nJ) Your highlighting of low and high attention scores in Figure 2 seems arbitrary. The caption suggests you highlight all low/high attention scores. In addition to not explaining how you decide what is a low/high attention score, you also clearly miss some low/high scores, e.g. 0.00, 0.00, and 0.05 in row 1 should be marked as low if 0.11 is, and 0.20 in row3 should be high if 0.15 is. It seems like you highlight only those scores that fit your narrative, e.g. the pizza in row three is not highlighted, but the train and bus matching the query 'train' are. I would suggest you highlight scores consistently. Alternatively, make clear you highlight only some examples: those which you discuss in the main text.\n\nK) Why is the performance on WILDS Camelyon17 the same for worst and average case? Does this mean the performance is exactly equal on all environments? Why is this the case?\n\n## Nits\n\nL) \"We attain significant out- of-distribution abilities in a multitude of specific tasks —such as writing poems\" --> Why are you sure that writing poems out of distribution for large language models – even if conditioned on some context?\n\nM) Figure 1b): It seems like it should be $x_i^e$ instead of $x_{i,1}^e$?\n\nN) The sentence before Eq. 5 does not make sense to me. Perhaps move the 'estimates the conditional expectation P (Y | X, C)' to after the definition in equation 5?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)