Causal Omnivore: Fusing Noisy Estimates of Spurious Correlations
Abstract
Spurious correlations are one of the biggest pain points for users of modern machine learning.
To handle this issue, many approaches attempt to learn features that are causally linked to the prediction variable.
Such techniques, however, suffer from various flaws—they are often prohibitively complex or based on heuristics and strong assumptions that may fail in practice.
There is no onesize-fits-all causal feature identification approach.
To address this challenge, we propose a simple way to fuse multiple noisy estimates of causal features.
Our approach treats the underlying causal structure as a latent variable and exploits recent developments in estimating latent structures without any access to ground truth.
In addition, our approach omnivorously integrates any source of causal signal.
We propose new sources, including an automated way to extract causal insights from existing ontologies or foundation models.
On multiple benchmark environmental shift datasets, our discovered features can train a model via vanilla empirical risk minimization that outperforms multiple baselines, including automated causal feature discovery techniques such as invariant risk minimization on three benchmark datasets.
1. Introduction
Standard training pipelines struggle to differentiate between features that are causally linked to the prediction target and those that are merely associations.
When measured in a new environment, such associations may no longer be predictive; they become spurious correlations.
This leads to models that are brittle: they may perform well in environments identical to those they were trained on, but fail to generalize to others.
The importance of this problem has spurred significant research in the hope of building tools to identify a set of causal features that transfer to any environment.
The holy grail is an algorithm that provably locates causal features from data without any additional signal.
In general, this is hopeless.
It is known that identifying causal relationships from observational data is not possible absent additional assumptions or knowledge.
Recent works attempt to use information from data drawn from multiple environments to discover a common set of causal features (Peters et al., 2016; Arjovsky et al., 2019; Krueger et al., 2021).
These techniques are promising but suffer from multiple flaws.
For example, invariant risk minimization (IRM) (Arjovsky et al., 2019) requires a vast number of environments to be guaranteed to learn causal features—and may perform worse than vanilla empirical risk minimization when this fails to happen (Rosenfeld et al., 2021).
Experimentally, none of these techniques are known to work in all cases (Dranker et al., 2021).
Furthermore, by targeting a full end-to-end solution usable in any scenario, they ignore the presence of easily-accessible sources of causal knowledge in many specific scenarios.
Given the substantial challenge of a single technique that always finds spurious correlations, an alternative is to build an omnivorous method that can flexibly take advantage of any kind of causal signal.
We propose CO MNIVORE , a system that takes a step towards satisfying this property by two key aspects.
First , it enables the use of multiple sources to generate potential candidate feature sets.
In particular, it allows for simple human-based specifications.
When not available, it enables for simple ways to automate such specifications.
Second , it extracts causal features from the resulting candidate feature sets by combining the outputs of multiple causal estimation approaches.
Using principles similar to those in weak supervision (Ratner et al., 2018), it estimates the reliability of each causal estimate output, without ground truth .
It then provides a higher-quality fused set of estimated features.
CO MNIVORE is compatible with any pre-existing approach to causal feature estimation.
It has the benefit of simplicity— not requiring any specialized loss functions or difficult bilevel optimization.
Effectively, CO MNIVORE simply asks as many sources of signal as possible for causal information, weights this information, and trains a downstream model on the detected features with vanilla empirical risk minimization.
We validate CO MNIVORE empirically, showing that it outperforms competing end-to-end baselines like IRM, while improving on ERM by 37.9% on three benchmark datasets.
2. Background and Problem Setting
We first describe some of the tools we will use and then detail the problem setting.
Identifying Causal Features Discovering causal features is an active area of research under which approach can take many forms.
We briefly describe the approach we use.
First, an important problem in causal inference is learning causal structures from observational data, interventions, structural assumptions, or heuristics.
Naturally, such approaches do not work in every setting; violations of their underlying assumptions can be thought of as noise.
For instance, Peter-Clark (PC) (Spirtes et al., 2000), FCI (Spirtes et al., 2000), and Greedy Equivalence Search (Chickering, 2003) assume the absence of certain conditions on the latent confounders between features.
Grow-Shrink (GS) (Margaritis, 2003), Incremental Association Markov Blanket (IAMB) (Tsamardinos et al., 2003), Interleaved IAMB (Yaramakala & Margaritis, 2005), and Exact Search (Silander & Myllymaki, 2012) require the underlying model to have a certain Bayesian structure.
More recent optimization-based methods (Zheng et al., 2018) (Zheng et al., 2020) are limited by optimization constraints.
These assumptions thus limit their accuracy when applied to complex and high-dimensional data.
Weak Supervision Weak supervision is a set of techniques that are used to construct labeled training sets (Ratner et al., 2016; 2018; Fu et al., 2020) from unlabeled training data.
The idea is that even though no labels are available, multiple noisy estimates of each label are observed.
These are the outputs of labeling functions λ 1 , . . .
, λ m .
The challenge is to determine the reliability of these functions and to use this information to fuse their outputs into a pseudolabel of higher quality than each of their individual votes.
We use similar principles to fuse noisy causal estimates.
Problem Setting We have access to a dataset { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } of samples drawn from some distribution D .
Potentially, we have access to additional training distribution sets from D 2 , D 3 , . . .
, D k .
We refer to these as the k training environments.
Different environments represent the different data collection settings (e.g., hospital, time periods, countries, demographics).
Our goal is to learn a model f ( x ) that performs well in new scenarios.
Typically this means that it generalizes well to test distributions D test , where D test and D train are nonoverlapping (i.e., D test ∪ D train = ∅ ).
However, we might also be interested in f ( x ) performing well in subpopulation shift scenarios.
In this case, our goal is to maximize f ( x ) ’s performance across all domains seen during training (i.e., D test ⊆
D train ), but the proportions of samples from each domain can change.
We also assume we have access to pretrained model or foundation model (FM) embeddings.
These embeddings are the outputs of a mapping g : X → Z from input space to latent embedding space.
This mapping is fixed and obtained from off-the-shelf models.
3. Approach
Algorithm
Input: Training dataset drawn from a distribution D orig = { ( x 1 , y 1 ) , . .
, ( x n , y n ) } , causal feature selection functions λ , embedding mapping f
Generate candidate sets C for C i ∈ C do for λ i ∈ λ do Generate causal structure estimates G i for G i ∈ G do if fusing method == Graph-based WS then
Compute λ j ’s weight w j while Annealing iteration do
Minimize weighted objective as in (1) to get ˆ
G collect all z i ∈
z that has a causal edge to label node if fusing method ==
Vanilla WS then for z
i in z do
Construct label matrix
L
Get causal predictions from WS system
ˆ Y
i ±
Result:
Singleton set of configurations .
We present our method to train models robust to spurious correlations: CO MNIVORE .
At high level, we break up the task into two parts.
Our first goal is to obtain many sources of potentially causal features and group them into a set of distinct feature transformations.
These might include the raw features, embeddings from pretrained models like a ResNet or foundation models like CLIP, the result from performing a manually-chosen transformation/augmentation on the dataset, or their combinations.
We call the resulting sets of feature transformations the candidate sets .
Directly relying on the features in these candidate sets may not be sufficient, however—they may also be affected by spurious correlations.
To further refine our estimated features, we run a suite of causal estimation approaches for each set.
We refer to these as the causal feature selection functions .
We will estimate the reliability of each of the selection functions and produce an improved overall estimate of the causal features.
Finally, we combine the resulting estimated features from each of the candidate sets and train the end model.
Generating candidate sets Our first task is to generate candidate sets { C 1 , C 2 , . .
, C b } .
These are transformed versions of the original features that ideally have some reduced spuriousness.
Potential choices of these include,
• Embeddings from off-the-shelf models, such as foundation models like CLIP, • Existing end-to-end invariant feature construction methods, like IRM, when suitable, • Manually-selected transformations/augmentations, • Automated transformations/augmentations.
We describe the latter two possibilities.
First, we observe that humans can often identify causal features with ease.
As a running example, consider the Waterbirds dataset (Wah et al., 2011)
(Zhou et al., 2018).
The goal is to classify birds as being water-based or terrestrial, and the background in the images of these birds (bodies of water versus land) acts as a spurious feature.
This is challenging for training algorithms to discern, but nearly trivial for humans.
It is easy to encode this human insight into transformations that can be built with off-the-shelf tools.
In this case, running a standard segmentation algorithm to isolate the bird acts as such a transformation, as shown in Figure 2.
Manually-selected transformations help translate easilyacquired human insights into high-quality candidate sets.
However, we do not always have access to such information.
In Appendix C, we show how to automate the process of encoding human insights into causal versus spurious features.
This will enable us to get the best-of-both worlds.
The term Candidate Set loosely applies to any result of transformations that can possibly reduce spuriousness in the original sample.
When using manually-selected transformation, this term refers to the transformed samples.
When using invariant feature construction like IRM, Candidate
Set refers to the latent feature extracted from the penultimate layer of an IRM model.
For the first case, we then get the latent representation using foundation model or any pretrained model θ : C → { z 11 , . . .
, z 1 d } , . . , { z n 1 , . . . , z nd } . Generating causal feature selection functions Next, we use the suite of causal inference algorithms listed in Section 2 to obtain the estimated causal structures for each candidate set.
These algorithms take the sets of features paired with the corresponding labels { ( { z 11 , . . .
, z 1 d } , y 1 ) , . . .
, ( { z n 1 , . .
, z nd } , y n ) } for each available training environment and output the estimates of causal structures that govern the relationships among individual features and with the label.
These causal structures are represented in form of DAGs in G .
Formally, given m causal algorithms, the output of each algorithm λ a is described by
[Equation]
One challenge is that such algorithms often have high complexity, sometimes superexponential in the number of fea- tures.
We use a simple way to address this difficulty.
We map the features into a lower-dimensional space, perform estimation in this space, and then return to the original space.
We use Feature Agglomeration (Steinbach et al., 2000), but our method is compatible with any dimensionality reduction technique.
At the end of this step, we have m DAGs { G 1 , . . .
, G m } per candidate set and environment.
Our approach treats the causal inference algorithms like labeling functions in weak supervision (section 2).
Since these algorithms assume certain conditions or heuristics, we expect their causal estimate outputs to be noisy.
Weak supervision nicely complements our need to combine these noisy estimates since it fuses noisy labels to build a labeled dataset.
Fusing noisy causal estimates Our final task is to obtain a fused estimate G from the DAGs.
Our goal is to obtain a better set of features compared to the noisy DAG estimates { G 1 , . . .
, G m } .
We employ two weak supervision-based techniques to combine the G a ’s into ˆ G :
Graph-based Weak Supervision
With this approach
, we learn weights w a for each estimate G a .
These weights correspond to average distances to a true G ∗ which we do not observe.
To estimate the weights, we use the algorithm in (Shin et al., 2022), described below.
We embed the graphs into R d , producing an embedding r ( G a ) . We set up the following system of equations for triplets ( a, b, c ) chosen from { 1 , . . . , m } : ∥ r ( G a ) − r ( G b ) ∥ 2 = ∥ r ( G a ) − r
( G ∗ ) ∥ 2 + ∥ g ( G b ) − r ( G ∗ ) ∥∥ r
( G a ) − r
G c ) ∥ 2 = ∥ r ( G a ) − r
( G ∗ )
∥ 2 + ∥ g ( G c ) − r
( G ∗ ) ∥∥ r ( G b ) − r ( G c ) ∥ 2 = ∥ r ( G b ) − r
( G ∗ )
∥ 2 + ∥
r ( G c ) − r ( G ∗ ) ∥ 2 .
To obtain ∥ r ( G a ) − r ( G ∗ ) ∥ 2 , we add the first two equations, subtract the third, and divide by two.
This is an estimate of the average distance between our (embeddings of) graphs; the weights w are just the reciprocals, so that w a = 1 / ( ∥ r ( G a ) − r
( G ∗ ) ∥ 2 ) .
Once we have estimated ˆ w a , we perform the following optimization,
[Equation]
Note that we compute the mean in the original DAG space, not in the embedding space.
We use simulated annealing (Kirkpatrick et al., 1983), an iterative global search optimization method, to obtain ˆ G .
Next we take all features z i
that have a causal path in ˆ G to the label node as the causal feature subset.
Vanilla Weak Supervision.
Alternatively, instead of searching for the best overall causal structure, we can try
O r i g i n a l T r a n s f o r m e d WILDS Waterbirds O r i g i n a l
T r a n s f o r m e d
WILDS CelebA to solve a perhaps more manageable problem: is feature z i causally related to the label y ?
We treat this problem for each feature
z i as a simple classification problem Y i ∈ {± 1 } where +1 means a causal edge present between z i and y in DAG G j , 0 indicates no relationship, and − 1 an anti-causal edge present between z i and y .
Inspired by (Fu et al., 2020), for each z
i , we first construct a k × m label matrix L , where k is the number of environments we have access to and m is the number of causal estimation functions.
Note that L is constructed for each z i in each candidate set separately.
We encode the predictions output by each estimation algorithm into L and pass it as input to any weak supervision approach, e.g., (Ratner et al., 2018; 2019; Fu et al., 2020).
4. Experiments
This section validates the following claims about CO MNIVORE :
• Performance (Section 4.1): We show that CO MNIVORE outperforms baseline end-to-end approaches (IRM (Arjovsky et al., 2019) and REx (Krueger et al., 2021)) on unseen environment D new with comparable performance on the original environment D orig on both subpopulation shift and domain generalization datasets.
• Ablations (Section 4.2): The source of CO MNIVORE ’s lift on D new while retaining good performance on D orig is all components of the approach.
We show this by comparing CO MNIVORE with training vanilla ERM and baseline approaches using the foundation model embeddings of the candidate sets.
Datasets We evaluate CO MNIVORE on three datasets in the
WILDS benchmark (Koh et al., 2020).
In subpopulation shift, we use Waterbirds , which combines bird images from the CUB dataset (Wah et al., 2011) with backgrounds from the Places dataset (Zhou et al., 2018); and CelebA celebrity faces dataset (Liu et al., 2015);
In domain generalization, we use Camelyon17 (Bándi et al., 2019), which task is to identify tumor in medical images.
We also evaluate on
[Table] [Table] ColorMNIST , where spurious correlations between digits and color are artificially created, similar to the synthetic task used in IRM and REx.
We list the complete dataset details in Appendix F.1.
Pre-trained embeddings We use pre-trained CLIP embeddings (Radford et al., 2021).
4.1. Performance Comparisons
We compare CO MNIVORE with IRM and REx, measuring accuracy on the original train distribution D orig and the new test distribution D new .
For CO MNIVORE , we train a simple 2-layer MLP using ERM on the sets of causal features acquired using both graph-based WS and vanilla
WS.
For IRM and REx, we experiment with 2-layer MLPs using two choices of feature extractors: CLIP and ResNet50.
The latter follows the choice of architecture used in the
WILDS benchmark (Koh et al., 2020).
We report the best of the two results.
Table 1 shows the results.
CO MNIVORE outperforms IRM and REx on D new across all datasets.
For D orig , CO MNIVORE ’s performance is comparable to the best baseline on Waterbirds (by 4.6%) and achieved the best accuracies on ColorMNIST (tie with IRM), CelebA and Camelyon17.
This reflects our method’s ability to ingest and refine a large number of causal features.
We also note that we construct our implementation of ColorMNIST.
In contrast with IRM and REx, we do not collapse the classes and directly use the digits 0 vs 1 .
In IRM’s version of ColorMNIST, IRM achieves 70.8% accuracy on
D orig and 66.9% D new ; and in REx’s implementation, REx achieves 71.5% accuracy on D orig and 68.7% on D new .
Our main experimental table contains the values we obtained on our version of the dataset.
As an additional point of comparison, we note that EIIL (Creager et al., 2021) achieves a 69.7% accuracy on D orig and 78.8% D new in Waterbirds dataset.
4.2. Ablations
We train models using vanilla ERM on the extracted features of the candidate sets and original images, without performing the causal estimate step.
We report the results in table 2.
On D new , CO MNIVORE outperforms both vanilla ERM trained using original images and candidate sets.
CO MNIVORE ’s performance on D orig is within relatively comparable accuracy with the best of vanilla ERM on Waterbirds,
CelebA, and ColorMNIST (by 5.6%, 5.2%, and 0.3%) and performed the best on Camelyon17.
This result demonstrates that CO MNIVORE ’s performance lift is produced by both candidate set generation and causal feature selection.
Appendix
We discuss related work (Appendix A), provide a glossary containing key terminology (Appendix B) and algorithm table (Appendix ?? ), introduce additional details into our algorithm claims (Appendix E), then give extra experimental details and results (Appendix F).
We also provide theoretical analysis considering simple setups where estimates of causal features are useful (Appendix D).
Finally, we show a simple approach to automate the process of encoding human insights into causal versus spurious features mentioned in section 3 (Appendix C).
A. Related Work
This section presents discussion of related work and connections to our work.
Invariant learning methods such as IRM (Arjovsky et al., 2019), REx (Krueger et al., 2021), and a multitude of similar works (Ahmed et al., 2020; Parascandolo et al., 2020; Ahuja et al., 2020; Creager et al., 2021) share a similar goal with our work.
The aim is finding feature representations that are invariant across domains or environments.
We can think of these invariant features as similar to our goal causal feature subset.
This is achieved mainly by minimizing specialized loss functions.
IRM and REx minimize the sum of loss terms across environments and thus require environment labels.
Environment Inference for Invariant Learning (EIIL) (Creager et al., 2021) and Predictive Group Invariance (PGI) (Ahmed et al., 2020) train an initial ERM model to infer environment labels and later on train another model with invariant learning objectives.
In contrast, CO MNIVORE estimates the causal features in-prior to training a model and thus circumvents the need for any specialized loss functions.
Improving robustness to spurious correlations and distribution shift is an extensive line of work that can be divided into two aspects, based on whether access to group/domain information is given or not.
In the line that requires group information a priori, one popular work is group distributionally robust optimization (GDRO) (Sagawa et al., 2019), which divides the data into explicit groups and then trains them to directly minimize the worst group-level error among these groups.
Similarly, Fish (Shi et al., 2021) and Inter-environment Gradient Alignment (IGA) (Koyama & Yamaguchi, 2020) aim to improve domain generalization performance by maximizing inter-domain gradient terms in their loss functions.
More similar to our approach are methods that do not assume group information at training time.
For instance, distributionally robust optimization (DRO) minimizes worst group loss within a ball centered around the training distribution (Ben-Tal et al., 2013; Duchi et al., 2016).
More recent methods (Nam et al., 2020; Liu et al., 2021; Zhang et al., 2022) train two ERM models: the first one is to estimate which data points play a crucial role in their subsequent steps (e.g., which points belong to minority groups, which samples come from the same class but has different spurious features, etc.).
Next, such methods train another ERM model with specialized objectives (e.g., to up-weight minority groups, using contrastive loss to learn invariant features, etc.).
Note that all of these works are compatible with our approach as well.
Causal inference algorithms (Spirtes et al., 2000; Margaritis, 2003; Chickering, 2003; Tsamardinos et al., 2003;
Yaramakala & Margaritis, 2005; Silander & Myllymaki, 2012; Zheng et al., 2018;?) seek to discover the structure that governs relationship between set of features in the data.
Ideally, for our purpose, if we feed the sets of features and labels into these algorithms, we hope to be able to extract the features that have a causal link to the label.
Unfortunately, this problem is statistically and computationally hard (Chickering, 1996; Chickering et al., 2004).
As a result, these methods resort to local heuristics and assumptions, thus limiting their accuracy when applied to complex high-dimensional data.
Our approach fuses these noisy estimates of causal structures to get the estimated set of causal features on which training an end model will be robust to spurious correlation and domain shift.
Weak supervision is a set of techniques that use noisy sources of labels to construct labeled training sets without access to ground truth labels (Ratner et al., 2016; 2018; Fu et al., 2020).
This technique is vastly explored for binary classification problems.
Recently, (Shin et al., 2022) enables weak supervision over broader sets of problems, which also serves as a basis for our graph-based weak supervision fusion method.
B. Glossary
The glossary is given in Table 3 below.
[Table]
C. Automating Candidate Set Transformations
There are many situations where a human user may not be aware of a spurious pattern in the data.
Had CelebA (Liu et al., 2015) not contained the appropriate annotation, a machine learning practitioner wishing to predict hair colors may have overlooked this feature’s spurious correlation with gender.
More generally, it is not always certain that users may have sufficient domain expertise to design hand-crafted transformations for candidate feature sets.
We describe a simple method to fully automate the candidate set transformation.
An illustrative example is provided for the
Waterbirds dataset (Sagawa et al., 2019).
To discover patterns in the training images, we generate a caption for each image using a CLIP-based captioner (Mokady et al., 2021), then extract captions’ keywords.
We search each label ( waterbird and landbird ) on Wikipedia (Wikipedia contributors) and extract the keywords from the first resulting article’s introduction section.
Spurious words are considered to be the top m most common caption keywords that do not occur in the article keywords.
We next break each training image into non-overlapping p × p patches.
If zero-shot CLIP (Radford et al., 2021) predicts any of the spurious words in a given patch with confidence greater than τ , the patch is covered.
A resulting image from this process is depicted in Figure 3. out
We note that there are many potential ways to fully automate the candidate set transforms by taking advantage of ontologies and pretrained models.
The proposed procedure requires only the label names and some form of task description, for instance that the dataset is comprised of images, allowing it to also be implemented in other settings outside the specific example described above.
C.1. Result with Automated Candidate Set
We evaluate CO MNIVORE when using the candidate set built from automated transformations described previously.
We use a patch size p of 75 , spurious word list length m of 10 and threshold τ of 0 .
6 on Waterbirds (Sagawa et al., 2019).
As shown in Table 4, this configuration yields a similar performance to CO MNIVORE with human-supervised LFs for D orig .
AutoLF’s score of 66 .
2 on D new also improves by 7 . 6% on the baseline D new Waterbirds results of Table 4.
We observed that the result is typically sensitive to the choices of threshold.
We hypothesize that expanding the approach to larger ontologies will further close the gap to manual performance.
[Table]
D. Theoretical Analysis
While the idea of fusing multiple noisy causal estimates is intuitively appealing, it is not clear whether we can expect this to work and under what conditions.
This section is dedicated to showing how, in certain simple scenarios, the resulting estimates of the causal structure are useful.
Setup and Noise Distributions
We will consider two scenarios.
In the first, we have some candidate set of features z 1 , . . .
, z k , and we are interested in determining whether z i is causal for output y .
In other words, we are predicting a set
D ⊆ { 1 , . . . , k } .
In the second scenario, we additionally take into account the causal structure, i.e., a directed acyclic graph
G over the nodes z 1 , . . .
, z k .
We denote the causal recovery techniques by λ a for a = 1 , . . .
, m , so that λ a : R k × n → 2 { 1 ,...,k } in the first case, or λ a : R k
→ G in the second case.
Recall that G is the set of DAGs on k nodes.
Finally, we have a access to k environments, where for each environment, we observe n samples of the features z 1 , . . .
, z k .
Each causal estimation function λ a may fail in a variety of ways; this may be because the underlying assumptions are not met, or because of noise, or for some other reason.
The outcome of such noise is either a predicted set D not equal to the true D ∗ , or a predicted graph G not equal to G ∗ .
We will model the noise in the estimation approaches with the following model inspired by (Shin et al., 2022):
[Equation]
where Z is the normalizing partition function, d H is the Hamming distance, and θ =
[ θ 1 , . .
, θ m ] T is a vector of parameters.
For sets, d H is simply the size of the symmetric diference.
We can also operate in the second scenario by switching D ∗ to
G ∗ on both sides of (2).
In this case, the Hamming distance over graphs counts the number of differences in edges.
Note how the model works: if θ a is large, then the probability mass is significantly reduced even for a small distance between the prediction and the true causal model; this implies that the quality of the approach λ a is high.
If θ a is small, then even a large distance does not significantly reduce the probability, so λ a is low-quality.
Note that richer models are possible; for example, we could replace the graph Hamming distance with the interventional distance as in (Peters & Bühlmann, 2015).
The advantage of the exponential family model above is that it is tractable without requiring significant specifications on the underlying causal model.
[Figure]
Estimating Qualities and Performing Fusion The main challenge is how to estimate θ 1 , . .
, θ m .
The two techniques in
Algorithm 1 work for these two scenarios.
We show that the second approach has consistent estimation of θ in terms of the number of environments k .
Theorem D.1.
Suppose λ a , . . .
, λ m are distributed according to (2) and we have access to k training environments.
Using vanilla weak supervision to estimate ˆ θ , we have that E [ ∥ ˆ θ − θ ∗ ∥ ] ≤ O (1 / √ k ) .
This implies that, given sufficiently many environments, the weights we learn for use in Algorithm 1 reflect the underlying quality of the causal estimation functions.
D.1. Theoretical Analysis Experiment
We evaluate a key claim from our theoretical characterization outlined above.
We expect that as the values of the θ parameter vector are larger, the quality of the causal estimation functions improves, and that our resulting algorithm produces causal features that perform well in a new environment.
We validate this notion using a synthetic dataset reflecting a simple linear regression setup.
In the original environment
D orig , the label is a function of all of the features, while in the new environment D new , the label is a function of only a subset of features—and the remaining features have a significantly different distribution from their counterparts in D orig .
The results are shown in Fig.
We swept the average magnitude of θ , used our two approaches based on graph-based WS (left) and vanilla WS (right), trained a linear regression end model, and measured the root mean squared error (RMSE).
As expected, using all features results in very good error in D orig (green curve) and very poor error in D new (red curve).
Applying our causal approaches resulted in nearly-as-good D orig performance (blue curve), and vastly improved D new performance (yellow curve).
As we hoped, the error of this curve generally decreases with improved quality estimates (i.e., larger θ ).
Additionally, we note that the vanilla WS approach, while slightly noisier, produces a smaller final error.
This suggests a closer analysis of the two approaches would be useful.
E. Additional Algorithmic Details E.1. Projection to Low-Dimensional Space
We use sklearn’s
(Buitinck et al., 2013) implementation of feature agglomeration, an unsupervised dimensionality reduction technique that uses agglomerative clustering to group together features that look very similar (Steinbach et al., 2000).
Our method also works with any dimensionality reduction technique like PCA (F.R.S., 1901).
We chose feature agglomeration because it provides an automatic mapping from higher to lower dimensional space, thus eliminating the need to manually set thresholds for the components.
E.2. L Matrix Computation
In Vanilla WS fusion method, for each z i , we construct a k × m label matrix L , where k is the number of environments we have access to and m is the number of causal estimation functions.
Formally, let A ( b,j ) be the d × d adjacency matrix representation of G j from b th environment, and the label node is the d th node in G j , each entry of L is defined by: L bj =

 1 , if A ( b,j ) id =0 , if A ( b,j ) id =− 1 , if A ( b,j ) id = −, b ∈ 1 . . . k, j ∈ 1 . . . m, i ∈ 1 , . . . d −(3)
