Invariant Learning with Annotation-free Environments

Authors: phuong quynh le
christin  seifert
jörg  schlötterer
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-24 01:47:38

Abstract: Invariant learning across environments is a promising approach to improve domain generalization compared to Empirical Risk Minimization (ERM). However, most invariant learning methods rely on the assumption that training examples are pre-partitioned into different known environments. We instead infer environments without the need for additional annotations, motivated by observations of the properties within the representation space of a trained ERM model. We show the preliminary effectiveness of our approach on the ColoredMNIST benchmark, achieving performance comparable to methods requiring explicit environment labels and on par with an annotation-free method that poses strong restrictions on the ERM reference model.

Main paper: 
 introduction: 
 Empirical Risk Minimization (ERM) is known to generalize poorly when the test data has a different distribution than the training data. When features have a statistical but not a causal relationship with class labels, we call this a spurious correlation and the corresponding features spurious features. For example, the Waterbirds dataset [Sagawa et al., 2020] has background as a spurious feature. During training, 95% of waterbirds have water background and 95% of landbirds have land background, whereas land and water background are equally distributed in the test set, causing a distribution shift. This shift results in poor predictive accuracy of ERM models, particularly driven by the inferior performance of groups that lack the spurious correlation (e.g., waterbirds on land). Focusing on individual groups (formed by combinations of class labels and spurious features), Sagawa et al. [2020] and Srivastava et al. [2020] proposed distributionally robust optimization methods to improve the worst group accuracy. Generalizing the setting, methods aim to learn invariant features across different so-called environments. These environments are partitionings of the data into subsets with varying presence of the spurious correlation (e.g., only 80% of waterbirds with water background in one environment), simulating data collection under varying circumstances. Approaches that focus on finding invariant representations of data across different training environments encompass optimization by a gradient penalty on classifier weights [Arjovsky et al., 2019] or by min-max frameworks [Krueger et al., 2021] . DecAug [Bai et al., 2021] learns disentangled features that separately capture category and context information and enhances generalization through a data augmentation method.However, these approaches have limitations in realistic scenarios. Consider for example the prediction of skin cancer (benign vs. malignant lesions) in the ISIC dataset [Codella et al., 2019] . In this dataset, almost half of the benign cases contain a colored patch that had been applied to the patient's skin in the hospital. None of the malignant cases contain a patch. Thus, half of the benign cases can be easily identified by the colored patch alone [Nauta et al., 2021] . In general, clinical markers, such as ruler measurements, surgical marks, or other artifacts, tend to become spurious features and reduce the generalization ability of ERM models [Mishra and Celebi, 2016; Pewton and Yap, 2022] . To generate training environments for invariant learning methods, it would be necessary to (i) identify all potential spurious features and (ii) artificially generate samples for the missing groups, i.e., groups with spurious features and uncorrelated class labels. Since not all artifacts may be known in advance and human annotation, especially by domain experts, is costly, the requirement to create suitable environments is unrealistic in such settings. EIIL [Creager et al., 2021] eliminates this requirement and instead infers appropriate environments from a reference model. Specifically, EIIL seeks to construct environments with maximal (anti)correlations, i.e., a feature is highly correlated with one class in one environment and another class in another environment, thereby forcing the prediction model to rely only on invariant features (causal for the prediction). However, this approach depends on a reference model that relies heavily on spurious correlations and requires estimating the distribution of assignments of instances to environments.Inspired by the observation by Le et al. [2024] that ERM induces clusters of samples with the same spurious features in the representation space and that these clusters can be exploited to find samples where the spurious correlation is missing, we adopt their strategy to infer suitable environments for invariant learning. We show that the adoption succeeds in finding counter-spurious correlation samples (conflict samples) from which we can construct environments that allow to learn invariant features. Our approach does not pose any requirements on the reference model and accordingly leaves the training process untouched. Clustering on the representation space of a trained model introduces little overhead, which makes our method an efficient approach to infer environments for invariant learning without the need for annotations.Our code is available at: https://github.com/aix-group/prusc. 
