Paper: Context is Environment (ICLR-2024)

Reviewer comment (Novelty):
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 (ANPs) (Kim et al., 2019), 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.

Author response:
> On Neural processes and attentive neural processes

There are a few important points worth emphasizing:
- The key challenge of domain generalization is to build learners that succeed at test time in new situations without any access to the labeled data. If we have access to labeled data, it falls in the realm of few-shot learning, an important, related, and complementary area of research. We hope that the reviewer appreciates the challenge of domain generalization. 

- We agree that the existing works on NPs, ANPs, which have now been cited in the paper, and the more recent proposals studying ICL have similar architectures. However, as rightly pointed out in the review, there is a critical distinction – they all rely on labeled data in the context. Our work explores the potential of these context-based transformer architectures in the unlabeled data setting. We believe this is an important departure that embraces the essence of domain generalization. As mentioned in the review, this endeavor is deemed sufficiently novel. That said, our method is not constrained to operate on solely $x$’s. As shown below and in Section D.4, ICRM can be easily adapted to training on labeled sequences. However, this approach (Supervised ICRM) is unsuitable for domain generalization settings where data labels are unavailable at inference.

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**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** |

*Here “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. 

- Our work also lays down the theory that shows when unlabeled data can help the trained models perform better than standard ERM in different settings – short and long contexts and in and out-of-distribution evaluations. The theoretical analysis in this setting poses a different challenge than labeled data. This is another reason we strongly believe our work should not be viewed as a special case of ANPs. The theory from existing works would not translate and explain the success of ICRM.

Reviewer comment (Novelty):
I would strongly encourage the authors to highlight ANPs (and related work) as important prior work, and dial down claims of novelty when appropriate. I 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.

Author response:
> On Neural processes and attentive neural processes

There are a few important points worth emphasizing:
- The key challenge of domain generalization is to build learners that succeed at test time in new situations without any access to the labeled data. If we have access to labeled data, it falls in the realm of few-shot learning, an important, related, and complementary area of research. We hope that the reviewer appreciates the challenge of domain generalization. 

- We agree that the existing works on NPs, ANPs, which have now been cited in the paper, and the more recent proposals studying ICL have similar architectures. However, as rightly pointed out in the review, there is a critical distinction – they all rely on labeled data in the context. Our work explores the potential of these context-based transformer architectures in the unlabeled data setting. We believe this is an important departure that embraces the essence of domain generalization. As mentioned in the review, this endeavor is deemed sufficiently novel. That said, our method is not constrained to operate on solely $x$’s. As shown below and in Section D.4, ICRM can be easily adapted to training on labeled sequences. However, this approach (Supervised ICRM) is unsuitable for domain generalization settings where data labels are unavailable at inference.

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**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** |

*Here “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. 

- Our work also lays down the theory that shows when unlabeled data can help the trained models perform better than standard ERM in different settings – short and long contexts and in and out-of-distribution evaluations. The theoretical analysis in this setting poses a different challenge than labeled data. This is another reason we strongly believe our work should not be viewed as a special case of ANPs. The theory from existing works would not translate and explain the success of ICRM.

Reviewer comment (Novelty):
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.

Author response:
> On Neural processes and attentive neural processes

There are a few important points worth emphasizing:
- The key challenge of domain generalization is to build learners that succeed at test time in new situations without any access to the labeled data. If we have access to labeled data, it falls in the realm of few-shot learning, an important, related, and complementary area of research. We hope that the reviewer appreciates the challenge of domain generalization. 

- We agree that the existing works on NPs, ANPs, which have now been cited in the paper, and the more recent proposals studying ICL have similar architectures. However, as rightly pointed out in the review, there is a critical distinction – they all rely on labeled data in the context. Our work explores the potential of these context-based transformer architectures in the unlabeled data setting. We believe this is an important departure that embraces the essence of domain generalization. As mentioned in the review, this endeavor is deemed sufficiently novel. That said, our method is not constrained to operate on solely $x$’s. As shown below and in Section D.4, ICRM can be easily adapted to training on labeled sequences. However, this approach (Supervised ICRM) is unsuitable for domain generalization settings where data labels are unavailable at inference.

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**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** |

*Here “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. 

- Our work also lays down the theory that shows when unlabeled data can help the trained models perform better than standard ERM in different settings – short and long contexts and in and out-of-distribution evaluations. The theoretical analysis in this setting poses a different challenge than labeled data. This is another reason we strongly believe our work should not be viewed as a special case of ANPs. The theory from existing works would not translate and explain the success of ICRM.

Reviewer comment (Theory):
Does 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.

Author response:
> On Neural processes and attentive neural processes

There are a few important points worth emphasizing:
- The key challenge of domain generalization is to build learners that succeed at test time in new situations without any access to the labeled data. If we have access to labeled data, it falls in the realm of few-shot learning, an important, related, and complementary area of research. We hope that the reviewer appreciates the challenge of domain generalization. 

- We agree that the existing works on NPs, ANPs, which have now been cited in the paper, and the more recent proposals studying ICL have similar architectures. However, as rightly pointed out in the review, there is a critical distinction – they all rely on labeled data in the context. Our work explores the potential of these context-based transformer architectures in the unlabeled data setting. We believe this is an important departure that embraces the essence of domain generalization. As mentioned in the review, this endeavor is deemed sufficiently novel. That said, our method is not constrained to operate on solely $x$’s. As shown below and in Section D.4, ICRM can be easily adapted to training on labeled sequences. However, this approach (Supervised ICRM) is unsuitable for domain generalization settings where data labels are unavailable at inference.

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**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** |

