Context is Environment

Performance decreases with context in WILDS Camelyon17.

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.

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.

On the definition of in-context learning

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”. While it is common to leverage additional information through labeled data, it is also possible to use unlabeled data to provide additional information, as demonstrated by our work. As shown in [2], labels do not play a significant role in several in-context learning tasks.

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

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

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

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.

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

We thank the reviewer for their detailed, thoughtful, and critical review. We are glad you appreciated the novelty of our framework and its benefits, and we believe this could form the basis for a positive review. We believe that there are a few confusions that have resulted in the low rating. We have tried our best to clarify them below and in the paper. The relevant changes are highlighted in magenta in the main body and the Appendix.

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.

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

On the gains achieved using unlabeled examples

We would like to highlight three key points here. Firstly, in three of the four datasets, we observe notable improvements through adaptation. Please also refer to the worst-group accuracy in Table 2. These datasets are particularly challenging, and even a 1-2 percent gain is significant. Secondly, Theorem 2 characterizes scenarios under which increasing unlabeled context can monotonically improve the performance of the trained model. Thirdly, we present results with comprehensive hyperparameter tuning for all baselines for two additional datasets, Imagenet-R and CIFAR 10 C, demonstrating the advantages of incorporating context.

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

We thank the reviewer for their continued discussion. We hope that our clarifications below can help with the concerns they raised and help see our work in a positive light.

Regarding misleading motivations

In their response, they state “I find this to be misleading given the current motivation and presentation of the paper focuses heavily on few-shot in-context learning gains”. We are not clear what they mean here. Both zero-shot and few-shot gains are at the heart of domain generalization, which is the primary focus of our paper. We do not see how the paper is misleading. Nowhere in our paper do we assert that gains are exclusively expected with large context lengths and not attainable with small context lengths. In fact, achieving gains at small context lengths is practically significant. Our theoretical results focus on in-context gains in the limit of both large and small context lengths (even one context length), as demonstrated in Theorems 1 and 2, respectively.

Regarding the issue of performance drop in CIFAR10-C

We use a toy example to explain the source of the performance drop. Consider a sequence $x_1, \cdots, x_n$ and corresponding labels $y_1, \cdots, y_n$. Suppose the label are generated as $y_i = f(x_{i},x_{i-1})$, where $x_i$ is the current query, $x_{i-1}$ is the previous query and $y_i$ is the label for the current query. We contrast learning a function with context length 1 and context length 2. The functions with context length 1 take as input $x_i, x_{i-1}$ and the functions with context length 2 take as input $x_i, x_{i-1}, x_{i-2}$. Observe that the function class of context length 2 contains the function class of context length 1. As a result, the optimal predictor (in terms of the prediction error over the training distribution) in length 2 class is at least as good as that in length 1 class. However, when it comes to learning the ideal function $f(x_{i},x_{i-1})$ with a finite number of samples, learning the ideal function from length 1 class should require fewer samples than learning the ideal function from length 2 class. In other words, with a fixed number of samples, the generalization error of the length 1 class should be smaller than length 2 class. In our experiments, these finite sample-driven learning-based issues could lead to performance drops.

Regarding the performance of CIFAR 10 C at intermediate points

We would be happy to add performance at all intermediate points for CIFAR 10 C and Imagenet R. And at the risk of repetition even if the performance was the same for all $k=1, 2, …, 25$, it is acceptable and not in contradiction with desiderata or the theory.

Regarding zero shot gains

You stated, “I remain not convinced why zero-shot predictions with ICRM should be better”. We have provided empirical evidence that zero-shot gains are achievable. We have provided a solid theoretical analysis for context lengths 1 and beyond. We also provide initial insights into zero-shot gains through examples in equation (7) and the selective augmentation perspective. Transforming these initial insights into a solid full-blown theory to explain “why ICRM improves zero-shot predictions” is an important future work, which we also stated in our previous response. We hope that the reviewer appreciates that it is not possible to explain all aspects of the experiments through the lens of theory in one paper.

Regarding the comparison of ERM vs ICRM at zero shot in the example from equation (7)

In Proposition 3 in the Appendix, we showed that for the example under consideration, the error of ERM grows in variance of the first feature. In the paragraph that follows Proposition 3, we show that the error of ICRM without any context is independent of the variance of the first feature. In the absence of any context, we assume ICRM uses a fixed value of zero as an estimate for the mean. Zero is an arbitrary choice and qualitatively the result stays the same even if some other fixed value is used. If the variance of the first feature is sufficiently high in the test environment, then the error of ERM is larger than the error of ICRM (without context). On the point about adding a fixed constant to ERM, this will also not work as the mean $mu_e^2$ varies across environments.

We thank the reviewer for their thought provoking review and encouraging words of appreciation for our work. We have incorporated the changes which can be found in magenta in the main body and the Appendix.

On the analogy between LLMs and ICRM

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.

On the nature of distribution shifts

We agree that ICRM cannot tackle arbitrary shifts in $P(Y|X)$, and the Voronoi cells are precisely meant to constrain that. However, there is an interesting side to it. The constraint imposed by the Voronoi cell is more relaxed than the standard covariate shift assumption where $P(Y|X)$ does not vary. This is afforded to ICRM by the fact that it operates on an extended feature space that takes both current input $X$ and distribution $p(X)$ from the test environment as input, while standard learners (based on ERM, IRM) operate solely on $X$.

On the fairness of comparison between ICRM and ERM in Section 5

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.

On the issue of ICRM exploiting “natural order”.

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.

On the drop in Camelyon-17.

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.

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.

On the difference w.r.t averaging based methods

Our experiments show that a long line of work starting from the seminal works on marginal transfer learning that first appeared in NeurIPS 2011, later improved by adaptive risk minimization (NeurIPS 2021) have adopted a suboptimal choice for capturing contextual information. Our proposal simplifies such literature in light of recently developed transformer and attention architectures. In a nutshell, we show that state-of-the-art performance is obtained by presenting all of the data as a sequence directly to the model, which can later decide to average in various non-linear ways. Further, as discussed in the paper, the size of the representation $\phi$ would have to grow linearly with the size of the training data to ensure storing all of the information available in the data

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.

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

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.

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.

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

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

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.

Clarifying the line on swapping two components

Yes, we mean it swaps parameter tuples corresponding to $y=0$ and $y=1$.

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.

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

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.

True labels for each image in Figure 2

We thank the reviewer for this suggestion. Using the label mapping here, 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.

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.

First of all, we would like to thank the reviewer for their diligent and insightful review and their very encouraging words. The changes in the manuscript can be found highlighted in magenta. We share our thoughts on the questions asked below.

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

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

On weakening claims about “What ICRM has learned”

This is indeed a great point—we have weakened our exposition to explicitly emphasize that it is our hypothesis that ICRM exhibits a propensity to learn more invariant predictors. The post-hoc analysis of attention maps in Section 6.3 complements it empirically.

To substantiate this further, we extract embeddings from the penultimate layer of the model trained using our approach of the data pooled together from training environments. Then, we use these embeddings to train a linear classifier that predicts the corresponding environment index. The results (details in Section D.4) demonstrate our approach's success, achieving accuracy rates exceeding 75% on FEMNIST and 98% on Rotated MNIST. This shows that context can help us infer the environment and hence, reinforces the notion of context is environment.

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

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.

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.

Thank you for your prompt and detailed response to each of my questions and concerns! Thank you also for updating your draft for clarity, and for also including the additional experimental results on ImageNet-R and CIFAR10-C. The authors' responses have sufficiently addressed my concerns, and after reviewing the larger discussion at this time I still believe this paper should be accepted.