Paper: Discovering Environments with XRM (ICLR-2024)

Reviewer comment (Theory):
If I understand correctly, the method seems to rely on the fact that misclassified examples are such because they do not contain a "spurious correlation" that a model would learn by default. The twin training serves to reinforce the tendency of one of the trained models to capture this spurious correlation. If this is indeed the case, then the overall method seems to depend on the (common) heuristic that models learn spurious correlations by default (a.k.a. shortcut learning). I think this is the same heuristic that is used in the existing methods criticised in Section 3. Critically, this heuristic relies on the fact that we know that the chosen architecture/training set lead to learning undesirable spurious correlations by default. What if one applies the method to a situation where a perfectly-fine, "robust" model is learned by default? I'm guessing that the method would then be detrimental. I'm not suggesting that we should be able to do better without additional knowledge (in fact [1] seems to show it's not possible) but the authors here do claim to do so, hence the need to point out this possible limitation (see also W5).

Author response:
> What if one applies the method to a situation where a perfectly-fine, "robust" model is learned by default?

Appendix B now contains a tour on four versions of the ColoredMNIST dataset, providing intuition about those cases where XRM works, and those instances where it may fail. As you suggest, XRM is very helpful when the spurious correlation is learned first, whereas ERM is a better alternative in those cases where the invariant correlation is learned first. In sum, we recommend running both ERM and XRM when the spurious correlation type is unknown.

Reviewer comment (Novelty):
If my understanding of the method is correct, the method is also very similar to the following works:  
- Works in the debiasing literature (e.g. LfF) that train a pair of models that respectively rely/do not rely on spurious features. These works are discussed in Sect. 3, and I do understand that they rely to some extent on the tuning of model capacity to ensure that it captures the spurious feature, but I am not convinced that the proposed parallel training (which seems to be the essential difference) leads to the discovery of something fundamentally different.  
- Works in the "model diversity" literature that train a pair of models that differ in their predictions [5,6]. These also proceed to train models in parallel, in a way that seems conceptually very similar to the step 1 proposed here (implementation details aside).

Author response:
> [Comparison to LfF]

We now clarify (in blue). The crucial ingredient in XRM is the mechanism of label-flipping, which allows us to (i) run our environment discovery method safely until convergence, and (ii) count label flips as a model selection criteria. This is a major advancement with respect to previous methods, which (i) required surgical early stopping, and (ii) didn't provide a model selection criteria. Consequently, previous work resorted to a validation set with environment annotations, the very information subject to discovery.

> [Comparison to model diversity]

We now clarify (in blue). Once again, the difference is in the label-flipping mechanism, which implements an "echo chamber" that biases the twins to increasingly rely on spurious correlations, and *agree* on their predictions. This is in fact the opposite of what is pursued in the model diversity literature—training multiple accurate models that differ in their predictions.

Reviewer comment (Novelty):
The proposed method only partitions the data into 2 "environments". I don't think this is really in line with the literature on DG (with which this work is supposed to connect) that is mostly based on invariance learning and needs a large number of training environments (e.g. IRM). This work therefore seems much more related to the simpler setting of "debiasing" methods (e.g. LfF) that aim at removing the reliance of a model from one precise biased feature. The methods used for the second phase are indeed simple baselines for debiasing, and not really DG methods. These are very strong baselines in these settings (and with the datasets considered), but I'm not sure this is what the reader would expect given all the mentions about DG.

Author response:
> The proposed method only partitions the data into 2 "environments"

Whether two environments suffice—when constructed appropriately—is in fact a fascinating question that we have pondered about for a while. We do not believe that XRM is identifying "one spurious correlation" but, more generally, "the direction of spuriousness”. We lack the theory we would love to have here, but we suspect that an Oracle with knowledge about what the test domain will be would be able to split the training data as to maximally reveal the train-test discrepancy. Thank you for pointing this out, we now discuss this stimulating topic (in blue).

> I'm not sure this is what the reader would expect given all the mentions about DG

We are currently conducting experiments on DomainBed, a key benchmark in domain generalization. We will include the findings in Appendix C in coming hours.

Reviewer comment (Novelty):
Absence of a comprehensive review of related work. Some directly-related methods are correctly cited/discussed throughout the paper, but there are other connected areas that are not really discussed. Examples include:  
- [1,2] discuss conditions under which environment discovery is possible. I think the theoretical statements in [1] are particularly important to discuss (I am not sure how the proposed method overcomes the impossibility stated in that paper; see also W5).  
- [3] was an early method that also proposed to "unshuffle" data (a term used in this paper) by simply clustering the data. Looking at the visualizations of discovered "environments" in Fig. 3, one wonders if these could also be discovered with such a simple clustering baseline.  
- [4] is another recent method that also seems to claim discovering partitions in the data (I suspect it has similar flaws to those discussed in the paper; it has appeared at ICCV 2023 after the ICLR deadline so it's totally fine to dismiss it though).

Author response:
> Absence of a comprehensive review of related work

Appendix A now contains an exhaustive literature review on domain generalization, subpopulation shift, and their relations to invariance and causal inference. We survey methods for environment discovery, as well as the relevant impossibility results and other interesting pieces of work you kindly mention.

Reviewer comment (Theory):
No discussion or empirical exploration of the limitations of the methods. No precise statement of the assumptions on which the method relies.

Author response:
> No discussion or empirical exploration of the limitations of the methods

Again, Appendix B now presents a ColoredMNIST dataset analysis, showing XRM's effectiveness and limitations in certain scenarios and recommending a combined use of ERM and XRM when the correlation type is unclear.

Due to no-free-lunch theorems, no learning method can work across the board, and must necessarily rest on some data assumptions, which are in many cases difficult to verify. Our new Appendix B (soon in the main text) provides the main intuition: XRM is effective when the spurious features are simpler than the invariant patterns of interest. In any case, to cover more ground, we recommend practitioners to add ERM runs to their sweeps.

Reviewer comment (Theory):
The paper should provide a clear discussion on the identifiability challenges presented in [1], which demonstrate that learning invariance without domain partition can be generally impossible. It is crucial to address the need for imposing inductive bias, additional assumptions, conditions, or auxiliary information to ensure the effectiveness of the proposed method. A thorough exploration of these aspects would enhance the paper's theoretical foundation and its practical applicability.

Author response:
Thank you for the very interesting reference [ZIN]. Appendix B now contains a tour on four versions of the ColoredMNIST dataset, providing intuition about those cases where XRM works, and those instances where it may fail. We frame this analysis within the impossibility results described in the ZIN paper.

Reviewer comment (Experiments):
According to [2, 3], spurious features are defined as any nodes in the causal graph other than the direct causes of the label. However, I have concerns about the evaluations conducted on datasets like waterbird, which explicitly contain only one dominating spurious feature. These datasets may not fully reflect the implications of the proposed methods on more realistic and high-dimensional datasets such as ImageNet variants. Moreover, the paper relies on the assumption that Empirical Risk Minimization (ERM) learns spurious features first, but it may not hold true for all types of spurious features as discussed in [2, 3]. It would be valuable to address these concerns and provide further insights into the generalizability of the method to diverse real-world datasets.

Author response:
Our main evaluation was carried out on seven datasets, where we would argue that Metashift and ImageNetBG are "ImageNet-like". We added as many datasets as we could within our constraints in an attempt to cover a wide range of spurious correlations.

Reviewer comment (Theory):
If a large number of spurious features are present, [1] demonstrates that there are necessary and sufficient conditions for learning invariance without explicit domain partitions, which can be quite restrictive. I have concerns about whether the proposed two-stage method can effectively address this problem given the limitations imposed by these conditions. It would be valuable for the authors to discuss how their method overcomes or accommodates these restrictions and whether it can achieve satisfactory results in scenarios with a significant number of spurious features.

Author response:
Appendix B now contains a tour on four versions of the ColoredMNIST dataset, providing intuition about those cases where XRM works, and those instances where it may fail. As you suggest, XRM is very helpful when the spurious correlation is learned first, whereas ERM is a better alternative in those cases where the invariant correlation is learned first. In sum, we recommend running both ERM and XRM when the spurious correlation type is unknown.

Reviewer comment (Experiments):
Several studies [5, 6] have highlighted the challenges associated with learning invariance in the presence of many spurious features. In a recent paper, [4] discovered that when dealing with a large number of spurious features, each ERM model tends to learn a subset of these features. [4] further demonstrates that rather than exclusively focusing on learning invariant features, it is beneficial for OOD performance to diversify the learned spurious features (referred to as spurious feature diversification). Spurious feature diversification is shown to explain the effectiveness of empirically strong methods like SWAD and Model soup. It would be valuable to investigate whether the proposed method (XRM) can enhance spurious feature diversification and demonstrate effective performance on a broader range of real-world datasets, such as PACS, OfficeHome, DomainNet, or ImageNet variants.

Author response:
Thank you for pointing us to feature diversification, we now discuss it in our manuscript (in blue).  
Additionally, per your suggestion, we are currently conducting experiments on DomainBed, and we will include the findings in Appendix C in coming hours.

Reviewer comment (Experiments):
The paper's claims may be slightly overstated. While the focus on subpopulation shift in distribution shift is indeed important, it might be more appropriate to avoid claiming to solve a long-standing problem in out-of-distribution generalization without further empirical studies on widely recognized DG benchmarks such as DomainBed and Wilds. These additional experiments could provide more convincing evidence of the proposed approach's effectiveness.

Author response:
In line with your suggestion, we are currently conducting experiments on DomainBed. We will share the findings in Appendix C in coming hours.

Reviewer comment (Novelty):
The paper lacks a comprehensive discussion of important related works concerning data splitting strategies for improved DG performance and subpopulation shift, such as references [1], [2], [3] and [4]. Notably, in [1], the authors have theoretically demonstrated the challenges of learning the invariant correlation between samples and labels in the absence of prior information. Including these relevant works would enhance the paper's literature review and contextualize the proposed approach.

Author response:
Appendix A now contains an exhaustive literature review on domain generalization, subpopulation shift, and their relations to invariance and causal inference. We also survey methods for environment discovery, as well as the relevant impossibility results you kindly mentioned.

Reviewer comment (Theory):
As highlighted in [1], it is crucial to understand the specific scenarios where XRM is expected to be effective. Therefore, it would be beneficial for the authors to provide further insights into the data distribution settings in which XRM is likely to perform well. Alternatively, the authors could explore providing theoretical guarantees to enhance the understanding of XRM's strengths and limitations.

Author response:
Appendix B now contains a tour on four versions of the ColoredMNIST dataset, providing intuition about those cases where XRM works, and those instances where it may fail.

Reviewer comment (Evaluation):
The observation that XRM outperforms Human-annotation methods is intriguing and warrants further explanation.

Author response:
Our experiments on Waterbirds, CelebA, MultiNLI, and CivilComments show that environment annotations discovered by XRM afford the same or better state-of-the-art test worst-group-accuracy than those environment annotations included in the original datasets by humans. Our hypothesis is that the biases identified by humans may not always align with the biases that models learn, and vice versa [1].