*Here “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. 

- Our work also lays down the theory that shows when unlabeled data can help the trained models perform better than standard ERM in different settings – short and long contexts and in and out-of-distribution evaluations. The theoretical analysis in this setting poses a different challenge than labeled data. This is another reason we strongly believe our work should not be viewed as a special case of ANPs. The theory from existing works would not translate and explain the success of ICRM.

Reviewer comment (Evaluation):
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.

Author response:
> **On adaptation with increasing context**

- Learning with few (even one) in-context samples is a feature of ICRM and not a flaw. In Theorem 1 in the paper, we show the benefit of ICRM in the limit of large context lengths. In Theorem 2, we show that ICRM can benefit merely from one unlabeled in-context example! So in short, it is possible (as we observe empirically) that very few examples provide most of the gain, with minimal gains as the number of examples increase more.

- The reviewer's evaluation seems to have missed our results from two additional datasets: ImageNet-R and CIFAR 10-C. Table 8 in Appendix explicitly shows in-context gains (difference in worst group performance at 100 context length and zero context length) of **20 percent** for FEMNIST, **6 percent** for Rotated MNIST, **0.6 percent** for TinyImagenet, **5 percent** for Imagenet R, **1.3** percent for CIFAR10 C. While FEMNIST is largest, others are useful gains as well.

Reviewer comment (Evaluation):
It seems that 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. There is no discussion of this in the paper, and one should be provided.

Author response:
> **On adaptation with increasing context**

- Learning with few (even one) in-context samples is a feature of ICRM and not a flaw. In Theorem 1 in the paper, we show the benefit of ICRM in the limit of large context lengths. In Theorem 2, we show that ICRM can benefit merely from one unlabeled in-context example! So in short, it is possible (as we observe empirically) that very few examples provide most of the gain, with minimal gains as the number of examples increase more.

- The reviewer's evaluation seems to have missed our results from two additional datasets: ImageNet-R and CIFAR 10-C. Table 8 in Appendix explicitly shows in-context gains (difference in worst group performance at 100 context length and zero context length) of **20 percent** for FEMNIST, **6 percent** for Rotated MNIST, **0.6 percent** for TinyImagenet, **5 percent** for Imagenet R, **1.3** percent for CIFAR10 C. While FEMNIST is largest, others are useful gains as well.

Reviewer comment (Experiments):
In comparing the ERM+ baseline (which uses "an identical architecture to ICRM, but without context") with ICRM at context size 0 across Tables 2 and 3, I was surprised to see that ICRM significantly outperforms ERM+ across tasks. What differences are left between ICRM and ERM+ that could explain this?

Author response:
> On the ERM+ baseline

ERM+ only operates on the current query and there is no context input to it. However, ICRM accepts the current query and context as input. At test time during zero shot evaluation, it is true that both ICRM and ERM+ do not have any context. However, during train time ICRM is trained with contexts and that can lead to a model that is better than ERM. The context allows the model to selectively augment the current query and train a better featurizer as also explained in the previous answer.

Reviewer comment (Evaluation):
The statement regarding ICRM outperforming baselines in zero-shot despite not leveraging any context from the test environment — and attributing this to training benefiting from contexts to find contextual features that ERM ignores — is unclear. How can ICRM benefit from context if there is no context?

Author response:
> **On the zero-shot gains of ICRM**

- The most natural source of gain for our method is to utilize additional information in the context. However, it can offer zero-shot benefits. We'll elaborate on this with two related descriptions.
	
- Observe that while ICRM models can predict without context at test time, these have always been trained to observe contexts during training. Therefore, when the ICRM model is trained, it has the opportunity to learn to extract features from the current image while selectively paying attention to other images in the context. By attending to images that the model finds useful for prediction, ICRM can end up learning “in weights” a higher-quality featurizer. That is, ICRM models may “depend” on context to learn environmental spurious features while burning “in-weights” the invariant features across environments.  Our current theoretical analysis is mainly geared towards explaining the benefit of contextual information. We believe it is a fruitful future work to develop a theory to explain the above selective augmentation perspective.

- The example in equation (7) and its extensions in Appendix A.6 explicitly constructs a linear example for which _ICRM exhibits better out-of-distribution performance than ERM without any contextual information at test time_. Since the model is trained over an extended feature space, it can learn the right invariances for inputs in the original feature space.

Reviewer comment (Novelty):
The claim that ICRM relies on an "emergent ability" of next-token predictors, namely in-context learning, is misleading. For ICRM, the model is explicitly trained to rely on context, similar to Neural Processes that predate ICL. Moreover, ICRM does not perform auto-regressive next-token prediction — the inputs are images and outputs are labels — and there are no labelled in-context examples, so it may not meet common expectations for in-context learning.

Author response:
> On the role of emergence

There appears to be some confusion here. We are referring to the emergent ability of LLMs, and nowhere do we claim that this ability is emergent for our method

Reviewer comment (Theory):
The repeated claim that ICRM can "fully exploit data in natural order" is questionable. If causal attention or positional embeddings are not used, the Transformer should be equivariant to input order and predictions invariant to context order, making the ‘nature does not shuffle data’ statement misleading. If such mechanisms are used (as in GPT-2), they might negatively affect performance and should be ablated.

Author response:
> On the role of natural order