background on invariant risk minimization: 
 Empirical Risk Minimization (ERM) [Vapnik, 1991] minimizes the error across training data, using all featuresmin f :X →Y 1 n n i=1 L(f (x i ), y i ), (ERM)where L is the loss function, f maps from input space X to the output space Y, f (x i ) is the predicted value and y i is the true value.Invariant Risk Minimization [Arjovsky et al., 2019] aims to extract environment-invariant features from input data to enable consistent predictions across environments.min Φ:X → Ĥ,ω: Ĥ→Y e∈Etrain R e (ω • Φ) subject to ω ∈ arg min ω: Ĥ→ Ŷ R e (ω • Φ), ∀e ∈ E train , (IRM)Ĥ is the invariant feature space, Y is the output space, R e (Φ) is the risk under a known environment e ∈ E train . IRM learns the function f = ω • Φ where Φ learns the invariant features from multiple environments. The final prediction is made by ω based on the extracted invariant feature space Ĥ.For practical reasons, Arjovsky et al. [2019] simplify the dual-objective optimization tomin Φ:X →Y e∈Etr R e (Φ) + λ • ∇ w w=1.0 R e (w • Φ) 2 (IRMv1)3 DatasetWe use ColoredMNIST [Arjovsky et al., 2019] , a synthetic dataset constructed from MNIST where colors are strongly correlated with the class labels. The dataset is constructed as follows:1. Set label ỹ = 1 for images with digits from 0 to 4, otherwise ỹ = 0.2. The final label y is defined by randomly flipping 25% of the labels ỹ (noise level n y = 0.25).3. The color id z is defined by flipping y with probability p e and images are colored red if z = 0 and green if z = 1 (see Fig. 1a for examples)  
method: 
 We first present initial observations on the ERM representation space trained in a spurious setting, and subsequently derive our annotation-free sampling to construct environments for IRM training. A key contribution of our method is the identification of conflict samples, i.e., instances with a relation between spurious attributes and class labels that is opposite to the spurious correlation present in the training data. Our method does not pose any restrictions on the model or training process, but solely relies on properties of the learned representation space.We train an ERM model on ColoredMNIST and extract the embedding of the penultimate layer.We perform k-means clustering2 on the embedding of the training set and obtain the clusters as visualized in Fig. 1c . Interestingly, the purity of the clusters is higher w.r.t. spurious features than w.r.t. class labels (99.99 vs. 85.13, cf. Tab. 1d, p e = 0.15). Therefore, we assume that each cluster defines a spurious feature. Since spurious features have strong correlations with a particular class, a cluster's purity is usually high with respect to both, the spurious feature and the class label. For our binary classification task, we define minority cases as samples from the class with fewer samples within a cluster. These samples share similar spurious features but have class labels that are not correlated with these features. In other words, minority cases in a cluster conflict with the spurious correlations introduced in the training set (in short, conflict samples). We base our environment construction on this observation.Discussion. While our initial analysis in this paper focuses on the ColoredMNIST dataset, which is simple and has strong spurious correlations (with a spurious ratio of 0.9 in the training data), related studies have shown that spurious features are typically learned early in training or are easier for models to learn, leading models to rely on these features [Shah et al., 2020] . Therefore, we hypothesize that the observation that representations cluster stronger w.r.t. to spurious features than w.r.t. class labels also holds in scenarios with more realistic or complex spurious correlations (e.g., multiple spurious correlations as in the CelebA dataset [Liu et al., 2015] ).We use the observation that spurious correlations introduce clusters to relax the requirements of hand-crafted environments for IRM. We define the minority set D m as the union of all minority cases over all clusters. Thus, the minority set contains all conflict samples, so D m is expected to have an inverse correlation between colors and labels compared to the training set. This also means that D m contains a spurious correlation that is opposite to the one in the original training set. To obtain a more balanced sample, we additionally sample dominant cases, i.e., cases that belong to the majority class in this cluster. Overall, we sample as many dominant cases as minority cases. Our balanced set D balance is the union of minority and dominant cases, and is designed to be likely class-balanced and balanced between conflict and non-conflict samples. To account for IRM's multiple environments, we train IRM with D m and D balance . 
predictive performance: 
 We follow the training architecture and data partitioning of previous work [Arjovsky et al., 2019] , using a simple CNN and the official test set with color correlation p e = 0.9. First, we train the baseline ERM model on the training set with p e = 0.15. We compare our approach to IRM [Arjovsky et al., 2019] and DecAug [Bai et al., 2021] , both of which require annotations, and to EIIL [Creager et al., 2021] , which infers suitable environments without annotations.The results in Tab. 2a show that IRM training with our sampling approach D m and D balance is competitive to methods that require annotations (slightly better than standard IRM and slightly worse than DecAug) and on par with the annotation-free method EIIL. In contrast to EIIL, our approach neither poses restrictions on the reference model, nor does it require to train additional parameters.We analyze the effectiveness of our sampling approach on two models: ERM trained with D balance and IRM trained with the two environments D m and D balance . The results are shown in Tab. 2a (highlighted rows). Our sampling of an approximately balanced set D balance improves the accuracy of the ERM model from 17% to 65%. This result is even close to that of invariant learning methods such as IRM. The performance is also comparable to the performance of ERM trained with a hand-crafted balanced set ERM spurious-free as shown in Fig. 2b , indicating that our sampling approach is capable of constructing a truly balanced set. 