F. Experimental Details F.1. Dataset Details
Table 5 shows details on train/dev/test splits for each dataset, as well as the number of smallest group samples in distribution shift datasets.
All splits are following the default provided by WILDS benchmark (Koh et al., 2020).
[Table]
We evaluate CO MNIVORE on three datasets in the WILDS benchmark (Koh et al., 2020).
In subpopulation shift,
Waterbirds ’s spurious correlation occurs between label Y = { landbirds , waterbirds } and background attribute A = { land , water } .
; CelebA celebrity faces dataset (Liu et al., 2015) has spurious correlation between the hair color label
Y = { blond , dark } and the gender attribute A
= { male , female } ; In domain generalization, Camelyon17 (Bándi et al., 2019)’s domain shift is the different hospitals where training and test samples are collected.
F.2. ColorMNIST
We construct our implementation of synthetic colored version of the MNIST dataset (Deng, 2012).
In contrast with IRM and
REx, we do not collapse the classes (i.e., y = 0 for digits 0 − 4 and y = 1 for digits 5 − 9 ).
Instead, we directly use the digits 0 vs 1 .
More specifically, we take MNIST subsets of digits 0 and 1 , assign a color to each digit, and flip the color on
D test .
We use the default train/dev/test splits provided by MNIST.
We also note that in IRM’s version of ColorMNIST, IRM achieves 70.8% accuracy on D orig and 66.9% D new ; and in REx’s implementation, REx achieves 71.5% accuracy on D orig and 68.7% on D new .
Our main experimental table contains the values we obtained on our version of the dataset.
F.3. Manual Candidate Sets
This section details the construction of manual candidate sets used in our experiments.
The original and transformed images are shown in figure 5.
For Waterbirds and celebA, segmentation is done using Pytorch’s off-the-shelf DeepLabV3 model (Chen et al., 2017).
For Camelyon17, the candidate set generated is the gaussian blurred version of the original images, generated using PIL’s Gaussian Blur filter (Umesh, 2012).
For ColorMNIST, the candidate set used is the original images and the black and white version.
Table 6 details the candidate sets used for best numbers reported.
F.4. Hyperparameters and Model Selection F.4.1. E ND C LASSIFICATION M ODEL
Experiments were done three times, and we reported an average of three runs.
Models are selected based on the best performance on the dev set (and OOD dev set for Camelyon17).
Experiments are conducted using two NVIDIA RTX
A
GPUs.
For all datasets, we train a 2-layer MLP with 512 hidden dimensions.
Best hyperparameters are reported in table 7.
All models are trained using 0.9 momentum and 0.1 l 2 regularization penalty.
Training epochs are set until 500, and we picked the checkpoint with the
[Table] [Table]
O r i g i n a l S e g m e n t
[Table]
r i g i n a l G a u ss i a n B l u
[Table]
r i g i n a l S e g m e n t + C r o p B o tt o m S e g m e n t +
C r o p F a c e
[Table]
[Table]
Baseline Implementations
IRM A ResNet50 is trained using the IRM implementation from the WILDS benchmark (Koh et al., 2020).
Reported results are averaged across three runs, using the hyperparameters yielding highest average accuracy on the dev set in any epoch.
In real-world applications, the best strategy would often be to select a model that balances somewhere in between maximizing average and worst-group accuracy as determined by domain experts.
In this work, however, we choose to report epochs that maximize average accuracy without regard to worst-group accuracy in order to establish a uniform, unbiased method to select the “best" hyperparameters and performance metrics.
The maximum possible number of epochs is 200.
Momentum of 0 .
9 , IRM λ of 100 and penalty annealing iterations ofare used for all datasets.
Learning rate and batch size are reported in Table 8.
[Table]
We do not report the hyperparameters for Camelyon, because we report IRM result on Cameyon based on the WILDS leaderboard (WILDS).
REx We train a 2-layer MLP with 256 hidden dimensions using REx implementation for all datasets.
Experiments were done three times, and we reported an average of three runs.
The maximum possible number of epochs is 500, and we picked the checkpoint with the highest performance on dev set (and OOD dev set for Camelyon17).
Penalty annealing iterations of 100 are used for all the datasets.
Other best hyperparameters are reported in Table 9.
[Table]