We employ both causal attention and positional embeddings. Causal attention is used to allow the trained model to only make predictions by attending to the observed data thus far. It is also a popular choice to study in-context learning [1]. In our current setup, we randomly sample data sequences from a domain. However, there is a partial order imposed by the environments, which we have explained in greater detail in the revised manuscript. Our intention was to convey that if the data had been collected with a natural order, our approach can accommodate it by processing the inputs accordingly.

As the review correctly highlights, using positional encodings for our datasets, where predictions are invariant to the order of data in a sequence, may adversely affect performance. The table below illustrates the performance of ICRM with and without positional embeddings. Clearly, both the model achieves similar performance access benchmarks with an exception of FEMNIST, where the model without positional embeddings exhibits 5\% gains in the worst group accuracy. 

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**FEMNIST** | 0 | 25 | 50 | 75 | 100 | 0 | 25 | 50 | 75 | 100 |
ICRM (with pos) | 78.4 | 86.6 | **87.3** | 87.0 | 87.3 | **56.0** | **72.3** | 69.3 | 70.8 | 70.0
ICRM (w/o pos) | **79.1** | **86.9** | 86.6 | **87.4** | **88.0** | 55.0 | 70.0 | **71.0** | **72.0** | **75.0**
| | | | | | | | | |
**Rotated MNIST** | 0 | 25 | 50 | 75 | 100 | 0 | 25 | 50 | 75 | 100 |
ICRM (with pos) | **92.7** | **95.8** | **95.7** | **95.8** | **95.9** | **81.1** | **87.8** | **87.0** | **87.5** | **87.8**
ICRM (w/o pos) | 92.1 | 95.3 | 95.5 | 95.5 | 95.6 | 80.1 | 86.0 | 86.3 | 86.8 | 86.9
| | | | | | | | | |
**Tiny ImageNet-C** | 0 | 25 | 50 | 75 | 100 | 0 | 25 | 50 | 75 | 100 |
ICRM (with pos) | 38.3 |  39.2 |  39.2 |  39.2 | 39.2 | 18.8  | 19.2  |  19.5 | 19.5 | 19.4 
ICRM (w/o pos) | **38.8** | **39.7** | **39.7** | **39.7** | **39.7** | **19.3** | **20.5** | **20.5** | **20.5** | **20.4**

*Here “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. These reported numbers are not obtained through a rigorous hyper-parameter sweep. 

[1] Garg, Shivam, et al. "What can transformers learn in-context? a case study of simple function classes." Advances in Neural Information Processing Systems 35 (2022): 30583-30598.

Reviewer comment (Reproducibility):
Regarding the statement "The maximum context length, or support, is fixed at 100 for all algorithms" — is the sampling done uniformly between 0 and 100 examples during training, or are different ICRM models trained for each context length? Additionally, the use of 'auto-regressive loss': does this involve minimizing predictive loss for each training input (teacher forcing) as in LLM training?

Author response:
> On training of ICRM models with maximum context length 100

Thanks to the attention in transformer models, we can obtain predictions for all samples in a sequence of length 100 in a single forward pass. Consequently, during each iteration, we randomly select a training sequence (of length 100) and generate predictions for each example with just one forward pass through a single ICRM model.

Reviewer comment (Evaluation):
The performance on WILDS Camelyon17 is the same for worst and average case — does this mean performance is exactly equal on all environments? Why?

Author response:
> Same average and worst group test accuracy on WILDS Camelyon17

As mentioned in the original submission, WILDS Camelyon17 comprises data from 5 hospitals, with 3 for training, 1 for validation, and 1 for testing. With only one testing environment, the average performance is equivalent to the worst-case performance across testing domains.

Reviewer comment (Presentation):
In Figure 1b, it seems like it should be $x_i^e$ instead of $x_{i,1}^e$.

Author response:
2. _Typographical error in Figure 1b._ We thank the reviewer for their attention to detail. We have made the correction in the revised manuscript.

Reviewer comment (Writing):
The sentence before Eq. 5 does not make sense. Perhaps move the 'estimates the conditional expectation P(Y | X, C)' phrase to after the definition in equation 5.

Author response:
3. _Sentence before equation 5._ We have now corrected this in the revised manuscript by moving the sentence to after equation

Reviewer comment (Writing):
In "Define δe to be a permutation of γe that swaps its two components" — what are the two components (as it appears there are three), and what does it mean to swap them?

Author response:
>  Clarifying the line on swapping two components

We mean it swaps parameter tuples corresponding to $y=0$ and $y=1$.

Reviewer comment (Evaluation):
How do the computational costs of ICRM and the baselines compare?

Author response:
> Computational costs of ICRM

Training an ICRM model on FEMNIST and Rotated MNIST typically consumes around 6 GB and 10 GB of GPU memory, respectively, and takes roughly 1-2 hours for training. The exact time varies depending on the context length and batch size for our model architecture with 12 layers, 128 dimensions, and 4 heads.
In contrast, the ERM baseline on these datasets uses approximately 4 GB and 5 GB, respectively, and takes around 30 minutes - 1 hour to complete.  Time and space complexity of other baselines may vary. 

*Reason for time efficiency of ICRM.*  ICRM makes predictions for all samples in a sequence of length $K$ in one forward pass. Thanks to the efficient PyTorch implementation of transformers, caching intermediate hidden states, typically the outputs of attention layers, significantly accelerates inference.

Reviewer comment (Theory):
In "minimizing the worst risk across" — since environments are sampled at random during training, does that not mean the average risk is being minimized instead?

Author response:
> Clarification of the kind of error is being minimized, the worst group or average error?

ICRM is trained to minimize the average error across training environments. Equation 2 refers to the min-max formulation of the optimal predictor in the domain generalization setting.