analysis and discussion: 
 ERM and IRM. By construction, the training and test distributions of ColoredMNIST have an inverse distribution of colors and labels. The environments constructed by our method may have an advantage, because we choose counter-spurious samples, thereby implicitly forming an inverse distribution. Therefore, we verify the ability to learn invariant features by varying the amount of spurious correlations in the test set. The ultimate goal of invariant learning is to have consistent performance across different environments. Fig. 2b shows the accuracy over test sets with different spurious ratios p e from 0.1 to 0.9 (official test set). ERM Dm,Dbalance and IRM Dm,Dbalance are trained with the sampled environments defined in Sec. 4. IRM requires (at least) two different environments for invariant learning (D m , D balance ) and ERM is trained with the concatenation of D m and D balance . ERM spurious-free is trained with a hand-crafted balanced subset, which does not contain spurious correlations. With the concatenation of D m and D balance , the spurious correlation learned by ERM Dm,Dbalance is opposite to the correlation in the training set, but aligns with the correlation in the test set (not as strong as in D m alone though). Accordingly, ERM Dm,Dbalance shows a clear trend following the increase of the spurious ratio, while ERM spurious-free and IRM Dm,Dbalance are largely consistent over different test environments and IRM always out-performs ERM spurious-free . This implies that ERM is sensitive to the spurious correlation (ratio) in the training and test set while IRM with multiple environments is robust against particular correlation ratios in individual environments and generally less dependent on spurious features.Environments Optimization. As Arjovsky et al. [2019] discuss, IRM cannot prevent models from learning spurious correlations, even when trained with hand-crafted environments. This is also evidenced by a performance gap between IRM and the Oracle model (cf. Tab. 2a). In Fig. 2b we evaluate another choice of hand-crafted environments, IRM pe=0.9,pe=0.1 , with a training set comprised of two environments with p e = 0.9 and p e = 0.1 respectively. We observe an increase in performance (Fig. 2b , red graph) with this choice of environments, being even on par with the Oracle (73%). This observation highlights the importance of defining a good combination of environments for invariant learning in general. It sets an interesting future research direction to investigate the explicit condition for environment construction in invariant learning, and more importantly, how to obtain the most beneficial environments without additional annotation requirements.Environment Inference without Annotations. EIIL infers suitable environments by a learnable probability distribution over the assignments of instances to environments q i (e ′ ) := q(e ′ |x i , y i ), indicating the probability that the i-th sample belongs to environment e ′ . EIIL estimates q by maximizing the regularization term in IRMv1 w.r.t a reference model Φ. The goal of this maximization is to infer environments that are governed by (supposedly spurious) features that maximally violate the invariant principle. That is, a (supposedly spurious) feature should be correlated with a particular class in one environment and with a different class in the other. Subsequent IRM training on these environments then allows to extract invariant features. For this process to be effective, the reference model Φ is required to heavily rely on spurious correlations. In practice, an ERM model in an early training stage (early stopped) is chosen for Φ. However, the availability of such a model strongly depends on the dataset and careful tuning of early-stopped criteria. The reason is that the ERM model initially focuses on spurious features during early training, resulting in significant performance disparities across subgroups, which provides a signal for EIIL to infer effective environments. However, this learning signal weakens when using a well-trained model [Creager et al., 2021] .In contrast, our approach infers environments from representation clustering (cf. Sec.4.1), without posing restrictions on the reference model, and the clustering overhead is small compared to the estimation of q. 
conclusion: 
 In this paper, we introduced a novel strategy to identify conflict samples in the training dataset in the presence of spurious correlations, based on the observation that instances tend to cluster more strongly w.r.t. spurious features than w.r.t. class labels in the learned representation space. Our approach allows for the easy construction of sub-datasets with varying spurious correlation ratios without explicit annotations, forming multiple environments suitable for invariant risk minimization.In future work, we plan to validate whether the method successfully extends to more general and complex scenarios, e.g., multi-class classification tasks, varying strength of spurious correlations, and in the presence of multiple spurious correlations. We further aim to identify the limits of our approach, i.e., boundary cases where the method fails. 
 
 Appendix: Preprint.Arjovsky et al. [2019] use two environments with an equal amount of instances in each. pe = 0.1 in one and pe = 0.2 in the other, resulting in an overall pe = 0.15 for ERM training.We choose k = 8 for all experiments in this paper. 


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: This submission proposes an invariant learning method in annotation-free environments by first identifying counter-spurious correlations to construct environments that enable learning invariant features.
main_review: main_review: **Strengths:**
- Clear and well-defined problem statement.
- Thorough discussion of potential future work.

**Comments:**
The method section could benefit from improved clarity and structure to enhance understanding and better highlight the contributions.

**Limitations:**
The results remain limited compared to baseline models.
title: title: Reducing Reliance on Spurious Correlations
main_review: main_review: **Summary:** 

Given a dataset with large spurious correlations, the work proposes a way of constructing a curated dataset, 
which reduces the reliance on spurious correlations of a model trained on this new dataset. 
The reliance on spurious correlations is compared with the same model trained on the original dataset and with 
a model trained with Invariant Risk Minimization (IRM). 
The construction of the curated dataset depends on a model trained on the original dataset. 


