E
XCHANGING L ESSONS B ETWEEN A LGORITHMIC F AIRNESS AND D OMAIN G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Standard learning approaches are designed to perform well on average for the data distribution available at training time.
Developing learning approaches that are not overly sensitive to the training distribution is central to research on domainor out-of-distribution generalization, robust optimization and fairness.
In this work we focus on links between research on domain generalization and algorithmic fairness—where performance under a distinct but related test distributions is studied—and show how the two ﬁelds can be mutually beneﬁcial.
While domain generalization methods typically rely on knowledge of disjoint “domains” or “environments”, “sensitive” label information indicating which demographic groups are at risk of discrimination is often used in the fairness literature.
Drawing inspiration from recent fairness approaches that improve worst-case performance without knowledge of sensitive groups, we propose a novel domain generalization method that handles the more realistic scenario where environment partitions are not provided.
We then show theoretically and empirically how different partitioning schemes can lead to increased or decreased generalization performance, enabling us to outperform Invariant Risk Minimization with handcrafted environments in multiple cases.
We also show how a re-interpretation of IRMv allows us for the ﬁrst time to directly optimize a common fairness criterion, groupsufﬁciency and thereby improve performance on a fair prediction task
I NTRODUCTION
Machine learning achieves super-human performance on many tasks when the test data is drawn from the same distribution as the training data.
However, when the two distributions differ, model performance can severely degrade to even below chance predictions (Geirhos et al., 2020).
Tiny perturbations can derail classiﬁers, as shown by adversarial examples (Szegedy et al., 2014) and common-corruptions in image classiﬁcation (Hendrycks & Dietterich, 2019).
Even new test sets collected from the same data acquisition pipeline induce distribution shifts that signiﬁcantly harm performance (Recht et al., 2019; Engstrom et al., 2020).
Many approaches have been proposed to overcome model brittleness in the face of input distribution changes.
Robust optimization aims to achieve good performance on any distribution close to the training distribution (Goodfellow et al., 2015; Duchi et al., 2016; Madry et al., 2018).
Domain generalization on the other hand tries to go one step further, to generalize to distributions potentially far away from the training distribution.
The ﬁeld of algorithmic fairness meanwhile primarily focuses on developing metrics to track and mitigate performance differences between different sub-populations or across similar individuals (Dwork et al., 2012; Corbett-Davies & Goel, 2018; Chouldechova & Roth, 2018).
Like domain generalization, evaluation using data related to but distinct from the training set is needed to characterize model failure.
These evaluations are curated through the design of audits , which play a central role in revealing unfair algorithmic decision making (Buolamwini & Gebru, 2018; Obermeyer et al., 2019).
While the ultimate goals of domain generalization and algorithmic fairness are closely aligned, little research has focused on their similarities and how they can inform each other constructively.
One of their main common goals can be characterized as:
Learning algorithms robust to changes across domains or population groups.
[Table]
not only allows models to generalize to different and unobserved but related
- butions, also mitigates unequal treatment of individuals solely based on
In this work we explore independently developed concepts from the domain generalization and fairness literatures and exchange lessons between them to motivate new methodology for both ﬁelds.
Inspired by fairness approaches for unknown group memberships (Kim et al., 2019; H´ebert-Johnson et al., 2018; Lahoti et al., 2020), we develop a new domain generalization method that does not require domain identiﬁers and yet can outperform manual speciﬁcation of domains (Table 1).
Leveraging domain generalization insights in a fairness context, we show the regularizer from IRMv1 (Arjovsky et al., 2019) optimizes a fairness criterion termed “group-sufﬁciency” which for the ﬁrst time enables us to explicitly optimize this criterion for non-convex losses in fair classiﬁcation
The following contributions show how lessons can be exchanged from the two ﬁelds:
• We draw several connections between the goals of domain generalization and those of algorithmic fairness, suggesting fruitful research directions in both ﬁelds (Section 2).
• Drawing inspiration from recent methods on inferring worst-case sensitive groups from data, we propose a novel domain generalization algorithm—Environment Inference for Invariant Learning (EIIL)—for cases where training data does not include environment partition labels (Section 3).
Our method outperforms IRM on the domain generalization benchmark ColorMNIST without access to environment labels (Section 4).
• We also show that IRM, originally developed for domain generalization tasks, affords a differentiable regularizer for the fairness notion of group sufﬁciency, which was previously hard to optimize for non-convex losses.
On a variant of the UCI Adult dataset where confounding bias is introduced, we leverage this insight with our method EIIL to improve group sufﬁciency without knowledge of sensitive groups, ultimately improving generalization performance for large distribution shifts compared with a baseline robust optimization method (Section 4).
• We characterize both theoretically and empirically the limitations of our proposed method, concluding that while EIIL can correct a baseline ERM solution that uses a spurious feature or “shortcut” for prediction, it is not suitable for all settings (Sections 3 and 4).
D OMAIN G ENERALIZATION AND A LGORITHMIC F AIRNESS
Here we lay out some connections between the two ﬁelds.
Table 2 provides a high-level comparison of the objectives and assumptions of several relevant methods.
Loosely speaking, recent approaches from both areas share the goal of matching some chosen statistic across a conditioning variable e , representing sensitive group membership in algorithmic fairness or an environment/domain indicator in domain generalization.
The statistic in question informs the learning objective for the resulting model, and is motivated differently in each case.
In domain generalization, learning is informed by the properties of the test distribution where good generalization should be achieved.
In algorithmic fairness the choice of statistic is motivated by a context-speciﬁc fairness notion , that likewise encourages a particular solution that achieves “fair” outcomes (Chouldechova & Roth, 2018).
Empty spaces in Table 2 suggest areas for future work, and bold-faced entries suggest connections we show in this paper.
[Table]
Notation Let X be the input space, E the set of environments (a.k.a.
“domains”), Y the target space.
Let x, y, e ∼ p obs ( x, y, e ) be observational data, with x ∈ X , y ∈ Y , and e ∈ E .
H denotes a representation space, from which a classiﬁer w ◦ Φ (that maps to the pre-image of ∆( Y ) via a linear map w ) can be applied.
Φ :
X → H denotes the parameterized mapping or “model” that we optimize.
We refer to Φ( x ) ∈ H as the “representation” of example x .
ˆ y ∈ Y denotes a hard prediction derived from the classiﬁer by stochastic sampling or probability thresholding.
(cid:96) : H × Y → R denotes the scalar loss, which guides the learning.
The empirical risk minimization (ERM) solution is found by minimizing the global risk, expressed as the expected loss over the observational distribution:
[Equation]
Domain Generalization Domain generalization is concerned with achieving low error rates on unseen test distributions.
One way to achieve domain generalization is by casting it as a robust optimization problem (Ben-Tal et al., 2009) where one aims to minimize the worst-case loss for every subset of the training set, or other well-deﬁned perturbation sets around the data (Duchi et al., 2016;
Madry et al., 2018).
Other approaches tackle domain generalization by adversarially learning representations invariant (Zhang et al., 2017; Hoffman et al., 2018; Ganin et al., 2016) or conditionally invariant (Li et al., 2018) to the environment.
Distributionally Robust Optimization (DRO) (Duchi et al., 2016), seeks good performance for all nearby distributions by minimizing the worst-case loss sup q E q [ (cid:96) ] s.t. D ( q || p ) < (cid:15) , where D denotes similarity between two distributions (e.g. χ 2 divergence) and (cid:15) is a hyperparameter.
The objective can be computed as an expectation over p via per-example importance weights γ i
= q ( x i ,y i ) p ( x i ,y i ) .
Recently, Invariant Learning approaches such as Invariant Risk Minimization (IRM) (Arjovsky et al., 2019) and Risk Extrapolation (REx) (Krueger et al., 2020) were proposed to overcome the limitations of domain invariant representation learning (Zhao et al., 2019) by discovering invariant relationships between inputs and targets across domains.
Invariance serves as a proxy for causality, as features representing “causes” of target labels rather than effects will generalize well under intervention.
In IRM, a representation Φ( x ) is learned that performs optimally within each environment—and is thus invariant to the choice of environment e ∈ E —with the ultimate goal of generalizing to an unknown test dataset p ( x, y | e test ) .
Because optimal classiﬁers under standard loss functions can be realized via a conditional label distribution ( f ∗ ( x ) =
E [ y | x ] ), then an invariant representation Φ( x ) must satisfy the following Environment Invariance Constraint :
[Equation]
Intuitively, the representation Φ( x ) encodes features of the input x that induce the same conditional distribution over labels across each environment.
Because trivial representations such as mapping all x onto the same value may satisfy environment invariance, other objectives must be introduced to encourage the predictive utility of Φ .
Arjovsky et al.
(2019) propose IRM as a way to satisfy (EI-C ONSTR ) while achieving a good overall risk.
As a practical instantiation, the authors introduce IRMv1, a gradient-penalty regularized objective enforcing simultaneous optimality of the same classiﬁer w ◦ Φ in all environments.
1 Denoting by
R e =
E p obs ( x,y | e ) [ (cid:96) ] the per-environment risk, the objective to be minimized is
[Equation]
Krueger et al.
(2020) propose the related Risk Extrapolation (REx) principle, which dictates a stronger preference to exactly equalize R e ∀ e (e.g. by penalizing variance across e ), which is shown to improve generalization in several settings.
Fairness Early approaches to learning fair representations (Zemel et al., 2013; Edwards & Storkey, 2015;
Louizos et al., 2015; Zhang et al., 2018; Madras et al., 2018) leveraged statistical independence regularizers from domain adaptation (Ben-David et al., 2010; Ganin et al., 2016), noting that marginal or conditional independence from domain to prediction relates to the fairness notions of demographic parity
ˆ y ⊥ e (Dwork et al., 2012) and equal opportunity ˆ
y ⊥ e | y
(Hardt et al., 2016).
Recall that (EI-C ONSTR ) involves an environment-speciﬁc conditional label expectation given a data representation E
[ y | Φ( x ) = h, e ] .
Objects of this type have been closely studied in the fair machine learning literature, where e now denotes a “sensitive” indicating membership in a protected demographic group (age, race, gender, etc.), and the vector representation Φ( x ) is typically replaced by a scalar score S ( x ) ∈ R .
E [ y | S ( x ) , e ] can now be interpreted as a calibration curve that must be regulated according to some fairness constraint.
Chouldechova (2017) showed that equalizing this calibration curve across groups is often incompatible with a common fairness constraint, demographic parity, while Liu et al.
(2019) studied “group sufﬁciency’‘—satisﬁed when E [ y | S ( x ) , e ] =
E [ y | S ( x ) , e (cid:48) ] ∀ e, e (cid:48) —of classiﬁers with strongly convex losses, concluding that ERM naturally ﬁnds group sufﬁcient solutions without fairness constraints.
Because Liu et al.
(2019) consider convex losses, their theoretical results do not hold for neural network representations.
However, by noting the link between group sufﬁciency and the constraint from (EI-C ONSTR ), we observe that the IRMv1 regularizer (applicable to neural nets) in fact minimizes the group sufﬁciency gap in the case of a scalar representation Φ( x ) ⊆ R , and when e indicates sensitive group membership.
It is worth noting that Arjovsky et al.
(2019) brieﬂy discuss using groups as environments, but without specifying a particular fairness criterion.
Reliance on sensitive demographic information is cumbersome since it often cannot be collected without legal or ethical repercussions.
H ´ ebert-Johnson et al.
(2018) discussed the problem of mitigating subgroup unfairness when group labels are unknown, and proposed Multicalibration as a way of ensuring a classiﬁer’s calibration curve is invariant to efﬁciently computable environment splits.
Since the proposed algorithm requires brute force enumeration over all possible environments/groups, Kim et al.
(2019) suggested a more practical algorithm by relaxing the calibration constraint to an accuracy constraint, yielding a Multiaccurate classiﬁer.
The goal here is to boost the predictions of a pre-trained classiﬁer through multiple rounds of auditing (searching for worstcase subgroups using an auxiliary model) rather than learning an invariant representation.
A related line of work also leverages inferred subgroup information to improve worst-case model performance using the framework of DRO.
Hashimoto et al.
(2018) applied DRO to encourage long-term fairness in a dynamical setting where the average loss for a subpopulation inﬂuences their propensity to continue engaging with the model.
Lahoti et al.
(2020) proposed Adversarially
Reweighted Learning, which extends DRO using an auxiliary model to compute the importance weights γ i mentioned above.
Amortizing this computation mitigates the tendency of DRO to overﬁt its reweighting strategy to noisy outliers.
Wang et al.
(2020) proposed a group DRO method for adaptively estimating soft assignments to groups suitable for the setting where group labels are noisy.
Limitations of generalization-ﬁrst fairness One exciting direction for future work is to apply methods developed in the domain generalization literature to tasks where distribution shift is related to some societal harm that should be mitigated.
However, researchers should be wary of blind “solutionism”, which can be ineffectual or harmful when the societal context surrounding the machine learning system is ignored (Selbst et al., 2019).
Moreover, many aspects of algorithmic discrimination are not simply a matter of achieving few errors on unseen distributions.
Unfairness due to task deﬁnition or dataset collection, as discussed in the study of target variable selection by Obermeyer et al.
(2019), may not be reversible by novel algorithmic developments.
I NVARIANCE W ITHOUT D EMOGRAPHICS OR E NVIRONMENTS
In this section we draw inspiration from recent work on fair prediction without sensitive labels (Kearns et al., 2018; H´ebert-Johnson et al., 2018; Hashimoto et al., 2018;
Lahoti et al., 2020) to propose a novel domain generalization algorithm that does not require a priori domain/environment knowledge.
To motivate the study of this setting and show the fairness and invariance considerations at play, consider the task of using a high dimensional medical image x to predict a target label y ∈ { 0 , 1 } indicating the presence of COVID-19 in the imaged patient.
DeGrave et al.
(2020) describe the common use of a composite dataset for this task, where the process of aggregating data across two source hospitals e ∈ { H 1 , H 2 } leads to a brittle neural net classiﬁer
w ◦ Φ( x ) that ﬁxates on spurious low-level artifacts in x as predictive features.
Now we will consider a slightly different scenario.
Consider a single hospital serving two different demographic populations e ∈ { P 1 , P 2 } .
While P 1 has mostly sick patients at time t = 0 due to the prevalence of COVID-19 in this subpopulation, P 2 currently has mostly well patients.
Then p ( x, y | e = P 1 , t = 0) and p ( x, y | e = P 2 , t = 0) will differ considerably, and moreover a classiﬁer using a spurious feature indicative of subpopulation membership—either a low-level image artifact or attribute of the medical record—may achieve low average error on the available data.
Of course such a classiﬁer may generalize poorly.
Consider temporal distribution shifts: suppose at time t = 1 , due to the geographic density of the virus changing over time, P 1 has mostly well patients while patients from P 2 are now mostly sick.
Now the spurious classiﬁer may suffer worse-than-chance error rates and imply unfair outcomes for disadvantaged groups.
In reality the early onset and frequency of exposure to COVID-19 has been unequally distributed along many social dimensions (class, race, occupation, etc.)
that could constitute protected groups (Tai et al., 2020), raising concerns of additional algorithmic discrimination.
Learning to be invariant to spurious features encoding demographics would prevent errors due to such temporal shifts.
While loss reweighting as in DRO/ARL can upweight error cases, without an explicit invariance regularizer the model may still do best on average by making use of the spurious feature.
IRM can remove the spurious feature in this particular case, but a method for discovering environment partitions directly may occasionally be needed.
This need is clear when demographic makeup is not directly observed and a method to sort each example into the maximally separating the spurious feature, i.e. inferring populations { P 1 , P 2 } , is needed for effective invariant learning.
3. E NVIRONMENT I NFERENCE FOR I NVARIANT L EARNING
We now derive a principle for inferring environments from observational data.
Our exposition extends IRMv1 (Equation 2), but we emphasize that our method EIIL is applicable more broadly to any environment-based learning objective.
We begin by introducing u
i ( e (cid:48) )
= p obs ( e (cid:48)
| x i , y i )
= 1 ( e i =
e (cid:48) ) as an indicator of the hand-crafted environment assignment per-example.
Noting that N e := (cid:80) i u i ( e ) represents the number of examples in environment e , we can re-express this objective to make its dependence on environment labels explicit
[Equation] [Equation]
Our general strategy is to replace the binary indicator u
i ( e ) , with a probability distribution q ( e | x i , y i ) , representing a soft assignment of the i -th example to the e -th environment.
q ( e | x i , y i ) should capture worst-case environments w.r.t the invariant learning objective; rewriting q ( e | x i , y i ) as q i ( e ) for consistency with the above expression, we arrive at the following bi-level optimization:
[Equation]
We leave the full exploration of this bi-level optimization to future work, but for now propose the following practical sequential approach, which we call EIILv1 (See Appendix A for pseudocode):
Input reference model ˜Φ ; 2.
Fix Φ ← ˜Φ and fully optimize the inner loop of (EIIL) to infer environments ˜ q
i ( e ) =
˜ q ( e | x i , y i ) ; 3.
Fix q ← ˜ q and fully optimize the outer loop to yield the new model Φ .
Instead of requiring hand-crafted environments, we instead require a trained reference model ˜Φ , which is arguably easier to produce and could be found using ERM on p obs ( x, y ) , for example.
In our experiments we consider binary environments 5 and explicitly parameterize the q ( e | x, y ) as a vector of probabilities for each example in the training data.
3. A NALYZING THE EIIL S OLUTION
To characterize the ability of EIILv1 to generalize to unseen test data, we now examine the inductive bias for generalization provided by the reference model ˜Φ .
We state the main result here and defer the proofs to Appendix B. Consider a dataset with some feature(s)
z which are spurious, and other(s) v which are valuable/causal w.r.t.
the label y .
Our proof considers binary features/labels and two environments, but the same argument extends to other cases.
Our goal is to ﬁnd a model Φ whose representation Φ( v, z ) is invariant w.r.t.
z and focuses solely on v .
Theorem 1 Consider environments that differ in the degree to which the label y agrees with the spurious features z : P ( 1 ( y = z ) | e 1 ) (cid:54) =
P ( 1 ( y = z ) | e 2 ) : then a reference model
˜Φ
Spurious that is invariant to valuable features v and solely focusing on spurious features z maximally violates the Invariance Principle ( EI-C ONSTR ).
Likewise, consider the case with ﬁxed representation Φ that focuses on the spurious features: then a choice of environments that maximally violates ( EIC ONSTR ) is e
1 = { v, z, y | 1 ( y = z ) } and e 2 = { v, z, y |
( y (cid:54) =
z ) } .
If environments are split according to agreement of y and z , then the constraint from (EI-C ONSTR ) is satisﬁed by a representation that ignores z : Φ( x ) ⊥ z .
Unfortunately this requires a priori knowledge of either the spurious feature z or a reference model
˜Φ Spurious that extracts it.
When the wrong solution ˜Φ Spurious is not a priori known, it can sometimes be recovered directly from the training data; for example in CMNIST we ﬁnd that ˜Φ ERM approximates ˜Φ Color .
This allows EIIL to ﬁnd environment partitions providing the starkest possible contrast for invariant learning.
Even if environment partitions are available, it may be possible to improve performance by inferring new partitions from scratch.
It can be shown (see Appendix B.2) that the environments provided in the CMNIST dataset (Arjovsky et al., 2019) do not maximally violate (EI-C ONSTR ) for a reference model ˜Φ Color , and are thus not maximally informative for learning to ignore color.
Accordingly,
EIIL improves test accuracy for IRM compared with the hand-crafted environments (Table 1).
[Figure]
E XPERIMENTS
We defer a proof-of-concept synthetic regression experiment to Appendix E for lack of space.
We proceed with the established domain generalization benchmark ColorMNIST and then discuss a variant of the algorithmic fairness dataset UCI Adult.
We note that benchmarking model performance on a shifted test distribution without access to validation samples—especially during model selection—is a difﬁcult open problem, a solution to which is beyond the scope of this paper.
Accordingly we use the default IRM hyperparameters wherever appropriate, and otherwise follow a recently proposed model selection strategy (Gulrajani & Lopez-Paz, 2020)
(see Appendix D).
4. C OLOR MNIST
ColorMNIST (CMNIST) is a noisy digit recognition task 8 where color is a spurious feature that correlates with the label at train time but anti-correlates at test time, with the correlation strength at train time varying across two pre-speciﬁed environments (Arjovsky et al., 2019).
Crucially, label noise is applied by ﬂipping y with probability θ y ; the default setting ( θ y = 0 . 25 ) implies that shape (the correct feature) is marginally less reliable than color in the train set, so naive ERM ignores shape to focus on color and suffers from below-chance performance at test time.
We evaluated the performance of the following methods: ERM: A naive MLP that does not make use of environment labels e , but instead optimizes the average loss on the aggregated environments; IRM( e HC ): the method of Arjovsky et al.
(2019) using hand-crafted environment labels;
IRM( e EIIL ): our proposed method (a.k.a.
EIILv1) that infers useful environments (not using handcrafted environment labels)
based on the naive ERM, then applies IRM to the inferred environments.
After noting that EIILv1—denoted IRM( e EIIL ) above—outperforms IRM without access to environment labels in the default setting (See Tables 1 and 6), we examine how the various methods perform as a function of θ y .
This parameter inﬂuences the ERM solution since low θ y implies shape is more reliable than color in the aggregated training data (thus ERM generalizes well), while the opposite trend holds for high θ y .
Because EIILv1 relies on a reference model ˜Φ , its performance is also affected when ˜Φ
=
ERM (Figure 1).
We ﬁnd that IRM( e EIIL ) generalizes better than IRM( e HC ) with sufﬁciently high label noise θ y > . 2 , but generalizes poorly under low label noise.
This is precisely due to the success of ERM in this setting, where shape is a more reliable feature in the training data than color.
We verify this conclusion by evaluating IRM( e EIIL )
when ˜Φ =
Φ Color , i.e. a hand-coded color-based predictor as reference.
This does relatively well across all settings of θ y , approaching the performance of the (oracle) baseline that classiﬁes using grayscale inputs.
4. C ENSUS DATA
We now study a fair prediction problem using a variant of the UCI Adult dataset, 9 which comprises 48 , 842 individual census records collected from the United States in 1994.
The task commonly used as an algorithmic fairness benchmark is to predict a binarized income indicator (thresholded at $50 , 000 ) as the target label, possibly considering sensitive attributes such as age, sex, and race.
Because the original task measures in-distribution test performance, we instead construct a variant of this dataset suitable for measuring out-of-distribution test performance, which we call ConfoundedAdult.
[Table]
Lahoti et al.
(2020) demonstrate the beneﬁt of per-example loss reweighting on UCI Adult using their method ARL to improve predictive performance for undersampled subgroups.
Following Lahoti et al.
(2020), we consider the effect of four sensitive subgroups —deﬁned by composing binarized race and sex labels—on model performance, assuming the model does not know a priori which features are sensitive.
However, we focus on a distinct generalization problem where a pernicious dataset bias confounds the training data, making subgroup membership predictive of the label on the training data.
At test time these correlations are reversed, so a predictor that infers subgroup membership to make predictions will perform poorly at test time (see Appendix C for details).
Dwork et al.
(2012) described a similar motivating scenario where the conditional distribution mapping features to target labels varies across demographic groups due to cultural differences, so the most predictive predictor for one group may not generalize to the others.
The large distribution shift of our test set can be understood as a worst-case audit to determine whether the classiﬁer uses subgroup information in its predictions.
Using EIILv1—to ﬁrst infer worst-case environments then ensure invariance across them—performs favorably on the audit test set, compared with ARL and a baseline MLP (Table 3).
We also ﬁnd that, without access to sensitive group information, using the IRMv1 penalty on the EIIL environments improves subgroup sufﬁciency (Figure 2).
Appendix E.3 provides an ablation showing that all components of the EIILv1 approach are needed to achieve the best performance.
[Figure]
C ONCLUSION
We discussed the common goals of algorithmic fairness and domain generalization, compared related methods from each literature and suggested how lessons can be exchanged between the two ﬁelds to inform future research.
The most concrete outcome of this discussion was our novel domain generalization method, Environment Inference for Invariant Learning (EIIL).
Drawing inspiration from fairness methods that optimize worst-case performance without access to demographic information, EIIL improves the performance of IRM on CMNIST without requiring a priori knowledge of the environments.
On a variant of the UCI Adult dataset, EIIL makes use of the IRMv1 regularizer to improve group sufﬁciency—a fairness criterion previously difﬁcult to optimize for non-convex losses—without requiring knowledge of the sensitive groups.
A EIIL V 1 P SUEDOCODE
[Table]
Randomly initialize e ∈ R |D| as vectorized logit of posterior with σ ( e i ) :
= q ( e | x i , y i ) .
for n ∈ 1 . .
N steps do
[Equation] [Equation]
regularizer
[Equation]
// D2 risk
[Equation]
// D2 invariance
regularizer
[Equation]
λG e e ← OptimUpdate ( e , ∇ e L ) end ˆ e ∼ Bernoulli ( σ ( e )) ; // sample splits
D 1 ← { x i , y i
| ˆ e i = 1 } , D 2 ← {
[Table]
i ,
[Table]
i |
[Table]
e i
[Table]
0 } ; // split data
B P ROOFS B. P ROOF OF T HEOREM
Consider a dataset with some feature(s) z which are spurious, and other(s) v which are valuable/causal w.r.t.
the label y .
This includes data generated by models where v → y → z , such that P ( y | v, z ) =
P ( y | v ) .
Assume further that the observations x are functions of both spurious and valuable features: x := f ( v, z ) .
The aim of invariant learning is to form a classiﬁer that predicts y from x that focuses solely on the causal features, i.e., is invariant to z and focuses solely on v .
Consider a classiﬁer that produces a score S ( x ) for example x .
In the binary classiﬁcation setting S is analogous to the model Φ , while the score S ( x ) is analogous to the representation Φ( x ) .
To quantify the degree to which the constraint in the Invariant Principle (EI-C ONSTR ) holds, we introduce a measure called the group sufﬁciency gap 10 :
[Equation]
Now consider the notion of an environment: some setting in which the x → y relationship varies (based on spurious features).
Assume a single binary spurious feature z .
We restate Theorem 1 as follows:
Claim: If environments are deﬁned based on the agreement of the spurious feature z and the label y , then a classiﬁer that predicts based on z alone maximizes the group-sufﬁciency gap (and vice versa – if a classiﬁer predicts y directly by predicting z , then deﬁning two environments based on agreement of label and spurious feature— e 1 = { v, z, y | 1 ( y = z ) } and e 2 = { v, z, y | 1 ( y (cid:54) =
z ) } —maximizes the gap).
We can show this by ﬁrst noting that if the environment is based on spurious feature-label agreement, then with e ∈ { 0 , 1 } we have e = 1 ( y = z ) .
If the classiﬁer predicts z , i.e. S ( x ) =
z , then we have ∆
[Equation]
S, e
[Equation]
=
[Equation] [Equation]
[Equation]
x
[Equation]
( y
[Equation]
= z )]]
For each instance of x either z = 0 or z = 1 .
Now we note that when z = we have
E ( y | z,
[Equation]
y = z )
[Equation]
=and
E ( y | z, 1 ( y (cid:54) =
z ))
= 0 , while when z = 0 E ( y | z, I [ y =
=
z ]) =
and E [ y | z, 1 ( y (cid:54) =
z )] = 1 .
Therefore for each example | E ( y | z ( x ) , 1 ( y = z )) − E ( y | z ( x ) , 1 ( y (cid:54) =
z ) | = 1 , contributing to an overall ∆( S, e ) = 1 , which is the maximum value for the sufﬁciency gap.
B. G IVEN CMNIST ENVIRONMENTS ARE SUBOPTIMAL W . R . T . SUFFICIENCY GAP
The regularizer from IRMv1 encourages a representation for which sufﬁciency gap is minimized between the available environments.
Therefore when faced with a new task it is natural to measure the natural sufﬁciency gap between these environments, mediated through a naive or baseline method.
Here we show that for CMNIST, when considering a naive color-based classiﬁer as the reference model, the given environment splits are actually
suboptimal w.r.t.
sufﬁciency gap, which motivates the proposed EIIL approach for inferring environments that have a more sever sufﬁciency gap for the reference model.
We begin by computing ∆( S, e ) , the sufﬁciency gap for color-based classiﬁer g over the given train environments { e 1 , e 2 } .
We introduce an auxiliary color variable z , which is not observed but can be sampled from via the color based classiﬁer g :
[Equation]
Denote by G REEN and R ED the set of green and red images, respectively.
I.e. we have z ∈ G iff z = 1 and x ∈ G REEN iff z
( x ) = 1 .
The the sufﬁciency gap is expressed as
[Equation]
The regularizer in IRMv1 is trying to reduce the sufﬁciency gap, so in some sense we can think about this gap as a learning signal for the IRM learner.
A natural question would be whether a different set of environment partition { e } can be found such that this learning signal is stronger, i.e. the sufﬁciency gap is increased.
We ﬁnd
the answer is yes.
Consider an environment distribution q ( e | x, y, z ) that assigns each data point to one of two environments.
Any assignment sufﬁces so far as its marginal matches the observed data: (cid:82) z (cid:82) e q ( x, y, z, e ) =
p obs ( x, y ) .
We can now express the sufﬁciency gap (given a color-based classiﬁer g ) as a function of the environment assignment q :
[Equation]
Where we use the same change of variables trick as above to replace g ( x ) with samples from p ( z | x ) (note that this is the color factor from the generative process p according with our assumption that g matches this distribution).
We want to show that there exists a q yielding a higher sufﬁciency gap than the given environments.
Consider q that yields the conditional label distribution
[Equation]
This can be realized by an encoder/auditor q ( e | x, y, z ) that ignores image features in x and partitions the example based on whether or not the label y and color z agree.
We also note that z is deterministically the color of the image in the generative process: p ( z | x ) = 1 ( x = R ED )
Now we can compute the sufﬁciency gap:
[Equation]
Note that 1 is the maximal sufﬁciency gap, meaning that the described environment partition maximizes the sufﬁciency gap w.r.t.
the color-based classiﬁer g .
C D ATASET DETAILS
Constructing the ConfoundedAdult dataset To create our semi-synthetic dataset, called ConfoundedAdult, we start by observing that the conditional distribution over labels varies across the subgroups, and in some cases subgroup membership is very predictive of the target label.
We construct a test set (a.k.a.
the audit set) where this relationship between subgroups and target label is reversed.
The four sensitive subgroups are deﬁned following the procedure of Lahoti et al.
(2020), with sex (recorded as binary:
Male/Female) and binarized race (Black/non-Black) attributes compose to make four possible subgroups: Non-Black Males (SG1), Non-Black Females (S2), Black Males (SG3), and Black Females (SG4).
We start with the observation that each subgroup has a different correlation strength with the target label, and in some cases subgroup membership alone can be used to achieve relatively low error rates in prediction.
As these correlations should be considered “spurious” to mitigate unequal treatment across groups, we create a semi-synthetic variant of the UCI Adult dataset, which we call ConfoundedAdult, where these spurious correlations are exaggerated.
Table 4 shows various conditional label distributions for the original dataset and our proposed variant.
The test set for
ConfoundedAdult revereses the correlation strengths, which can be thought of as a worst-case audit to ensure the model is not relying on subgroup membership alone in its predictions.
We generate samples for ConfoundedAdult using importance sampling, keeping the original train/test splits from
UCI Adult as well as the subgroup sizes, but sampling individual examples under/over-sampled according to importance weights p ConfoundedAdult p UCIAdult .
D E XPERIMENTAL DETAILS
Model selection Krueger et al.
(2020) discussed the pitfalls of achieving good test performance on CMNIST by using test data to tune hyperparameters.
Because our primary interest is in the properties of the inferred environment rather than the ﬁnal test performance, we sidestep this issue in the Synthetic Regression and CMNIST experiments by using the default parameters of IRM without further tuning.
However for the ConfoundedAdult dataset a speciﬁc strategy for model selection is needed.
[Table]
We refer the interested reader to Gulrajani & Lopez-Paz (2020) for an extensive discussion of possible model selection strategies.
They also provide a large empirical study showing that ERM is difﬁcult baseline to beat when all methods are put on equal footing w.r.t.
model selection.
In our case, we use the most relaxed model selection method proposed by Gulrajani & Lopez-Paz (2020), which amounts to allowing each method a 20 test evaluations using hyperparameter chosen at random from a reasonable range, with the best hyperparameter setting selected for each method.
While none of the methods is given an unfair advantage in the search over hyperparameters, the basic model selection premise does not translate to real-world applications, since information about the test-time distribution is required to select hyperparameters.
Thus these results can be understood as being overly optimistic for each method, although the relative ordering between the methods can still be compared.
CMNIST IRM is trained on these two environments and tested on a holdout environment constructed from 10 , 000 test images in the same way as the training environments, where colour is predictive of the noisy label 10% of the time.
So using color as a feature to predict the label will lead to an accuracy of roughly 10% on the test environment, while it yields 80% and 90% accuracy respectively on the training environments.
To evaluate IRM( e EIIL ) we remove the environment identiﬁer from the training set and thus have one training set comprised of 50 , 000 images from both original training environments.
We then train an MLP with binary cross-entropy loss on the training environments, freeze its weights and use the obtained model to learn environment splits that maximally violate the IRM penalty.
When optimizing the inner loop of EIIL, we use Adam with learning rate 0.001 for 10 , 000 steps with full data batches used to computed gradients.
The obtained environment partitions are then used to train a new model from scratch with IRM.
Following Arjovsky et al.
(2019), we allow the representation to train for several hundred annealing steps before applying the IRMv1 penalty.
Census data Following Lahoti et al.
(2020), we use a two-hidden-layer MLP architecture for all methods, with 64 and 32 hidden units respectively, and a linear adversary for ARL.
We optimize all methods using Adagrad; learning rates, number of steps, and batch sizes chosen by the model selection strategy described above (with 20 test evaluations per method), as are penalty weights for
IRMv1 regularizer and standard weight decay.
For the inner loop of EIIL (inferring the environments), we use the same settings as in CMNIST.
We ﬁnd that the performance of EIIL is somewhat sensitive to the number of steps taken with the IRMv1 penalty applied.
To limit the number of test queries needed during model selection, we use an early stopping heuristic by enforcing the IRMv penalty only during the ﬁnal 500 steps of training, with the previous steps serving as annealing period to learn a baseline representation to be regularized.
Unlike the previous datasets, here we use minibatches to compute gradients during IRM training (for consistency with the ARL method, which uses minibatches).
However, full batch gradients are still used for inferring environments in
EIIL.
[Table]
E A DDITIONAL E MPIRICAL R ESULTS E. S YNTHETIC D ATA
We begin with a regression setting originally used as a toy dataset for evaluating IRM (Arjovsky et al., 2019).
The features x ∈ R N comprise a “causal” feature v ∈ R N/ 2 concatenated with a “non-causal” feature z ∈ R N/ 2 : x = [ v , z ] .
Noise varies across hand-crafted environments e :
[Equation]
We evaluated the performance of the following methods:
• ERM: A naive regressor that does not make use of environment labels e , but instead optimizes the average loss on the aggregated environments; • IRM( e HC ): the method of Arjovsky et al.
(2019) using hand-crafted environment labels; • ICP( e HC ): the method of Peters et al.
(2016) using hand-crafted environment labels; • IRM( e EIIL ): our proposed method (which does use hand-crafted environment labels) that infers useful environments based on the naive ERM, then applies IRM to the inferred environments
The regression methods ﬁt a scalar target y = 1 T y via a regression model
ˆ y ≈ w T x to minimize || y −
ˆ y || w.r.t.
w , plus an invariance penalty as needed.
The optimal (causally correct) solution is w ∗ =
Given a solution [ ˆ w v , ˆ w z ] from one of the methods, we report the mean squared error for the causal and non-causal dimensions as ||
ˆ w v − 1 ||
and ||
ˆ w z − 0 ||
(Table 5).
Because v is marginally noisier than z , ERM focuses on the spurious z .
IRM using hand-crafted environments, denoted IRM( e HC ), exploits variability in noise level in the non-causal feature (which depends on the variability of σ y ) to achieve lower error.
Using EIILv1 instead of hand crafted environments yields an improvement on the resulting IRM solution (denoted IRM( e EIIL )) by learning worst-case environments for invariant training.
We show in a follow-up experiment that the EIILv1 solution is indeed sensitive to the choice of reference representation, and in fact, can only discover useful environments (environments that allow
IRM( e EIIL )to learn the correct causal representation) when the reference representation encodes the incorrect inductive bias by focusing on the spurious feature.
We can explore this dependence of
EIILv1 on the mix of spurious and non-spurious features in the reference model by constructing a ˜Φ that varies in the degree it focuses on the spurious feature, according to convex mixing parameter α ∈ [0 , 1] .
α = 0 indicates focusing entirely on the correct causal feature , while α = 1 indicates focusing on the spurious feature .
We refer to this variant as IRM( e EIIL | ˜Φ =
Φ α − S PURIOUS ), and measure its performance as a function of α (Figure 3).
Environment inference only yields good testtime performance for high values of α , where the reference model captures the incorrect inductive bias.
[Figure] [Figure]
[Table]
E. C OLOR MNIST
Table 6 expands on the results from Table 1 by adding additional methods discussed in Section 3.
In Table 7 we measure the performance of some alternative strategies for optimizing the bi-level problem from Equation (EIIL).
In particular, we consider alternating updates to the representation Φ and environment assignments q , as well as solving the inner/outer loop of EIIL multiple times.
On the CMNIST dataset, none of these variants offers a performance beneﬁt above the method used in
Section 4.
EIIL loops =
k indicates that the inner and outer loops of the EIIL objective in Equation (EIIL) are successively optimized k times, with k = 1 corresponding to IRM( e EIIL ), the method studied in the main experiments section.
In other words, Φ loops =1 is solved using IRM( e EIIL ), then this representation is used as a reference classiﬁer to ﬁnd Φ loops = k +1 =
IRM ( e EIIL | ˜Φ =
Φ loops = k ) in the
[Table] [Table]
next “loop” of computation.
This also means that the training time needed is k times the training time of IRM( e EIIL ).
As we expect from our theoretical analysis, using the IRM( e EIIL ) solution as a reference classiﬁer for another round of EIIL is detrimental: since the reference classiﬁer already relies on the correct shape feature, environments that encourage invariance to this feature are found in the second round, so the EIIL loops =2 classiﬁer uses color rather than shape.
EIIL AltUpdates consists of optimizing Equation 1 using alternating steps to Φ and q .
Unforatuntely, whereas this strategy works well for other bi-level optimization problems such as GANs, it seems to do poorly in this setting.
This method outperforms ERM but does not exceed chance-level predictions on the test set.
E. C ENSUS D ATA
Ablation Here we provide an ablation study extending the results from Section 4.2 to demonstrate that both ingredients in the EIILv1 solution—ﬁnding worst-case environment splits and regularizing using the IRMv1 penalty—are necessary to achieve good test-time performance on the
ConfoundedAdult dataset.
From Lahoti et al.
we see that ARL can perform favorably compared with DRO (Hashimoto et al., 2018) in adaptively computing how much each example should contribute to the overall loss, i.e. computing the per-example γ i in C =
E x i ,y i ∼ p [ γ i (cid:96) (Φ( x i ) , y i )] .
Because all per-environment risks in IRM are weighted equally (see Equation 2), and each per-environment risk comprises an average across per-example losses within the environment, each example contributes its loss to the overall objective in accordance with the size of its assigned environment.
For example with two environments e 1 and e 2 of sizes | e 1 | and | e 2 | , we implicitly have the per-example weights of γ i =| e 1 | for i ∈
e 1 and γ i = 1 | e 2 | for i ∈ e 2 , indicating that examples in the smaller environment count more towards the overall objective.
Because EIILv1 is known to discover worst-case environments of unequal sizes, we measure the performance of EIILv1 using only this reweighting, without adding the gradient-norm penalty typically used in IRM (i.e. setting λ = 0 ).
To determine the beneﬁt of worst-case environment discovery, we also measure IRM with random assignment of environments.
Table 8 shows the results, conﬁrming that both ingredients are required to attain good performance using EIILv1.
F A DDITIONAL T HEORETICAL R ESULTS F. O PTIMAL SOFT PARTITIONS MAXIMIMALLY VIOLATE THE INVARIANCE PRINCPLE
We want to show that ﬁnding environment assignments that maximize the violation of the softened version of the regularizer from Equation 3 also maximally violates the invariance princple.
Because the invaraince principle E [ Y | Φ( X ) , e ] =
E
[ Y | Φ( X ) , e (cid:48) ] ∀ e, e (cid:48) is difﬁcult to quantify for continuous Φ( X ) , we consider a binned version of the representation, with b denoting the discrete index of the bin in representation space.
Let q i ∈
[0 , 1] denote the soft assignment of example i to environment 1, and 1 − q i denote its converse, the assignment of example i to environment 2.
Denote by y i ∈ { 0 , 1 } the binary target for example i , and ˆ y ∈ [0 , 1] as the model prediction on this example.
Assume that (cid:96) represents a cross entropy or squared error loss so that ∇ w (cid:96) (ˆ y, y ) = (ˆ y − y ) Φ( x ) .
Consider the IRMv1 regularizer with soft assignment, expressed as
[Equation]
Now consider that the space of Φ( X ) is discretized into disjoint bins b over its support, using z i,b ∈ { 0 , 1 } to indicate whether example i falls into bin b according to its mapping Φ( x i ) .
Thus we have
D ( q ) = (cid:88)
[Equation]
The important point is that within a bin, all examples have roughly the same Φ( x i ) value
[Equation]
and the same value for ˆ y i as well.
So denoting K (1) b
[Equation]
as the relevant constant within-bin summations
[Equation]
we have the following objective to be maximized by
EIIL:
[Equation]
One way to maximize this is to assign all y i
= 1 values to environment 1 ( q i = 1 for these examples) and all y i = 0 to the other environment ( q i = 0) .
We can show this is maximized by considering all of the examples except the i -th one have been assigned this way, and then that the loss is maximized by assigneing the i -th example according to this rule.
Now we want to show that the same assignment maximially violates the invariance principle (showing that this soft EIIL solution provides maximal non-invariance).
Intuitively within each bin the difference between E [ y | e = 1] and E [ y | e = 2] is maximized (within the bin)
if one of these expected label distributions is 1 while the other is 0 .
This can be achieved by assigning all the y i
=values to the ﬁrst environment and the y i = 0 values to the second.
Thus a global optimum for the relaxed version of EIIL (using the IRMv1 regularizer) also maximally violates the invariance principle.