Reviewer comment (Theory):
In "to minimize the auto-regressive loss" — why is this loss auto-regressive? It does not depend on the model's past predictions.

Author response:
> Clarification about autoregressive loss

We apologize for the incorrect usage of the word ‘autoregressive’ here. By autoregressive loss, we meant to refer to the loss of our autoregressive model. We have corrected this in the revised manuscript. As discussed in Section 4, we minimize the expected loss over predictions using all the context lengths. Denote context of length $j$ as $c^e_{j+1} = (x^e_1, \ldots, x^e_{j})$, for all $j = 1, \ldots, t-1$, and $c^e_1 = \emptyset$. Then mathematically, the loss function is given by $\sum_{j=1}^t \ell(h(x^e_j; c^e_j), y^e_j)$.

Reviewer comment (Theory):
I find myself somewhat confused by the analogy between LLMs and ICRM - the paper makes it seem as though 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 with respect to what follows them). I think ICRM is interesting nonetheless but I have trouble getting the intuition straight because of this.

Author response:
We agree that there is not an exact one-to-one correspondence between all aspects of an LLM and ICRM. However, a few important aspects bear a strong parallel in our implementation. Firstly, both LLM and ICRM condition on the entire sequence of experience so far. Secondly, both typically exploit attention mechanisms to select the examples to pay attention to. However, there are differences, as rightly pointed out in the review. For us, the examples in a context are unlabeled inputs sampled at random from the environment, and thus the order of the examples does not matter. This is a facet specific to datasets we consider and, hence is different from standard language modeling. However, the proposal can handle datasets with some natural ordering based on time (e.g., video) and space (e.g. text). In such a scenario, we preserve the natural ordering instead of shuffling it.

Reviewer comment (Experiments):
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 μ 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.

Author response:
This example is meant to bring to light how a key difference between the nature of ICRM and ERM translates to an important difference in the learned model. ERM operates on the current query $x$, and ICRM in its standard form operates on $x$ and $p(x)$. In this example, we make an additional assumption that ICRM operates on $x$ and $\mu(p(x))$, which we have stated. Under this assumption, we show that ICRM learns the right invariances. It is a pedagogical example to show that “if one could use distributional features as additional inputs, one could reveal the invariance of interest”. Further, our experiments are empirical support for ICRM having such capability for distributional feature discovery.

Reviewer comment (Theory):
I don't understand how ICRM allows models to take advantage of "natural order", since the examples are drawn randomly at training time, correct?

Author response:
This point is absolutely correct. In the current setup, we randomly sequence data and do not exploit natural order. We were alluding to possible future work where, if the data had some natural order, namely, say a video where you classify each frame, our ICRM  proposal extends without change and exploits natural order. In sum, our remark about “natural order” is made as an attempt to motivate future data collectors to retain this precious time-metadata, so methods such as ICRM can be deployed in full force.

Reviewer comment (Evaluation):
I find it a little puzzling that the performance of ICRM decreases with added context in Camelyon17 - is there an explanation for this?

Author response:
In the Camelyon dataset, standard ERM often encounters a failure mode where the model exploits shortcut features such as luminosity. Consider a model that only takes the current query as input but has an effective featurizer that eliminates luminosity. Such a model can perform well on Camelyon-17. We posit that ICRM learns a similar featurizer by observing context during training. Consequently, the ICRM model has to learn to ignore additional queries provided during test time.  In some cases, these additional queries may have acted as distractors, causing a 1.2 percent drop in performance.

Reviewer comment (Evaluation):
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?

Author response:
The most natural source of gain for our method is to utilize additional information in the context. However, it can offer zero-shot benefits. We'll elaborate on this with two related descriptions.
	
- Observe that while ICRM models can predict without context at test time, these have always been trained to observe contexts during training. Therefore, when the ICRM model is trained, it has the opportunity to learn to extract features from the current image while paying attention to other images in the context. By attending to images that the model finds useful for prediction, ICRM can end up learning “in weights” a higher-quality featurizer. That is, ICRM models may “depend” on context to learn environmental spurious features while burning “in-weights” the invariant features across environments. Our current theoretical analysis is mainly geared towards explaining the benefit of contextual information. We believe it is a fruitful future work to develop a theory to explain the selective augmentation perspective.
- The example in equation (7) and its extensions in Appendix A.6 show that the extended feature space of ICRM can lead to better out-of-distribution performance even when it does not have any contextual information at test time. Since the model is trained over an extended feature space, it can learn the right invariances for inputs in the original feature space.

Reviewer comment (Theory):
Thm 3: "there exists an ICL algorithm" - what is an ICL algorithm? I don't think this is defined as a mathematical object.

Author response:
An ICL algorithm receives the data from multiple environments as input and outputs a predictor that takes the current query and the context prepended to it as input. We have clarified it in the manuscript.

Reviewer comment (Writing):
End of Sec 5: I think you mean that lung cancer is "deterministic" in the single smoker case, rather than it "invariably follows".

Author response:
That’s correct, for one given smoker, the outcome of disease given its entire vector of characteristics is deterministic.

Reviewer comment (Writing):
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.

Author response:
In the part, where we referred to OOD generalization capabilities of  LLMs, we have prefixed the claims with “may” as it is a hypothesis that is still being debated.

Reviewer comment (Writing):
In general, I find the two concrete examples in the paper more confusing than 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.