**Questions and Suggestions:**

Q1: You use k-means on some model embeddings to choose your training dataset composition. 
Since k-means usually does not work very well in high dimension, have you considered what one should do in the 
case of a high-dimensional embedding space? Or do you have a reason why it should still work?

Q2: In section 2, I suggest being more clear on the difference between Empirical Risk Minimization and Invariant Risk Minimization. 
For example by explicitly stating the optimization problem for Empirical Risk Minimization.   

Q3: I suggest rewriting the first sentence of your abstract: 
"Invariant learning across environments is a promising approach to improve domain generalization compared to ERM." 
It is not clear what you mean by "environments" here. Also, please use the full name and not an abbreviation the first time you mention a concept, 
in this case Empirical Risk Minimization (ERM).

Q4: I suggest rewriting the paragraph beginning with "However, these approaches require explicit..." on line 28, since it is difficult to read in its current form. 
Instead of repeating the difficulties ((i) and (ii)), I suggest beginning with your example, e.g. 
"A realistic example of spurious features ..." and only clarifying the difficulties after. 

Q5: When you write: ERM_{D_m, D_balance}, do you mean that the model was trained on a concatenation of D_m and D_balance? 
So ERM_{D_m, D_balance} is actually trained on a dataset which has the inverse spurious correlation of the original dataset, 
just less strong than in the original data? This is what it looks like in figure 2 (b). Please be more clear about it.

Q6: How would you extend this method to classification with more than two classes? If you don't know how to extend the method, it should be mentioned as a limitation. 


**Strengths:**

S1: The accuracy results presented are averages over 10 runs, which gives a good feel for the amount of variation which can be expected 
in this experimental setting. 

S2: Once the first model has been trained, the new curated dataset can be constructed automatically. There is no need for hand-picking samples.  


**Weaknesses:** 

W1: There is no clearly marked conclusion.

W2: There are several places where the writing is unclear, see Q2-Q5.   

W3: The experiments are only on a binary classification problem. The result would be stronger, 
if it could be shown that the procedure also works in a setting with more classes. 


**Justification:**

I recommend this paper to be accepted to the workshop with the added clarifications.

The main weaknesses can be fixed with rewriting.
title: title: Interesting idea but with weak experiments
main_review: main_review: The paper introduces an approach to infer environments for invariant learning without the need for explicit annotations, which is better than prior works that rely on pre-defined environments. Here are some issues.
1. The experiments focus on ColoredMNIST, which is simple. It would be beneficial to extend the analysis to more real-world datasets to validate the generalizability of the approach.
2. The method depends on strong spurious correlations, e.g., the reliance on clusters formed by spurious features. But how the method performs in scenarios with weaker or more complex spurious correlations.
3. Including a section analyzing situations where the method fails or performs suboptimally is better.

Comments: comment: We thank all reviewers for their insightful comments and have provided additional details and incorporated changes in individual responses.
This paper is an initial investigation of the method on the simple ColoredMNIST dataset. We added a discussion of why we deem the method to work in more complex scenarios as well to section 4 and plan to investigate those (and also boundary cases, where the method fails) in future work.
**Answers to Questions**

A1: “k-means usually does not work very well in high dimension” is rooted in the curse of dimensionality (most instances have near identical distance, rendering distance-based comparison meaningless). Whether this curse is present in machine learning applications (where data generation processes are not independent per dimension) has been called into question at least (see e.g., [1]). We expect classification models to learn patterns in the training data and, therefore, representations to cluster in dense (enough) regions that allow meaningful distance-based comparison.

A2 - A4: We adjusted the writing according to your suggestions.

A5: Correct. We added a more explicit elaboration to Section 6 ERM and IRM.

A6: The goal of the method is to construct (at least) 2 different environments with different levels of spurious correlations, i.e. in the paper we choose D_m containing all conflict samples and D_balance containing (probably) zero spurious correlation. The goal of IRM is to learn invariant features across different environments, so even in a k-class (k > 2) classification task, having 2 environments is enough to run the method.

We added the conclusion (Section 7)

[1] Lin, W. Y., Liu, S., Ren, C., Cheung, N. M., Li, H., & Matsushita, Y. (2021). Shell theory: A statistical model of reality. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(10), 6438-6453.
[Limitations]

While some baselines in previous studies show similar performance, our method offers distinct advantages over other approaches. 
- Compared to DecAug or IRM, our method does not require environment labels or spurious feature annotations. 
- Compared to EIIL, our approach reduces the computational complexity and eliminates the need for a heavily spurious-dependent reference model, as discussed in Section 6 (Environment Inference without Annotations). 

We highlighted our key contribution in the method section of the camera-ready version.