Author response:
In this example, we only meant to show that the model can adapt its output based on the context (e.g., a teenager vs a physics graduate). In standard examples of capital prediction tasks of in-context learning shown in the introduction, the model also adapts its behavior based on contextual information. We think the confusion stems from the fact that in-context learning is often associated with (input, label) pair demonstrations. However, a model can use other sources of information in the context and not just the labels. As shown in [1]. labels do not play a significant role in several in-context learning tasks.

[1] Work, What Makes In-Context Learning. "Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?."

Reviewer comment (Theory):
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.

Author response:
We thank the reviewer for their critical feedback that helps us improve the work and also for appreciating the ambitious idea that we undertake in this work. We have highlighted the changes in magenta in both the main body and the Appendix of the revised manuscript.

---

> On integration of the two concepts of domain generalization and in-context learning: 

We will clarify that, in order to benefit from in-context learning in domain generalization, the context itself must be a sequence of unlabeled inputs. The first and most popular conception of In-context learning [1]  involves providing the model with contextual information, typically a few sample $(x, y)$ pairs that represent a specific “task”. In contrast, we introduce an alternative perspective on in-context learning, where unlabeled $x$ inputs act as the contextual backdrop for a task, also known as an 'environment.'  Below, we emphasize the significance of this approach along two key dimensions.

- **Why unlabelled $x$’s:** In domain generalization, at test-time, the learner only has access to the unlabeled x’s from the test environment and not both $(x,y)$ pairs. We embrace this challenge and develop a proposal inspired by in-context learning that shows we can continue benefiting from information in the unlabeled $x$’s. Our theory qualifies the conditions under which the information from unlabelled $x$’s only can help improve upon standard ERM in both in-distribution and out-of-distribution settings.

- **Example of when only $x$’s can provide useful context:** Consider the following example – there are two types of users, one who has a more cursive handwriting and the other type uses straight edges.  At test time, the classifier may be faced with samples from cursive users. By observing the images ($x$’s) only of the samples drawn by the user without the label, the model can quickly realize that it is dealing with a cursive type user. These intuitions are made formal by our theory. Our method capitalizes on these intuitions to show the empirical benefits. 

Further, our method can be easily adapted to deal with $(x,y)$ pairs as input. We refer to this approach as Supervised ICRM. As shown in the table below and Section D.3 and as anticipated, Supervised ICRM outperforms ICRM but is not suitable for domain generalization settings where data labels are unavailable at inference.

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**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** |

To further the evidence that our proposal of using unlabeled $x$’s is powerful, we provide experiments on two more datasets Imagenet-R (consisting of renditions of Imagenet classes) and CIFAR10-C dataset. Please see the table below or refer to Section D.2. 

Dataset |  |  | Avg.* |  |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
CIFAR10-C | 0 | 25 | 50 | 75 | 100 | | 0 | 25 | 50 | 75 | 100 |
ARM   | 65.9 | 66.0 | 66.0 | 66.0 | 66.0  | | 39.3 | 39.3 | 39.4 | 39.3 | 39.4
ERM   | 66.1 | 66.1 | 66.1 | 66.1 | 66.1 | | 39.8 | 39.8  | 39.8  | 39.8  | 39.8
ICRM  | **70.6** | **71.9** |  **71.9** | **71.9** | **71.9** | | **54.6** | **56.0** | **55.8** | **55.8** | **55.9**
|  |  |  |  |  |  |  |  |  |  |
ImageNet R | 0 | 25 | 50 | 75 | 100 | | 0 | 25 | 50 | 75 | 100 |
ARM   | 56.3 | 58.1 | 58.8 | **59.8** | 59.0  | | 47.4 | 45.3 | 47.2 | **49.8** | 47.4
ERM   | **58.9** | 58.9 | 58.9 | 58.9 | 58.9 | |**48.0** | **48.0**  | **48.0**  | 48.0  | 48.0 
ICRM   | 57.4 | **59.7** |  **59.6** | 59.4 | **60.5** | | 45.4 | **48.0** | 47.2 | 46.9 | **50.6**

*Here, “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. 

[1] Brown, Tom, et al. "Language models are few-shot learners." Advances in neural information processing systems 33 (2020): 1877-1901.

Reviewer comment (Experiments):
While 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).

Author response:
> On the assumption that testing data sequences are sampled from the same domain

While we did mention this point in the previous draft, we acknowledge that it may not have been sufficiently emphasized. In the revised manuscript, we have reiterated this aspect multiple times. We agree that this is an assumption, if true, shows the success of the proposal. Additionally, we consider extending our approach to settings where this assumption does not hold as an exciting avenue for future research. Nevertheless, it's important to note that the current assumption is practical for many scenarios. In particular, all work in domain generalization (see DomainBed [1], WILDS [2]) follows the same fixed-test-domain protocol.

[1] Gulrajani, Ishaan, and David Lopez-Paz. "In search of lost domain generalization." arXiv preprint arXiv:2007.01434 (2020).

[2] Koh, Pang Wei, et al. "Wilds: A benchmark of in-the-wild distribution shifts." International Conference on Machine Learning. PMLR, 2021.

Reviewer comment (Writing):
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.

Author response:
> On the relationship to in-context learning in LLMs

Our whole journey and thought process in this paper was inspired by LLMs and their ability to tackle OOD tasks through in-context learning. Our approach is inspired by it and is not an exact emulation of the workings of a LLM. We have added additional clarifications on this in the paper.

Reviewer comment (Theory):
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?

Author response:
> On the concern regarding Theorem 1

We are uncertain about which inequality is being referred to here."
- If the reference is to $I(Y; E|X)>0$, then we emphasize that this simply states that X does not carry all the relevant information to predict Y. There is some extra information to be gained from knowing the environment one is operating in. For instance, if one is driving, the decision to speed or not should depend on the current image $X$ in front of the driver and the operating conditions (crowded city, weather, etc., that form $E$). This assumption is made in all marginal transfer learning papers.

- Or perhaps the reference is to the inequality in equation (10). Perhaps your concern refers to the condition above equation 10 in the Appendix, which states $I(Y; C_t | X, E)=0$. Below, we show how to derive it. Equation 10 rests on the assumption that $X_i, Y_i$’s are independent conditional on $E$. This is a fairly standard data generation assumption in most domain generalization papers (e.g., invariant risk minimization, its precursors, and many works that follow up on it). The assumption just means that within an environment, $X_i, Y_i$ samples are independent of one another. Conditional on the user, if the first digit drawn by a user is 9, that does not impact the choice of the digit it draws next.

Consider two samples $(X_1,Y_1,X_2,Y_2)$ from the environment $E=e$. From conditional independence of the pair it follows that 
$P(X_1=x_1, Y_1=y_1, X_2=x_2, Y_2=y_2 | E=e) = P(X_1=x_1, Y=y_1|E=e) P(X_2=x_2, Y_2=y_2 |E=e)$

We now take the sum over $y_1$’s on both sides of the above equality to obtain. 
$\sum_{y_1} P(X_1=x_1, Y_1=y_1, X_2=x_2, Y_2=y_2 | E=e) = \sum_{y_1}P(X_1=x_1,Y=y_1|E=e) P(X_2=x_2,Y_2=y_2 |E=e)$

The LHS simplifies to $P(X_1=x_1,X_2=x_2, Y_2=y_2|E=e)$ while the RHS simplifies to $P(X_1=x_1|E=e) P(X_2=x_2,Y_2=y_2|E=e)$. Thus we obtain 
$P(X_1=x_1,X_2=x_2, Y_2=y_2|E=e) = P(X_1=x_1|E=e) P(X_2=x_2, Y_2=y_2|E=e)$ 

The RHS can be simplified further and we obtain 
$P(X_1=x_1,X_2=x_2, Y_2=y_2|E=e) =  P(X_1=x_1|E=e)P(X_2=x_2|E=e)P(Y_2=y_2|X_2=x_2,E=e)$ 

From the above equality, we obtain that $Y_2$ is independent of $X_1$  conditional on $X_2, E$. Since $X_1$ alone forms the context, $Y_2$ is independent of the context conditional on $X_2, E$. Hence, in this example $I(Y_2;C| X_2,E)=0$. The above argument readily extends to longer sequences, and thus we get $I(Y; C_t | X, E)=0.$

Reviewer comment (Reproducibility):
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?

Author response:
> Regarding $h(.)$ and parameters for different environments

$h(.)$ possesses different parameters for  different environments. Theorem 1 deals with in-distribution  evaluations. Hence, the parameter of the test environment for Theorem 1 is set to be one of the training environments. In Theorem 3, we deal with out-of-distribution environments. In this case, the parameters of the test environment are constrained to be in the Voronoi cell of the training environment. Note that Voronoi cells do not constrain the test parameters to be arbitrarily close ,and quite remarkably, the parameter vector can be quite far. In our experiments, since we conduct out-of-distribution evaluations, in some cases these parameters can be quite different from those seen at train time.

Reviewer comment (Theory):
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.

Author response:
> On the assumption that testing data sequences are sampled from the same domain

While we did mention this point in the previous draft, we acknowledge that it may not have been sufficiently emphasized. In the revised manuscript, we have reiterated this aspect multiple times. We agree that this is an assumption, if true, shows the success of the proposal. Additionally, we consider extending our approach to settings where this assumption does not hold as an exciting avenue for future research. Nevertheless, it's important to note that the current assumption is practical for many scenarios. In particular, all work in domain generalization (see DomainBed [1], WILDS [2]) follows the same fixed-test-domain protocol.

[1] Gulrajani, Ishaan, and David Lopez-Paz. "In search of lost domain generalization." arXiv preprint arXiv:2007.01434 (2020).

[2] Koh, Pang Wei, et al. "Wilds: A benchmark of in-the-wild distribution shifts." International Conference on Machine Learning. PMLR, 2021.

Reviewer comment (Experiments):
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.

Author response:
> On additional baselines

We have incorporated six additional recent and popular baselines into our study. These baselines either extend the invariance-based paradigm (Fish, IB-IRM), leverage contextual information differently (BN-adapt, Bayesian BN-adapt), or apply classic regularizations (Mixup, IB-ERM). You can find the detailed comparisons in the table below and Section D.2.

Dataset |  |  | Avg.* |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
**FEMNIST** | 0 | 25 | 50 | 75 | 100 | 0 | 25 | 50 | 75 | 100 |
BN Adapt  | 78.3 | 76.9 | 80.3 | 81.5 | 82.4 | 54.3 | 56.2 | 61.9 | 64.7 | 65.3
Bayesian BN Adapt  | 78.3 | 79.6 | 81.3 | 82.2 | 82.9 | 54.3 | 60.4 | 64.7 | 65.5 | 66.3 |
Fish | 74.2 | 74.2| 74.2| 74.2| 74.2| 47.3| 47.3|47.3 | 47.3 |47.3| 
IB-ERM | 74.3 |74.3 |74.3 |74.3 |74.3 |53.0  |53.0  |53.0  | 53.0 |53.0 | 
IB-IRM | 68.9 |68.9 |68.9 |68.9 |68.9 |46.5 | 46.5|46.5 | 46.5 |46.5 
Mixup | 80.5 |80.5 |80.5 |80.5 |80.5 | 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 | 0 | 25 | 50 | 75 | 100 |
BN Adapt  | **94.6** | 87.0 | 91.5 | 93.7 | 94.3 | 80.5 | 70.9 | 76.9 | 79.8 | 80.9
Bayesian BN Adapt  | **94.6** |  91.2 | 93.4 | 94.3 | 94.7 | 80.5 | 75.4 | 79.2 | 80.7 | 81.3
Fish | 93.5 | 93.5| 93.5|93.5 |93.5 | 76.1| 76.1| 76.1| 76.1 | 76.1 |
IB-ERM | 92.4 | 92.4| 92.4| 92.4| 92.4|77.7 |77.7 |77.7 |77.7| 77.7 | 
IB-IRM | 84.8 |84.8 |84.8 |84.8 |84.8 | 65.1| 65.1| 65.1| 65.1|65.1| 
Mixup | 93.2 | 93.2|93.2 |93.2 |93.2 |76.7 | 76.7 |76.7 |76.7 |76.7| 
ICRM       |  93.6 |   **96.1**|  **96.2**| **96.2** | **96.2**|  **82.5**|  **88.5**|  **88.5**| **88.8**| **88.8**
| | | | | | | | | |
**WILDS Camelyon17** | 0 | 25 | 50 | 75 | 100 | 0 | 25 | 50 | 75 | 100
BN Adapt | 67.5 | 82.0 | 87.4 | 89.7 | 89.9 | - | - | - | - | - |
Bayesian BN Adapt  | 67.5 | 82.0 | 87.3 | 89.6 | 89.7 |- | - | - | - | - |
Fish | 53.8 |53.8 | 53.8| 53.8| 53.8| - | - | - | - | - |
IB-ERM | 50.3 | 50.3| 50.3| 50.3|50.3 |- | - | - | - | - |
IB-IRM | 51.1 | 51.1| 51.1| 51.1| 51.1|- | - | - | - | - |
Mixup | 36.2 |36.2  |36.2  |36.2  |36.2  |- | - | - | - | - |
ICRM |  **92.0** | **90.7**  | **90.8**  | **90.8** | **90.8** | - | - | - | - | - |
| | | | | | | | | |
**Tiny ImageNet-C** | 0 | 25 | 50 | 75 | 100 | 0 | 25 | 50 | 75 | 100  
BN Adapt  | 31.7 | 1.7 | 1.7 | 1.9 | 2.1 | 9.4 | 1.3 | 1.4 | 1.6 | 1.7
Bayesian BN Adapt  | 31.7  | 2.2 | 2.1 | 2.3 | 2.4 | 9.4 | 1.6 | 1.6 | 1.8 | 1.8 
Fish | 32.5 |32.5 |32.5 |32.5 |32.5 | 11.5 |11.5  | 11.5 |11.5  | 11.5 
IB-ERM | 34.7 | 34.7| 34.7| 34.7| 34.7| 14.7|14.7 |14.7 |14.7| 14.7 | 
IB-IRM | 35.2 |35.2 |35.2 | 35.2| 35.2| 16.0|16.0 |16.0 | 16.0 |16.0| 
Mixup | 36.2 |36.2 | 36.2| 36.2| 36.2 |11.6 |11.6 |11.6 | 11.6| 11.6 | 
ICRM       |  **38.3** | **39.2**  | **39.2**  | **39.2** | **39.2** | **18.8**  |  **19.2** | **19.5**  | **19.5** | **19.4** |

Reviewer comment (Experiments):
**Choice of datasets for evaluation + including experiments on "harder" datasets with spurious correlations that vary across environments**.  
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".

Author response:
> On the choice of evaluation and including experiments on “harder” datasets with spurious correlations

- In the literature, distribution shifts are often divided into two categories: diversity shift and correlation shift [1]. Our work primarily focuses on diversity shift, deferring the exploration of correlation shift to future research. Handling datasets like Waterbirds is challenging, as the typical evaluation assumes test data sampled i.i.d from all environments (subpopulations), making a 3:1 training-to-testing split nonstandard. Further, adapting ICRM to such a setting requires significant modifications to the training pipeline and additional analysis. Therefore, our focus is on domain generalization benchmarks, where all environments except one are used for training, aligning with our assumption of _only one environment during testing_.
- To test the model in other hard settings, we have included the ImageNet R dataset [2], which consists of challenging renditions of Imagenet R classes, and the CIFAR10-C dataset. In both datasets, we observe consistent improvements offered by ICRM (see table below and also in the Appendix of the edited draft). 


Dataset |  |  | Avg.* |  |  |  |  |  | WG.*  |  |  |
|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|----------|
CIFAR10-C | 0 | 25 | 50 | 75 | 100 | | 0 | 25 | 50 | 75 | 100 |
ARM   | 65.9 | 66.0 | 66.0 | 66.0 | 66.0  | | 39.3 | 39.3 | 39.4 | 39.3 | 39.4
ERM   | 66.1 | 66.1 | 66.1 | 66.1 | 66.1 | | 39.8 | 39.8  | 39.8  | 39.8  | 39.8
ICRM  | **70.6** | **71.9** |  **71.9** | **71.9** | **71.9** | | **54.6** | **56.0** | **55.8** | **55.8** | **55.9**
|  |  |  |  |  |  |  |  |  |  |
ImageNet R | 0 | 25 | 50 | 75 | 100 | | 0 | 25 | 50 | 75 | 100 |
ARM   | 56.3 | 58.1 | 58.8 | **59.8** | 59.0  | | 47.4 | 45.3 | 47.2 | **49.8** | 47.4
ERM   | **58.9** | 58.9 | 58.9 | 58.9 | 58.9 | |**48.0** | **48.0**  | **48.0**  | 48.0  | 48.0 
ICRM   | 57.4 | **59.7** |  **59.6** | 59.4 | **60.5** | | 45.4 | **48.0** | 47.2 | 46.9 | **50.6**

*Here “Avg.” and “WG.” respectively denote the average and the worst group accuracy across test environments. \

[1] Ye, Nanyang, et al. "Ood-bench: Quantifying and understanding two dimensions of out-of-distribution generalization." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[2] Hendrycks, Dan, et al. "The many faces of robustness: A critical analysis of out-of-distribution generalization." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021.

Reviewer comment (Experiments):
**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.  
In my opinion, this approach is perfect for domain generalization settings where the "environment" is literally the institution in which the algorithm is being deployed (e.g., 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.

Author response:
> More discussion on how to use ICRM in practice

You make an excellent point on the type of distribution shifts our method is suited for. In datasets such as Waterbirds, the test domain is a mixture of hidden groups/domains and not separated into domains based on the groups. Extending our method to tackle distribution shifts in slowly-changing test environments and changes in subpopulation such as Waterbirds, CelebA type dataset is an exciting future work. We have added this discussion in the the revised manuscript

Reviewer comment (Writing):
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.

Author response:
> On “Context is Environment” vs. “Environment is Context''  

We consider the environment to be a more fundamental primitive. We believe that the data generation process in nature begins with sampling of an environment, followed by the generation of samples from this environment, forming the context. Due to this, it is the context that helps us infer the environment and hence context is environment.

Reviewer comment (Writing):
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?

Author response:
> On clarification of ERM 

In our paper, Empirical Risk Minimization (ERM) refers to the training of a single model by minimizing the error across the entire pooled dataset. We have incorporated the recommended clarification in the manuscript.

Reviewer comment (Writing):
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?

Author response:
> Clarification on “reveal” an “invariance” 

By “invariance” we mean a feature representation inducing a classifier that performs well across environments. By “revealing” we mean that the learner can find and use such an invariance..

Reviewer comment (Writing):
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?

Author response:
> On the linear growth in the size of the representation

In this particular case, we aimed to highlight the limitations of marginal transfer learning proposals that use non-linear averaging on the inputs. By size, we mean the dimensionality of the representation $\phi(x)$. Appendix Section A.7 presents a family of examples where such marginal transfer learning proposals fail. In particular, we give an example of function classes for which the size of $\phi(x)$ has to grow in the context length to perfectly match the function class. The overarching intuition here is that one would need an $\theta(n)$-dimensional $\phi(x)$ to ensure storing all of the information available in $n$ examples.

Reviewer comment (Theory):
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 phenomenon (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".

Author response:
> On linear least squares example

We assume in this example that ICRM directly operates on the features summarizing the means $\mu$. ERM in the standard form only operates on the current query and not on an extended feature space. This example illustrates how ICRM by operating on an extended feature space can discover invariances for features in the original input space. This is as opposed to a global ERM across environments, or any ERM on a separate environment.

Reviewer comment (Writing):
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?

Author response:
> On the example of smoking and invariance

The smoking example attempts to illustrate the fact that, as we consider growing collections of examples (smokers), we need to mine more and more features to find an invariant map—for instance, consider the enormous amount of inputs that we would need to invariantly predict the outcome of disease for all smokers in the world with one single model. On the other extreme, if considering one smoker only, the invariant rule is constant and equal to the observed disease outcome. This illustrates a tension between how many environments one desires to be invariant across ($|E|$) and how many features are necessary to guarantee such invariance ($|C|$).

Reviewer comment (Experiments):
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?

Author response:
> Clarification about images in the input sequences in Figure 2

Figure 2 displays four random sequences, each containing images randomly sampled from _real data_ in a test environment. It's important to note that none of the images in these sequences were manually augmented. The key takeaway from this visualization is that ICRM effectively learns to attend to a select few samples in an input sequence. These samples either belong to the same class or exhibit similar features, despite potentially belonging to different classes.

Reviewer comment (Presentation):
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?)

Author response:
> True labels for each image in Figure 2
 
We thank the reviewer for this suggestion. Using the label mapping [here](https://gist.github.com/aaronpolhamus/964a4411c0906315deb9f4a3723aac57), we have replaced the class indices with their true labels in Figure 2. Clearly, the third row shows that the model, when presented with a query image of a "bullet train", attends not only on other trains but also on a "school bus"---indicating a semantic understanding of similarity. Further, in the last row, since the associated image from class "volleyball” in the query, also features a volleyball player, the model learns to pay attention to other images with individuals, including those wearing academic gowns.

Reviewer comment (Experiments):
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.

Author response:
> On the natural order of data point in the conclusion

Existing datasets in domain generalization only carry environment-level separation, and the proposal is designed to handle the partial ordering of data dictated by this separation. However, if the data had more refined information (e.g., time-based ordering in video), the proposal can exploit that too. We have clarified this in the revised manuscript. In sum, our remark about “natural order” is made to motivate future data collectors to retain this precious time metadata, so methods such as ICRM can be deployed in full force.